File size: 1,902 Bytes
b22c324 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | # 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()
|