mtyrrell/chatfed_retriever

1from abc import ABC, abstractmethod2from typing import List, Dict, Any, Optional3from gradio_client import Client4import logging5import os6import time7 8class VectorStoreInterface(ABC):9    """Abstract interface for different vector store implementations."""10    11    @abstractmethod12    def search(self, query: str, top_k: int, **kwargs) -> List[Dict[str, Any]]:13        """Search for similar documents."""14        pass15 16class HuggingFaceSpacesVectorStore(VectorStoreInterface):17    """Vector store implementation for Hugging Face Spaces with MCP endpoints."""18    19    def __init__(self, space_url: str, collection_name: str, hf_token: Optional[str] = None):20        token = os.getenv("HF_TOKEN")21        repo_id = space_url22            23        logging.info(f"Connecting to Hugging Face Space: {repo_id}")24        25        if token:26            self.client = Client(repo_id, hf_token=token)27        else:28            self.client = Client(repo_id)29            30        self.collection_name = collection_name31        32    def search(self, query: str, top_k: int, **kwargs) -> List[Dict[str, Any]]:33        """Search using Hugging Face Spaces MCP API."""34        try:35            # Use the /search_text endpoint as documented in the API36            result = self.client.predict(37                query=query,38                collection_name=self.collection_name,39                model_name=kwargs.get('model_name'),40                top_k=top_k,41                api_name="/search_text"42            )43            44            logging.info(f"Successfully retrieved {len(result) if result else 0} documents")45            return result46            47        except Exception as e:48            logging.error(f"Error searching Hugging Face Spaces: {str(e)}")49            raise e50 51# class QdrantVectorStore(VectorStoreInterface):52#     """Vector store implementation for direct Qdrant connection."""53#     # needs to be generalized for other vector stores (or add a new class for each vector store)54#     def __init__(self, host: str, port: int, collection_name: str, api_key: Optional[str] = None):55#         from qdrant_client import QdrantClient56#         from langchain_community.vectorstores import Qdrant57        58#         self.client = QdrantClient(59#             host=host,60#             port=port,61#             api_key=api_key62#         )63#         self.collection_name = collection_name64#          # Embedding model not implemented 65        66#     def search(self, query: str, top_k: int, **kwargs) -> List[Dict[str, Any]]:67#         """Search using direct Qdrant connection."""68#         # Embedding model not implemented 69#         raise NotImplementedError("Direct Qdrant search needs embedding model configuration")70 71def create_vectorstore(config: Any) -> VectorStoreInterface:72    """Factory function to create appropriate vector store based on configuration."""73    vectorstore_type = config.get("vectorstore", "TYPE")74    75    if vectorstore_type.lower() == "huggingface_spaces":76        space_url = config.get("vectorstore", "SPACE_URL")77        collection_name = config.get("vectorstore", "COLLECTION_NAME")78        hf_token = config.get("vectorstore", "HF_TOKEN", fallback=None)79        return HuggingFaceSpacesVectorStore(space_url, collection_name, hf_token)80    81    elif vectorstore_type.lower() == "qdrant":82        host = config.get("vectorstore", "HOST")83        port = int(config.get("vectorstore", "PORT"))84        collection_name = config.get("vectorstore", "COLLECTION_NAME")85        api_key = config.get("vectorstore", "API_KEY", fallback=None)86        return QdrantVectorStore(host, port, collection_name, api_key)87    88    else:89        raise ValueError(f"Unsupported vector store type: {vectorstore_type}")