kagent-adk 0.6.7__py3-none-any.whl → 0.6.9__py3-none-any.whl

{kagent_adk → kagent/adk}/__init__.py

@@ -1,7 +1,7 @@
 import importlib.metadata
 
-from .a2a import KAgentApp
-from .models import AgentConfig
+from ._a2a import KAgentApp
+from .types import AgentConfig
 
 __version__ = importlib.metadata.version("kagent_adk")
 

kagent_adk/a2a.py → kagent/adk/_a2a.py

@@ -1,20 +1,13 @@
 #! /usr/bin/env python3
 import faulthandler
-import inspect
 import logging
 import os
-import sys
-from contextlib import asynccontextmanager
-from typing import Awaitable, Callable, override
+from typing import Callable
 
 import httpx
-from a2a.auth.user import User
-from a2a.server.agent_execution import RequestContext, SimpleRequestContextBuilder
-from a2a.server.apps import A2AStarletteApplication
-from a2a.server.context import ServerCallContext
+from a2a.server.apps import A2AFastAPIApplication
 from a2a.server.request_handlers import DefaultRequestHandler
-from a2a.server.tasks import TaskStore
-from a2a.types import AgentCard, MessageSendParams, Task
+from a2a.types import AgentCard
 from fastapi import FastAPI, Request
 from fastapi.responses import PlainTextResponse
 from google.adk.agents import BaseAgent
@@ -22,54 +15,16 @@ from google.adk.runners import Runner
 from google.adk.sessions import InMemorySessionService
 from google.genai import types
 
+from kagent.core.a2a import KAgentRequestContextBuilder, KAgentTaskStore
+
 from ._agent_executor import A2aAgentExecutor
 from ._session_service import KAgentSessionService
-from ._task_store import KAgentTaskStore
 from ._token import KAgentTokenService
 
 # --- Configure Logging ---
 logger = logging.getLogger(__name__)
 
 
-class KAgentUser(User):
-    def __init__(self, user_id: str):
-        self.user_id = user_id
-
-    @property
-    def is_authenticated(self) -> bool:
-        return True
-
-    @property
-    def user_name(self) -> str:
-        return self.user_id
-
-
-class KAgentRequestContextBuilder(SimpleRequestContextBuilder):
-    """
-    A request context builder that will be used to hack in the user_id for now.
-    """
-
-    def __init__(self, task_store: TaskStore):
-        super().__init__(task_store=task_store)
-
-    async def build(
-        self,
-        params: MessageSendParams | None = None,
-        task_id: str | None = None,
-        context_id: str | None = None,
-        task: Task | None = None,
-        context: ServerCallContext | None = None,
-    ) -> RequestContext:
-        if context:
-            # grab the user id from the header
-            headers = context.state.get("headers", {})
-            user_id = headers.get("x-user-id", None)
-            if user_id:
-                context.user = KAgentUser(user_id=user_id)
-        request_context = await super().build(params, task_id, context_id, task, context)
-        return request_context
-
-
 def health_check(request: Request) -> PlainTextResponse:
     return PlainTextResponse("OK")
 
@@ -126,7 +81,7 @@ class KAgentApp:
             request_context_builder=request_context_builder,
         )
 
-        a2a_app = A2AStarletteApplication(
+        a2a_app = A2AFastAPIApplication(
             agent_card=self.agent_card,
             http_handler=request_handler,
         )

{kagent_adk → kagent/adk}/_agent_executor.py

@@ -13,6 +13,7 @@ from a2a.server.events.event_queue import EventQueue
 from a2a.types import (
     Artifact,
     Message,
+    Part,
     Role,
     TaskArtifactUpdateEvent,
     TaskState,
@@ -20,15 +21,16 @@ from a2a.types import (
     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 opentelemetry import trace
 from pydantic import BaseModel
 from typing_extensions import override
 
+from kagent.core.a2a import TaskResultAggregator, get_kagent_metadata_key
+
+from .converters.event_converter import convert_event_to_a2a_events
+from .converters.request_converter import convert_a2a_request_to_adk_run_args
+
 logger = logging.getLogger("google_adk." + __name__)
 
 
@@ -134,7 +136,7 @@ class A2aAgentExecutor(AgentExecutor):
                             message=Message(
                                 message_id=str(uuid.uuid4()),
                                 role=Role.agent,
-                                parts=[TextPart(text=str(e))],
+                                parts=[Part(TextPart(text=str(e)))],
                             ),
                         ),
                         context_id=context.context_id,
@@ -193,9 +195,9 @@ class A2aAgentExecutor(AgentExecutor):
                 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"],
+                    get_kagent_metadata_key("app_name"): runner.app_name,
+                    get_kagent_metadata_key("user_id"): run_args["user_id"],
+                    get_kagent_metadata_key("session_id"): run_args["session_id"],
                 },
             )
         )

{kagent_adk → kagent/adk}/_token.py

@@ -3,6 +3,8 @@ import asyncio
 from contextlib import asynccontextmanager
 from typing import Any, Optional
 
+import httpx
+
 KAGENT_TOKEN_PATH = "/var/run/secrets/tokens/kagent-token"
 logger = logging.getLogger(__name__)
 
@@ -59,7 +61,7 @@ class KAgentTokenService:
                 async with self.update_lock:
                     self.token = token
 
-    async def _add_bearer_token(self, request):
+    async def _add_bearer_token(self, request: httpx.Request):
         # Your function to generate headers dynamically
         token = await self._get_token()
         headers = {"X-Agent-Name": self.app_name}

kagent/adk/cli.py

@@ -0,0 +1,111 @@
+import asyncio
+import json
+import logging
+import os
+from typing import Annotated
+
+import typer
+import uvicorn
+from a2a.types import AgentCard
+from google.adk.cli.utils.agent_loader import AgentLoader
+
+from kagent.core import KAgentConfig, configure_tracing
+
+from . import AgentConfig, KAgentApp
+
+logger = logging.getLogger(__name__)
+
+app = typer.Typer()
+
+
+@app.command()
+def static(
+    host: str = "127.0.0.1",
+    port: int = 8080,
+    workers: int = 1,
+    filepath: str = "/config",
+    reload: Annotated[bool, typer.Option("--reload")] = False,
+):
+    app_cfg = KAgentConfig()
+
+    with open(os.path.join(filepath, "config.json"), "r") as f:
+        config = json.load(f)
+    agent_config = AgentConfig.model_validate(config)
+    with open(os.path.join(filepath, "agent-card.json"), "r") as f:
+        agent_card = json.load(f)
+    agent_card = AgentCard.model_validate(agent_card)
+    root_agent = agent_config.to_agent(app_cfg.name)
+
+    kagent_app = KAgentApp(root_agent, agent_card, app_cfg.url, app_cfg.app_name)
+
+    server = kagent_app.build()
+    configure_tracing(server)
+
+    uvicorn.run(
+        server,
+        host=host,
+        port=port,
+        workers=workers,
+        reload=reload,
+    )
+
+
+@app.command()
+def run(
+    name: Annotated[str, typer.Argument(help="The name of the agent to run")],
+    working_dir: str = ".",
+    host: str = "127.0.0.1",
+    port: int = 8080,
+    workers: int = 1,
+):
+    app_cfg = KAgentConfig()
+
+    agent_loader = AgentLoader(agents_dir=working_dir)
+    root_agent = agent_loader.load_agent(name)
+
+    with open(os.path.join(working_dir, name, "agent-card.json"), "r") as f:
+        agent_card = json.load(f)
+    agent_card = AgentCard.model_validate(agent_card)
+    kagent_app = KAgentApp(root_agent, agent_card, app_cfg.url, app_cfg.app_name)
+    server = kagent_app.build()
+    configure_tracing(server)
+
+    uvicorn.run(
+        server,
+        host=host,
+        port=port,
+        workers=workers,
+    )
+
+
+async def test_agent(agent_config: AgentConfig, agent_card: AgentCard, task: str):
+    app_cfg = KAgentConfig()
+    agent = agent_config.to_agent(app_cfg.name)
+    app = KAgentApp(agent, agent_card, app_cfg.url, app_cfg.app_name)
+    await app.test(task)
+
+
+@app.command()
+def test(
+    task: Annotated[str, typer.Option("--task", help="The task to test the agent with")],
+    filepath: Annotated[str, typer.Option("--filepath", help="The path to the agent config file")],
+):
+    with open(filepath, "r") as f:
+        content = f.read()
+        config = json.loads(content)
+
+    with open(os.path.join(filepath, "agent-card.json"), "r") as f:
+        agent_card = json.load(f)
+    agent_card = AgentCard.model_validate(agent_card)
+    agent_config = AgentConfig.model_validate(config)
+    asyncio.run(test_agent(agent_config, agent_card, task))
+
+
+def run_cli():
+    logging.basicConfig(level=logging.INFO)
+    logging.info("Starting KAgent")
+    app()
+
+
+if __name__ == "__main__":
+    run_cli()

kagent/adk/converters/event_converter.py

@@ -0,0 +1,315 @@
+from __future__ import annotations
+
+import logging
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional
+
+from a2a.server.events import Event as A2AEvent
+from a2a.types import DataPart, Message, Role, Task, TaskState, TaskStatus, TaskStatusUpdateEvent, TextPart
+from a2a.types import Part as A2APart
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows.functions import REQUEST_EUC_FUNCTION_CALL_NAME
+from google.genai import types as genai_types
+
+from kagent.core.a2a import (
+    A2A_DATA_PART_METADATA_IS_LONG_RUNNING_KEY,
+    A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL,
+    A2A_DATA_PART_METADATA_TYPE_KEY,
+    get_kagent_metadata_key,
+)
+
+from .part_converter import (
+    convert_genai_part_to_a2a_part,
+)
+
+# Constants
+
+ARTIFACT_ID_SEPARATOR = "-"
+DEFAULT_ERROR_MESSAGE = "An error occurred during processing"
+
+# Logger
+logger = logging.getLogger("kagent_adk." + __name__)
+
+
+def _serialize_metadata_value(value: Any) -> str:
+    """Safely serializes metadata values to string format.
+
+    Args:
+      value: The value to serialize.
+
+    Returns:
+      String representation of the value.
+    """
+    if hasattr(value, "model_dump"):
+        try:
+            return value.model_dump(exclude_none=True, by_alias=True)
+        except Exception as e:
+            logger.warning("Failed to serialize metadata value: %s", e)
+            return str(value)
+    return str(value)
+
+
+def _get_context_metadata(event: Event, invocation_context: InvocationContext) -> Dict[str, str]:
+    """Gets the context metadata for the event.
+
+    Args:
+      event: The ADK event to extract metadata from.
+      invocation_context: The invocation context containing session information.
+
+    Returns:
+      A dictionary containing the context metadata.
+
+    Raises:
+      ValueError: If required fields are missing from event or context.
+    """
+    if not event:
+        raise ValueError("Event cannot be None")
+    if not invocation_context:
+        raise ValueError("Invocation context cannot be None")
+
+    try:
+        metadata = {
+            get_kagent_metadata_key("app_name"): invocation_context.app_name,
+            get_kagent_metadata_key("user_id"): invocation_context.user_id,
+            get_kagent_metadata_key("session_id"): invocation_context.session.id,
+            get_kagent_metadata_key("invocation_id"): event.invocation_id,
+            get_kagent_metadata_key("author"): event.author,
+        }
+
+        # Add optional metadata fields if present
+        optional_fields = [
+            ("branch", event.branch),
+            ("grounding_metadata", event.grounding_metadata),
+            ("custom_metadata", event.custom_metadata),
+            ("usage_metadata", event.usage_metadata),
+            ("error_code", event.error_code),
+        ]
+
+        for field_name, field_value in optional_fields:
+            if field_value is not None:
+                metadata[get_kagent_metadata_key(field_name)] = _serialize_metadata_value(field_value)
+
+        return metadata
+
+    except Exception as e:
+        logger.error("Failed to create context metadata: %s", e)
+        raise
+
+
+def _create_artifact_id(app_name: str, user_id: str, session_id: str, filename: str, version: int) -> str:
+    """Creates a unique artifact ID.
+
+    Args:
+      app_name: The application name.
+      user_id: The user ID.
+      session_id: The session ID.
+      filename: The artifact filename.
+      version: The artifact version.
+
+    Returns:
+      A unique artifact ID string.
+    """
+    components = [app_name, user_id, session_id, filename, str(version)]
+    return ARTIFACT_ID_SEPARATOR.join(components)
+
+
+def _process_long_running_tool(a2a_part: A2APart, event: Event) -> None:
+    """Processes long-running tool metadata for an A2A part.
+
+    Args:
+      a2a_part: The A2A part to potentially mark as long-running.
+      event: The ADK event containing long-running tool information.
+    """
+    if (
+        isinstance(a2a_part.root, DataPart)
+        and event.long_running_tool_ids
+        and a2a_part.root.metadata
+        and a2a_part.root.metadata.get(get_kagent_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY))
+        == A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
+        and a2a_part.root.data.get("id") in event.long_running_tool_ids
+    ):
+        a2a_part.root.metadata[get_kagent_metadata_key(A2A_DATA_PART_METADATA_IS_LONG_RUNNING_KEY)] = True
+
+
+def convert_event_to_a2a_message(
+    event: Event, invocation_context: InvocationContext, role: Role = Role.agent
+) -> Optional[Message]:
+    """Converts an ADK event to an A2A message.
+
+    Args:
+      event: The ADK event to convert.
+      invocation_context: The invocation context.
+
+    Returns:
+      An A2A Message if the event has content, None otherwise.
+
+    Raises:
+      ValueError: If required parameters are invalid.
+    """
+    if not event:
+        raise ValueError("Event cannot be None")
+    if not invocation_context:
+        raise ValueError("Invocation context cannot be None")
+
+    if not event.content or not event.content.parts:
+        return None
+
+    try:
+        a2a_parts = []
+        for part in event.content.parts:
+            a2a_part = convert_genai_part_to_a2a_part(part)
+            if a2a_part:
+                a2a_parts.append(a2a_part)
+                _process_long_running_tool(a2a_part, event)
+
+        if a2a_parts:
+            return Message(message_id=str(uuid.uuid4()), role=role, parts=a2a_parts)
+
+    except Exception as e:
+        logger.error("Failed to convert event to status message: %s", e)
+        raise
+
+    return None
+
+
+def _create_error_status_event(
+    event: Event,
+    invocation_context: InvocationContext,
+    task_id: Optional[str] = None,
+    context_id: Optional[str] = None,
+) -> TaskStatusUpdateEvent:
+    """Creates a TaskStatusUpdateEvent for error scenarios.
+
+    Args:
+      event: The ADK event containing error information.
+      invocation_context: The invocation context.
+      task_id: Optional task ID to use for generated events.
+      context_id: Optional Context ID to use for generated events.
+
+    Returns:
+      A TaskStatusUpdateEvent with FAILED state.
+    """
+    error_message = getattr(event, "error_message", None) or DEFAULT_ERROR_MESSAGE
+
+    # Get context metadata and add error code
+    event_metadata = _get_context_metadata(event, invocation_context)
+    if event.error_code:
+        event_metadata[get_kagent_metadata_key("error_code")] = str(event.error_code)
+
+    return TaskStatusUpdateEvent(
+        task_id=task_id,
+        context_id=context_id,
+        metadata=event_metadata,
+        status=TaskStatus(
+            state=TaskState.failed,
+            message=Message(
+                message_id=str(uuid.uuid4()),
+                role=Role.agent,
+                parts=[A2APart(TextPart(text=error_message))],
+                metadata={get_kagent_metadata_key("error_code"): str(event.error_code)} if event.error_code else {},
+            ),
+            timestamp=datetime.now(timezone.utc).isoformat(),
+        ),
+        final=False,
+    )
+
+
+def _create_status_update_event(
+    message: Message,
+    invocation_context: InvocationContext,
+    event: Event,
+    task_id: Optional[str] = None,
+    context_id: Optional[str] = None,
+) -> TaskStatusUpdateEvent:
+    """Creates a TaskStatusUpdateEvent for running scenarios.
+
+    Args:
+      message: The A2A message to include.
+      invocation_context: The invocation context.
+      event: The ADK event.
+      task_id: Optional task ID to use for generated events.
+      context_id: Optional Context ID to use for generated events.
+
+
+    Returns:
+      A TaskStatusUpdateEvent with RUNNING state.
+    """
+    status = TaskStatus(
+        state=TaskState.working,
+        message=message,
+        timestamp=datetime.now(timezone.utc).isoformat(),
+    )
+
+    if any(
+        part.root.metadata.get(get_kagent_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY))
+        == A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
+        and part.root.metadata.get(get_kagent_metadata_key(A2A_DATA_PART_METADATA_IS_LONG_RUNNING_KEY)) is True
+        and part.root.data.get("name") == REQUEST_EUC_FUNCTION_CALL_NAME
+        for part in message.parts
+        if part.root.metadata
+    ):
+        status.state = TaskState.auth_required
+    elif any(
+        part.root.metadata.get(get_kagent_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY))
+        == A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
+        and part.root.metadata.get(get_kagent_metadata_key(A2A_DATA_PART_METADATA_IS_LONG_RUNNING_KEY)) is True
+        for part in message.parts
+        if part.root.metadata
+    ):
+        status.state = TaskState.input_required
+
+    return TaskStatusUpdateEvent(
+        task_id=task_id,
+        context_id=context_id,
+        status=status,
+        metadata=_get_context_metadata(event, invocation_context),
+        final=False,
+    )
+
+
+def convert_event_to_a2a_events(
+    event: Event,
+    invocation_context: InvocationContext,
+    task_id: Optional[str] = None,
+    context_id: Optional[str] = None,
+) -> List[A2AEvent]:
+    """Converts a GenAI event to a list of A2A events.
+
+    Args:
+      event: The ADK event to convert.
+      invocation_context: The invocation context.
+      task_id: Optional task ID to use for generated events.
+      context_id: Optional Context ID to use for generated events.
+
+    Returns:
+      A list of A2A events representing the converted ADK event.
+
+    Raises:
+      ValueError: If required parameters are invalid.
+    """
+    if not event:
+        raise ValueError("Event cannot be None")
+    if not invocation_context:
+        raise ValueError("Invocation context cannot be None")
+
+    a2a_events = []
+
+    try:
+        # Handle error scenarios
+        if event.error_code:
+            error_event = _create_error_status_event(event, invocation_context, task_id, context_id)
+            a2a_events.append(error_event)
+
+        # Handle regular message content
+        message = convert_event_to_a2a_message(event, invocation_context)
+        if message:
+            running_event = _create_status_update_event(message, invocation_context, event, task_id, context_id)
+            a2a_events.append(running_event)
+
+    except Exception as e:
+        logger.error("Failed to convert event to A2A events: %s", e)
+        raise
+
+    return a2a_events