agentengine-sdk-python 0.6.2__py3-none-any.whl → 0.6.9__py3-none-any.whl

ksadk/a2a/executor.py

@@ -0,0 +1,201 @@
+"""A2A executor that adapts a KsADK runner to the A2A server contract."""
+
+from __future__ import annotations
+
+import json
+import logging
+from collections.abc import AsyncIterator
+from typing import Any
+
+from a2a.server.agent_execution import AgentExecutor, RequestContext
+from a2a.server.events import EventQueue
+from a2a.server.tasks import TaskUpdater
+from a2a.types import Part, TextPart
+
+logger = logging.getLogger(__name__)
+
+
+class KsAgentExecutor(AgentExecutor):
+    """Execute a KsADK runner inside the A2A request lifecycle."""
+
+    def __init__(self, runner: Any, prefer_stream: bool = True) -> None:
+        self.runner = runner
+        self.prefer_stream = prefer_stream
+
+    async def execute(
+        self,
+        context: RequestContext,
+        event_queue: EventQueue,
+    ) -> None:
+        updater = TaskUpdater(
+            event_queue=event_queue,
+            task_id=context.task_id or "unknown-task",
+            context_id=context.context_id or "unknown-context",
+        )
+
+        try:
+            await updater.start_work()
+            runner_input = self._build_runner_input(context)
+            output = await self._run_runner(context, updater, runner_input)
+            completion_message = (
+                updater.new_agent_message(parts=[Part(root=TextPart(text=output))])
+                if output
+                else None
+            )
+            await updater.complete(message=completion_message)
+        except Exception as exc:
+            logger.exception("A2A task execution failed")
+            error_text = self._coerce_text(exc)
+            await updater.failed(
+                message=updater.new_agent_message(
+                    parts=[Part(root=TextPart(text=error_text))]
+                )
+            )
+
+    async def cancel(
+        self,
+        context: RequestContext,
+        event_queue: EventQueue,
+    ) -> None:
+        updater = TaskUpdater(
+            event_queue=event_queue,
+            task_id=context.task_id or "unknown-task",
+            context_id=context.context_id or "unknown-context",
+        )
+        await updater.cancel(
+            message=updater.new_agent_message(
+                parts=[Part(root=TextPart(text="Request canceled"))]
+            )
+        )
+
+    async def _run_runner(
+        self,
+        context: RequestContext,
+        updater: TaskUpdater,
+        runner_input: dict[str, Any],
+    ) -> str:
+        stream = getattr(self.runner, "stream", None)
+        if self.prefer_stream and callable(stream):
+            return await self._run_streaming(context, updater, stream, runner_input)
+
+        result = await self.runner.invoke(runner_input)
+        text = self._coerce_text(result)
+        if text:
+            await updater.add_artifact(
+                parts=[Part(root=TextPart(text=text))],
+                artifact_id=f"{context.task_id}-response",
+                name="response",
+                last_chunk=True,
+            )
+        return text
+
+    async def _run_streaming(
+        self,
+        context: RequestContext,
+        updater: TaskUpdater,
+        stream: Any,
+        runner_input: dict[str, Any],
+    ) -> str:
+        output_text = ""
+        emitted_chunks = 0
+        artifact_id = f"{context.task_id}-response"
+
+        async for chunk, last_chunk in self._with_last_flag(stream(runner_input)):
+            chunk_type = chunk.get("type") if isinstance(chunk, dict) else None
+            if chunk_type == "final":
+                final_text = (
+                    self._coerce_text(chunk.get("output"))
+                    if isinstance(chunk, dict)
+                    else ""
+                )
+                if not final_text:
+                    continue
+                if not output_text:
+                    output_text = final_text
+                elif final_text.startswith(output_text):
+                    suffix = final_text[len(output_text) :]
+                    output_text = final_text
+                    if suffix:
+                        emitted_chunks += 1
+                        await updater.add_artifact(
+                            parts=[Part(root=TextPart(text=suffix))],
+                            artifact_id=artifact_id,
+                            name="response",
+                            append=bool(emitted_chunks > 1),
+                            last_chunk=last_chunk,
+                        )
+                    continue
+                else:
+                    output_text = final_text
+                    await updater.add_artifact(
+                        parts=[Part(root=TextPart(text=final_text))],
+                        artifact_id=artifact_id,
+                        name="response",
+                        append=False,
+                        last_chunk=last_chunk,
+                    )
+                continue
+
+            text = self._coerce_text(chunk)
+            if not text:
+                continue
+            output_text += text
+            emitted_chunks += 1
+            await updater.add_artifact(
+                parts=[Part(root=TextPart(text=text))],
+                artifact_id=artifact_id,
+                name="response",
+                append=bool(emitted_chunks > 1),
+                last_chunk=last_chunk,
+            )
+
+        return output_text
+
+    async def _with_last_flag(
+        self,
+        iterator: AsyncIterator[Any],
+    ) -> AsyncIterator[tuple[Any, bool]]:
+        try:
+            previous = await anext(iterator)
+        except StopAsyncIteration:
+            return
+
+        while True:
+            try:
+                current = await anext(iterator)
+            except StopAsyncIteration:
+                yield previous, True
+                return
+            yield previous, False
+            previous = current
+
+    @staticmethod
+    def _build_runner_input(context: RequestContext) -> dict[str, Any]:
+        metadata = dict(context.metadata)
+        state = metadata.get("state", {})
+        if not isinstance(state, dict):
+            state = {}
+
+        return {
+            "input": context.get_user_input(),
+            "task_id": context.task_id,
+            "context_id": context.context_id,
+            "session_id": context.context_id,
+            "state": dict(state),
+            "branch": metadata.get("branch", ""),
+            "metadata": metadata,
+        }
+
+    @classmethod
+    def _coerce_text(cls, payload: Any) -> str:
+        if payload is None:
+            return ""
+        if isinstance(payload, str):
+            return payload
+        if isinstance(payload, dict):
+            for key in ("output", "delta", "data"):
+                value = payload.get(key)
+                if value is not None:
+                    return cls._coerce_text(value)
+            return json.dumps(payload, ensure_ascii=False, default=str)
+        return str(payload)

ksadk/a2a/server.py

@@ -0,0 +1,79 @@
+"""High-level server helpers for exposing a KsADK runner over A2A."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+from a2a.server.apps import A2AStarletteApplication
+from a2a.server.request_handlers import DefaultRequestHandler
+from a2a.server.tasks import InMemoryTaskStore
+from starlette.applications import Starlette
+
+from ksadk.a2a.card_builder import AgentCardBuilder
+from ksadk.a2a.executor import KsAgentExecutor
+
+
+class KsA2AServer:
+    """Expose a KsADK runner as an A2A-compatible ASGI app."""
+
+    def __init__(
+        self,
+        runner: Any,
+        app_name: str,
+        url: str,
+        description: str = "",
+        skills: Sequence[str] | None = None,
+        version: str = "1.0.0",
+    ) -> None:
+        self.runner = runner
+        self.app_name = app_name
+        self.url = url
+        self.description = description
+        self.skills = list(skills or [])
+        self.version = version
+
+        self.agent_card = AgentCardBuilder(
+            name=app_name,
+            url=url,
+            description=description,
+            skills=self.skills,
+            version=version,
+        ).build()
+        self.executor = KsAgentExecutor(runner=runner)
+        self.task_store = InMemoryTaskStore()
+        self.request_handler = DefaultRequestHandler(
+            agent_executor=self.executor,
+            task_store=self.task_store,
+        )
+        self.application = A2AStarletteApplication(
+            agent_card=self.agent_card,
+            http_handler=self.request_handler,
+        )
+
+    def build(self, **kwargs: Any) -> Starlette:
+        """Build the Starlette application with the standard A2A routes."""
+        return self.application.build(**kwargs)
+
+    def add_routes_to_app(self, app: Starlette) -> None:
+        """Attach the A2A routes to an existing Starlette application."""
+        self.application.add_routes_to_app(app)
+
+
+def to_a2a(
+    runner: Any,
+    app_name: str,
+    url: str,
+    description: str = "",
+    skills: Sequence[str] | None = None,
+    version: str = "1.0.0",
+) -> KsA2AServer:
+    """Convenience helper for building a KsA2AServer."""
+    return KsA2AServer(
+        runner=runner,
+        app_name=app_name,
+        url=url,
+        description=description,
+        skills=skills,
+        version=version,
+    )

ksadk/agents/__init__.py

@@ -0,0 +1,20 @@
+"""原生多 Agent 编排层。"""
+
+from ksadk.agents.base import BaseOrchestrationAgent
+from ksadk.agents.context import OrchestrationContext
+from ksadk.agents.event import AgentEvent, EventType
+from ksadk.agents.loop_agent import LoopAgent
+from ksadk.agents.parallel_agent import ParallelAgent
+from ksadk.agents.runner_adapter import RunnerAgent
+from ksadk.agents.sequential_agent import SequentialAgent
+
+__all__ = [
+    "AgentEvent",
+    "BaseOrchestrationAgent",
+    "EventType",
+    "LoopAgent",
+    "OrchestrationContext",
+    "ParallelAgent",
+    "RunnerAgent",
+    "SequentialAgent",
+]

ksadk/agents/base.py

@@ -0,0 +1,370 @@
+"""编排 agent 基类。"""
+
+from __future__ import annotations
+
+import inspect
+import keyword
+import logging
+import re
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator, AsyncIterator, Callable
+from dataclasses import replace
+from typing import Any
+
+from ksadk.agents.context import OrchestrationContext
+from ksadk.agents.event import AgentEvent, EventType
+
+logger = logging.getLogger(__name__)
+
+SubAgentType = Any
+_CAMEL_BOUNDARY_RE = re.compile(r"(?<!^)(?=[A-Z])")
+_TOP_LEVEL_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_-]*$")
+
+
+class BaseOrchestrationAgent(ABC):
+    """多 agent 编排基类。"""
+
+    def __init__(
+        self,
+        name: str,
+        sub_agents: list[SubAgentType] | None = None,
+        description: str = "",
+        before_run_hook: Callable[[OrchestrationContext], Any] | None = None,
+        after_run_hook: Callable[[OrchestrationContext], Any] | None = None,
+    ):
+        self._validate_top_level_name(name)
+        self.name = name
+        self.description = description
+        self.sub_agents = sub_agents or []
+        self.before_run_hook = before_run_hook
+        self.after_run_hook = after_run_hook
+        self._validate_sub_agents(self.sub_agents)
+
+    async def run_async(
+        self, context: OrchestrationContext | None = None
+    ) -> AsyncGenerator[AgentEvent, None]:
+        run_context = context or OrchestrationContext()
+        logger.info("[%s] orchestration started", self.name)
+        yield AgentEvent(
+            agent_name=self.name,
+            event_type=EventType.AGENT_START,
+            branch=run_context.branch,
+        )
+
+        try:
+            async for hook_event in self._run_hook(self.before_run_hook, run_context):
+                yield hook_event
+
+            async for event in self._run_impl(run_context):
+                yield self._ensure_event_branch(event, run_context.branch)
+
+            async for hook_event in self._run_hook(self.after_run_hook, run_context):
+                yield hook_event
+        except Exception as exc:
+            yield AgentEvent(
+                agent_name=self.name,
+                event_type=EventType.ERROR,
+                data=str(exc),
+                branch=run_context.branch,
+                metadata={"exception_type": type(exc).__name__},
+            )
+            raise
+
+        logger.info("[%s] orchestration finished", self.name)
+        yield AgentEvent(
+            agent_name=self.name,
+            event_type=EventType.AGENT_END,
+            branch=run_context.branch,
+        )
+
+    @abstractmethod
+    async def _run_impl(self, context: OrchestrationContext) -> AsyncGenerator[AgentEvent, None]:
+        if False:
+            yield AgentEvent(agent_name=self.name, event_type=EventType.TEXT_OUTPUT)
+
+    async def _run_sub_agent(
+        self, agent: SubAgentType, context: OrchestrationContext
+    ) -> AsyncGenerator[AgentEvent, None]:
+        agent_name = self._get_agent_name(agent)
+
+        if isinstance(agent, BaseOrchestrationAgent):
+            async for event in agent.run_async(context):
+                yield self._apply_event_state(event, context)
+            return
+
+        if self._is_runner_like(agent):
+            async for event in self._run_runner_like(agent_name, agent, context):
+                yield event
+            return
+
+        if callable(agent):
+            async for event in self._run_callable(agent_name, agent, context):
+                yield event
+            return
+
+        raise TypeError(f"Unsupported sub-agent type: {type(agent)!r}")
+
+    @classmethod
+    def _get_agent_name(cls, agent: SubAgentType) -> str:
+        if isinstance(agent, BaseOrchestrationAgent):
+            cls._validate_identifier_name(agent.name)
+            return agent.name
+
+        explicit_name = getattr(agent, "name", None)
+        if isinstance(explicit_name, str) and explicit_name:
+            cls._validate_identifier_name(explicit_name)
+            return explicit_name
+
+        if inspect.isfunction(agent) or inspect.ismethod(agent):
+            inferred_name = getattr(agent, "__name__", "")
+            cls._validate_identifier_name(inferred_name)
+            return inferred_name
+
+        class_name = agent.__class__.__name__
+        inferred_name = _CAMEL_BOUNDARY_RE.sub("_", class_name).lower().lstrip("_")
+        cls._validate_identifier_name(inferred_name)
+        return inferred_name
+
+    @classmethod
+    def _validate_identifier_name(cls, name: str) -> None:
+        if not isinstance(name, str) or not name:
+            raise ValueError("Agent name must be a non-empty string")
+        if not name.isidentifier() or keyword.iskeyword(name):
+            raise ValueError(f"Agent name must be a valid identifier: {name}")
+
+    @classmethod
+    def _validate_top_level_name(cls, name: str) -> None:
+        if not isinstance(name, str) or not name:
+            raise ValueError("Agent name must be a non-empty string")
+        if not _TOP_LEVEL_NAME_RE.match(name):
+            raise ValueError(f"Agent name must be a valid identifier: {name}")
+
+    @classmethod
+    def _validate_sub_agents(cls, sub_agents: list[SubAgentType]) -> None:
+        names: list[str] = []
+        for agent in sub_agents:
+            if cls._is_runner_like(agent):
+                names.append(cls._get_agent_name(agent))
+                continue
+
+            if callable(agent):
+                names.append(cls._get_agent_name(agent))
+                if not (inspect.iscoroutinefunction(agent) or inspect.isasyncgenfunction(agent)):
+                    raise TypeError("Custom callable sub-agents must be async")
+                continue
+
+            if isinstance(agent, BaseOrchestrationAgent):
+                names.append(cls._get_agent_name(agent))
+                continue
+
+            raise TypeError(f"Unsupported sub-agent type: {type(agent)!r}")
+
+        if len(names) != len(set(names)):
+            raise ValueError(f"Sub-agent names must be unique: {names}")
+
+    @staticmethod
+    def _is_runner_like(agent: Any) -> bool:
+        invoke = getattr(agent, "invoke", None)
+        return callable(invoke)
+
+    async def _run_callable(
+        self,
+        agent_name: str,
+        agent: Callable[[OrchestrationContext], Any],
+        context: OrchestrationContext,
+    ) -> AsyncGenerator[AgentEvent, None]:
+        try:
+            result = agent(context)
+            if inspect.isasyncgen(result):
+                async for item in result:
+                    yield self._apply_event_state(
+                        self._coerce_event(agent_name, item, context.branch),
+                        context,
+                    )
+                return
+
+            if not inspect.isawaitable(result):
+                raise TypeError("Custom callable sub-agents must be async")
+
+            value = await result
+        except Exception as exc:
+            yield AgentEvent(
+                agent_name=agent_name,
+                event_type=EventType.ERROR,
+                data=str(exc),
+                branch=context.branch,
+                metadata={"exception_type": type(exc).__name__},
+            )
+            raise
+
+        if value is None:
+            return
+
+        if isinstance(value, (list, tuple)):
+            for item in value:
+                yield self._apply_event_state(
+                    self._coerce_event(agent_name, item, context.branch),
+                    context,
+                )
+            return
+
+        yield self._apply_event_state(
+            self._coerce_event(agent_name, value, context.branch),
+            context,
+        )
+
+    async def _run_runner_like(
+        self,
+        agent_name: str,
+        runner: Any,
+        context: OrchestrationContext,
+    ) -> AsyncGenerator[AgentEvent, None]:
+        output_parts: list[Any] = []
+        runner_input = self._build_runner_input(context)
+
+        try:
+            stream = getattr(runner, "stream", None)
+            if callable(stream):
+                async for chunk in stream(runner_input):
+                    event = self._apply_event_state(
+                        self._coerce_event(agent_name, chunk, context.branch),
+                        context,
+                    )
+                    if event.event_type == EventType.TEXT_OUTPUT:
+                        output_parts.append(event.data)
+                    yield event
+            else:
+                result = await runner.invoke(runner_input)
+                event = self._apply_event_state(
+                    self._coerce_event(agent_name, result, context.branch),
+                    context,
+                )
+                if event.event_type == EventType.TEXT_OUTPUT:
+                    output_parts.append(event.data)
+                yield event
+        except Exception as exc:
+            yield AgentEvent(
+                agent_name=agent_name,
+                event_type=EventType.ERROR,
+                data=str(exc),
+                branch=context.branch,
+                metadata={"exception_type": type(exc).__name__},
+            )
+            raise
+
+        if output_parts:
+            context.set(f"{agent_name}_output", self._merge_output_parts(output_parts))
+
+    async def _run_hook(
+        self,
+        hook: Callable[[OrchestrationContext], Any] | None,
+        context: OrchestrationContext,
+    ) -> AsyncIterator[AgentEvent]:
+        if hook is None:
+            return
+
+        result = hook(context)
+        if inspect.isawaitable(result):
+            result = await result
+
+        if result is None:
+            return
+
+        if isinstance(result, (list, tuple)):
+            for item in result:
+                yield self._coerce_event(self.name, item, context.branch)
+            return
+
+        yield self._coerce_event(self.name, result, context.branch)
+
+    @staticmethod
+    def _build_runner_input(context: OrchestrationContext) -> dict[str, Any]:
+        return {
+            "input": context.get("input", ""),
+            "state": dict(context.state),
+            "session_id": context.session_id,
+            "branch": context.branch,
+        }
+
+    @staticmethod
+    def _merge_output_parts(parts: list[Any]) -> Any:
+        if all(isinstance(part, str) for part in parts):
+            return "".join(parts)
+        if len(parts) == 1:
+            return parts[0]
+        return parts
+
+    @classmethod
+    def _coerce_event(
+        cls,
+        agent_name: str,
+        payload: Any,
+        branch: str,
+    ) -> AgentEvent:
+        if isinstance(payload, AgentEvent):
+            event = payload
+            if not event.agent_name:
+                event = replace(event, agent_name=agent_name)
+            if not event.branch:
+                event = replace(event, branch=branch)
+            return event
+
+        if isinstance(payload, dict):
+            raw_event_type = payload.get("event_type")
+            event_type = cls._coerce_event_type(raw_event_type)
+            data = payload.get("data")
+            if data is None:
+                data = payload.get("output", payload.get("delta"))
+            state_delta = dict(payload.get("state_delta", {}))
+            metadata = dict(payload.get("metadata", {}))
+
+            if "type" in payload and "chunk_type" not in metadata:
+                metadata["chunk_type"] = payload["type"]
+
+            if event_type is None:
+                event_type = (
+                    EventType.STATE_CHANGE
+                    if state_delta and data is None
+                    else EventType.TEXT_OUTPUT
+                )
+
+            if data is None and event_type == EventType.STATE_CHANGE and state_delta:
+                data = state_delta
+
+            escalate = bool(payload.get("escalate", False) or event_type == EventType.ESCALATE)
+            return AgentEvent(
+                agent_name=agent_name,
+                event_type=event_type,
+                data=data,
+                branch=branch,
+                metadata=metadata,
+                state_delta=state_delta,
+                escalate=escalate,
+            )
+
+        return AgentEvent(
+            agent_name=agent_name,
+            event_type=EventType.TEXT_OUTPUT,
+            data=payload,
+            branch=branch,
+        )
+
+    @staticmethod
+    def _coerce_event_type(raw_event_type: Any) -> EventType | None:
+        if raw_event_type is None:
+            return None
+        if isinstance(raw_event_type, EventType):
+            return raw_event_type
+        return EventType(str(raw_event_type))
+
+    @staticmethod
+    def _ensure_event_branch(event: AgentEvent, branch: str) -> AgentEvent:
+        if event.branch or not branch:
+            return event
+        return replace(event, branch=branch)
+
+    @staticmethod
+    def _apply_event_state(event: AgentEvent, context: OrchestrationContext) -> AgentEvent:
+        if event.state_delta:
+            context.update(event.state_delta)
+        return event