Refactoring to address comments

Change-Id: I5eca290bd814eb7480175d14822d7833f815bd3a

Author: kthota-g

examples/demo/agent_proxy.py

@@ -6,23 +6,23 @@
 from a2a.types import (
     Artifact,
     CancelTaskRequest,
-    CancelTaskSuccessResponse,
+    CancelTaskResponse,
     JSONRPCErrorResponse,
     SendTaskRequest,
+    SendTaskResponse,
     SendTaskStreamingRequest,
     SendTaskStreamingResponse,
     SendTaskStreamingSuccessResponse,
     SendTaskSuccessResponse,
     Task,
     TaskArtifactUpdateEvent,
-    TaskNotCancelableError,
     TaskResubscriptionRequest,
     TaskState,
     TaskStatus,
     TaskStatusUpdateEvent,
     UnsupportedOperationError,
 )
-from a2a.utils import get_text_artifact
+from a2a.utils import build_text_artifact
 
 
 class FakeAgent:
@@ -45,17 +45,19 @@ def __init__(self):
 
     async def on_send(
         self, task: Task, request: SendTaskRequest
-    ) -> SendTaskSuccessResponse | JSONRPCErrorResponse:
+    ) -> SendTaskResponse:
         result = await self.agent.invoke()
 
         if not task.artifacts:
             task.artifacts = []
 
-        artifact: Artifact = get_text_artifact(result, len(task.artifacts))
+        artifact: Artifact = build_text_artifact(result, len(task.artifacts))
         task.artifacts.append(artifact)
         task.status.state = TaskState.completed
 
-        return SendTaskSuccessResponse(id=request.id, result=task)
+        return SendTaskResponse(
+            root=SendTaskSuccessResponse(id=request.id, result=task)
+        )
 
     async def on_send_subscribe(  # type: ignore
         self, task: Task, request: SendTaskStreamingRequest
@@ -67,7 +69,7 @@ async def on_send_subscribe(  # type: ignore
         async for chunk in self.agent.stream():
             artifact_update = TaskArtifactUpdateEvent(
                 id=task.id,
-                artifact=get_text_artifact(chunk, new_index),
+                artifact=build_text_artifact(chunk, new_index),
                 append=i > 0,
                 lastChunk=False,  # TODO: set this value, but is this needed?
             )
@@ -91,9 +93,11 @@ async def on_send_subscribe(  # type: ignore
 
     async def on_cancel(
         self, task: Task, request: CancelTaskRequest
-    ) -> CancelTaskSuccessResponse | JSONRPCErrorResponse:
-        return JSONRPCErrorResponse(
-            id=request.id, error=TaskNotCancelableError()
+    ) -> CancelTaskResponse:
+        return CancelTaskResponse(
+            root=JSONRPCErrorResponse(
+                id=request.id, error=UnsupportedOperationError()
+            )
         )
 
     async def on_resubscribe(  # type: ignore

examples/langgraph/agent.py

@@ -1,12 +1,20 @@
+import logging
+
+from collections.abc import AsyncIterable
+from typing import Any, Literal
+
+import httpx
+
+from langchain_core.messages import AIMessage, ToolMessage
+from langchain_core.runnables.config import (
+    RunnableConfig,
+)
+from langchain_core.tools import tool  # type: ignore
 from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_core.tools import tool
-from langgraph.prebuilt import create_react_agent
 from langgraph.checkpoint.memory import MemorySaver
-from langchain_core.messages import AIMessage, ToolMessage
-import httpx
-from typing import Any, AsyncIterable, Literal
+from langgraph.prebuilt import create_react_agent  # type: ignore
 from pydantic import BaseModel
-import logging
+
 
 logger = logging.getLogger(__name__)
 
@@ -15,9 +23,9 @@
 
 @tool
 def get_exchange_rate(
-    currency_from: str = "USD",
-    currency_to: str = "EUR",
-    currency_date: str = "latest",
+    currency_from: str = 'USD',
+    currency_to: str = 'EUR',
+    currency_date: str = 'latest',
 ):
     """Use this to get current exchange rate.
 
@@ -28,118 +36,120 @@ def get_exchange_rate(
 
     Returns:
         A dictionary containing the exchange rate data, or an error message if the request fails.
-    """    
+    """
     try:
         response = httpx.get(
-            f"https://api.frankfurter.app/{currency_date}",
-            params={"from": currency_from, "to": currency_to},
+            f'https://api.frankfurter.app/{currency_date}',
+            params={'from': currency_from, 'to': currency_to},
         )
         response.raise_for_status()
 
         data = response.json()
-        if "rates" not in data:
-            logger.error(f"rates not found in response: {data}")
-            return {"error": "Invalid API response format."}
-        logger.info(f"API response: {data}")
+        if 'rates' not in data:
+            logger.error(f'rates not found in response: {data}')
+            return {'error': 'Invalid API response format.'}
+        logger.info(f'API response: {data}')
         return data
     except httpx.HTTPError as e:
-        logger.error(f"API request failed: {e}")
-        return {"error": f"API request failed: {e}"}
+        logger.error(f'API request failed: {e}')
+        return {'error': f'API request failed: {e}'}
     except ValueError:
-        logger.error("Invalid JSON response from API")
-        return {"error": "Invalid JSON response from API."}
+        logger.error('Invalid JSON response from API')
+        return {'error': 'Invalid JSON response from API.'}
 
 
 class ResponseFormat(BaseModel):
     """Respond to the user in this format."""
-    status: Literal["input_required", "completed", "error"] = "input_required"
+
+    status: Literal['input_required', 'completed', 'error'] = 'input_required'
     message: str
 
+
 class CurrencyAgent:
+    """Currency Conversion Agent Example."""
 
     SYSTEM_INSTRUCTION = (
-        "You are a specialized assistant for currency conversions. "
+        'You are a specialized assistant for currency conversions. '
         "Your sole purpose is to use the 'get_exchange_rate' tool to answer questions about currency exchange rates. "
-        "If the user asks about anything other than currency conversion or exchange rates, "
-        "politely state that you cannot help with that topic and can only assist with currency-related queries. "
-        "Do not attempt to answer unrelated questions or use tools for other purposes."
+        'If the user asks about anything other than currency conversion or exchange rates, '
+        'politely state that you cannot help with that topic and can only assist with currency-related queries. '
+        'Do not attempt to answer unrelated questions or use tools for other purposes.'
     )
 
     RESPONSE_FORMAT_INSTRUCTION: str = (
-        "Select status as completed if the request is complete"
-        "Select status as input_required if the input is a question to the user"
-        "Set response status to error if the input indicates an error"
+        'Select status as completed if the request is complete'
+        'Select status as input_required if the input is a question to the user'
+        'Set response status to error if the input indicates an error'
     )
-     
+
     def __init__(self):
-        self.model = ChatGoogleGenerativeAI(model="gemini-2.0-flash")
+        self.model = ChatGoogleGenerativeAI(model='gemini-2.0-flash')
         self.tools = [get_exchange_rate]
 
         self.graph = create_react_agent(
-            self.model, tools=self.tools, checkpointer=memory, prompt = self.SYSTEM_INSTRUCTION, response_format=(self.RESPONSE_FORMAT_INSTRUCTION, ResponseFormat)
+            self.model,
+            tools=self.tools,
+            checkpointer=memory,
+            prompt=self.SYSTEM_INSTRUCTION,
+            response_format=(self.RESPONSE_FORMAT_INSTRUCTION, ResponseFormat),
         )
 
     def invoke(self, query: str, sessionId: str) -> dict[str, Any]:
-        config: dict[str, Any] = {"configurable": {"thread_id": sessionId}}
-        self.graph.invoke({"messages": [("user", query)]}, config)        
-        response = self.get_agent_response(config)        
-        return response
-        
-
-    async def stream(self, query: str, sessionId: str) -> AsyncIterable[dict[str, Any]]:
-        inputs: dict[str, Any] = {"messages": [("user", query)]}
-        config: dict[str, Any] = {"configurable": {"thread_id": sessionId}}
-
-        for item in self.graph.stream(inputs, config, stream_mode="values"):
-            message = item["messages"][-1]
+        config: RunnableConfig = {'configurable': {'thread_id': sessionId}}
+        self.graph.invoke({'messages': [('user', query)]}, config)
+        return self.get_agent_response(config)
+
+    async def stream(
+        self, query: str, sessionId: str
+    ) -> AsyncIterable[dict[str, Any]]:
+        inputs: dict[str, Any] = {'messages': [('user', query)]}
+        config: RunnableConfig = {'configurable': {'thread_id': sessionId}}
+
+        for item in self.graph.stream(inputs, config, stream_mode='values'):
+            message = item['messages'][-1]
             if (
                 isinstance(message, AIMessage)
                 and message.tool_calls
                 and len(message.tool_calls) > 0
             ):
                 yield {
-                    "is_task_complete": False,
-                    "require_user_input": False,
-                    "content": "Looking up the exchange rates...",
+                    'is_task_complete': False,
+                    'require_user_input': False,
+                    'content': 'Looking up the exchange rates...',
                 }
             elif isinstance(message, ToolMessage):
                 yield {
-                    "is_task_complete": False,
-                    "require_user_input": False,
-                    "content": "Processing the exchange rates..",
-                }            
-        
+                    'is_task_complete': False,
+                    'require_user_input': False,
+                    'content': 'Processing the exchange rates..',
+                }
+
         yield self.get_agent_response(config)
 
-        
-    def get_agent_response(self, config: dict[str, Any]) -> dict[str, Any]:
-        current_state = self.graph.get_state(config)  
-        
-        structured_response = current_state.values.get('structured_response')        
-        if structured_response and isinstance(structured_response, ResponseFormat): 
-            if structured_response.status == "input_required":
-                return {
-                    "is_task_complete": False,
-                    "require_user_input": True,
-                    "content": structured_response.message
-                }
-            elif structured_response.status == "error":
+    def get_agent_response(self, config: RunnableConfig) -> dict[str, Any]:
+        current_state = self.graph.get_state(config)
+
+        structured_response = current_state.values.get('structured_response')
+        if structured_response and isinstance(
+            structured_response, ResponseFormat
+        ):
+            if structured_response.status in {'input_required', 'error'}:
                 return {
-                    "is_task_complete": False,
-                    "require_user_input": True,
-                    "content": structured_response.message
+                    'is_task_complete': False,
+                    'require_user_input': True,
+                    'content': structured_response.message,
                 }
-            elif structured_response.status == "completed":
+            if structured_response.status == 'completed':
                 return {
-                    "is_task_complete": True,
-                    "require_user_input": False,
-                    "content": structured_response.message
+                    'is_task_complete': True,
+                    'require_user_input': False,
+                    'content': structured_response.message,
                 }
 
         return {
-            "is_task_complete": False,
-            "require_user_input": True,
-            "content": "We are unable to process your request at the moment. Please try again.",
+            'is_task_complete': False,
+            'require_user_input': True,
+            'content': 'We are unable to process your request at the moment. Please try again.',
         }
 
-    SUPPORTED_CONTENT_TYPES = ["text", "text/plain"]
+    SUPPORTED_CONTENT_TYPES = ['text', 'text/plain']