@@ -139,7 +139,39 @@ def test_add_file_with_metadata(self):
139139 doc = cursor .fetchone ()
140140 assert doc
141141 assert doc [0 ] == "This is a test document with metadata."
142- assert doc [1 ] == json .dumps (metadata )
142+ assert doc [1 ] == json .dumps (
143+ {
144+ ** metadata ,
145+ "generated" : {"title" : "This is a test document with metadata." },
146+ }
147+ )
148+
149+ def test_add_documents_with_generated_title (self ):
150+ with tempfile .NamedTemporaryFile (mode = "w" , suffix = ".txt" , delete = False ) as doc1 :
151+ doc1 .write ("# Title 1\n This is the second test document." )
152+ with tempfile .NamedTemporaryFile (mode = "w" , suffix = ".txt" , delete = False ) as doc2 :
153+ doc2 .write ("# Title 2\n This is the second test document." )
154+
155+ doc3 = "# Title 3\n This is the third test document."
156+ doc4 = "# Title 4\n This is the fourth test document."
157+
158+ rag = SQLiteRag .create (db_path = ":memory:" )
159+
160+ rag .add (doc1 .name )
161+ rag .add (doc2 .name )
162+ rag .add_text (doc3 )
163+ rag .add_text (doc4 )
164+
165+ conn = rag ._conn
166+ cursor = conn .execute ("SELECT metadata FROM documents" )
167+ docs = cursor .fetchall ()
168+ assert len (docs ) == 4
169+
170+ titles = [json .loads (doc [0 ]).get ("generated" , {}).get ("title" ) for doc in docs ]
171+ assert "Title 1" in titles
172+ assert "Title 2" in titles
173+ assert "Title 3" in titles
174+ assert "Title 4" in titles
143175
144176 def test_add_empty_file (self ):
145177 with tempfile .NamedTemporaryFile (mode = "w" , suffix = ".txt" , delete = False ) as f :
@@ -229,7 +261,14 @@ def test_add_text_with_metadata(self):
229261 assert doc
230262 assert doc [0 ] == "This is a test document content with metadata."
231263 assert doc [1 ] == "test_doc_with_metadata.txt"
232- assert doc [2 ] == json .dumps (metadata )
264+ assert doc [2 ] == json .dumps (
265+ {
266+ ** metadata ,
267+ "generated" : {
268+ "title" : "This is a test document content with metadata."
269+ },
270+ }
271+ )
233272
234273 def test_list_documents (self ):
235274 rag = SQLiteRag .create (":memory:" )
0 commit comments