core/rag/store.py

import chromadb
from chromadb.config import Settings
import os

# Path to save the database inside ai-service/data
BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
DB_PATH = os.path.join(BASE_DIR, "data", "chroma_db")

class VectorStore:
    def __init__(self, collection_name="platform_knowledge"):
        """Initialize persistent ChromaDB client."""
        # Folder create agar nahi hai to
        os.makedirs(DB_PATH, exist_ok=True)
        
        self.client = chromadb.PersistentClient(path=DB_PATH)
        
        # Create or get collection
        self.collection = self.client.get_or_create_collection(name=collection_name)

    def add_text(self, text_chunks, metadatas, ids):
        """Text data ko DB mein save karna."""
        try:
            self.collection.upsert(
                documents=text_chunks,
                metadatas=metadatas,
                ids=ids
            )
            return True
        except Exception as e:
            print(f"[RAG Error] Failed to add text: {str(e)}")
            return False

    def search(self, query, n_results=2):
        """Question ke hisaab se matching data lana."""
        try:
            results = self.collection.query(
                query_texts=[query],
                n_results=n_results
            )
            # Thoda sa formatting taaki clean data mile
            if results['documents']:
                return results['documents'][0] # Return list of matching texts
            return []
        except Exception as e:
            print(f"[RAG Error] Search failed: {str(e)}")
            return []