agentic_workflow/retriever.py

# from chroma.chroma_client import knowledge_col, corrections_col
# from chatbot_embed import embed
# from agentic_workflow.config import TOP_K
from sentence_transformers import SentenceTransformer
import chromadb
from pathlib import Path

base_dir = Path(__file__).resolve().parent.parent
embed_model = SentenceTransformer("multi-qa-mpnet-base-dot-v1")
db_path = base_dir / "chroma" / "chroma_db"
client     = chromadb.PersistentClient(path= str(db_path))

print(client.list_collections()) 

def retrieve_from_vector_db(query: str, collection_name: str = "bank_faq", k: int = 3) -> list[dict]:
    """
    Embeds the query using the same model used during ingestion,
    queries ChromaDB, and returns chunks with their similarity scores.
    """

    collection = client.get_collection(name=collection_name)

    # Embed the query with the same model used at ingestion
    query_embedding = embed_model.encode(query).tolist()

    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=k,
        include=["documents", "metadatas", "distances"]
    )

    # ChromaDB returns distances (lower = more similar), convert to similarity score
    chunks = []
    for doc, metadata, distance in zip(
        results["documents"][0],
        results["metadatas"][0],
        results["distances"][0]
    ):
        chunks.append({
            "text": doc,
            "metadata": metadata,
            "similarity_score": round(1 - distance, 4)   # convert distance → similarity
        })

    return chunks

results = retrieve_from_vector_db(query="What are the requirements to open a bank account?", k=3)

# print("Test 1: ")
# for i, chunk in enumerate(results, 1):
#     print(f"  Chunk {i}:")
#     print(f"    Score:    {chunk['similarity_score']}")
#     print(f"    Metadata: {chunk['metadata']}")
#     print(f"    Text:     {chunk['text'][:100]}...")
# print()