| from typing import Any, Dict, List | |
| from haystack.schema import Document | |
| from fastrag.retrievers import QuantizedBiEncoderRetriever | |
| class EndpointHandler: | |
| def __init__(self, path=""): | |
| model_id = "Intel/bge-small-en-v1.5-rag-int8-static" | |
| self.retriever = QuantizedBiEncoderRetriever(embedding_model=model_id) | |
| def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: | |
| queries = data.get("queries", None) | |
| documents = data.get("documents", None) | |
| 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) | |
| 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") | |