kagent-adk 0.0.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

kagent_adk/__init__.py

@@ -1,10 +1,8 @@
 import importlib.metadata
 
 from .a2a import KAgentApp
-from .kagent_session_service import KAgentSessionService
-from .kagent_task_store import KAgentTaskStore
 from .models import AgentConfig
 
 __version__ = importlib.metadata.version("kagent_adk")
 
-__all__ = ["KAgentSessionService", "KAgentTaskStore", "KAgentApp", "AgentConfig"]
+__all__ = ["KAgentApp", "AgentConfig"]

kagent_adk/_agent_executor.py

@@ -0,0 +1,259 @@
+from __future__ import annotations
+
+import inspect
+import logging
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Awaitable, Callable, Optional
+
+from a2a.server.agent_execution import AgentExecutor
+from a2a.server.agent_execution.context import RequestContext
+from a2a.server.events.event_queue import EventQueue
+from a2a.types import (
+    Artifact,
+    Message,
+    Role,
+    TaskArtifactUpdateEvent,
+    TaskState,
+    TaskStatus,
+    TaskStatusUpdateEvent,
+    TextPart,
+)
+from google.adk.a2a.converters.event_converter import convert_event_to_a2a_events
+from google.adk.a2a.converters.request_converter import convert_a2a_request_to_adk_run_args
+from google.adk.a2a.converters.utils import _get_adk_metadata_key
+from google.adk.a2a.executor.task_result_aggregator import TaskResultAggregator
+from google.adk.runners import Runner
+from google.adk.utils.feature_decorator import experimental
+from pydantic import BaseModel
+from typing_extensions import override
+
+logger = logging.getLogger("google_adk." + __name__)
+
+
+@experimental
+class A2aAgentExecutorConfig(BaseModel):
+    """Configuration for the A2aAgentExecutor."""
+
+    pass
+
+
+# This class is a copy of the A2aAgentExecutor class in the ADK sdk,
+# with the following changes:
+# - The runner is ALWAYS a callable that returns a Runner instance
+# - The runner is cleaned up at the end of the execution
+@experimental
+class A2aAgentExecutor(AgentExecutor):
+    """An AgentExecutor that runs an ADK Agent against an A2A request and
+    publishes updates to an event queue.
+    """
+
+    def __init__(
+        self,
+        *,
+        runner: Callable[..., Runner | Awaitable[Runner]],
+        config: Optional[A2aAgentExecutorConfig] = None,
+    ):
+        super().__init__()
+        self._runner = runner
+        self._config = config
+
+    async def _resolve_runner(self) -> Runner:
+        """Resolve the runner, handling cases where it's a callable that returns a Runner."""
+        if callable(self._runner):
+            # Call the function to get the runner
+            result = self._runner()
+
+            # Handle async callables
+            if inspect.iscoroutine(result):
+                resolved_runner = await result
+            else:
+                resolved_runner = result
+
+            # Ensure we got a Runner instance
+            if not isinstance(resolved_runner, Runner):
+                raise TypeError(f"Callable must return a Runner instance, got {type(resolved_runner)}")
+
+            return resolved_runner
+
+        raise TypeError(
+            f"Runner must be a Runner instance or a callable that returns a Runner, got {type(self._runner)}"
+        )
+
+    @override
+    async def cancel(self, context: RequestContext, event_queue: EventQueue):
+        """Cancel the execution."""
+        # TODO: Implement proper cancellation logic if needed
+        raise NotImplementedError("Cancellation is not supported")
+
+    @override
+    async def execute(
+        self,
+        context: RequestContext,
+        event_queue: EventQueue,
+    ):
+        """Executes an A2A request and publishes updates to the event queue
+        specified. It runs as following:
+        * Takes the input from the A2A request
+        * Convert the input to ADK input content, and runs the ADK agent
+        * Collects output events of the underlying ADK Agent
+        * Converts the ADK output events into A2A task updates
+        * Publishes the updates back to A2A server via event queue
+        """
+        if not context.message:
+            raise ValueError("A2A request must have a message")
+
+        # for new task, create a task submitted event
+        if not context.current_task:
+            await event_queue.enqueue_event(
+                TaskStatusUpdateEvent(
+                    task_id=context.task_id,
+                    status=TaskStatus(
+                        state=TaskState.submitted,
+                        message=context.message,
+                        timestamp=datetime.now(timezone.utc).isoformat(),
+                    ),
+                    context_id=context.context_id,
+                    final=False,
+                )
+            )
+
+        # Handle the request and publish updates to the event queue
+        runner = await self._resolve_runner()
+        try:
+            await self._handle_request(context, event_queue, runner)
+        except Exception as e:
+            logger.error("Error handling A2A request: %s", e, exc_info=True)
+            # Publish failure event
+            try:
+                await event_queue.enqueue_event(
+                    TaskStatusUpdateEvent(
+                        task_id=context.task_id,
+                        status=TaskStatus(
+                            state=TaskState.failed,
+                            timestamp=datetime.now(timezone.utc).isoformat(),
+                            message=Message(
+                                message_id=str(uuid.uuid4()),
+                                role=Role.agent,
+                                parts=[TextPart(text=str(e))],
+                            ),
+                        ),
+                        context_id=context.context_id,
+                        final=True,
+                    )
+                )
+            except Exception as enqueue_error:
+                logger.error("Failed to publish failure event: %s", enqueue_error, exc_info=True)
+        finally:
+            await runner.close()
+
+    async def _handle_request(
+        self,
+        context: RequestContext,
+        event_queue: EventQueue,
+        runner: Runner,
+    ):
+        # Convert the a2a request to ADK run args
+        run_args = convert_a2a_request_to_adk_run_args(context)
+
+        # ensure the session exists
+        session = await self._prepare_session(context, run_args, runner)
+
+        # create invocation context
+        invocation_context = runner._new_invocation_context(
+            session=session,
+            new_message=run_args["new_message"],
+            run_config=run_args["run_config"],
+        )
+
+        # publish the task working event
+        await event_queue.enqueue_event(
+            TaskStatusUpdateEvent(
+                task_id=context.task_id,
+                status=TaskStatus(
+                    state=TaskState.working,
+                    timestamp=datetime.now(timezone.utc).isoformat(),
+                ),
+                context_id=context.context_id,
+                final=False,
+                metadata={
+                    _get_adk_metadata_key("app_name"): runner.app_name,
+                    _get_adk_metadata_key("user_id"): run_args["user_id"],
+                    _get_adk_metadata_key("session_id"): run_args["session_id"],
+                },
+            )
+        )
+
+        task_result_aggregator = TaskResultAggregator()
+        async for adk_event in runner.run_async(**run_args):
+            for a2a_event in convert_event_to_a2a_events(
+                adk_event, invocation_context, context.task_id, context.context_id
+            ):
+                task_result_aggregator.process_event(a2a_event)
+                await event_queue.enqueue_event(a2a_event)
+
+        # publish the task result event - this is final
+        if (
+            task_result_aggregator.task_state == TaskState.working
+            and task_result_aggregator.task_status_message is not None
+            and task_result_aggregator.task_status_message.parts
+        ):
+            # if task is still working properly, publish the artifact update event as
+            # the final result according to a2a protocol.
+            await event_queue.enqueue_event(
+                TaskArtifactUpdateEvent(
+                    task_id=context.task_id,
+                    last_chunk=True,
+                    context_id=context.context_id,
+                    artifact=Artifact(
+                        artifact_id=str(uuid.uuid4()),
+                        parts=task_result_aggregator.task_status_message.parts,
+                    ),
+                )
+            )
+            # public the final status update event
+            await event_queue.enqueue_event(
+                TaskStatusUpdateEvent(
+                    task_id=context.task_id,
+                    status=TaskStatus(
+                        state=TaskState.completed,
+                        timestamp=datetime.now(timezone.utc).isoformat(),
+                    ),
+                    context_id=context.context_id,
+                    final=True,
+                )
+            )
+        else:
+            await event_queue.enqueue_event(
+                TaskStatusUpdateEvent(
+                    task_id=context.task_id,
+                    status=TaskStatus(
+                        state=task_result_aggregator.task_state,
+                        timestamp=datetime.now(timezone.utc).isoformat(),
+                        message=task_result_aggregator.task_status_message,
+                    ),
+                    context_id=context.context_id,
+                    final=True,
+                )
+            )
+
+    async def _prepare_session(self, context: RequestContext, run_args: dict[str, Any], runner: Runner):
+        session_id = run_args["session_id"]
+        # create a new session if not exists
+        user_id = run_args["user_id"]
+        session = await runner.session_service.get_session(
+            app_name=runner.app_name,
+            user_id=user_id,
+            session_id=session_id,
+        )
+        if session is None:
+            session = await runner.session_service.create_session(
+                app_name=runner.app_name,
+                user_id=user_id,
+                state={},
+                session_id=session_id,
+            )
+            # Update run_args with the new session_id
+            run_args["session_id"] = session.id
+
+        return session

kagent_adk/a2a.py

@@ -1,10 +1,11 @@
 #! /usr/bin/env python3
 import faulthandler
+import inspect
 import logging
 import os
 import sys
 from contextlib import asynccontextmanager
-from typing import Callable
+from typing import Awaitable, Callable, override
 
 import httpx
 from a2a.auth.user import User
@@ -16,14 +17,14 @@ from a2a.server.tasks import TaskStore
 from a2a.types import AgentCard, MessageSendParams, Task
 from fastapi import FastAPI, Request
 from fastapi.responses import PlainTextResponse
-from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor
 from google.adk.agents import BaseAgent
 from google.adk.runners import Runner
 from google.adk.sessions import InMemorySessionService
 from google.genai import types
 
-from .kagent_session_service import KAgentSessionService
-from .kagent_task_store import KAgentTaskStore
+from ._agent_executor import A2aAgentExecutor
+from ._session_service import KAgentSessionService
+from ._task_store import KAgentTaskStore
 
 # --- Constants ---
 USER_ID = "admin@kagent.dev"
@@ -89,7 +90,7 @@ kagent_url_override = os.getenv("KAGENT_URL")
 class KAgentApp:
     def __init__(
         self,
-        root_agent: BaseAgent | Callable[[], BaseAgent],
+        root_agent: BaseAgent,
         agent_card: AgentCard,
         kagent_url: str,
         app_name: str,
@@ -103,33 +104,15 @@ class KAgentApp:
         http_client = httpx.AsyncClient(base_url=kagent_url_override or self.kagent_url)
         session_service = KAgentSessionService(http_client)
 
-        if isinstance(self.root_agent, Callable):
-            agent_factory = self.root_agent
-
-            def create_runner() -> Runner:
-                return Runner(
-                    agent=agent_factory(),
-                    app_name=self.app_name,
-                    session_service=session_service,
-                )
-
-            runner = create_runner
-        elif isinstance(self.root_agent, BaseAgent):
-            agent_instance = self.root_agent
-
-            def create_runner() -> Runner:
-                return Runner(
-                    agent=agent_instance,
-                    app_name=self.app_name,
-                    session_service=session_service,
-                )
-
-            runner = create_runner
-        else:
-            raise ValueError(f"Invalid root agent: {self.root_agent}")
+        def create_runner() -> Runner:
+            return Runner(
+                agent=self.root_agent,
+                app_name=self.app_name,
+                session_service=session_service,
+            )
 
         agent_executor = A2aAgentExecutor(
-            runner=runner,
+            runner=create_runner,
         )
 
         kagent_task_store = KAgentTaskStore(http_client)
@@ -146,12 +129,6 @@ class KAgentApp:
             http_handler=request_handler,
         )
 
-        # @asynccontextmanager
-        # async def agent_lifespan(app: FastAPI):
-        #     yield
-        #     if isinstance(runner, Runner):
-        #         await runner.close()
-
         faulthandler.enable()
         app = FastAPI()
 

{kagent_adk-0.0.1.dist-info → kagent_adk-0.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kagent-adk
-Version: 0.0.1
+Version: 0.5.0
 Summary: kagent-adk is an sdk for integrating adk agents with kagent
 Requires-Python: >=3.12.11
 Requires-Dist: a2a-sdk>=0.2.16
@@ -23,6 +23,8 @@ Requires-Dist: opentelemetry-sdk>=1.32.0
 Requires-Dist: protobuf>=6.31.1
 Requires-Dist: pydantic>=2.5.0
 Requires-Dist: typing-extensions>=4.8.0
+Provides-Extra: memory
+Requires-Dist: psutil>=6.1.0; extra == 'memory'
 Provides-Extra: test
 Requires-Dist: pytest-asyncio>=0.25.3; extra == 'test'
 Requires-Dist: pytest>=8.3.5; extra == 'test'

kagent_adk-0.5.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+kagent_adk/__init__.py,sha256=3oB8gbSzsvVmqV8w-BGKkRlH3JPfK6o27AfOD6wAd8o,182
+kagent_adk/_agent_executor.py,sha256=L1ReZ8VPseMcy3-HGE-kbo97M9G8NqSWDDlMplE7Kgw,9798
+kagent_adk/_session_service.py,sha256=A47gsfDVp8jITzeW987AHTJLEhcU_mU3ik_SFptFGIc,5815
+kagent_adk/_task_store.py,sha256=3ApKbFfcDZmcEnwef6bCDhBhoGY9ZYwwyP671B1DHFo,889
+kagent_adk/a2a.py,sha256=HBEdGq4gPr78AD88GhygsPhntnctY5dw5pt-BN7FpaI,5647
+kagent_adk/models.py,sha256=dtwUQny2r5K1JtCtl_mGk3dNMW-XnPPjWJT8JbyWE0E,3654
+kagent_adk-0.5.0.dist-info/METADATA,sha256=k5WvtF1jPEEuhBvZXVGbGqgqx7X9WiyoHtYR2CdVzYg,1092
+kagent_adk-0.5.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+kagent_adk-0.5.0.dist-info/RECORD,,

kagent_adk-0.0.1.dist-info/RECORD

@@ -1,8 +0,0 @@
-kagent_adk/__init__.py,sha256=S0SQAy9KBIX8Vm55TO5d4Jj94AvmTDXtG8YurUKrSj4,329
-kagent_adk/a2a.py,sha256=d49IIoC2r__Zmjf5iuD_pFp8XKevrDn1Yh3sH0FRZKY,6460
-kagent_adk/kagent_session_service.py,sha256=A47gsfDVp8jITzeW987AHTJLEhcU_mU3ik_SFptFGIc,5815
-kagent_adk/kagent_task_store.py,sha256=3ApKbFfcDZmcEnwef6bCDhBhoGY9ZYwwyP671B1DHFo,889
-kagent_adk/models.py,sha256=dtwUQny2r5K1JtCtl_mGk3dNMW-XnPPjWJT8JbyWE0E,3654
-kagent_adk-0.0.1.dist-info/METADATA,sha256=FumdkqL9hAtChYuWbWkvMoZ31e5cJ1hdtZw5DbgaPjw,1021
-kagent_adk-0.0.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-kagent_adk-0.0.1.dist-info/RECORD,,