add agent history

core/agent.py

@@ -197,14 +197,22 @@ async def _execute_tool(self, name: str, args: dict, auth: AuthContext, source_m
             case _:
                 raise ValueError(f"Unknown tool: {name}")
 
-    async def run(self, query: str, auth: AuthContext) -> str:
+    async def run(self, query: str, auth: AuthContext, conversation_history: list = None) -> str:
         """Synchronously run the agent and return the final answer."""
         # Per-run state to avoid cross-request leakage
         source_map: dict = {}
         messages = [
             {"role": "system", "content": self.system_prompt},
-            {"role": "user", "content": query},
         ]
+
+        # Add conversation history if provided
+        if conversation_history:
+            for msg in conversation_history[:-1]:  # Exclude the last message (current user query)
+                messages.append({"role": msg["role"], "content": msg["content"]})
+
+        # Add the current user query
+        messages.append({"role": "user", "content": query})
+
         tool_history = []  # Initialize tool history list
         # Get the full model name from the registered models config
         settings = get_settings()
@@ -370,8 +378,28 @@ async def run(self, query: str, auth: AuthContext) -> str:
 
                 # Return final content, tool history, display objects and sources
                 display_objects = crop_images_in_display_objects(display_objects)
+
+                # Generate a user-friendly response text from display objects
+                response_text = ""
+                if display_objects:
+                    # Extract text content from display objects for a clean response
+                    text_contents = []
+                    for obj in display_objects:
+                        if obj.get("type") == "text" and obj.get("content"):
+                            text_contents.append(obj["content"])
+
+                    if text_contents:
+                        # Join text contents with proper spacing
+                        response_text = "\n\n".join(text_contents)
+                    else:
+                        # If no text objects, provide a generic response
+                        response_text = "I've found relevant information in the documents. Please see the display objects above for details."
+                else:
+                    # Fallback to original content if no display objects
+                    response_text = msg.content
+
                 return {
-                    "response": msg.content,
+                    "response": response_text,
                     "tool_history": tool_history,
                     "display_objects": display_objects,
                     "sources": sources,

core/api.py

@@ -626,21 +626,75 @@ async def get_chat_history(
 
 @app.post("/agent", response_model=Dict[str, Any])
 @telemetry.track(operation_type="agent_query")
-async def agent_query(request: AgentQueryRequest, auth: AuthContext = Depends(verify_token)):
+async def agent_query(
+    request: AgentQueryRequest,
+    auth: AuthContext = Depends(verify_token),
+    redis: arq.ArqRedis = Depends(get_redis_pool),
+):
     """Execute an agent-style query using the :class:`MorphikAgent`.
 
     Args:
-        request: The query payload containing the natural language question.
+        request: The query payload containing the natural language question and optional chat_id.
         auth: Authentication context used to enforce limits and access control.
+        redis: Redis connection for chat history storage.
 
     Returns:
         A dictionary with the agent's full response.
     """
+    # Chat history retrieval
+    history_key = None
+    history: List[Dict[str, Any]] = []
+    if request.chat_id:
+        history_key = f"chat:{request.chat_id}"
+        stored = await redis.get(history_key)
+        if stored:
+            try:
+                history = json.loads(stored)
+            except Exception:
+                history = []
+        else:
+            db_hist = await document_service.db.get_chat_history(request.chat_id, auth.user_id, auth.app_id)
+            if db_hist:
+                history = db_hist
+
+        history.append(
+            {
+                "role": "user",
+                "content": request.query,
+                "timestamp": datetime.now(UTC).isoformat(),
+            }
+        )
+
     # Check free-tier agent call limits in cloud mode
     if settings.MODE == "cloud" and auth.user_id:
         await check_and_increment_limits(auth, "agent", 1)
+
     # Use the shared MorphikAgent instance; per-run state is now isolated internally
-    response = await morphik_agent.run(request.query, auth)
+    response = await morphik_agent.run(request.query, auth, history)
+
+    # Chat history storage
+    if history_key:
+        # Store the full agent response with structured data
+        agent_message = {
+            "role": "assistant",
+            "content": response.get("response", ""),
+            "timestamp": datetime.now(UTC).isoformat(),
+            # Store agent-specific structured data
+            "agent_data": {
+                "display_objects": response.get("display_objects", []),
+                "tool_history": response.get("tool_history", []),
+                "sources": response.get("sources", []),
+            },
+        }
+        history.append(agent_message)
+        await redis.set(history_key, json.dumps(history))
+        await document_service.db.upsert_chat_history(
+            request.chat_id,
+            auth.user_id,
+            auth.app_id,
+            history,
+        )
+
     # Return the complete response dictionary
     return response
 

core/models/chat.py

@@ -1,5 +1,5 @@
 from datetime import UTC, datetime
-from typing import List, Literal, Optional
+from typing import Any, Dict, List, Literal, Optional
 
 from pydantic import BaseModel, Field
 
@@ -10,6 +10,7 @@ class ChatMessage(BaseModel):
     role: Literal["user", "assistant"]
     content: str
     timestamp: str = Field(default_factory=lambda: datetime.now(UTC).isoformat())
+    agent_data: Optional[Dict[str, Any]] = None  # For agent-specific data like display_objects, tool_history, sources
 
 
 class ChatConversation(BaseModel):

core/models/request.py

@@ -151,3 +151,7 @@ class AgentQueryRequest(BaseModel):
     """Request model for agent queries"""
 
     query: str = Field(..., description="Natural language query for the Morphik agent")
+    chat_id: Optional[str] = Field(
+        None,
+        description="Optional chat session ID for persisting conversation history",
+    )

core/tests/integration/test_api.py

@@ -3868,3 +3868,44 @@ async def test_chat_persistence(client: AsyncClient):
     assert len(history) >= 4
     assert history[0]["role"] == "user"
     assert history[1]["role"] == "assistant"
+
+
+@pytest.mark.asyncio
+async def test_agent_conversation_persistence(client: AsyncClient):
+    """Ensure agent conversation history is persisted across queries."""
+    headers = create_auth_header()
+    chat_id = str(uuid.uuid4())
+
+    # First agent message
+    resp1 = await client.post(
+        "/agent",
+        json={"query": "Hello, my name is Bob", "chat_id": chat_id},
+        headers=headers,
+    )
+    assert resp1.status_code == 200
+    data1 = resp1.json()
+    assert "response" in data1
+
+    # Second agent message that references the first
+    resp2 = await client.post(
+        "/agent",
+        json={"query": "What is my name?", "chat_id": chat_id},
+        headers=headers,
+    )
+    assert resp2.status_code == 200
+    data2 = resp2.json()
+    assert "response" in data2
+    # The agent should remember the name "Bob" from the conversation history
+    assert "Bob" in data2["response"] or "bob" in data2["response"].lower()
+
+    # Verify conversation history is stored
+    hist = await client.get(f"/chat/{chat_id}", headers=headers)
+    assert hist.status_code == 200
+    history = hist.json()
+    assert len(history) >= 4  # 2 user messages + 2 assistant responses
+    assert history[0]["role"] == "user"
+    assert history[0]["content"] == "Hello, my name is Bob"
+    assert history[1]["role"] == "assistant"
+    assert history[2]["role"] == "user"
+    assert history[2]["content"] == "What is my name?"
+    assert history[3]["role"] == "assistant"

ee/ui-component/components/chat/AgentChatSection.tsx

@@ -16,6 +16,7 @@ interface AgentChatSectionProps {
   initialMessages?: ChatMessage[];
   isReadonly?: boolean;
   onAgentSubmit?: (query: string) => void;
+  chatId?: string;
 }
 
 /**
@@ -27,6 +28,7 @@ const AgentChatSection: React.FC<AgentChatSectionProps> = ({
   initialMessages = [],
   isReadonly = false,
   onAgentSubmit,
+  chatId,
 }) => {
   // State for managing chat
   const [messages, setMessages] = useState<AgentUIMessage[]>(
@@ -45,6 +47,54 @@ const AgentChatSection: React.FC<AgentChatSectionProps> = ({
   const messagesContainerRef = useRef<HTMLDivElement>(null);
   const messagesEndRef = useRef<HTMLDivElement>(null);
 
+  // Load agent messages from chat history when component mounts or chatId changes
+  useEffect(() => {
+    const loadAgentHistory = async () => {
+      if (chatId && apiBaseUrl && (authToken || apiBaseUrl.includes("localhost"))) {
+        try {
+          const response = await fetch(`${apiBaseUrl}/chat/${chatId}`, {
+            headers: {
+              ...(authToken ? { Authorization: `Bearer ${authToken}` } : {}),
+            },
+          });
+          if (response.ok) {
+            const data = await response.json();
+            const agentMessagesFromHistory = data.map((m: any) => {
+              const baseMessage = {
+                id: generateUUID(),
+                role: m.role,
+                content: m.content,
+                createdAt: new Date(m.timestamp),
+              };
+
+              // If this is an assistant message with agent_data, reconstruct experimental_agentData
+              if (m.role === "assistant" && m.agent_data) {
+                return {
+                  ...baseMessage,
+                  experimental_agentData: {
+                    tool_history: m.agent_data.tool_history || [],
+                    displayObjects: m.agent_data.display_objects || [],
+                    sources: m.agent_data.sources || [],
+                  },
+                };
+              }
+
+              return baseMessage;
+            });
+            setMessages(agentMessagesFromHistory);
+          }
+        } catch (err) {
+          console.error("Failed to load agent chat history", err);
+        }
+      }
+    };
+
+    // Only load if we don't have initial messages
+    if (initialMessages.length === 0) {
+      loadAgentHistory();
+    }
+  }, [chatId, apiBaseUrl, authToken, initialMessages.length]);
+
   // Function to handle form submission
   const handleSubmit = async (e?: React.FormEvent) => {
     e?.preventDefault();
@@ -92,6 +142,7 @@ const AgentChatSection: React.FC<AgentChatSectionProps> = ({
         },
         body: JSON.stringify({
           query: userMessage.content,
+          chat_id: chatId,
         }),
       });
 

ee/ui-component/components/chat/ChatSection.tsx

@@ -120,6 +120,54 @@ const ChatSection: React.FC<ChatSectionProps> = ({
   const [agentMessages, setAgentMessages] = useState<AgentUIMessage[]>([]);
   const [agentStatus, setAgentStatus] = useState<"idle" | "submitted" | "completed">("idle");
 
+  // Load agent messages from chat history when switching to agent mode
+  useEffect(() => {
+    const loadAgentHistory = async () => {
+      if (isAgentMode && chatId && apiBaseUrl && (authToken || apiBaseUrl.includes("localhost"))) {
+        try {
+          const response = await fetch(`${apiBaseUrl}/chat/${chatId}`, {
+            headers: {
+              ...(authToken ? { Authorization: `Bearer ${authToken}` } : {}),
+            },
+          });
+          if (response.ok) {
+            const data = await response.json();
+            const agentMessagesFromHistory = data.map((m: any) => {
+              const baseMessage = {
+                id: generateUUID(),
+                role: m.role,
+                content: m.content,
+                createdAt: new Date(m.timestamp),
+              };
+
+              // If this is an assistant message with agent_data, reconstruct experimental_agentData
+              if (m.role === "assistant" && m.agent_data) {
+                return {
+                  ...baseMessage,
+                  experimental_agentData: {
+                    tool_history: m.agent_data.tool_history || [],
+                    displayObjects: m.agent_data.display_objects || [],
+                    sources: m.agent_data.sources || [],
+                  },
+                };
+              }
+
+              return baseMessage;
+            });
+            setAgentMessages(agentMessagesFromHistory);
+          }
+        } catch (err) {
+          console.error("Failed to load agent chat history", err);
+        }
+      } else if (!isAgentMode) {
+        // Clear agent messages when switching back to regular chat mode
+        setAgentMessages([]);
+      }
+    };
+
+    loadAgentHistory();
+  }, [isAgentMode, chatId, apiBaseUrl, authToken]);
+
   // Fetch available graphs for dropdown
   const fetchGraphs = useCallback(async () => {
     if (!apiBaseUrl) return;
@@ -308,7 +356,10 @@ const ChatSection: React.FC<ChatSectionProps> = ({
           "Content-Type": "application/json",
           ...(authToken ? { Authorization: `Bearer ${authToken}` } : {}),
         },
-        body: JSON.stringify({ query: userMessage.content }),
+        body: JSON.stringify({
+          query: userMessage.content,
+          chat_id: chatId,
+        }),
       });
 
       if (!response.ok) {