aip-agents-binary 0.5.25b8__py3-none-any.whl → 0.6.0__py3-none-any.whl

aip_agents/agent/base_langgraph_agent.py

@@ -21,13 +21,13 @@ from concurrent.futures import Future
 from contextlib import suppress
 from contextvars import ContextVar
 from dataclasses import dataclass
-from typing import Annotated, Any
+from typing import Annotated, Any, cast
 
 from a2a.types import AgentCard
 from aiostream import stream as astream
-from gllm_core.event import EventEmitter
-from gllm_core.event.handler import StreamEventHandler
-from gllm_core.schema import Chunk
+from gllm_core.event import EventEmitter  # type: ignore[import-untyped]
+from gllm_core.event.handler import StreamEventHandler  # type: ignore[import-untyped]
+from gllm_core.schema import Chunk  # type: ignore[import-untyped]
 from langchain_core.messages import AIMessage, BaseMessage, ToolMessage
 from langchain_core.tools import BaseTool
 from langgraph.graph import StateGraph
@@ -197,6 +197,7 @@ class BaseLangGraphAgent(BaseAgent):
         self.enable_a2a_token_streaming = enable_a2a_token_streaming
         self.event_emitter = event_emitter
         self.checkpointer = checkpointer
+        self.tool_output_manager = None
 
         self._mem0_client: Any | None = None
         self.memory: BaseMemory | None = None
@@ -384,10 +385,13 @@ class BaseLangGraphAgent(BaseAgent):
         try:
             user_id = override_user_id or self.memory_agent_id
             if hasattr(self.memory, MemoryMethod.SEARCH):
-                results = self.memory.search(query=query, user_id=user_id, limit=self.memory_retrieval_limit)  # type: ignore[attr-defined]
+                results = self.memory.search(  # type: ignore[attr-defined]
+                    query=query,
+                    user_id=user_id,
+                    limit=self.memory_retrieval_limit,
+                )
                 return results if isinstance(results, list) else []
-            else:
-                return []
+            return []
         except Exception as e:  # noqa: BLE001
             logger.debug(f"Memory: search ignored error: {e}")
         return []
@@ -415,7 +419,11 @@ class BaseLangGraphAgent(BaseAgent):
                 future = save_async(user_text=str(user_text), ai_text=str(ai_text), user_id=user_id)
                 self._watch_memory_future(future, user_id)
             elif hasattr(self.memory, MemoryMethod.SAVE_INTERACTION):
-                self.memory.save_interaction(user_text=str(user_text), ai_text=str(ai_text), user_id=user_id)  # type: ignore[attr-defined]
+                self.memory.save_interaction(  # type: ignore[attr-defined]
+                    user_text=str(user_text),
+                    ai_text=str(ai_text),
+                    user_id=user_id,
+                )
             else:
                 logger.warning(
                     "Memory: save_interaction method NOT available on memory adapter "
@@ -560,7 +568,11 @@ class BaseLangGraphAgent(BaseAgent):
             return
 
         try:
-            tool.set_tool_config(tool_config_data)
+            set_tool_config = getattr(tool, "set_tool_config", None)
+            if callable(set_tool_config):
+                set_tool_config(tool_config_data)
+            else:
+                raise AttributeError("set_tool_config not available")
             logger.info(f"Agent '{self.name}': Configured tool '{tool.name}' with agent defaults: {tool_config_data}")
         except Exception as e:
             logger.warning(f"Agent '{self.name}': Failed to configure tool '{tool.name}': {e}")
@@ -598,7 +610,7 @@ class BaseLangGraphAgent(BaseAgent):
         self._sanitize_tool_names()
         try:
             if self.state_schema:
-                graph_builder = StateGraph(self.state_schema)
+                graph_builder: StateGraph = StateGraph(self.state_schema)
             else:
 
                 class DefaultAgentState(TypedDict):
@@ -715,7 +727,7 @@ class BaseLangGraphAgent(BaseAgent):
             return None
         last_item = list_state[-1]
         if isinstance(last_item, AIMessage) and getattr(last_item, "content", None) is not None:
-            output_content = last_item.content
+            output_content = self._normalize_event_content(last_item.content)
         elif isinstance(last_item, str):
             output_content = last_item
         return output_content
@@ -995,7 +1007,7 @@ class BaseLangGraphAgent(BaseAgent):
         key = self.thread_id_key or "thread_id"
         return configurable.get(key)
 
-    def _process_langgraph_event(self, event: dict[str, Any]) -> str | dict[str, Any] | None:
+    def _process_langgraph_event(self, event: Any) -> str | dict[str, Any] | A2AEvent | None:
         """Process a single LangGraph streaming event.
 
         Args:
@@ -1045,6 +1057,36 @@ class BaseLangGraphAgent(BaseAgent):
 
         return True
 
+    @staticmethod
+    def _normalize_usage_metadata(usage: Any | None) -> dict[str, Any] | None:
+        """Normalize usage metadata to a dictionary when possible.
+
+        Args:
+            usage: Usage metadata from LangChain messages.
+
+        Returns:
+            A dictionary copy when usage is available, otherwise None.
+        """
+        if usage is None:
+            return None
+        if isinstance(usage, dict):
+            return dict(usage)
+        return cast(dict[str, Any], usage)
+
+    @staticmethod
+    def _normalize_event_content(content: Any) -> str:
+        """Normalize event content to a string payload.
+
+        Args:
+            content: Raw content payload from LangChain/LangGraph.
+
+        Returns:
+            String representation suitable for A2A events.
+        """
+        if isinstance(content, str):
+            return content
+        return json.dumps(content)
+
     async def _stream_with_langgraph(self, query: str, **kwargs: Any) -> AsyncGenerator[str | dict[str, Any], None]:
         """Handle streaming for LangChain models using LangGraph's native streaming.
 
@@ -1135,9 +1177,13 @@ class BaseLangGraphAgent(BaseAgent):
             logger.info(f"Agent '{self.name}': Initializing MCP tools with persistent sessions.")
 
             # Add timeout for initialization to prevent hanging
-            await asyncio.wait_for(self.mcp_client.initialize(), timeout=30.0)
+            mcp_client = self.mcp_client
+            if mcp_client is None:
+                return
+
+            await asyncio.wait_for(mcp_client.initialize(), timeout=30.0)
 
-            mcp_tools = await self.mcp_client.get_tools()
+            mcp_tools = await mcp_client.get_tools()
 
             if not mcp_tools:
                 logger.warning(f"Agent '{self.name}': No MCP tools retrieved from configured servers.")
@@ -1169,7 +1215,7 @@ class BaseLangGraphAgent(BaseAgent):
                 logger.warning(f"Agent '{self.name}': Error during MCP client cleanup: {e}")
                 # Don't re-raise - cleanup should be best-effort
 
-    async def arun_a2a_stream(self, query: str, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
+    async def arun_a2a_stream(self, query: str, **kwargs: Any) -> AsyncGenerator[A2AEvent, None]:
         """Asynchronously streams the agent's response in A2A format.
 
         Args:
@@ -1190,7 +1236,7 @@ class BaseLangGraphAgent(BaseAgent):
         task_id: str | None = None,
         context_id: str | None = None,
         **kwargs: Any,
-    ) -> AsyncGenerator[dict[str, Any], None]:
+    ) -> AsyncGenerator[A2AEvent, None]:
         """Stream agent response as SSE-compatible chunks.
 
         This method wraps arun_a2a_stream and transforms output to the normalized
@@ -1222,7 +1268,8 @@ class BaseLangGraphAgent(BaseAgent):
         pii_mapping = kwargs.get("pii_mapping")
         transformer = SSEChunkTransformer(task_id=task_id, context_id=context_id, pii_mapping=pii_mapping)
         try:
-            async for chunk in transformer.transform_stream(self.arun_a2a_stream(query, **kwargs)):
+            stream = self.arun_a2a_stream(query, **kwargs)
+            async for chunk in transformer.transform_stream(stream):
                 yield chunk
         except Exception as e:
             # Lazy import to support optional guardrails dependency
@@ -1358,14 +1405,20 @@ class BaseLangGraphAgent(BaseAgent):
         Returns:
             A2AEvent with TOOL_CALL event type and structured tool information.
         """
-        tool_calls_details = [
-            {
-                "id": tool_call.get("id"),
-                "name": tool_call["name"],
-                "args": tool_call["args"],
-            }
-            for tool_call in message.tool_calls
-        ]
+        tool_calls_details: list[dict[str, Any]] = []
+        manager = getattr(self, "tool_output_manager", None)
+        thread_id = _THREAD_ID_CVAR.get()
+        for tool_call in message.tool_calls:
+            args = tool_call["args"]
+            if manager and thread_id and isinstance(args, dict):
+                args = manager.rewrite_args_with_latest_reference(args, thread_id)
+            tool_calls_details.append(
+                {
+                    "id": tool_call.get("id"),
+                    "name": tool_call["name"],
+                    "args": args,
+                }
+            )
         tool_names = [details["name"] for details in tool_calls_details]
 
         event = self._create_a2a_event(
@@ -1374,7 +1427,7 @@ class BaseLangGraphAgent(BaseAgent):
             tool_info={"tool_calls": tool_calls_details, "status": "running"},
             metadata={"status": Status.RUNNING},
             is_final=False,
-            step_usage=message.usage_metadata,
+            step_usage=self._normalize_usage_metadata(message.usage_metadata),
         )
 
         self._record_emitted_tool_calls(tool_calls_details)
@@ -1594,8 +1647,9 @@ class BaseLangGraphAgent(BaseAgent):
         """
         is_final_response = self._is_final_response(message)
         metadata = self._build_metadata_for_final_response(is_final_response, state)
+        raw_content = message.content
         content = deanonymize_final_response_content(
-            content=message.content,
+            content=raw_content if isinstance(raw_content, str) else json.dumps(raw_content),
             is_final_response=is_final_response,
             metadata=metadata,
         )
@@ -1605,7 +1659,7 @@ class BaseLangGraphAgent(BaseAgent):
             tool_info=None,
             metadata=metadata,
             is_final=is_final_response,
-            step_usage=message.usage_metadata,
+            step_usage=self._normalize_usage_metadata(message.usage_metadata),
         )
         return event, is_final_response
 
@@ -1882,7 +1936,7 @@ class BaseLangGraphAgent(BaseAgent):
         """
         current_thread_id: str | None = None
         try:
-            configurable = config.get("configurable", {})  # type: ignore[assignment]
+            configurable = config.get("configurable", {})
             thread_key = self.thread_id_key or "thread_id"
             current_thread_id = str(configurable.get(thread_key)) if configurable.get(thread_key) else None
         except Exception:
@@ -1964,7 +2018,7 @@ class BaseLangGraphAgent(BaseAgent):
         )
         return events, is_final, updated_message_count
 
-    async def _arun_a2a_stream(self, query: str, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
+    async def _arun_a2a_stream(self, query: str, **kwargs: Any) -> AsyncGenerator[A2AEvent, None]:
         """Internal implementation of arun_a2a_stream without MCP handling.
 
         Args:
@@ -2031,7 +2085,7 @@ class BaseLangGraphAgent(BaseAgent):
             enable_token_streaming=self.enable_a2a_token_streaming,
         )
 
-    async def _handle_streaming_process(self, context: "_StreamingContext") -> AsyncGenerator[dict[str, Any], None]:
+    async def _handle_streaming_process(self, context: "_StreamingContext") -> AsyncGenerator[A2AEvent, None]:
         """Handle the main streaming process including initial status and event processing.
 
         Args:
@@ -2048,7 +2102,7 @@ class BaseLangGraphAgent(BaseAgent):
             self._log_streaming_event_debug("process_stream_item", event)
             yield event
 
-    def _create_initial_status_event(self) -> dict[str, Any]:
+    def _create_initial_status_event(self) -> A2AEvent:
         """Create and setup the initial status event."""
         initial_status_event = self._create_a2a_event(
             event_type=A2AStreamEventType.STATUS_UPDATE, content=DefaultStepMessages.EN.value
@@ -2065,7 +2119,7 @@ class BaseLangGraphAgent(BaseAgent):
 
         return initial_status_event
 
-    async def _process_streaming_items(self, context: "_StreamingContext") -> AsyncGenerator[dict[str, Any], None]:
+    async def _process_streaming_items(self, context: "_StreamingContext") -> AsyncGenerator[A2AEvent, None]:
         """Process individual streaming items from the LangGraph execution.
 
         Handles the core streaming logic by iterating through items produced by
@@ -2098,9 +2152,7 @@ class BaseLangGraphAgent(BaseAgent):
             async for event in self._create_graph_stream_events(enhanced_input, context):
                 yield event
 
-    async def _process_a2a_streaming_with_tokens(
-        self, context: "_StreamingContext"
-    ) -> AsyncGenerator[dict[str, Any], None]:
+    async def _process_a2a_streaming_with_tokens(self, context: "_StreamingContext") -> AsyncGenerator[A2AEvent, None]:
         """Process A2A streaming with token streaming support using aiostream.
 
         Supports both LM Invoker and LangChain models by detecting the appropriate
@@ -2133,6 +2185,9 @@ class BaseLangGraphAgent(BaseAgent):
                 token_stream, enhanced_input = self._create_token_stream(context)
                 graph_stream = self._create_graph_stream_events(enhanced_input, context)
 
+                if token_stream is None:
+                    raise RuntimeError(f"Agent '{self.name}': Token stream not available for LM invoker.")
+
                 merged = astream.merge(token_stream, graph_stream)
                 async with merged.stream() as merged_stream:
                     async for event in merged_stream:
@@ -2148,7 +2203,7 @@ class BaseLangGraphAgent(BaseAgent):
             logger.error(f"Agent '{self.name}': Error during A2A token streaming: {e}")
             raise
 
-    async def _create_lm_invoker_token_stream(self) -> AsyncGenerator[dict[str, Any], None]:
+    async def _create_lm_invoker_token_stream(self) -> AsyncGenerator[A2AEvent, None]:
         """Generate A2A events from LM Invoker token stream.
 
         Uses StreamEventHandler to capture tokens emitted by LM Invoker.
@@ -2160,6 +2215,8 @@ class BaseLangGraphAgent(BaseAgent):
             RuntimeError: If no StreamEventHandler is found in event_emitter.
         """
         stream_handler = self._get_stream_handler()
+        if stream_handler is None:
+            raise RuntimeError(f"Agent '{self.name}': StreamEventHandler is required for token streaming.")
 
         try:
             async for event in stream_handler.stream():
@@ -2175,7 +2232,7 @@ class BaseLangGraphAgent(BaseAgent):
     def _create_token_stream(
         self,
         context: "_StreamingContext",
-    ) -> tuple[AsyncGenerator[dict[str, Any], None], dict[str, Any]]:
+    ) -> tuple[AsyncGenerator[A2AEvent, None] | None, dict[str, Any]]:
         """Create appropriate token stream and enhanced input for the active model backend.
 
         Args:
@@ -2197,7 +2254,7 @@ class BaseLangGraphAgent(BaseAgent):
 
     async def _create_graph_stream_events(
         self, enhanced_input: dict[str, Any], context: "_StreamingContext"
-    ) -> AsyncGenerator[dict[str, Any], None]:
+    ) -> AsyncGenerator[A2AEvent, None]:
         """Generate A2A events from graph execution.
 
         Args:
@@ -2216,8 +2273,9 @@ class BaseLangGraphAgent(BaseAgent):
             async for item in graph_execution:
                 stream_mode, stream_data = item
 
-                if stream_mode == StreamMode.MESSAGES:
-                    async for token_event in self._process_message_stream_item(stream_data):
+                if stream_mode == StreamMode.MESSAGES.value:
+                    message_data = cast(tuple[Any, dict[str, Any]], stream_data)
+                    async for token_event in self._process_message_stream_item(message_data):
                         yield token_event
                     continue
 
@@ -2236,10 +2294,10 @@ class BaseLangGraphAgent(BaseAgent):
         Returns:
             List of stream modes to use for graph execution.
         """
-        stream_modes = [StreamMode.VALUES, StreamMode.CUSTOM]
+        stream_modes = [StreamMode.VALUES.value, StreamMode.CUSTOM.value]
 
         if context.enable_token_streaming and not self._has_lm_invoker():
-            stream_modes.append(StreamMode.MESSAGES)
+            stream_modes.append(StreamMode.MESSAGES.value)
 
         return stream_modes
 
@@ -2249,7 +2307,7 @@ class BaseLangGraphAgent(BaseAgent):
         stream_mode: str,
         stream_data: Any,
         context: "_StreamingContext",
-    ) -> AsyncGenerator[dict[str, Any], None]:
+    ) -> AsyncGenerator[A2AEvent, None]:
         """Process a single graph stream item and yield A2A events.
 
         Args:
@@ -2261,10 +2319,12 @@ class BaseLangGraphAgent(BaseAgent):
         Yields:
             A2A events generated from the stream item.
         """
-        context.final_state = copy.copy(stream_data) if stream_mode == StreamMode.VALUES else context.final_state
+        context.final_state = copy.copy(stream_data) if stream_mode == StreamMode.VALUES.value else context.final_state
 
+        pending_artifacts = context.pending_artifacts if context.pending_artifacts is not None else []
+        seen_artifact_hashes = context.seen_artifact_hashes if context.seen_artifact_hashes is not None else set()
         events, is_final, context.processed_message_count = self._handle_stream_item(
-            item, context.pending_artifacts, context.seen_artifact_hashes, context.processed_message_count
+            item, pending_artifacts, seen_artifact_hashes, context.processed_message_count
         )
 
         if is_final:
@@ -2277,7 +2337,7 @@ class BaseLangGraphAgent(BaseAgent):
 
     async def _process_message_stream_item(
         self, message_data: tuple[Any, dict[str, Any]]
-    ) -> AsyncGenerator[dict[str, Any], None]:
+    ) -> AsyncGenerator[A2AEvent, None]:
         """Process message stream items to extract token events.
 
         The "messages" stream mode yields tuples of (AIMessageChunk, metadata).
@@ -2314,9 +2374,7 @@ class BaseLangGraphAgent(BaseAgent):
         except Exception as e:
             logger.error(f"Agent '{self.name}': Error processing message stream item: {e}")
 
-    def _update_final_response_for_streaming(
-        self, context: "_StreamingContext", event: dict[str, Any]
-    ) -> dict[str, Any]:
+    def _update_final_response_for_streaming(self, context: "_StreamingContext", event: A2AEvent) -> A2AEvent:
         """Update final response events with appropriate streaming configuration.
 
         For FINAL_RESPONSE events, this method updates the metadata and optionally clears
@@ -2330,13 +2388,17 @@ class BaseLangGraphAgent(BaseAgent):
             The processed event dictionary with updated metadata and content
         """
         if event.get("event_type") == A2AStreamEventType.FINAL_RESPONSE:
-            event["metadata"][MetadataFieldKeys.TOKEN_STREAMING] = False
+            metadata = event.get("metadata")
+            if not isinstance(metadata, dict):
+                metadata = {}
+                event["metadata"] = metadata
+            metadata[MetadataFieldKeys.TOKEN_STREAMING] = False
             if context.enable_token_streaming:
                 event["content"] = ""
-                event["metadata"][MetadataFieldKeys.TOKEN_STREAMING] = True
+                metadata[MetadataFieldKeys.TOKEN_STREAMING] = True
         return event
 
-    def _convert_raw_token_to_a2a_event(self, raw_event: str) -> dict[str, Any] | None:
+    def _convert_raw_token_to_a2a_event(self, raw_event: str) -> A2AEvent | None:
         """Parse raw token event into A2A event.
 
         Args:
@@ -2359,7 +2421,7 @@ class BaseLangGraphAgent(BaseAgent):
             logger.debug(f"Agent '{self.name}': Error parsing token event: {e}")
         return None
 
-    def _capture_final_content_if_needed(self, context: "_StreamingContext", event: dict[str, Any]) -> None:
+    def _capture_final_content_if_needed(self, context: "_StreamingContext", event: A2AEvent) -> None:
         """Capture final content from A2A events for memory persistence.
 
         Monitors A2A events for final response content and triggers early memory
@@ -2437,7 +2499,7 @@ class BaseLangGraphAgent(BaseAgent):
         except Exception:
             pass
 
-    async def _ensure_final_completion(self, context: "_StreamingContext") -> AsyncGenerator[dict[str, Any], None]:
+    async def _ensure_final_completion(self, context: "_StreamingContext") -> AsyncGenerator[A2AEvent, None]:
         """Ensure final completion events are yielded if not already done.
 
         Args:
@@ -2448,7 +2510,9 @@ class BaseLangGraphAgent(BaseAgent):
             dict[str, Any]: The final completion event.
         """
         if not context.final_event_yielded:
-            completion_event = self._create_completion_event(context.pending_artifacts, context.final_state)
+            pending_artifacts = context.pending_artifacts if context.pending_artifacts is not None else []
+            final_state = context.final_state or {}
+            completion_event = self._create_completion_event(pending_artifacts, final_state)
             self._log_streaming_event_debug("final_completion", completion_event)
             yield completion_event
 
@@ -2456,7 +2520,7 @@ class BaseLangGraphAgent(BaseAgent):
         self,
         context: "_StreamingContext",
         error: Exception,
-    ) -> AsyncGenerator[dict[str, Any], None]:
+    ) -> AsyncGenerator[A2AEvent, None]:
         """Handle streaming errors gracefully.
 
         Provides error handling for the A2A streaming process, ensuring errors
@@ -2473,11 +2537,14 @@ class BaseLangGraphAgent(BaseAgent):
                 error message, optionally including any pending artifacts.
         """
         logger.error(f"Error in agent stream: {error}", exc_info=True)
-        error_event = {"status": "failed", "content": f"Error: {str(error)}"}
-
-        if context.pending_artifacts:
-            error_event["artifacts"] = context.pending_artifacts
-
+        error_event = self._create_a2a_event(
+            event_type=A2AStreamEventType.ERROR,
+            content=f"Error: {str(error)}",
+            metadata={"status": "failed"},
+            artifacts=context.pending_artifacts,
+            is_final=True,
+        )
+        error_event["status"] = "failed"
         self._log_streaming_event_debug("error_event", error_event)
         yield error_event
 
@@ -2533,7 +2600,7 @@ class BaseLangGraphAgent(BaseAgent):
 
         return metadata
 
-    def _create_completion_event(self, pending_artifacts: list, final_state: dict[str, Any]):
+    def _create_completion_event(self, pending_artifacts: list, final_state: dict[str, Any]) -> A2AEvent:
         """Helper to create the completion event with artifacts and references if available.
 
         This method is used to create the completion event with artifacts and references if available.
@@ -2587,7 +2654,7 @@ class BaseLangGraphAgent(BaseAgent):
         else:
             return tool_name[:4]
 
-    def _generate_tool_call_step_id(self, tool_info: dict[str, Any], counter: int) -> str:
+    def _generate_tool_call_step_id(self, tool_info: dict[str, Any] | None, counter: int) -> str:
         """Generate step_id for tool call events.
 
         Args:
@@ -2623,7 +2690,7 @@ class BaseLangGraphAgent(BaseAgent):
 
         return f"{category}_{combined_name}_parent_{counter:03d}"
 
-    def _generate_tool_result_step_id(self, tool_info: dict[str, Any], counter: int) -> str:
+    def _generate_tool_result_step_id(self, tool_info: dict[str, Any] | None, counter: int) -> str:
         """Generate step_id for tool result events.
 
         Args:
@@ -2749,7 +2816,7 @@ class BaseLangGraphAgent(BaseAgent):
     def _create_a2a_event(  # noqa: PLR0913
         self,
         event_type: A2AStreamEventType,
-        content: str,
+        content: Any,
         metadata: dict[str, Any] | None = None,
         tool_info: dict[str, Any] | None = None,
         thinking_and_activity_info: dict[str, Any] | None = None,
@@ -2787,9 +2854,11 @@ class BaseLangGraphAgent(BaseAgent):
         # Inject cumulative time since the first STATUS_UPDATE for this thread
         # Do not set cumulative time here; server executor enforces it for all SSE events
 
+        normalized_content = self._normalize_event_content(content)
+
         event = {
             "event_type": event_type,
-            "content": content,
+            "content": normalized_content,
             "metadata": enriched_metadata,
             "tool_info": tool_info,
             "is_final": is_final,
@@ -2803,7 +2872,7 @@ class BaseLangGraphAgent(BaseAgent):
             event["thinking_and_activity_info"] = thinking_and_activity_info
 
         try:
-            content_preview = content if isinstance(content, str) else str(content)
+            content_preview = normalized_content
             logger.info(
                 "A2A emitting event: type=%s step_id=%s final=%s preview=%s",
                 getattr(event_type, "value", event_type),