Ameer1606/PostgresPro-Support

1from qdrant_client import QdrantClient2from langchain_community.embeddings import HuggingFaceEmbeddings3from langchain_qdrant import QdrantVectorStore4from src.config import QDRANT_COLLECTION_NAME, EMBEDDING_MODEL, QDRANT_URL, QDRANT_API_KEY5 6class PostgresRetriever:7    def __init__(self):8        self.embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)9        self.client = client = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)10        11        # Check if collection exists to avoid errors12        try:13            collections = [c.name for c in self.client.get_collections().collections]14            if QDRANT_COLLECTION_NAME not in collections:15                print("Warning: Vector database collection not found. Please run ingest.py first.")16        except Exception as e:17            print(f"Warning: Could not connect to Qdrant: {e}")18 19        self.vector_store = QdrantVectorStore(20            client=self.client,21            collection_name=QDRANT_COLLECTION_NAME,22            embedding=self.embeddings,23        )24        # return top 5 chunks25        self.retriever = self.vector_store.as_retriever(search_kwargs={"k": 5})26 27    def get_relevant_context(self, query: str) -> str:28        """Fetches relevant documents and formats them with explicit citations."""29        try:30            docs = self.retriever.invoke(query)31            if not docs:32                return "No relevant documentation found."33            34            formatted_context = []35            for i, doc in enumerate(docs):36                source = doc.metadata.get("source", "Unknown Source")37                page = doc.metadata.get("page", "Unknown Page")38                39                # Format required by citation guardrails40                citation = f"Source: {source}, Page: {page}"41                chunk_text = f"[Document {i+1}] ({citation})\n{doc.page_content}"42                formatted_context.append(chunk_text)43                44            return "\n\n".join(formatted_context)45        except Exception as e:46            return f"Error retrieving context: {str(e)}"47