Update app.py

app.py

@@ -6,6 +6,7 @@ from qdrant_client import QdrantClient
 from langchain.agents import Tool, AgentExecutor, create_openai_tools_agent
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.tools import BaseTool
+from langchain.memory import ConversationBufferMemory
 from typing import Type, Optional
 import os
 import warnings
@@ -52,6 +53,19 @@ class AgentQuery(BaseModel):
 
 # === UTILITY FUNCTIONS ===
 
+
+# Store chat histories keyed by user_id
+user_memories = {}
+
+def get_user_memory(user_id: str):
+    """Fetch or create memory for a user."""
+    if user_id not in user_memories:
+        user_memories[user_id] = ConversationBufferMemory(
+            memory_key="chat_history",
+            return_messages=True
+        )
+    return user_memories[user_id]
+
 def get_encoded_auth_token(user: int, token: str) -> str:
     auth_string = f"{user}:{token}"
     return base64.b64encode(auth_string.encode("utf-8")).decode("utf-8")
@@ -263,35 +277,68 @@ agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
 
 # === API ENDPOINTS ===
 
-@app.post("/bot")
-def chat_with_agent(query: AgentQuery):
-    """Main agent endpoint - handles both document search and project queries intelligently"""
-    try:
-        # Prepare the input for the agent
-        agent_input = query.message
+# @app.post("/bot")
+# def chat_with_agent(query: AgentQuery):
+#     """Main agent endpoint - handles both document search and project queries intelligently"""
+#     try:
+#         # Prepare the input for the agent
+#         agent_input = query.message
+        
+#         # If user provided credentials, add them to the context
         
-        # If user provided credentials, add them to the context
+#         if query.userLoginId is not None:
+#             agent_input += f" [UserLoginId: {query.userLoginId}"
+#             if query.orgId is not None:
+#                 agent_input += f", OrgId: {query.orgId}"
+#             agent_input += "]"
         
-        if query.userLoginId is not None:
-            agent_input += f" [UserLoginId: {query.userLoginId}"
-            if query.orgId is not None:
-                agent_input += f", OrgId: {query.orgId}"
-            agent_input += "]"
+#         # Use the agent to process the query
+#         result = agent_executor.invoke({"input": agent_input})
         
-        # Use the agent to process the query
-        result = agent_executor.invoke({"input": agent_input})
+#         return {
+#             "message": query.message,
+#             "answer": result["output"],
+#             "agent_used": True
+#         }
         
+#     except Exception as e:
+#         return {
+#             "message": query.message,
+#             "answer": f"An error occurred: {str(e)}",
+#             "agent_used": True
+#         }
+@app.post("/bot")
+def chat_with_agent(query: AgentQuery):
+    """Main agent endpoint with per-user chat history."""
+    try:
+        user_id = query.userLoginId or "anonymous"  # fallback if no ID
+        memory = get_user_memory(user_id)
+
+        agent = initialize_agent(
+            tools,
+            llm,
+            agent="chat-conversational-react-description",
+            memory=memory,
+            verbose=True
+        )
+
+        # Build input
+        agent_input = query.message
+        if query.orgId is not None:
+            agent_input += f" [OrgId: {query.orgId}]"
+
+        # Run agent
+        result = agent.invoke({"input": agent_input})
+
         return {
             "message": query.message,
             "answer": result["output"],
-            "agent_used": True
+            "chat_history": memory.load_memory_variables({})["chat_history"]
         }
-        
     except Exception as e:
         return {
             "message": query.message,
-            "answer": f"An error occurred: {str(e)}",
-            "agent_used": True
+            "answer": f"An error occurred: {str(e)}"
         }
 
 @app.post("/chat-documents")
@@ -311,6 +358,26 @@ def chat_documents_only(query: Query):
             "tool_used": "document_search"
         }
 
+@app.get("/get-chat-history/{user_id}")
+def get_chat_history(user_id: str, n: int = 10):
+    """
+    Fetch the last N chat messages for a user.
+    Default: 10
+    """
+    memory = user_memories.get(user_id)
+
+    if not memory:
+        return {"user_id": user_id, "chat_history": []}
+
+    history = memory.load_memory_variables({})["chat_history"]
+
+    # Only return last N messages
+    return {
+        "user_id": user_id,
+        "chat_history": history[-n:]
+    }
+
+
 @app.post("/list-projects")
 def list_projects(request: ProjectRequest):
     """Direct project listing without agent"""