Create chatbot.py

chatbot.py

@@ -0,0 +1,168 @@
+import uuid
+from datetime import datetime
+from urllib.parse import quote_plus
+
+from pymongo import MongoClient
+from langchain.prompts import ChatPromptTemplate
+from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory
+from langchain.chains import ConversationalRetrievalChain
+
+from llm_provider import llm
+from vectorstore_manager import get_user_retriever
+
+# === Prompt Template ===
+quiz_solving_prompt = '''
+You are an assistant specialized in solving quizzes. Your goal is to provide accurate, concise, and contextually relevant answers.
+Use the following retrieved context to answer the user's question.
+If the context lacks sufficient information, respond with "I don't know." Do not make up answers or provide unverified information.
+
+Guidelines:
+1. Extract key information from the context to form a coherent response.
+2. Maintain a clear and professional tone.
+3. If the question requires clarification, specify it politely.
+
+Retrieved context:
+{context}
+
+User's question:
+{question}
+
+Your response:
+'''
+
+user_prompt = ChatPromptTemplate.from_messages([
+    ("system", quiz_solving_prompt),
+    ("human", "{question}")
+])
+
+# === MongoDB Configuration ===
+PASSWORD = quote_plus("momimaad@123")
+MONGO_URI = f"mongodb+srv://hammad:{PASSWORD}@cluster0.2a9yu.mongodb.net/"
+DB_NAME = "Education_chatbot"
+HISTORY_COLLECTION = "chat_histories"       # used by MongoDBChatMessageHistory
+SESSIONS_COLLECTION = "chat_sessions"       # to track chat metadata
+CHAINS_COLLECTION = "user_chains"           # to track per-user vectorstore paths
+
+# Initialize MongoDB client and collections
+client = MongoClient(MONGO_URI)
+db = client[DB_NAME]
+sessions_collection = db[SESSIONS_COLLECTION]
+chains_collection = db[CHAINS_COLLECTION]
+
+# === Core Functions ===
+
+def create_new_chat(user_id: str) -> str:
+    """
+    Create a new chat session for the given user, persist metadata in MongoDB,
+    and ensure a vectorstore path is registered for that user.
+    Returns the new chat_id.
+    """
+    chat_id = f"{user_id}-{uuid.uuid4()}"
+    created_at = datetime.utcnow()
+
+    # Persist chat session metadata
+    sessions_collection.insert_one({
+        "chat_id": chat_id,
+        "user_id": user_id,
+        "created_at": created_at
+    })
+
+    # Initialize chat history storage in its own collection via LangChain helper
+    MongoDBChatMessageHistory(
+        session_id=chat_id,
+        connection_string=MONGO_URI,
+        database_name=DB_NAME,
+        collection_name=HISTORY_COLLECTION,
+    )
+
+    # If the user has no chain/vectorstore registered yet, register it
+    if chains_collection.count_documents({"user_id": user_id}, limit=1) == 0:
+        # This also creates the vectorstore on disk via vectorstore_manager.ingest_report
+        # you should call ingest_report first elsewhere before chat
+        chains_collection.insert_one({
+            "user_id": user_id,
+            "vectorstore_path": f"user_vectorstores/{user_id}_faiss"
+        })
+
+    return chat_id
+
+def get_chain_for_user(user_id: str, chat_id: str) -> ConversationalRetrievalChain:
+    """
+    Reconstructs (or creates) the user's ConversationalRetrievalChain
+    using their vectorstore and the chat-specific memory object.
+    """
+    # Load chat history memory
+    chat_history = MongoDBChatMessageHistory(
+        session_id=chat_id,
+        connection_string=MONGO_URI,
+        database_name=DB_NAME,
+        collection_name=HISTORY_COLLECTION,
+    )
+
+    # Look up vectorstore path
+    chain_doc = chains_collection.find_one({"user_id": user_id})
+    if not chain_doc:
+        raise ValueError(f"No vectorstore registered for user {user_id}")
+
+    # Initialize retriever from vectorstore
+    retriever = get_user_retriever(user_id)
+
+    # Create and return the chain
+    return ConversationalRetrievalChain.from_llm(
+        llm=llm,
+        retriever=retriever,
+        return_source_documents=True,
+        chain_type="stuff",
+        combine_docs_chain_kwargs={"prompt": user_prompt},
+        memory=chat_history,
+        verbose=False,
+    )
+
+def summarize_messages(chat_history: MongoDBChatMessageHistory) -> bool:
+    """
+    If the chat history grows too long, summarize it to keep the memory concise.
+    Returns True if a summary was performed.
+    """
+    messages = chat_history.messages
+    if not messages:
+        return False
+
+    summarization_prompt = ChatPromptTemplate.from_messages([
+        ("system", "Summarize the following conversation into a concise message:"),
+        ("human", "{chat_history}")
+    ])
+    summarization_chain = summarization_prompt | llm
+    summary = summarization_chain.invoke({"chat_history": messages})
+
+    chat_history.clear()
+    chat_history.add_ai_message(summary.content)
+    return True
+
+def stream_chat_response(user_id: str, chat_id: str, query: str):
+    """
+    Given a user_id, chat_id, and a query string, streams back the AI response
+    while persisting both user and AI messages to MongoDB.
+    """
+    # Ensure the chain and memory are set up
+    chain = get_chain_for_user(user_id, chat_id)
+    chat_history = chain.memory  # the MongoDBChatMessageHistory instance
+
+    # Optionally summarize if too many messages
+    summarize_messages(chat_history)
+
+    # Add the user message to history
+    chat_history.add_user_message(query)
+
+    # Stream the response
+    response_accum = ""
+    for chunk in chain.stream({"question": query, "chat_history": chat_history.messages}):
+        if "answer" in chunk:
+            print(chunk["answer"], end="", flush=True)
+            response_accum += chunk["answer"]
+        else:
+            # Unexpected chunk format
+            print(f"[Unexpected chunk]: {chunk}")
+
+    # Persist the AI's final message
+    if response_accum:
+        chat_history.add_ai_message(response_accum)