This good for now merge it

Author: ryaneggz

backend/src/flows/__init__.py

@@ -21,6 +21,7 @@
 from src.utils.format import init_system_prompt
 from src.schemas.contexts import ContextSchema
 from src.schemas.entities.a2a import A2AServers
+from src.utils.middleware import add_ai_message_metadata, pii_middleware
 
 
 async def add_memories_to_system():
@@ -51,7 +52,7 @@ def graph_builder(
     subagents: list[SubAgent] = [],
     prompt: str = "You are a helpful assistant.",
     model: str = "openai:gpt-5-nano",
-    context_schema: Type[Any] | None = None,
+    context_schema: Type[ContextSchema] | None = None,
     checkpointer: BaseCheckpointSaver | None = None,
     store: BaseStore | None = None,
     graph_id: Literal[
@@ -75,6 +76,7 @@ def graph_builder(
         system_prompt=prompt,
         checkpointer=checkpointer,
         context_schema=context_schema,
+        middleware=[add_ai_message_metadata] + pii_middleware(),
         store=store,
     )
     return deep_agent

backend/src/schemas/contexts/__init__.py

@@ -3,4 +3,5 @@
 
 @dataclass
 class ContextSchema:
-    user_id: str
+    model: str = None
+    user_id: str = None

backend/src/utils/middleware.py

@@ -0,0 +1,38 @@
+from langchain.agents.middleware import after_model
+from langchain.agents import AgentState
+from langchain_core.messages import AIMessage
+from langgraph.runtime import Runtime
+from src.schemas.contexts import ContextSchema
+from langchain.agents.middleware import PIIMiddleware
+
+@after_model
+def add_ai_message_metadata(state: AgentState, runtime: Runtime[ContextSchema]) -> dict | None:
+    """Attach AI message metadata to final response."""
+    if state["messages"]:
+        last_msg = state["messages"][-1]
+        if isinstance(last_msg, AIMessage) and not last_msg.tool_calls:
+            last_msg.model = runtime.context.model
+    return None
+
+def pii_middleware() -> dict | None:
+     return [
+		# Redact email addresses
+		# PIIMiddleware(
+        #     "email",
+        #     strategy="redact",
+        #     apply_to_input=True,
+        # ),
+		# Mask credit card numbers
+   		PIIMiddleware(
+            "credit_card",
+            strategy="mask",
+            apply_to_input=True,
+        ),
+     	# Block API keys - raise error if detected
+        PIIMiddleware(
+			"api_key",
+			detector=r"sk-[A-Za-z0-9]+",
+			strategy="block",
+			apply_to_input=True,
+		),
+	 ]

backend/src/utils/stream.py

@@ -1,3 +1,4 @@
+from langchain.agents.middleware import PIIDetectionError
 import ujson
 from langchain_core.language_models import BaseChatModel
 from langchain_core.runnables import RunnableConfig
@@ -7,16 +8,16 @@
 from langgraph.types import StreamMode
 from deepagents import SubAgent
 
+from src.schemas.contexts import ContextSchema
 from src.contexts.service import ServiceContext
-from src.schemas.entities import LLMInput, LLMRequest
+from src.schemas.entities import LLMInput
 from src.constants import APP_LOG_LEVEL
 from src.flows import construct_agent
 from src.services.db import get_checkpoint_db
 from src.utils.messages import from_message_to_dict
 from langchain_core.messages import (
-    AIMessage,
     AIMessageChunk,
-    BaseMessage,
+    HumanMessage,
     ToolMessage,
 )
 from src.utils.logger import log_to_file, logger
@@ -189,9 +190,9 @@ async def stream_generator(
                 {"messages": input.messages}, 
                 stream_mode=["messages", "values"], 
                 config=config,
-                context=None
+                context=ContextSchema(model=agent.model)
             ):
-                # Serialize and yield each chunk as SSEq
+                # Serialize and yield each chunk as SSE
                 stream_chunk = handle_multi_mode(chunk)
                 if stream_chunk:
                     stream_type = stream_chunk[0]
@@ -202,10 +203,16 @@ async def stream_generator(
                     log_to_file(str(data), agent.model) and APP_LOG_LEVEL == "DEBUG"
                     logger.debug(f"data: {str(data)}")
                     yield f"data: {data}\n\n"
+        except PIIDetectionError as e:
+            # Yield error as SSE if streaming fails
+            logger.warning(f"Sensitive data detected in the query: {e}")
+            # raise HTTPException(status_code=500, detail=str(e))
+            error_msg = ujson.dumps(("error", str(e)))
+            yield f"data: {error_msg}\n\n"
 
         except Exception as e:
             # Yield error as SSE if streaming fails
-            logger.exception("Error in event_generator: %s", e)
+            logger.exception("Error in stream_generator: %s", e)
             # raise HTTPException(status_code=500, detail=str(e))
             error_msg = ujson.dumps(("error", str(e)))
             yield f"data: {error_msg}\n\n"
@@ -214,12 +221,20 @@ async def stream_generator(
                 final_state = await agent.graph.aget_state(config)
                 configurable = final_state.config.get("configurable", {})
                 messages = final_state.values.get('messages', [])
+                
+                # Get the last HumanMessage
+                last_human_message = None
+                for message in reversed(messages):
+                    if isinstance(message, HumanMessage):
+                        last_human_message = message
+                        break
+                
                 await service_context.thread_service.update(
                     thread_id=configurable.get("thread_id"),
                     data={
                         "thread_id": configurable.get("thread_id"),
                         "checkpoint_id": configurable.get("checkpoint_id"),
-                        "messages": [messages[-1].model_dump()],
+                        "messages": [last_human_message.model_dump()] if last_human_message else [],
                         "files": files_map,
                         "updated_at": get_time(),
                     }

frontend/src/components/lists/ChatMessages.tsx

@@ -12,6 +12,7 @@ import SearchEngineTool from "../tools/SearchEngine";
 import ChartRenderWidget from "../tools/ChartRenderWidget";
 import CopyTextButton from "../buttons/CopyTextButton";
 import FileViewer from "../viewers/FileViewer";
+import { latestHumanMessage } from "@/lib/utils/message";
 
 const MAX_LENGTH = 1000;
 
@@ -62,9 +63,11 @@ function ToolAction({
 export function Message({
 	message,
 	isLatest = false,
+	messages,
 }: {
 	message: any;
 	isLatest?: boolean;
+	messages: any[];
 }) {
 	const ICON_SIZE = 4;
 	const [isEditing, setIsEditing] = useState(false);
@@ -256,14 +259,14 @@ export function Message({
 
 						<div className="flex items-center gap-2">
 							<button className="text-sm text-muted-foreground">
-								{message.model}
+								{message.model || latestHumanMessage(messages)?.model || "Unknown model"}
 							</button>
 
 							{isLatest && streamingRate?.rate && (
 								<span
 									className={`text-sm text-muted-foreground/70 ${loading ? "animate-pulse" : ""}`}
 								>
-									{streamingRate.rate} tok/s
+									{streamingRate.rate} tok/s • {streamingRate.count} tokens
 								</span>
 							)}
 						</div>
@@ -315,6 +318,7 @@ const ChatMessages = ({ messages }: { messages: any[] }) => {
 								key={message.id}
 								message={message}
 								isLatest={index === messages.length - 1}
+								messages={messages}
 							/>
 						))
 					) : (

frontend/src/lib/utils/message.ts

@@ -0,0 +1,14 @@
+
+
+export function latestHumanMessage(messages: any[] | undefined | null) {
+	if (!Array.isArray(messages) || messages.length === 0) {
+		return null;
+	}
+	for (let i = messages.length - 1; i >= 0; i--) {
+		const msg = messages[i];
+		if (msg && msg.type === "human") {
+			return msg;
+		}
+	}
+	return null;
+}