Update app.py

app.py

@@ -4,7 +4,6 @@ from langchain.vectorstores import Chroma
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationalRetrievalChain
-from langchain.memory.chat_message_histories import ChatMessageHistory
 from langchain_openai import ChatOpenAI
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain.embeddings import HuggingFaceEmbeddings
@@ -14,7 +13,7 @@ embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-Mi
 persist_directory = 'vec_db'
 
 vectordb = Chroma(persist_directory=persist_directory,
-                           embedding_function=embedding_model)
+                  embedding_function=embedding_model)
 
 vectordb_retriever = vectordb.as_retriever(search_kwargs={'k':5})
 
@@ -23,56 +22,62 @@ llm = ChatOpenAI(model="gpt-4.1-nano", temperature=0.7)
 with open("instructions.txt", 'r') as file:
     instructions = file.read()
 
-
-# Custom prompt
 custom_prompt = ChatPromptTemplate.from_messages([
     ("system", instructions),
     MessagesPlaceholder(variable_name="chat_history"),
     ("user", "Question: {input}\nContext: {context}")
 ])
 
-# Memory
-memory = ConversationBufferMemory(
-    memory_key="chat_history",
-    return_messages=True
-)
-
 question_answer_chain = create_stuff_documents_chain(llm, custom_prompt)
-
 chain = create_retrieval_chain(vectordb_retriever, question_answer_chain)
 
-def conversate_assistant(query, history):
-    greetings = {"hey", "hi", "hello"}
-    normalized_query = query.strip().lower()
+greetings = {"hey", "hi", "hello"}
 
-    if len(memory.load_memory_variables({})["chat_history"]) >=6:
-        chat_history = memory.load_memory_variables({})["chat_history"][-6::]
-    else:
-        chat_history = memory.load_memory_variables({})["chat_history"]
+def conversate_assistant(query, history, memory):
+    # Load limited chat history from memory
+    history_data = memory.load_memory_variables({})["chat_history"]
+    chat_history = history_data[-6:] if len(history_data) >= 6 else history_data
+
+    normalized_query = query.strip().lower()
 
-    # If greeting, skip retrieval and context
     if normalized_query in greetings:
-        response = question_answer_chain.invoke({
+        response = question_answer_chain({
             "input": query,
-            "context": [],  # empty context for greetings
+            "context": [],
             "chat_history": chat_history
         })
-        answer = response
+        answer = response.get("output") or str(response)
     else:
-        response = chain.invoke({
+        response = chain({
             "input": query,
             "chat_history": chat_history
         })
-        answer = response['answer']
-    
-    # Save to memory
+        answer = response.get('answer') or str(response)
+
+    # Save context in this session's memory
     memory.save_context({"input": query}, {"output": answer})
 
-    return answer
+    return answer, history + [(query, answer)], memory
+
+# Gradio interface with state for memory and chat history
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox()
+    state = gr.State([])   # to keep chat history visible in UI
+    memory_state = gr.State(None)  # to keep ConversationBufferMemory per session
+
+    def init_memory():
+        return ConversationBufferMemory(
+            memory_key="chat_history",
+            return_messages=True
+        )
+
+    def respond(user_message, chat_history, memory):
+        if memory is None:
+            memory = init_memory()
+        answer, chat_history, memory = conversate_assistant(user_message, chat_history, memory)
+        return chat_history, memory, ""
 
-demo = gr.ChatInterface(
-    conversate_assistant,
-    type="messages"
-)
+    msg.submit(respond, inputs=[msg, state, memory_state], outputs=[chatbot, memory_state, msg])
 
-demo.launch()
+demo.launch()