Merge pull request #8 from satendrakumar/develop

Code refactoring

Author: satendrakumar

a2a_client_app.py

@@ -1,11 +1,11 @@
 import asyncio
 import uuid
 
-from src.a2a.a2a_client import A2ASimpleClient
+from src.a2a.a2a_client import A2AClient
 
 
 async def main():
-    a2a_client: A2ASimpleClient = A2ASimpleClient()
+    a2a_client: A2AClient = A2AClient()
     agent_host_url = "http://localhost:8000/a2a"
 
     trending_task = a2a_client.create_task(
@@ -23,10 +23,7 @@ async def main():
     print(trending_topics)
     print(analysis)
 
-    print(
-        "###############################################################################################################"
-    )
-    print("Multi-turn conversation with an agent............................")
+    print("Multi-turn conversation with an agent...................")
     context_id = str(uuid.uuid4())
     print(f"Starting conversation with context_id: {context_id}")
 

src/a2a/a2a_client.py

@@ -9,7 +9,7 @@
 AGENT_CARD_PATH = "/agent-card.json"
 
 
-class A2ASimpleClient:
+class A2AClient:
     def __init__(self, default_timeout: float = 240.0):
         # Cache for agent metadata
         self._agent_info_cache: dict[str, dict[str, Any] | None] = {}

src/a2a/a2a_fastapi_app.py

@@ -1,47 +1,25 @@
 from collections.abc import Callable
 from typing import Any
 
-from fastapi import APIRouter, FastAPI
-from google.adk import Runner
-from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor, A2aAgentExecutorConfig
-from google.adk.agents import LlmAgent
-from google.adk.artifacts import InMemoryArtifactService
-from google.adk.memory import InMemoryMemoryService
-from google.adk.sessions import InMemorySessionService
-from starlette.applications import Starlette
-
 from a2a.server.apps.jsonrpc.jsonrpc_app import CallContextBuilder, JSONRPCApplication
 from a2a.server.context import ServerCallContext
-from a2a.server.request_handlers import DefaultRequestHandler
 from a2a.server.request_handlers.request_handler import RequestHandler
-from a2a.server.tasks import InMemoryTaskStore
 from a2a.types import AgentCard
 from a2a.utils.constants import AGENT_CARD_WELL_KNOWN_PATH, DEFAULT_RPC_URL, EXTENDED_AGENT_CARD_PATH
-
-
-def get_agent_request_handler(agent: LlmAgent):
-    runner = Runner(
-        app_name=agent.name,
-        agent=agent,
-        artifact_service=InMemoryArtifactService(),
-        session_service=InMemorySessionService(),
-        memory_service=InMemoryMemoryService(),
-    )
-    config = A2aAgentExecutorConfig()
-    executor = A2aAgentExecutor(runner=runner, config=config)
-    return DefaultRequestHandler(agent_executor=executor, task_store=InMemoryTaskStore())
+from fastapi import APIRouter, FastAPI
+from starlette.applications import Starlette
 
 
 class A2AFastApiApp(JSONRPCApplication):
     def __init__(
-        self,
-        fastapi_app: FastAPI,
-        agent_card: AgentCard,
-        http_handler: RequestHandler,
-        extended_agent_card: AgentCard | None = None,
-        context_builder: CallContextBuilder | None = None,
-        card_modifier: Callable[[AgentCard], AgentCard] | None = None,
-        extended_card_modifier: Callable[[AgentCard, ServerCallContext], AgentCard] | None = None,
+            self,
+            fastapi_app: FastAPI,
+            agent_card: AgentCard,
+            http_handler: RequestHandler,
+            extended_agent_card: AgentCard | None = None,
+            context_builder: CallContextBuilder | None = None,
+            card_modifier: Callable[[AgentCard], AgentCard] | None = None,
+            extended_card_modifier: Callable[[AgentCard, ServerCallContext], AgentCard] | None = None,
     ):
         super().__init__(
             agent_card=agent_card,
@@ -54,11 +32,11 @@ def __init__(
         self.fastapi_app = fastapi_app
 
     def build(
-        self,
-        agent_card_url: str = AGENT_CARD_WELL_KNOWN_PATH,
-        rpc_url: str = DEFAULT_RPC_URL,
-        extended_agent_card_url: str = EXTENDED_AGENT_CARD_PATH,
-        **kwargs: Any,
+            self,
+            agent_card_url: str = AGENT_CARD_WELL_KNOWN_PATH,
+            rpc_url: str = DEFAULT_RPC_URL,
+            extended_agent_card_url: str = EXTENDED_AGENT_CARD_PATH,
+            **kwargs: Any,
     ) -> Starlette:
         name_prefix = rpc_url.replace("/", "")
         router = APIRouter()

src/a2a/a2a_request_handler.py

@@ -0,0 +1,23 @@
+from a2a.server.request_handlers import DefaultRequestHandler
+from a2a.server.tasks import InMemoryTaskStore
+from google.adk import Runner
+from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor, A2aAgentExecutorConfig
+from google.adk.agents import LlmAgent
+from google.adk.artifacts import InMemoryArtifactService
+from google.adk.memory import InMemoryMemoryService
+from google.adk.sessions import InMemorySessionService
+
+
+class A2ARequestHandler:
+    @staticmethod
+    def get_request_handler(agent: LlmAgent):
+        runner = Runner(
+            app_name=agent.name,
+            agent=agent,
+            artifact_service=InMemoryArtifactService(),
+            session_service=InMemorySessionService(),
+            memory_service=InMemoryMemoryService(),
+        )
+        config = A2aAgentExecutorConfig()
+        executor = A2aAgentExecutor(runner=runner, config=config)
+        return DefaultRequestHandler(agent_executor=executor, task_store=InMemoryTaskStore())

src/a2a/a2a_utils.py

@@ -1,24 +1,34 @@
 from typing import Callable
 
+from a2a.types import AgentCard
 from fastapi import FastAPI
 from google.adk.agents import LlmAgent
 
-from a2a.types import AgentCard
-from src.a2a.a2a_fastapi_app import A2AFastApiApp, get_agent_request_handler
+from src.a2a.a2a_fastapi_app import A2AFastApiApp
+from src.a2a.a2a_request_handler import A2ARequestHandler
 
 
 class A2AUtils:
+    """
+    Utility class for A2A (Agent-to-Agent) communication.
+
+    This class provides static methods to assist in constructing and managing agent
+    communication processes in the A2A framework. It is specifically designed to help
+    configure and set up agents, their request handlers, and their associated API endpoints.
+
+    """
     @staticmethod
     def build(
-        name: str,
-        get_agent: Callable[[str], LlmAgent],
-        get_agent_card: Callable[[str], AgentCard],
-        model_name: str,
-        agent_base_url: str,
-        app: FastAPI,
+            name: str,
+            get_agent: Callable[[str], LlmAgent],
+            get_agent_card: Callable[[str], AgentCard],
+            model_name: str,
+            agent_base_url: str,
+            app: FastAPI,
     ) -> None:
         agent = get_agent(model_name)
-        agent_request_handler = get_agent_request_handler(agent)
+        agent_request_handler = A2ARequestHandler.get_request_handler(agent)
         agent_card = get_agent_card(f"{agent_base_url}/{name}/")
         agent_server = A2AFastApiApp(fastapi_app=app, agent_card=agent_card, http_handler=agent_request_handler)
         agent_server.build(rpc_url=f"/{name}/", agent_card_url=f"/{name}/{{path:path}}")
+