File tree Expand file tree Collapse file tree
Expand file tree Collapse file tree Original file line number Diff line number Diff line change 22
33# SQLite RAG
44
5- [ ![ Run Tests] ( https://github.com/sqliteai/sqlite-rag/actions/workflows/test.yaml/badge.svg?branch=main&event=release )] ( https://github.com/sqliteai/sqlite-rag/actions/workflows/test.yaml )
5+ [ ![ Run Tests] ( https://github.com/sqliteai/sqlite-rag/actions/workflows/test.yaml/badge.svg )] ( https://github.com/sqliteai/sqlite-rag/actions/workflows/test.yaml )
66[ ![ codecov] ( https://codecov.io/github/sqliteai/sqlite-rag/graph/badge.svg?token=30KYPY7864 )] ( https://codecov.io/github/sqliteai/sqlite-rag )
77![ PyPI - Version] ( https://img.shields.io/pypi/v/sqlite-rag?link=https%3A%2F%2Fpypi.org%2Fproject%2Fsqlite-rag%2F )
88![ PyPI - Python Version] ( https://img.shields.io/pypi/pyversions/sqlite-rag?link=https%3A%2F%2Fpypi.org%2Fproject%2Fsqlite-rag )
Original file line number Diff line number Diff line change @@ -17,7 +17,6 @@ classifiers = [
1717 " Operating System :: OS Independent" ,
1818]
1919dependencies = [
20- " attrs" ,
2120 " typer" ,
2221 " huggingface_hub[hf_transfer]" ,
2322 " markitdown[docx]" ,
Original file line number Diff line number Diff line change 1- from attr import dataclass
1+ from dataclasses import dataclass
22
33
44@dataclass
Original file line number Diff line number Diff line change 11import hashlib
22import re
3+ from dataclasses import dataclass , field
34from datetime import datetime
45from typing import Optional
56
6- from attr import dataclass
7-
87from .chunk import Chunk
98
109
@@ -15,11 +14,11 @@ class Document:
1514 id : str | None = None
1615 content : str = ""
1716 uri : str | None = None
18- metadata : dict = {}
17+ metadata : dict = field ( default_factory = dict )
1918 created_at : datetime | None = None
2019 updated_at : datetime | None = None
2120
22- chunks : list ["Chunk" ] = []
21+ chunks : list ["Chunk" ] = field ( default_factory = list )
2322
2423 def hash (self ) -> str :
2524 """Generate a hash for the document content using SHA-3 for maximum collision resistance"""
Original file line number Diff line number Diff line change @@ -103,7 +103,7 @@ def add(
103103 if use_relative_paths
104104 else str (file_path .absolute ())
105105 )
106- document = Document (content = content , uri = uri , metadata = metadata )
106+ document = Document (content = content , uri = uri , metadata = metadata . copy () )
107107
108108 exists = self ._repository .document_exists_by_hash (document .hash ())
109109 if exists :
@@ -132,7 +132,7 @@ def add_text(
132132 """Add a text content into the database"""
133133 self ._ensure_initialized ()
134134
135- document = Document (content = text , uri = uri , metadata = metadata )
135+ document = Document (content = text , uri = uri , metadata = metadata . copy () )
136136
137137 self ._engine .create_new_context ()
138138 document = self ._engine .process (document )
Original file line number Diff line number Diff line change @@ -141,6 +141,33 @@ def test_add_file_with_metadata(self):
141141 assert doc [0 ] == "This is a test document with metadata."
142142 assert doc [1 ] == json .dumps (metadata )
143143
144+ def test_add_documents_with_generated_title (self ):
145+ with tempfile .NamedTemporaryFile (mode = "w" , suffix = ".txt" , delete = False ) as doc1 :
146+ doc1 .write ("# Title 1\n This is the second test document." )
147+ with tempfile .NamedTemporaryFile (mode = "w" , suffix = ".txt" , delete = False ) as doc2 :
148+ doc2 .write ("# Title 2\n This is the second test document." )
149+
150+ doc3 = "# Title 3\n This is the third test document."
151+ doc4 = "# Title 4\n This is the fourth test document."
152+
153+ rag = SQLiteRag .create (db_path = ":memory:" )
154+
155+ rag .add (doc1 .name )
156+ rag .add (doc2 .name )
157+ rag .add_text (doc3 )
158+ rag .add_text (doc4 )
159+
160+ conn = rag ._conn
161+ cursor = conn .execute ("SELECT metadata FROM documents" )
162+ docs = cursor .fetchall ()
163+ assert len (docs ) == 4
164+
165+ titles = [json .loads (doc [0 ]).get ("generated" , {}).get ("title" ) for doc in docs ]
166+ assert "Title 1" in titles
167+ assert "Title 2" in titles
168+ assert "Title 3" in titles
169+ assert "Title 4" in titles
170+
144171 def test_add_empty_file (self ):
145172 with tempfile .NamedTemporaryFile (mode = "w" , suffix = ".txt" , delete = False ) as f :
146173 temp_file_path = f .name
You can’t perform that action at this time.
0 commit comments