Update app.py

app.py

@@ -4,80 +4,75 @@ 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
 
-embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+ embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
-persist_directory = 'vec_db'
+ persist_directory = 'vec_db'
 
-vectordb = Chroma(persist_directory=persist_directory,
-                  embedding_function=embedding_model)
+ vectordb = Chroma(persist_directory=persist_directory,
+                           embedding_function=embedding_model)
 
-vectordb_retriever = vectordb.as_retriever(search_kwargs={'k':5})
+ vectordb_retriever = vectordb.as_retriever(search_kwargs={'k':5})
 
-llm = ChatOpenAI(model="gpt-4.1-nano", temperature=0.7)
+ llm = ChatOpenAI(model="gpt-4.1-nano", temperature=0.7)
 
-with open("instructions.txt", 'r') as file:
+ 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}")
 ])
 
-question_answer_chain = create_stuff_documents_chain(llm, custom_prompt)
-chain = create_retrieval_chain(vectordb_retriever, question_answer_chain)
+ # Memory
+memory = ConversationBufferMemory(
+    memory_key="chat_history",
+    return_messages=True
+)
 
-greetings = {"hey", "hi", "hello"}
+ question_answer_chain = create_stuff_documents_chain(llm, custom_prompt)
 
-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
+ chain = create_retrieval_chain(vectordb_retriever, question_answer_chain)
 
+ def conversate_assistant(query, history):
+    greetings = {"hey", "hi", "hello"}
     normalized_query = query.strip().lower()
 
+     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"]
+
+     # If greeting, skip retrieval and context
     if normalized_query in greetings:
-        response = question_answer_chain({
+        response = question_answer_chain.invoke({
             "input": query,
-            "context": [],
+            "context": [],  # empty context for greetings
             "chat_history": chat_history
         })
-        answer = response.get("output") or str(response)
+        answer = response
     else:
-        response = chain({
+        response = chain.invoke({
             "input": query,
             "chat_history": chat_history
         })
-        answer = response.get('answer') or str(response)
-
-    # Save context in this session's memory
+        answer = response['answer']
+    
+    # Save to memory
     memory.save_context({"input": query}, {"output": 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, ""
+     return answer
 
-    msg.submit(respond, inputs=[msg, state, memory_state], outputs=[chatbot, memory_state, msg])
+ demo = gr.ChatInterface(
+    conversate_assistant,
+    type="messages"
+)
 
-demo.launch()
+ demo.launch()