Update retrieval mode

backend/server/retrieval.py

@@ -69,4 +69,4 @@ async def extract_from_content(
             "instructions": extractor.instruction,
         }
     )
-    return ExtractResponse(extracted=[result.extracted])
+    return result

backend/tests/unit_tests/api/test_api_extract.py

@@ -4,6 +4,7 @@
 from uuid import UUID
 
 from langchain.text_splitter import CharacterTextSplitter
+from langchain_community.embeddings import FakeEmbeddings
 from langchain_core.runnables import RunnableLambda
 
 from tests.db import get_async_client
@@ -23,11 +24,20 @@ def mock_text_splitter(*args, **kwargs):
     return CharacterTextSplitter()
 
 
+def mock_embeddings(*args, **kwargs):
+    return FakeEmbeddings(size=10)
+
+
 @patch(
     "server.extraction_runnable.extraction_runnable",
     new=RunnableLambda(mock_extraction_runnable),
 )
+@patch(
+    "server.retrieval.extraction_runnable",
+    new=RunnableLambda(mock_extraction_runnable),
+)
 @patch("server.extraction_runnable.TokenTextSplitter", mock_text_splitter)
+@patch("server.retrieval.OpenAIEmbeddings", mock_embeddings)
 async def test_extract_from_file() -> None:
     """Test extract from file API."""
     async with get_async_client() as client:
@@ -67,6 +77,18 @@ async def test_extract_from_file() -> None:
         assert response.status_code == 200
         assert response.json() == {"data": ["Test Conte"]}
 
+        # Test retrieval
+        response = await client.post(
+            "/extract",
+            data={
+                "extractor_id": extractor_id,
+                "text": "Test Content",
+                "mode": "retrieval",
+            },
+        )
+        assert response.status_code == 200
+        assert response.json() == {"data": ["Test Conte"]}
+
         # We'll use multi-form data here.
         # Create a named temporary file
         with tempfile.NamedTemporaryFile(mode="w+t", delete=False) as f: