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