kagent-adk 0.0.1__py3-none-any.whl

kagent_adk/__init__.py

@@ -0,0 +1,10 @@
+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"]

kagent_adk/a2a.py

@@ -0,0 +1,202 @@
+#! /usr/bin/env python3
+import faulthandler
+import logging
+import os
+import sys
+from contextlib import asynccontextmanager
+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.request_handlers import DefaultRequestHandler
+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
+
+# --- Constants ---
+USER_ID = "admin@kagent.dev"
+
+# --- 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 False
+
+    @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, user_id: str, task_store: TaskStore):
+        super().__init__(task_store=task_store)
+        self.user_id = user_id
+
+    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 not context:
+            context = ServerCallContext(user=KAgentUser(user_id=self.user_id))
+        else:
+            context.user = KAgentUser(user_id=self.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")
+
+
+def thread_dump(request: Request) -> PlainTextResponse:
+    import io
+
+    buf = io.StringIO()
+    faulthandler.dump_traceback(file=buf)
+    buf.seek(0)
+    return PlainTextResponse(buf.read())
+
+
+kagent_url_override = os.getenv("KAGENT_URL")
+
+
+class KAgentApp:
+    def __init__(
+        self,
+        root_agent: BaseAgent | Callable[[], BaseAgent],
+        agent_card: AgentCard,
+        kagent_url: str,
+        app_name: str,
+    ):
+        self.root_agent = root_agent
+        self.kagent_url = kagent_url
+        self.app_name = app_name
+        self.agent_card = agent_card
+
+    def build(self) -> FastAPI:
+        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}")
+
+        agent_executor = A2aAgentExecutor(
+            runner=runner,
+        )
+
+        kagent_task_store = KAgentTaskStore(http_client)
+
+        request_context_builder = KAgentRequestContextBuilder(user_id=USER_ID, task_store=kagent_task_store)
+        request_handler = DefaultRequestHandler(
+            agent_executor=agent_executor,
+            task_store=kagent_task_store,
+            request_context_builder=request_context_builder,
+        )
+
+        a2a_app = A2AStarletteApplication(
+            agent_card=self.agent_card,
+            http_handler=request_handler,
+        )
+
+        # @asynccontextmanager
+        # async def agent_lifespan(app: FastAPI):
+        #     yield
+        #     if isinstance(runner, Runner):
+        #         await runner.close()
+
+        faulthandler.enable()
+        app = FastAPI()
+
+        # Health check/readiness probe
+        app.add_route("/health", methods=["GET"], route=health_check)
+        app.add_route("/thread_dump", methods=["GET"], route=thread_dump)
+        a2a_app.add_routes_to_app(app)
+
+        return app
+
+    async def test(self, task: str):
+        session_service = InMemorySessionService()
+        SESSION_ID = "12345"
+        USER_ID = "admin"
+        await session_service.create_session(
+            app_name=self.app_name,
+            session_id=SESSION_ID,
+            user_id=USER_ID,
+        )
+        if isinstance(self.root_agent, Callable):
+            agent_factory = self.root_agent
+            root_agent = agent_factory()
+        else:
+            root_agent = self.root_agent
+
+        runner = Runner(
+            agent=root_agent,
+            app_name=self.app_name,
+            session_service=session_service,
+        )
+
+        logger.info(f"\n>>> User Query: {task}")
+
+        # Prepare the user's message in ADK format
+        content = types.Content(role="user", parts=[types.Part(text=task)])
+        # Key Concept: run_async executes the agent logic and yields Events.
+        # We iterate through events to find the final answer.
+        async for event in runner.run_async(
+            user_id=USER_ID,
+            session_id=SESSION_ID,
+            new_message=content,
+        ):
+            # You can uncomment the line below to see *all* events during execution
+            # print(f"  [Event] Author: {event.author}, Type: {type(event).__name__}, Final: {event.is_final_response()}, Content: {event.content}")
+
+            # Key Concept: is_final_response() marks the concluding message for the turn.
+            jsn = event.model_dump_json()
+            logger.info(f"  [Event] {jsn}")

kagent_adk/kagent_session_service.py

@@ -0,0 +1,175 @@
+import logging
+from typing import Any, Optional
+
+import httpx
+from google.adk.events.event import Event
+from google.adk.sessions import Session
+from google.adk.sessions.base_session_service import (
+    BaseSessionService,
+    GetSessionConfig,
+    ListSessionsResponse,
+)
+from typing_extensions import override
+
+logger = logging.getLogger("kagent." + __name__)
+
+
+class KAgentSessionService(BaseSessionService):
+    """A session service implementation that uses the Kagent API.
+    This service integrates with the Kagent server to manage session state
+    and persistence through HTTP API calls.
+    """
+
+    def __init__(self, client: httpx.AsyncClient):
+        super().__init__()
+        self.client = client
+
+    async def _get_user_id(self) -> str:
+        """Get the default user ID. Override this method to implement custom user ID logic."""
+        return "admin@kagent.dev"
+
+    @override
+    async def create_session(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        state: Optional[dict[str, Any]] = None,
+        session_id: Optional[str] = None,
+    ) -> Session:
+        # Prepare request data
+        request_data = {
+            "user_id": user_id,
+            "agent_ref": app_name,  # Use app_name as agent reference
+        }
+        if session_id:
+            request_data["id"] = session_id
+
+        # Make API call to create session
+        response = await self.client.post(
+            "/api/sessions",
+            json=request_data,
+            headers={"X-User-ID": user_id},
+        )
+        response.raise_for_status()
+
+        data = response.json()
+        if not data.get("data"):
+            raise RuntimeError(f"Failed to create session: {data.get('message', 'Unknown error')}")
+
+        session_data = data["data"]
+
+        # Convert to ADK Session format
+        return Session(id=session_data["id"], user_id=session_data["user_id"], state=state or {}, app_name=app_name)
+
+    @override
+    async def get_session(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        config: Optional[GetSessionConfig] = None,
+    ) -> Optional[Session]:
+        try:
+            url = f"/api/sessions/{session_id}?user_id={user_id}"
+            if config:
+                if config.after_timestamp:
+                    # TODO: implement
+                    # url += f"&after={config.after_timestamp}"
+                    pass
+                if config.num_recent_events:
+                    url += f"&limit={config.num_recent_events}"
+                else:
+                    url += "&limit=-1"
+            else:
+                # return all
+                url += "&limit=-1"
+
+            # Make API call to get session
+            response: httpx.Response = await self.client.get(
+                url,
+                headers={"X-User-ID": user_id},
+            )
+            if response.status_code == 404:
+                return None
+            response.raise_for_status()
+
+            data = response.json()
+            if not data.get("data"):
+                return None
+
+            if not data.get("data").get("session"):
+                return None
+            session_data = data["data"]["session"]
+
+            events_data = data["data"]["events"]
+
+            events: list[Event] = []
+            for event_data in events_data:
+                events.append(Event.model_validate_json(event_data["data"]))
+
+            # Convert to ADK Session format
+            return Session(
+                id=session_data["id"],
+                user_id=session_data["user_id"],
+                events=events,
+                app_name=app_name,
+                state={},  # TODO: restore State
+            )
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 404:
+                return None
+            raise
+
+    @override
+    async def list_sessions(self, *, app_name: str, user_id: str) -> ListSessionsResponse:
+        # Make API call to list sessions
+        response = await self.client.get(f"/api/sessions?user_id={user_id}", headers={"X-User-ID": user_id})
+        response.raise_for_status()
+
+        data = response.json()
+        sessions_data = data.get("data", [])
+
+        # Convert to ADK Session format
+        sessions = []
+        for session_data in sessions_data:
+            session = Session(id=session_data["id"], user_id=session_data["user_id"], state={}, app_name=app_name)
+            sessions.append(session)
+
+        return ListSessionsResponse(sessions=sessions)
+
+    def list_sessions_sync(self, *, app_name: str, user_id: str) -> ListSessionsResponse:
+        raise NotImplementedError("not supported. use async")
+
+    @override
+    async def delete_session(self, *, app_name: str, user_id: str, session_id: str) -> None:
+        # Make API call to delete session
+        response = await self.client.delete(
+            f"/api/sessions/{session_id}?user_id={user_id}",
+            headers={"X-User-ID": user_id},
+        )
+        response.raise_for_status()
+
+    @override
+    async def append_event(self, session: Session, event: Event) -> Event:
+        # Convert ADK Event to JSON format
+        event_data = {
+            "id": event.id,
+            "data": event.model_dump_json(),
+        }
+
+        # Make API call to append event to session
+        response = await self.client.post(
+            f"/api/sessions/{session.id}/events?user_id={session.user_id}",
+            json=event_data,
+            headers={"X-User-ID": session.user_id},
+        )
+        response.raise_for_status()
+
+        # TODO: potentially pull and update the session from the server
+        # Update the in-memory session.
+        session.last_update_time = event.timestamp
+        await super().append_event(session=session, event=event)
+
+        return event

kagent_adk/kagent_task_store.py

@@ -0,0 +1,30 @@
+from typing import override
+
+import httpx
+from a2a.server.tasks import TaskStore
+from a2a.types import Task
+
+
+class KAgentTaskStore(TaskStore):
+    client: httpx.AsyncClient
+
+    def __init__(self, client: httpx.AsyncClient):
+        self.client = client
+
+    @override
+    async def save(self, task: Task) -> None:
+        response = await self.client.post("/api/tasks", json=task.model_dump())
+        response.raise_for_status()
+
+    @override
+    async def get(self, task_id: str) -> Task | None:
+        response = await self.client.get(f"/api/tasks/{task_id}")
+        if response.status_code == 404:
+            return None
+        response.raise_for_status()
+        return Task.model_validate(response.json())
+
+    @override
+    async def delete(self, task_id: str) -> None:
+        response = await self.client.delete(f"/api/tasks/{task_id}")
+        response.raise_for_status()

kagent_adk/models.py

@@ -0,0 +1,110 @@
+import logging
+from typing import Literal, Self, Union
+
+from a2a.types import AgentCard
+from google.adk.agents import Agent
+from google.adk.agents.llm_agent import ToolUnion
+from google.adk.agents.run_config import RunConfig, StreamingMode
+from google.adk.models.anthropic_llm import Claude as ClaudeLLM
+from google.adk.models.google_llm import Gemini as GeminiLLM
+from google.adk.models.lite_llm import LiteLlm
+from google.adk.tools.agent_tool import AgentTool
+from google.adk.tools.mcp_tool import MCPToolset, SseConnectionParams, StreamableHTTPConnectionParams
+from pydantic import BaseModel, Field
+
+logger = logging.getLogger(__name__)
+
+
+class HttpMcpServerConfig(BaseModel):
+    params: StreamableHTTPConnectionParams
+    tools: list[str] = Field(default_factory=list)
+
+
+class SseMcpServerConfig(BaseModel):
+    params: SseConnectionParams
+    tools: list[str] = Field(default_factory=list)
+
+
+class BaseLLM(BaseModel):
+    model: str
+
+
+class OpenAI(BaseLLM):
+    base_url: str | None = None
+
+    type: Literal["openai"]
+
+
+class AzureOpenAI(BaseLLM):
+    type: Literal["azure_openai"]
+
+
+class Anthropic(BaseLLM):
+    base_url: str | None = None
+
+    type: Literal["anthropic"]
+
+
+class GeminiVertexAI(BaseLLM):
+    type: Literal["gemini_vertex_ai"]
+
+
+class GeminiAnthropic(BaseLLM):
+    type: Literal["gemini_anthropic"]
+
+
+class Ollama(BaseLLM):
+    type: Literal["ollama"]
+
+
+class Gemini(BaseLLM):
+    type: Literal["gemini"]
+
+
+class AgentConfig(BaseModel):
+    kagent_url: str  # The URL of the KAgent server
+    agent_card: AgentCard
+    name: str
+    model: Union[OpenAI, Anthropic, GeminiVertexAI, GeminiAnthropic, Ollama, AzureOpenAI, Gemini] = Field(
+        discriminator="type"
+    )
+    description: str
+    instruction: str
+    http_tools: list[HttpMcpServerConfig] | None = None  # tools, always MCP
+    sse_tools: list[SseMcpServerConfig] | None = None  # tools, always MCP
+    agents: list[Self] | None = None  # agent names
+
+    def to_agent(self) -> Agent:
+        mcp_toolsets: list[ToolUnion] = []
+        if self.http_tools:
+            for http_tool in self.http_tools:  # add http tools
+                mcp_toolsets.append(MCPToolset(connection_params=http_tool.params, tool_filter=http_tool.tools))
+        if self.sse_tools:
+            for sse_tool in self.sse_tools:  # add stdio tools
+                mcp_toolsets.append(MCPToolset(connection_params=sse_tool.params, tool_filter=sse_tool.tools))
+        if self.agents:
+            for agent in self.agents:  # Add sub agents as tools
+                mcp_toolsets.append(AgentTool(agent.to_agent()))
+        if self.model.type == "openai":
+            model = LiteLlm(model=f"openai/{self.model.model}", base_url=self.model.base_url)
+        elif self.model.type == "anthropic":
+            model = LiteLlm(model=f"anthropic/{self.model.model}", base_url=self.model.base_url)
+        elif self.model.type == "gemini_vertex_ai":
+            model = GeminiLLM(model=self.model.model)
+        elif self.model.type == "gemini_anthropic":
+            model = ClaudeLLM(model=self.model.model)
+        elif self.model.type == "ollama":
+            model = LiteLlm(model=f"ollama_chat/{self.model.model}")
+        elif self.model.type == "azure_openai":
+            model = LiteLlm(model=f"azure/{self.model.model}")
+        elif self.model.type == "gemini":
+            model = self.model.model
+        else:
+            raise ValueError(f"Invalid model type: {self.model.type}")
+        return Agent(
+            name=self.name,
+            model=model,
+            description=self.description,
+            instruction=self.instruction,
+            tools=mcp_toolsets,
+        )

kagent_adk-0.0.1.dist-info/METADATA

@@ -0,0 +1,28 @@
+Metadata-Version: 2.4
+Name: kagent-adk
+Version: 0.0.1
+Summary: kagent-adk is an sdk for integrating adk agents with kagent
+Requires-Python: >=3.12.11
+Requires-Dist: a2a-sdk>=0.2.16
+Requires-Dist: anthropic[vertex]>=0.49.0
+Requires-Dist: anyio>=4.9.0
+Requires-Dist: fastapi>=0.115.1
+Requires-Dist: google-adk>=1.8.0
+Requires-Dist: google-auth>=2.40.2
+Requires-Dist: google-genai>=1.21.1
+Requires-Dist: httpx>=0.25.0
+Requires-Dist: jsonref>=1.1.0
+Requires-Dist: litellm>=1.74.3
+Requires-Dist: mcp>=1.12.0
+Requires-Dist: openai>=1.72.0
+Requires-Dist: opentelemetry-api>=1.32.0
+Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.32.0
+Requires-Dist: opentelemetry-instrumentation-httpx>=0.52.0
+Requires-Dist: opentelemetry-instrumentation-openai>=0.39.0
+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: test
+Requires-Dist: pytest-asyncio>=0.25.3; extra == 'test'
+Requires-Dist: pytest>=8.3.5; extra == 'test'

kagent_adk-0.0.1.dist-info/RECORD

@@ -0,0 +1,8 @@
+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,,

kagent_adk-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any