🚅 refactor: Implement Async Support for Embedding Routes (#160)

Author: sbruel

app/routes/document_routes.py

@@ -309,11 +309,13 @@ async def embed_local_file(
         user_id = entity_id if entity_id else request.state.user.get("id")
 
     try:
-        loader, known_type = get_loader(
+        loader, known_type, file_ext = get_loader(
             document.filename, document.file_content_type, document.filepath
         )
-        data = loader.load()
-        result = await store_data_in_vector_db(data, document.file_id, user_id)
+        data = await run_in_executor(None, loader.load)
+        result = await store_data_in_vector_db(
+            data, document.file_id, user_id, clean_content=file_ext == "pdf"
+        )
 
         if result:
             return {
@@ -388,7 +390,7 @@ async def embed_file(
         loader, known_type, file_ext = get_loader(
             file.filename, file.content_type, temp_file_path
         )
-        data = loader.load()
+        data = await run_in_executor(None, loader.load)
         result = await store_data_in_vector_db(
             data=data, file_id=file_id, user_id=user_id, clean_content=file_ext == "pdf"
         )
@@ -519,12 +521,14 @@ async def embed_file_upload(
         )
 
     try:
-        loader, known_type = get_loader(
+        loader, known_type, file_ext = get_loader(
             uploaded_file.filename, uploaded_file.content_type, temp_file_path
         )
 
-        data = loader.load()
-        result = await store_data_in_vector_db(data, file_id, user_id)
+        data = await run_in_executor(None, loader.load)
+        result = await store_data_in_vector_db(
+            data, file_id, user_id, clean_content=file_ext == "pdf"
+        )
 
         if not result:
             raise HTTPException(

tests/test_main.py

@@ -116,17 +116,6 @@ def test_embed_local_file(tmp_path, auth_headers, monkeypatch):
     test_file = tmp_path / "test.txt"
     test_file.write_text("This is a test document.")
 
-    # Override get_loader on the document_routes module so that /local/embed
-    # returns exactly 2 values (loader, known_type) and a loader that returns Document objects.
-    from app.routes import document_routes
-    def dummy_get_loader(filename, file_content_type, filepath):
-        class DummyLoader:
-            def load(self):
-                # Return a list of Document objects.
-                return [Document(page_content="Dummy document", metadata={})]
-        return DummyLoader(), True
-    monkeypatch.setattr(document_routes, "get_loader", dummy_get_loader)
-
     data = {
         "filepath": str(test_file),
         "filename": "test.txt",
@@ -166,16 +155,6 @@ def test_embed_file_upload(tmp_path, auth_headers, monkeypatch):
     test_file = tmp_path / "upload_test.txt"
     test_file.write_text(file_content)
 
-    # Override get_loader on the document_routes module so that /embed-upload
-    # returns exactly 2 values and a loader that returns Document objects.
-    from app.routes import document_routes
-    def dummy_get_loader(filename, file_content_type, filepath):
-        class DummyLoader:
-            def load(self):
-                return [Document(page_content="Dummy document", metadata={})]
-        return DummyLoader(), True
-    monkeypatch.setattr(document_routes, "get_loader", dummy_get_loader)
-
     with test_file.open("rb") as f:
         response = client.post(
             "/embed-upload",
@@ -200,4 +179,4 @@ def test_query_multiple(auth_headers):
     assert isinstance(json_data, list)
     if json_data:
         doc = json_data[0][0]
-        assert doc["page_content"] == "Queried content"
+        assert doc["page_content"] == "Queried content"