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

ksadk/agents/context.py

@@ -0,0 +1,38 @@
+"""编排上下文。"""
+
+from __future__ import annotations
+
+import copy
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class OrchestrationContext:
+    """用于在编排 agent 之间共享和分支状态。"""
+
+    session_id: str = ""
+    state: dict[str, Any] = field(default_factory=dict)
+    branch: str = ""
+    parent_context: "OrchestrationContext | None" = None
+
+    def get(self, key: str, default: Any = None) -> Any:
+        return self.state.get(key, default)
+
+    def set(self, key: str, value: Any) -> None:
+        self.state[key] = value
+
+    def update(self, values: dict[str, Any]) -> None:
+        self.state.update(values)
+
+    def snapshot_state(self) -> dict[str, Any]:
+        return copy.deepcopy(self.state)
+
+    def create_branch(self, branch_name: str) -> "OrchestrationContext":
+        branch = f"{self.branch}.{branch_name}" if self.branch else branch_name
+        return OrchestrationContext(
+            session_id=self.session_id,
+            state=self.snapshot_state(),
+            branch=branch,
+            parent_context=self,
+        )

ksadk/agents/event.py

@@ -0,0 +1,35 @@
+"""编排事件模型。"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class EventType(str, Enum):
+    """编排层事件类型。"""
+
+    AGENT_START = "agent_start"
+    AGENT_END = "agent_end"
+    TEXT_OUTPUT = "text_output"
+    TOOL_CALL = "tool_call"
+    TOOL_RESULT = "tool_result"
+    STATE_CHANGE = "state_change"
+    ESCALATE = "escalate"
+    ERROR = "error"
+
+
+@dataclass(slots=True)
+class AgentEvent:
+    """编排事件。"""
+
+    agent_name: str
+    event_type: EventType
+    data: Any = None
+    branch: str = ""
+    timestamp: float = field(default_factory=time.time)
+    metadata: dict[str, Any] = field(default_factory=dict)
+    state_delta: dict[str, Any] = field(default_factory=dict)
+    escalate: bool = False

ksadk/agents/loop_agent.py

@@ -0,0 +1,71 @@
+"""循环编排 agent。"""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import AsyncGenerator, Callable
+from typing import Any
+
+from ksadk.agents.base import BaseOrchestrationAgent
+from ksadk.agents.context import OrchestrationContext
+from ksadk.agents.event import AgentEvent, EventType
+
+
+class LoopAgent(BaseOrchestrationAgent):
+    """重复执行子 agent，直到命中退出条件。"""
+
+    def __init__(
+        self,
+        name: str,
+        sub_agents: list[Any] | None = None,
+        max_iterations: int | None = None,
+        exit_condition: Callable[[OrchestrationContext], Any] | None = None,
+        **kwargs: Any,
+    ):
+        if max_iterations is not None and max_iterations < 0:
+            raise ValueError("max_iterations must be non-negative")
+        super().__init__(name=name, sub_agents=sub_agents, **kwargs)
+        self.max_iterations = max_iterations
+        self.exit_condition = exit_condition
+
+    async def _run_impl(self, context: OrchestrationContext) -> AsyncGenerator[AgentEvent, None]:
+        iteration = 0
+
+        while self.max_iterations is None or iteration < self.max_iterations:
+            metadata = {"iteration": iteration}
+            yield AgentEvent(
+                agent_name=self.name,
+                event_type=EventType.AGENT_START,
+                branch=context.branch,
+                metadata=metadata,
+                data=metadata,
+            )
+
+            should_exit = False
+            for sub_agent in self.sub_agents:
+                async for event in self._run_sub_agent(sub_agent, context):
+                    yield event
+                    if event.escalate or event.event_type == EventType.ESCALATE:
+                        should_exit = True
+                        break
+                if should_exit:
+                    break
+
+            yield AgentEvent(
+                agent_name=self.name,
+                event_type=EventType.AGENT_END,
+                branch=context.branch,
+            )
+
+            iteration += 1
+            if should_exit:
+                break
+
+            if self.exit_condition is None:
+                continue
+
+            result = self.exit_condition(context)
+            if inspect.isawaitable(result):
+                result = await result
+            if result:
+                break

ksadk/agents/parallel_agent.py

@@ -0,0 +1,109 @@
+"""并行编排 agent。"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncGenerator
+from dataclasses import replace
+from typing import Any
+
+from ksadk.agents.base import BaseOrchestrationAgent
+from ksadk.agents.context import OrchestrationContext
+from ksadk.agents.event import AgentEvent, EventType
+
+_MISSING = object()
+
+
+class ParallelAgent(BaseOrchestrationAgent):
+    """在独立分支上下文中并发执行子 agent。"""
+
+    async def _run_impl(self, context: OrchestrationContext) -> AsyncGenerator[AgentEvent, None]:
+        if not self.sub_agents:
+            return
+
+        queue: asyncio.Queue[tuple[str, AgentEvent | None]] = asyncio.Queue()
+        branch_contexts: dict[str, OrchestrationContext] = {}
+        base_state = context.snapshot_state()
+        branch_errors: list[Exception] = []
+        tasks: list[asyncio.Task[None]] = []
+
+        async def run_branch(
+            agent_name: str,
+            sub_agent: Any,
+            branch_context: OrchestrationContext,
+        ) -> None:
+            try:
+                async for event in self._run_sub_agent(sub_agent, branch_context):
+                    if not event.branch:
+                        event = replace(event, branch=branch_context.branch)
+                    await queue.put((agent_name, event))
+            except Exception as exc:
+                branch_errors.append(exc)
+            finally:
+                await queue.put((agent_name, None))
+
+        for sub_agent in self.sub_agents:
+            agent_name = self._get_agent_name(sub_agent)
+            branch_context = context.create_branch(agent_name)
+            branch_contexts[agent_name] = branch_context
+            tasks.append(asyncio.create_task(run_branch(agent_name, sub_agent, branch_context)))
+
+        finished = 0
+        try:
+            while finished < len(tasks):
+                _, item = await queue.get()
+                if item is None:
+                    finished += 1
+                    continue
+                yield item
+        finally:
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+        if branch_errors:
+            raise branch_errors[0]
+
+        merge_delta = self._merge_branch_states(context, base_state, branch_contexts)
+        if merge_delta:
+            yield AgentEvent(
+                agent_name=self.name,
+                event_type=EventType.STATE_CHANGE,
+                branch=context.branch,
+                data=merge_delta,
+                state_delta=merge_delta,
+            )
+
+    def _merge_branch_states(
+        self,
+        context: OrchestrationContext,
+        base_state: dict[str, Any],
+        branch_contexts: dict[str, OrchestrationContext],
+    ) -> dict[str, Any]:
+        branch_deltas: dict[str, dict[str, Any]] = {}
+        by_key: dict[str, dict[str, Any]] = {}
+
+        for branch_name, branch_context in branch_contexts.items():
+            delta: dict[str, Any] = {}
+            for key, value in branch_context.state.items():
+                if base_state.get(key, _MISSING) != value:
+                    delta[key] = value
+                    by_key.setdefault(key, {})[branch_name] = value
+            branch_deltas[branch_name] = delta
+
+        merge_delta: dict[str, Any] = {f"{self.name}_results": branch_deltas}
+        context.set(f"{self.name}_results", branch_deltas)
+
+        conflicts: dict[str, dict[str, Any]] = {}
+        for key, branch_values in by_key.items():
+            values = list(branch_values.values())
+            first = values[0]
+            if all(value == first for value in values[1:]):
+                context.set(key, first)
+                merge_delta[key] = first
+                continue
+            conflicts[key] = branch_values
+
+        if conflicts:
+            context.set(f"{self.name}_conflicts", conflicts)
+            merge_delta[f"{self.name}_conflicts"] = conflicts
+
+        return merge_delta

ksadk/agents/runner_adapter.py

@@ -0,0 +1,24 @@
+"""Runner 适配器。"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from ksadk.agents.base import BaseOrchestrationAgent
+from ksadk.agents.context import OrchestrationContext
+from ksadk.agents.event import AgentEvent
+
+
+class RunnerAgent(BaseOrchestrationAgent):
+    """将 runner-like 对象包装为可编排子 agent。"""
+
+    def __init__(self, name: str, runner: Any, description: str = ""):
+        if not self._is_runner_like(runner):
+            raise TypeError("RunnerAgent requires a runner-like object with an invoke method")
+        super().__init__(name=name, sub_agents=None, description=description)
+        self.runner = runner
+
+    async def _run_impl(self, context: OrchestrationContext) -> AsyncGenerator[AgentEvent, None]:
+        async for event in self._run_runner_like(self.name, self.runner, context):
+            yield event

ksadk/agents/sequential_agent.py

@@ -0,0 +1,34 @@
+"""顺序编排 agent。"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+
+from ksadk.agents.base import BaseOrchestrationAgent
+from ksadk.agents.context import OrchestrationContext
+from ksadk.agents.event import AgentEvent, EventType
+
+
+class SequentialAgent(BaseOrchestrationAgent):
+    """按序依次执行子 agent。"""
+
+    async def _run_impl(self, context: OrchestrationContext) -> AsyncGenerator[AgentEvent, None]:
+        for index, sub_agent in enumerate(self.sub_agents):
+            agent_name = self._get_agent_name(sub_agent)
+            metadata = {"sub_agent": agent_name, "index": index}
+            yield AgentEvent(
+                agent_name=self.name,
+                event_type=EventType.AGENT_START,
+                branch=context.branch,
+                metadata=metadata,
+                data=metadata,
+            )
+            async for event in self._run_sub_agent(sub_agent, context):
+                yield event
+            yield AgentEvent(
+                agent_name=self.name,
+                event_type=EventType.AGENT_END,
+                branch=context.branch,
+                metadata=metadata,
+                data=metadata,
+            )

ksadk/api/__init__.py

@@ -0,0 +1,3 @@
+from .client import AgentEngineAPIError, AgentEngineClient, AttachmentContent, DryRunExit
+
+__all__ = ["AgentEngineClient", "AgentEngineAPIError", "DryRunExit", "AttachmentContent"]