vector_db.py

import os
import chromadb
from chromadb.config import Settings
from langchain_text_splitters import RecursiveCharacterTextSplitter

PERSIST_DIRECTORY = "./chroma_db"
COLLECTION_NAME = "knowledge_base"

_client = None
_collection = None

def get_client():
    global _client
    if _client is None:
        _client = chromadb.PersistentClient(path=PERSIST_DIRECTORY)
    return _client

def get_collection():
    global _collection
    if _collection is None:
        client = get_client()
        _collection = client.get_or_create_collection(
            name=COLLECTION_NAME,
            metadata={"hnsw:space": "cosine"}
        )
    return _collection

def delete_by_source_id(source_id: str) -> int:
    collection = get_collection()
    existing = collection.get(where={"source_id": source_id})
    count = len(existing["ids"]) if existing["ids"] else 0
    if count > 0:
        collection.delete(where={"source_id": source_id})
        print(f"Deleted {count} chunks for source_id: {source_id}")
    return count

def add_texts_to_db(texts: list, metadatas: list, source_id: str) -> int:
    collection = get_collection()
    
    existing = collection.get(where={"source_id": source_id})
    if existing["ids"]:
        print(f"Existing version found. Deleting {len(existing['ids'])} old vectors...")
        collection.delete(where={"source_id": source_id})
    
    if not texts:
        return 0
    
    import uuid
    ids = [str(uuid.uuid4()) for _ in texts]
    
    collection.add(
        documents=texts,
        metadatas=metadatas,
        ids=ids
    )
    
    return len(texts)

def get_all_source_ids() -> list:
    collection = get_collection()
    try:
        all_data = collection.get()
        if not all_data["metadatas"]:
            return []
        
        source_ids = set()
        for meta in all_data["metadatas"]:
            if meta and "source_id" in meta:
                source_ids.add(meta["source_id"])
        
        return sorted(list(source_ids))
    except Exception as e:
        print(f"Error fetching source IDs: {e}")
        return []

def get_source_metadata(source_id: str) -> dict:
    collection = get_collection()
    try:
        data = collection.get(where={"source_id": source_id})
        if data["metadatas"] and len(data["metadatas"]) > 0:
            meta = data["metadatas"][0]
            return {
                "source_id": source_id,
                "product": meta.get("product", "Unknown"),
                "feature": meta.get("feature", "Unknown"),
                "access_level": meta.get("access_level", "Unknown"),
                "chunk_count": len(data["ids"])
            }
    except Exception as e:
        print(f"Error fetching metadata for {source_id}: {e}")
    return {"source_id": source_id, "chunk_count": 0}

def query_by_access_level(query_text: str, access_level: str, n_results: int = 5) -> dict:
    collection = get_collection()
    try:
        results = collection.query(
            query_texts=[query_text],
            where={"access_level": access_level},
            n_results=n_results
        )
        return results
    except Exception as e:
        print(f"Error querying: {e}")
        return {"documents": [], "metadatas": [], "ids": []}

def query_with_filters(query_text: str, role: str, product: str = None, n_results: int = 5) -> dict:
    collection = get_collection()
    try:
        filters = []
        
        if role == "EXTERNAL":
            filters.append({"access_level": "EXTERNAL"})
        
        if product and product != "All Products":
            filters.append({"product": product})
        
        where_clause = None
        if len(filters) == 1:
            where_clause = filters[0]
        elif len(filters) > 1:
            where_clause = {"$and": filters}
        
        if where_clause:
            results = collection.query(
                query_texts=[query_text],
                where=where_clause,
                n_results=n_results
            )
        else:
            results = collection.query(
                query_texts=[query_text],
                n_results=n_results
            )
        
        return results
    except Exception as e:
        print(f"Error querying with filters: {e}")
        return {"documents": [], "metadatas": [], "ids": []}

def get_collection_count() -> int:
    collection = get_collection()
    return collection.count()

def add_manual_knowledge(question: str, answer: str, product: str, access_level: str) -> bool:
    """Add a manual Q&A fix to the knowledge base"""
    import time
    import uuid
    
    collection = get_collection()
    
    content = f"Q: {question}\n\nA: {answer}"
    source_id = f"manual_fix_{int(time.time())}"
    
    metadata = {
        "source_id": source_id,
        "type": "manual_fix",
        "product": product,
        "access_level": access_level,
        "feature": "Manual Fix"
    }
    
    try:
        collection.add(
            documents=[content],
            metadatas=[metadata],
            ids=[str(uuid.uuid4())]
        )
        return True
    except Exception as e:
        print(f"Error adding manual knowledge: {e}")
        return False

def get_all_manual_fixes():
    """Get all manual fixes from the knowledge base as a DataFrame"""
    import pandas as pd
    
    collection = get_collection()
    
    try:
        results = collection.get(where={"type": "manual_fix"})
        
        if not results["ids"]:
            return pd.DataFrame()
        
        data = []
        for i, doc_id in enumerate(results["ids"]):
            meta = results["metadatas"][i] if results["metadatas"] else {}
            doc = results["documents"][i] if results["documents"] else ""
            
            data.append({
                "ID": meta.get("source_id", doc_id),
                "Product": meta.get("product", "N/A"),
                "Access Level": meta.get("access_level", "N/A"),
                "Content": doc[:80] + "..." if len(doc) > 80 else doc
            })
        
        return pd.DataFrame(data)
    except Exception as e:
        print(f"Error getting manual fixes: {e}")
        return pd.DataFrame()

def delete_manual_knowledge(source_id: str) -> bool:
    """Delete a manual fix from the knowledge base"""
    collection = get_collection()
    
    try:
        collection.delete(where={"source_id": source_id})
        return True
    except Exception as e:
        print(f"Error deleting manual knowledge: {e}")
        return False

def chunk_text(text: str, chunk_size: int = 1000, chunk_overlap: int = 200) -> list:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=len,
        separators=["\n\n", "\n", " ", ""]
    )
    return splitter.split_text(text)

def ingest_text(
    text: str,
    source_id: str,
    product: str,
    feature: str,
    access_level: str,
    chunk_size: int = 1000
) -> int:
    chunks = chunk_text(text, chunk_size=chunk_size)
    
    metadatas = [
        {
            "source_id": source_id,
            "product": product,
            "feature": feature,
            "access_level": access_level
        }
        for _ in chunks
    ]
    
    count = add_texts_to_db(chunks, metadatas, source_id)
    return count