chroma/chatbot_retriever.py

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