Merge remote-tracking branch 'origin/main' into 112-selectai-integration

Author: gotsysdba

src/client/utils/st_common.py

@@ -7,7 +7,6 @@
 import os
 from io import BytesIO
 from typing import Union, get_args
-from uuid import uuid4
 import pandas as pd
 
 import streamlit as st
@@ -28,12 +27,6 @@
 #############################################################################
 # Common Helpers
 #############################################################################
-def client_gen_id() -> ClientIdType:
-    """Generate a new client ID; can be used to clear history"""
-    logger.info("Creating new client identifier")
-    return str(uuid4())
-
-
 def local_file_payload(uploaded_files: Union[BytesIO, list[BytesIO]]) -> list:
     """Upload Single file from Streamlit to the Server"""
     # If it's a single file, convert it to a list for consistent processing
@@ -124,8 +117,11 @@ def history_sidebar() -> None:
         on_change=update_user_settings("ll_model"),
     )
     if button_col.button("Clear", disabled=not chat_history_enable, use_container_width=True):
-        # Establish a new thread
-        state.user_settings["client"] = client_gen_id()
+        # Clean out history
+        try:
+            api_call.patch(endpoint="v1/chat/history")
+        except api_call.ApiError as ex:
+            logger.error("Clearing Chat History for %s failed: %s", state.user_settings["client"], ex)
         clear_state_key("user_client")
 
 

src/launch_client.py

@@ -8,12 +8,13 @@
 # spell-checker:ignore streamlit, scriptrunner
 
 import os
+from uuid import uuid4
 
 import streamlit as st
 from streamlit import session_state as state
 
 from client.utils import api_call
-from client.utils.st_common import set_server_state, client_gen_id
+from client.utils.st_common import set_server_state
 
 import common.logging_config as logging_config
 
@@ -62,7 +63,7 @@ def main() -> None:
     if "user_settings" not in state:
         try:
             state.user_settings = api_call.post(
-                endpoint="v1/settings", params={"client": client_gen_id()}, retries=10, backoff_factor=1.5
+                endpoint="v1/settings", params={"client": str(uuid4())}, retries=10, backoff_factor=1.5
             )
         except api_call.ApiError:
             logger.error("Unable to contact API Server; setting as Down!")

src/server/endpoints.py

@@ -18,7 +18,14 @@
 from pydantic import HttpUrl
 
 from langgraph.graph.state import CompiledStateGraph
-from langchain_core.messages import HumanMessage, AnyMessage, convert_to_openai_messages, ChatMessage
+from langgraph.graph.message import REMOVE_ALL_MESSAGES
+from langchain_core.messages import (
+    HumanMessage,
+    AnyMessage,
+    convert_to_openai_messages,
+    ChatMessage,
+    RemoveMessage,
+)
 from langchain_core.runnables import RunnableConfig
 from giskard.rag import evaluate, QATestset
 from giskard.rag.metrics import correctness_metric
@@ -751,12 +758,32 @@ async def chat_stream(request: schema.ChatRequest, client: schema.ClientIdType =
             media_type="application/octet-stream",
         )
 
+    @auth.patch(
+        "/v1/chat/history",
+        description="Delete Chat History",
+        response_model=list[schema.ChatMessage],
+    )
+    async def chat_history_clean(client: schema.ClientIdType = Header(...)) -> list[ChatMessage]:
+        agent: CompiledStateGraph = chatbot.chatbot_graph
+        try:
+            _ = agent.update_state(
+                config=RunnableConfig(
+                    configurable={
+                        "thread_id": client,
+                    }
+                ),
+                values={"messages": RemoveMessage(id=REMOVE_ALL_MESSAGES)},
+            )
+            return [ChatMessage(content="As requested, I've forgotten our conversation.", role="system")]
+        except KeyError:
+            return [ChatMessage(content="I'm sorry, I have no history of this conversation.", role="system")]
+
     @auth.get(
         "/v1/chat/history",
         description="Get Chat History",
         response_model=list[schema.ChatMessage],
     )
-    async def chat_history(client: schema.ClientIdType = Header(...)) -> list[ChatMessage]:
+    async def chat_history_return(client: schema.ClientIdType = Header(...)) -> list[ChatMessage]:
         """Return Chat History"""
         agent: CompiledStateGraph = chatbot.chatbot_graph
         try:
@@ -771,7 +798,7 @@ async def chat_history(client: schema.ClientIdType = Header(...)) -> list[ChatMe
             chat_messages = convert_to_openai_messages(messages)
             return chat_messages
         except KeyError:
-            return [ChatMessage(content="I'm sorry, I have no history of this conversation", role="system")]
+            return [ChatMessage(content="I'm sorry, I have no history of this conversation.", role="system")]
 
     #################################################
     # testbed Endpoints

tests/server/test_endpoints_chat.py

@@ -29,7 +29,8 @@ class TestInvalidAuthEndpoints:
         [
             pytest.param("/v1/chat/completions", "post", id="chat_post"),
             pytest.param("/v1/chat/streams", "post", id="chat_stream"),
-            pytest.param("/v1/chat/history", "get", id="chat_history"),
+            pytest.param("/v1/chat/history", "patch", id="chat_history_clean"),
+            pytest.param("/v1/chat/history", "get", id="chat_history_return"),
         ],
     )
     def test_endpoints(self, client, auth_headers, endpoint, api_method, auth_type, status_code):
@@ -125,14 +126,6 @@ def test_chat_stream_valid_mock(self, client, auth_headers):
             content = b"".join(response.iter_bytes())
             assert b"Test streaming response" in content
 
-    def test_chat_history_empty(self, client, auth_headers):
-        """Test no model chat completion request"""
-        response = client.get("/v1/chat/history", headers=auth_headers["valid_auth"])
-        assert response.status_code == 200
-        history = response.json()
-        assert history[0]["role"] == "system"
-        assert history[0]["content"] == "I'm sorry, I have no history of this conversation"
-
     def test_chat_history_valid_mock(self, client, auth_headers):
         """Test valid chat history request"""
         # Create the mock history response
@@ -152,3 +145,25 @@ def test_chat_history_valid_mock(self, client, auth_headers):
             assert len(history) == 2
             assert history[0]["role"] == "user"
             assert history[0]["content"] == "Hello"
+
+    def test_chat_history_clean(self, client, auth_headers):
+        """Test chat history with no history"""
+        with patch("server.agents.chatbot.chatbot_graph") as mock_graph:
+            mock_graph.get_state.side_effect = KeyError()
+            response = client.patch("/v1/chat/history", headers=auth_headers["valid_auth"])
+            assert response.status_code == 200
+            history = response.json()
+            assert len(history) == 1
+            assert history[0]["role"] == "system"
+            assert "forgotten" in history[0]["content"].lower()
+
+    def test_chat_history_empty(self, client, auth_headers):
+        """Test chat history with no history"""
+        with patch("server.agents.chatbot.chatbot_graph") as mock_graph:
+            mock_graph.get_state.side_effect = KeyError()
+            response = client.get("/v1/chat/history", headers=auth_headers["valid_auth"])
+            assert response.status_code == 200
+            history = response.json()
+            assert len(history) == 1
+            assert history[0]["role"] == "system"
+            assert "no history" in history[0]["content"].lower()