Update handler.py

handler.py

@@ -1,5 +1,6 @@
 from typing import Any, Dict, List
 
+from haystack.schema import Document
 from fastrag.retrievers import QuantizedBiEncoderRetriever
 
 
@@ -9,9 +10,23 @@ class EndpointHandler:
         self.retriever = QuantizedBiEncoderRetriever(embedding_model=model_id)
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        queries = data.get("queries")
+        queries = data.get("queries", None)
+        documents = data.get("documents", None)
 
-        assert isinstance(queries, list), "Expected a list of queries"
-        assert all(isinstance(query, str) for query in queries), "Expected a list of strings"
+        if queries is not None:
+            assert isinstance(queries, list), "Expected queries to be a list"
+            assert all(
+                isinstance(query, str) for query in queries
+            ), "Expected each query in queries to be a string"
 
-        return self.retriever.embed_queries(queries=queries)
+            return self.retriever.embed_queries(queries=queries)
+        elif documents is not None:
+            assert isinstance(documents, list), "Expected documents to be a list"
+            assert all(
+                isinstance(document, dict) for document in documents
+            ), "Expected each document in documents to be a dictionary"
+
+            documents = [Document.from_dict(document) for document in documents]
+            return self.retriever.embed_documents(documents=documents)
+        else:
+            raise ValueError("Expected either queries or documents")