added vectordb

.gitignore

@@ -3,5 +3,5 @@ __pycache__/
 .venv/
 
 models/
-chroma/
+
 *.sqlite3

chroma/chatbot_chroma_db.py

@@ -0,0 +1,186 @@
+import chromadb
+from sklearn.cluster import KMeans
+import pandas as pd
+import json
+import ast
+from itertools import groupby
+import uuid
+from datetime import datetime
+from pathlib import Path
+
+root_dir = Path(__file__).parent.parent
+data_dir = root_dir / "datasets"
+
+embed_faq_df = pd.read_csv(data_dir / "processed" / "bank_faq" / "embed_faq.csv")
+
+# chroma db
+client = chromadb.PersistentClient(path="./chroma_db")
+
+# client.delete_collection("bank_faq")
+# client.delete_collection("chat_history")
+
+# create collection
+collection = client.get_or_create_collection(
+    name="bank_faq",
+    metadata={"hnsw:space": "cosine"}  # cosine similarity for sentence transformers
+)
+
+embed_faq_df["embedding"] = embed_faq_df["embedding"].apply(ast.literal_eval)
+# prepare data
+# ids        = [str(i) for i in range(len(embed_faq_df))]
+ids = [
+    f"doc_{row['source_id']}_chunk_{row['chunk_num']}"
+    for _, row in embed_faq_df.iterrows()
+]
+embeddings = embed_faq_df["embedding"].tolist()
+documents  = embed_faq_df["text"].tolist()
+metadatas  = [
+    {
+        "doc_id": f"doc_{row['source_id']}_chunk_{row['chunk_num']}",
+        "domain": row["domain"],
+        "source_id": row["source_id"],
+        "chunk_num": row["chunk_num"],
+        "question": row["question"]
+    }
+    for _, row in embed_faq_df.iterrows()
+]
+
+# print(ids)
+# print(type(embeddings[0]))
+# print(documents)
+# print(metadatas)
+
+# store
+collection.add(
+    ids=ids,
+    embeddings=embeddings,
+    documents=documents,
+    metadatas=metadatas
+)
+
+print(f"Stored {collection.count()} documents")
+
+faq_collection = client.get_collection(name="bank_faq")
+results = faq_collection.get(
+    include=["documents", "metadatas"]
+)
+
+source_id_to_doc_ids = {}
+for doc_id, meta in zip(ids, metadatas):
+    source_id = meta["source_id"]
+    if source_id not in source_id_to_doc_ids:
+        source_id_to_doc_ids[source_id] = []
+    source_id_to_doc_ids[source_id].append(doc_id)
+
+# build a dict grouped by source_id
+groups = {}
+for doc_id, meta, text in zip(results["ids"], results["metadatas"], results["documents"]):
+    source_id = meta["source_id"]
+    if source_id not in groups:
+        groups[source_id] = []
+    groups[source_id].append({
+        "doc_id":    doc_id,
+        "chunk_num": int(meta["chunk_num"]),
+        "domain":    meta["domain"],
+        "question":  meta["question"],
+        "text":      text.split("A:", 1)[1].strip() if "A:" in text else text,
+    })
+
+# build viewable
+viewable = []
+for source_id, chunks in groups.items():
+    chunks = sorted(chunks, key=lambda x: x["chunk_num"])  # sort by chunk order
+    
+    entry = {
+        "id":       source_id,
+        "domain":   chunks[0]["domain"],
+        "question": chunks[0]["question"],
+    }
+
+    if len(chunks) == 1:
+        entry["answer"] = chunks[0]["text"]
+    else:
+        entry["answer"] = {"chunks": [{"chunk": c["chunk_num"] + 1, "doc_id":   c["doc_id"], "text": c["text"]} for c in chunks]}
+
+    viewable.append(entry)
+
+
+with open("chroma_preview.json", "w") as f:
+    json.dump(viewable, f, indent=2)
+
+print("Saved to chroma_preview.json")
+
+with open("chroma_preview.json", "r") as f:        # or test.json / val.json
+    faq_db = json.load(f)
+
+for item in faq_db:
+    source_id = item["id"]
+    item["relevant_doc_ids"] = source_id_to_doc_ids.get(source_id, [])
+
+with open("chroma_preview.json", "w") as f:
+    json.dump(faq_db, f, indent=2)
+
+# save chats in db
+chat_collection = client.get_or_create_collection(name="chat_history")
+
+def create_new_chat():
+    # creates a new chat session and returns its ID
+    chat_id = str(uuid.uuid4())
+    print(f"New chat created: {chat_id}")
+    return chat_id
+
+def get_chat_history(chat_id):
+    """Retrieve full history of a specific chat."""
+    results = chat_collection.get(
+        where={"chat_id": chat_id}
+    )
+
+    # Pair up queries and answers
+    history = []
+    for doc, meta in zip(results["documents"], results["metadatas"]):
+        history.append({
+            "query": doc,
+            "answer": meta["answer"],
+            "timestamp": meta["timestamp"]
+        })
+
+    # Sort by timestamp
+    history.sort(key=lambda x: x["timestamp"])
+    return history
+
+def load_chats_from_db():
+    # rebuild the chats dict from ChromaDB on page load
+    try:
+        results = chat_collection.get(include=["metadatas", "documents"])
+    except Exception:
+        return {}
+
+    chats = {}
+    for doc, meta in zip(results["documents"], results["metadatas"]):
+        chat_id   = meta.get("chat_id")
+        chat_name = meta.get("chat_name", "Chat 1")
+        timestamp = meta.get("timestamp", "")
+        query     = doc
+        answer    = meta.get("answer", "")
+
+        if chat_name not in chats:
+            chats[chat_name] = {"messages": [], "chat_id": chat_id}
+
+        chats[chat_name]["messages"].append({"role": "user",      "content": query, "timestamp": timestamp})
+        chats[chat_name]["messages"].append({"role": "assistant", "content": answer, "timestamp": timestamp})
+
+    # sort messages within each chat by timestamp
+    for chat_name in chats:
+        chats[chat_name]["messages"].sort(key=lambda x: x.get("timestamp", ""))
+
+    return chats
+
+def delete_chat(chat_id: str):
+    """Remove a chat from ChromaDB and the JSON log."""
+    # Remove from ChromaDB
+    results = chat_collection.get(where={"chat_id": chat_id})
+    if results["ids"]:
+        chat_collection.delete(ids=results["ids"])
+
+    print(f"Chat deleted: {chat_id}")
+

chroma/chatbot_retriever.py

@@ -0,0 +1,259 @@
+import chromadb
+from langchain_ollama import ChatOllama
+from langchain_classic.chains import create_history_aware_retriever, create_retrieval_chain
+from langchain_classic.chains.combine_documents import create_stuff_documents_chain
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_chroma import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_core.messages import HumanMessage, AIMessage
+from chroma.chatbot_chroma_db import create_new_chat, get_chat_history 
+from datetime import datetime
+import uuid
+from presidio_analyzer import AnalyzerEngine, PatternRecognizer, Pattern
+from presidio_anonymizer import AnonymizerEngine
+from presidio_analyzer import PatternRecognizer, Pattern
+from presidio_anonymizer.entities import OperatorConfig
+from presidio_analyzer.nlp_engine import NlpEngineProvider
+from sentence_transformers import CrossEncoder
+from langchain_classic.retrievers import ContextualCompressionRetriever
+from langchain_classic.retrievers.document_compressors import CrossEncoderReranker
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+from sentence_transformers import CrossEncoder
+from sentence_transformers import SentenceTransformer
+
+
+
+# load model and collection
+# reranking_model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
+# reranking_model.save("../models/reranker")
+# model = SentenceTransformer("multi-qa-mpnet-base-dot-v1")
+# model.save("../models/embedder")
+
+rerank_model = HuggingFaceCrossEncoder(model_name="./models/reranker")
+compressor = CrossEncoderReranker(model=rerank_model, top_n=3)
+model      = HuggingFaceEmbeddings(model_name="./models/embedder")
+client     = chromadb.PersistentClient(path="./chroma_db")
+chat_collection = client.get_collection("chat_history", embedding_function=None)
+
+print(client.list_collections()) 
+
+vectorstore = Chroma(
+    client = client,
+    collection_name = "bank_faq",
+    embedding_function = model
+)
+
+retriever = vectorstore.as_retriever(
+    search_type = "similarity",
+    search_kwargs = {"k": 3}
+)
+
+reranking_retriever = ContextualCompressionRetriever(
+    base_compressor=compressor,
+    base_retriever=vectorstore.as_retriever(search_kwargs={"k": 20})
+)   
+
+llm = ChatOllama(model="llama3.2", temperature=0)
+
+# contextualize qn
+def contextualize_question():
+    question_reformulation_prompt = """
+    Given a chat history and the latest user question
+    which might reference context in the chat history,
+    formulate a standalone question that can be understood
+    without the chat history. Do NOT answer the question,
+    just reformulate it if needed and otherwise return it as is."""
+
+    question_reformulation_template = ChatPromptTemplate.from_messages([
+        ("system", question_reformulation_prompt),
+        MessagesPlaceholder("chat_history"),
+        ("human", "{input}"),
+    ])
+
+    return create_history_aware_retriever(llm, reranking_retriever, question_reformulation_template)
+
+# answer qn
+def answer_question():
+    # If the answer is not in the context, say "I don't have information on that.
+    answer_prompt = """You are a helpful banking assistant.
+    Answer the user's question using ONLY the context below.
+    
+    Before answering, check whether the user's message:
+    1. Contains sensitive information (account numbers, passwords, PINs, OTPs, card numbers, etc.), OR
+    2. Requests a banking transaction or account-specific action (transfers, payments, account changes, balance checks, etc.).
+
+    If either condition is true:
+    - Do NOT process, verify, store, or repeat the information.
+    - Inform the user not to share sensitive information.
+    - Explain that account-specific actions cannot be performed through this chatbot.
+    - Do not ask for additional sensitive information.
+    - Do not continue the transaction.
+
+    Only after passing the security check should you use the context to answer the question.
+
+    If the answer is clearly in the context, answer it directly and MUST provide it.
+
+    If the answer is NOT in the context, first consider whether the question 
+    might be vague or incomplete. If so, ask a single clarifying question to 
+    help the user elaborate (e.g. "Could you tell me more about what you're 
+    looking for?", "Are you referring to a specific account type or product?").
+
+    Do not generate information on your own that is not within the context.
+
+    Do not say 'I don't have information' if the context contains relevant details.
+
+    Only say "I don't have information on that" if:
+    - The question is already specific and clear, AND
+    - The answer is genuinely not in the context."
+    
+    
+    Context: {context}"""
+    
+    answer_template = ChatPromptTemplate.from_messages([
+        ("system", answer_prompt),
+        MessagesPlaceholder("chat_history"),
+        ("human", "{input}"),
+    ])
+
+    return create_stuff_documents_chain(llm, answer_template)
+
+history_aware_retriever = contextualize_question()
+answer_chain            = answer_question()
+
+rag_chain = create_retrieval_chain(history_aware_retriever, answer_chain)
+
+chat_history = []
+
+# analyzer = AnalyzerEngine()
+anonymizer = AnonymizerEngine()
+
+provider = NlpEngineProvider(nlp_configuration={
+    "nlp_engine_name": "spacy",
+    "models": [{"lang_code": "en", "model_name": "en_core_web_lg"}]
+})
+
+analyzer = AnalyzerEngine(nlp_engine=provider.create_engine())
+
+account_recognizer = PatternRecognizer(
+    supported_entity="ACCOUNT_NUMBER",
+    patterns=[Pattern(name="account_number", regex=r"\d{3}-\d{3}-\d{3}", score=0.8)],
+    context=["account", "acct", "num", "number", "no.", "#"]
+)
+
+analyzer.registry.add_recognizer(account_recognizer)
+
+def _analyze(text: str):
+    return analyzer.analyze(text=text, language="en", entities=["ACCOUNT_NUMBER", "PERSON", "PHONE_NUMBER", "EMAIL_ADDRESS", "CREDIT_CARD"])
+
+def scrub_pii(text: str) -> str:
+    results = _analyze(text)
+    scrubbed = anonymizer.anonymize(text=text, analyzer_results=results)
+    return scrubbed.text
+
+SECURITY_MESSAGE = (
+    """ 
+    For your security and privacy, please do not share sensitive 
+    information such as account numbers, passwords, PINs, or OTPs 
+    in this chat. If you require account-related assistance, please 
+    use the bank's secure authenticated channels.
+    """
+)
+
+def contains_pii(text: str) -> bool:
+    results = _analyze(text)
+    return len(results) > 0
+
+
+def ask(query, chat_id=None, chat_name = None, persist=True):
+
+    if chat_id is None:
+        lc_history = []
+        past = []
+    else:
+        MAX_HISTORY_TURNS = 10
+        past = get_chat_history(chat_id)[-MAX_HISTORY_TURNS:]
+        # past = get_chat_history(chat_id) 
+        lc_history = []
+        for turn in past:
+            lc_history.append(HumanMessage(content=turn["query"]))
+            lc_history.append(AIMessage(content=turn["answer"]))
+
+    
+    # load long term history from db
+    clarification_phrases = ["could you", "can you elaborate", "what do you mean", "are you referring"]
+    clarification_count = sum(
+        1 for turn in past
+        if any(phrase in turn["answer"].lower() for phrase in clarification_phrases)
+    )
+
+    clean_query = scrub_pii(query)
+    # print(clean_query)
+
+    # docs = retriever.invoke(clean_query)
+    # print("Retrieved docs:")
+    # for i, doc in enumerate(docs):
+    #     print(f"\n--- Doc {i+1} ---")
+    #     print(doc.page_content[:500])
+    
+    # run rag chain with full history
+    response = rag_chain.invoke({
+        "input":        clean_query,
+        "chat_history": lc_history
+    })
+    answer = response["answer"]
+
+    is_clarifying = any(phrase in answer.lower() for phrase in clarification_phrases)
+    if is_clarifying and clarification_count >= 1:
+        answer = "I'm sorry, I don't have information on that."
+
+    # persist new turn back to db
+    if persist and chat_id is not None:
+        message_id = str(uuid.uuid4())
+        chat_collection.add(
+            ids=[message_id],
+            documents=[clean_query],
+            metadatas=[{
+                "chat_id":   chat_id,
+                "chat_name": chat_name,
+                "answer":    answer,
+                "timestamp": datetime.now().isoformat()
+            }]
+        )
+
+    return answer
+
+# Test questions
+# print(ask("What are the minimum and maximum age limits for investing in M&N Life Uday"))
+
+# print(ask("what documents do I need to open a bank account?"))
+# print(ask("what about for a savings account?")) 
+# print(ask("do I need KYC for that too?")) 
+
+def retrieve_with_rerank(query, top_k=20, top_n=5):
+    # Step 1: Retrieve
+    raw_chunks = vectorstore.similarity_search(query, k=top_k)
+    
+    # Step 2: Score each chunk against the query
+    pairs = [(query, chunk.page_content) for chunk in raw_chunks]
+    scores = rerank_model.predict(pairs)
+    
+    # Step 3: Sort by score and take top_n
+    ranked = sorted(zip(scores, raw_chunks), key=lambda x: x[0], reverse=True)
+    return [chunk for _, chunk in ranked[:top_n]]
+
+
+
+def answer_query(query):
+    # Before: chunks = vector_store.similarity_search(query, k=5)
+    # After:
+    chunks = retrieve_with_rerank(query, top_k=20, top_n=5)
+    
+    context = "\n\n".join([c.page_content for c in chunks])
+    
+    response = llm.invoke(f"""
+    Context: {context}
+    Question: {query}
+    Answer only from the context above.
+    """)
+    return response
+