kcs-agent 0.1.0__py3-none-any.whl

kcs_agent/__init__.py

@@ -0,0 +1,175 @@
+from .agent import Agent
+from .events import AgentEvent, AgentEventType, AgentResult
+from .exceptions import (
+    AgentCancelledError,
+    AgentError,
+    AgentIterationLimitError,
+    AgentProtocolError,
+    SessionVersionConflictError,
+    ToolDefinitionError,
+    ToolInvocationError,
+)
+from .extensions import (
+    AgentContext,
+    AgentExtension,
+    ExtensionFailurePolicy,
+    HistoryExtension,
+    StateExtension,
+    StaticContextExtension,
+    ToolExtension,
+    ToolPromptExtension,
+)
+from .messages import (
+    AnyMessage,
+    AssistantMessage,
+    AssistantMessageData,
+    AssistantMessageMetadata,
+    ImageAssetSource,
+    ImageBytesSource,
+    ImageContent,
+    ImageDetail,
+    ImageSource,
+    ImageUrlSource,
+    Message,
+    MessageRole,
+    SystemMessage,
+    SystemMessageData,
+    SystemMessageMetadata,
+    TextContent,
+    ToolCall,
+    ToolMessage,
+    ToolMessageData,
+    ToolMessageMetadata,
+    UserContent,
+    UserContentPart,
+    UserMessage,
+    UserMessageData,
+    UserMessageMetadata,
+)
+from .model import (
+    AgentModel,
+    ModelEvent,
+    ModelEventType,
+    ModelRequest,
+    ModelResponse,
+    ModelUsage,
+    ReasoningEffort,
+    ToolCallDelta,
+    ToolDefinition,
+)
+from .providers import (
+    AnthropicProvider,
+    DeepSeekProvider,
+    GoogleProvider,
+    OllamaProvider,
+    OpenAIProvider,
+    ProviderAuthError,
+    ProviderError,
+    ProviderResponseError,
+)
+from .schema import Parameter, annotation_schema, callable_schema
+from .state import (
+    AgentErrorInfo,
+    AgentState,
+    CancellationToken,
+    CheckpointReason,
+    ExtensionStateRegistry,
+    RunState,
+    RunStatus,
+    SessionState,
+    ToolCallStreamState,
+    UsageState,
+)
+from .tools import (
+    AgentTool,
+    ToolDocstring,
+    ToolParameterDocumentation,
+    parse_tool,
+    parse_tool_docstring,
+    render_tool_guidance,
+    tool,
+)
+
+__all__ = [
+    "Agent",
+    "AgentContext",
+    "AgentCancelledError",
+    "AgentError",
+    "AgentErrorInfo",
+    "AgentEvent",
+    "AgentEventType",
+    "AgentExtension",
+    "AgentIterationLimitError",
+    "AgentModel",
+    "AgentProtocolError",
+    "AgentResult",
+    "AgentState",
+    "AgentTool",
+    "AnyMessage",
+    "AssistantMessage",
+    "AssistantMessageData",
+    "AssistantMessageMetadata",
+    "CancellationToken",
+    "CheckpointReason",
+    "ExtensionFailurePolicy",
+    "ExtensionStateRegistry",
+    "HistoryExtension",
+    "ImageAssetSource",
+    "ImageBytesSource",
+    "ImageContent",
+    "ImageDetail",
+    "ImageSource",
+    "ImageUrlSource",
+    "Message",
+    "MessageRole",
+    "ModelEvent",
+    "ModelEventType",
+    "ModelRequest",
+    "ModelResponse",
+    "ModelUsage",
+    "Parameter",
+    "ReasoningEffort",
+    "RunState",
+    "RunStatus",
+    "SessionState",
+    "SessionVersionConflictError",
+    "StateExtension",
+    "StaticContextExtension",
+    "SystemMessage",
+    "SystemMessageData",
+    "SystemMessageMetadata",
+    "ToolCall",
+    "ToolCallDelta",
+    "ToolCallStreamState",
+    "ToolDefinition",
+    "ToolDefinitionError",
+    "ToolDocstring",
+    "ToolExtension",
+    "ToolPromptExtension",
+    "ToolInvocationError",
+    "ToolMessage",
+    "ToolMessageData",
+    "ToolMessageMetadata",
+    "ToolParameterDocumentation",
+    "TextContent",
+    "UsageState",
+    "annotation_schema",
+    "callable_schema",
+    "parse_tool",
+    "parse_tool_docstring",
+    "render_tool_guidance",
+    "AnthropicProvider",
+    "DeepSeekProvider",
+    "GoogleProvider",
+    "OllamaProvider",
+    "OpenAIProvider",
+    "ProviderAuthError",
+    "ProviderError",
+    "ProviderResponseError",
+    "tool",
+    "UserMessage",
+    "UserContent",
+    "UserContentPart",
+    "UserMessageData",
+    "UserMessageMetadata",
+]

kcs_agent/agent.py

@@ -0,0 +1,409 @@
+import asyncio
+import json
+from collections.abc import AsyncIterator, Callable, Sequence
+from contextlib import aclosing
+from typing import Any, cast
+
+from .events import AgentEvent, AgentEventType, AgentResult
+from .exceptions import AgentCancelledError, AgentIterationLimitError, AgentProtocolError, ToolDefinitionError
+from .extensions import AgentContext, AgentExtension, ExtensionFailurePolicy
+from .messages import (
+    AnyMessage,
+    Message,
+    SystemMessage,
+    SystemMessageData,
+    ToolCall,
+    ToolMessage,
+    ToolMessageData,
+    ToolMessageMetadata,
+)
+from .model import AgentModel, ModelEventType, ModelRequest, ModelResponse, ReasoningEffort
+from .state import (
+    AgentErrorInfo,
+    AgentState,
+    CancellationToken,
+    CheckpointReason,
+    RunState,
+    RunStatus,
+    SessionState,
+    utc_now,
+)
+from .tools import AgentTool, parse_tool
+
+
+class Agent[SessionDataT]:
+    """Minimal provider-neutral agent loop with typed state and extension-owned capabilities."""
+
+    def __init__(
+        self,
+        model: AgentModel,
+        *,
+        system_prompt: str = "You are a helpful assistant.",
+        extensions: Sequence[AgentExtension[SessionDataT]] = (),
+        max_iterations: int = 12,
+        default_reasoning_effort: ReasoningEffort = ReasoningEffort.MEDIUM,
+        state_factory: Callable[[], SessionDataT] | None = None,
+    ) -> None:
+        if max_iterations < 1:
+            raise ValueError("max_iterations must be at least one")
+        self._model = model
+        self._system_prompt = system_prompt.strip()
+        self._extensions = tuple(extensions)
+        self._max_iterations = max_iterations
+        self._default_reasoning_effort = default_reasoning_effort
+        self._state_factory = state_factory
+
+    async def run(
+        self,
+        messages: AnyMessage | Sequence[AnyMessage],
+        *,
+        session: SessionState[SessionDataT] | None = None,
+        metadata: dict[str, Any] | None = None,
+        reasoning_effort: ReasoningEffort | None = None,
+        cancellation_token: CancellationToken | None = None,
+    ) -> AgentResult[SessionDataT]:
+        """Run to completion and return the final assistant message with typed state."""
+        result: AgentResult[SessionDataT] | None = None
+        cancelled: AgentCancelledError | None = None
+        async for event in self.stream(
+            messages,
+            session=session,
+            metadata=metadata,
+            reasoning_effort=reasoning_effort,
+            cancellation_token=cancellation_token,
+        ):
+            match event.type:
+                case AgentEventType.RUN_COMPLETED:
+                    result = cast(AgentResult[SessionDataT], event.result)
+                case AgentEventType.RUN_CANCELLED:
+                    cancelled = cast(AgentCancelledError, event.error)
+                case _:
+                    pass
+        if cancelled is not None:
+            raise cancelled
+        if result is None:
+            raise AgentProtocolError("Agent stream ended without a final result")
+        return result
+
+    async def stream(
+        self,
+        messages: AnyMessage | Sequence[AnyMessage],
+        *,
+        session: SessionState[SessionDataT] | None = None,
+        metadata: dict[str, Any] | None = None,
+        reasoning_effort: ReasoningEffort | None = None,
+        cancellation_token: CancellationToken | None = None,
+    ) -> AsyncIterator[AgentEvent]:
+        """Execute model and tool steps while yielding typed lifecycle events in exact order."""
+        initial_messages = self._normalize_messages(messages)
+        context = AgentContext(
+            state=AgentState(
+                run=RunState(
+                    messages=initial_messages,
+                    reasoning_effort=(
+                        reasoning_effort if reasoning_effort is not None else self._default_reasoning_effort
+                    ),
+                ),
+                session=session or self._create_session(),
+            ),
+            metadata=dict(metadata or {}),
+        )
+        run = context.state.run
+        token = cancellation_token or CancellationToken()
+        try:
+            await self._call_extensions("load_state", context)
+            run.status = RunStatus.RUNNING
+            run.started_at = utc_now()
+            await self._call_extensions("on_run_start", context)
+            tools = await self._collect_tools(context)
+            injected = await self._collect_context(context)
+            system = [SystemMessage(data=SystemMessageData(self._system_prompt))] if self._system_prompt else []
+            run.messages = [*system, *injected, *run.messages]
+            yield self._event(context, AgentEventType.RUN_STARTED)
+            await self._checkpoint(context, CheckpointReason.RUN_STARTED)
+
+            for iteration in range(1, self._max_iterations + 1):
+                self._ensure_not_cancelled(run, token)
+                run.iteration = iteration
+                run.text_buffer = ""
+                run.reasoning_buffer = ""
+                run.tool_call_streams = {}
+                request = ModelRequest(
+                    messages=tuple(run.messages),
+                    tools=tuple(item.definition for item in tools.values()),
+                    reasoning_effort=run.reasoning_effort,
+                    metadata=context.metadata,
+                )
+                for extension in self._extensions:
+                    replacement = await self._invoke_extension(extension, "before_model", context, request)
+                    if replacement is not None:
+                        if not isinstance(replacement, ModelRequest):
+                            raise AgentProtocolError("before_model must return ModelRequest or None")
+                        request = replacement
+                run.status = RunStatus.WAITING_FOR_MODEL
+                yield self._event(context, AgentEventType.MODEL_STARTED)
+                response: ModelResponse | None = None
+                response_seen = False
+                async with aclosing(self._model.stream(request)) as model_events:
+                    async for model_event in model_events:
+                        self._ensure_not_cancelled(run, token)
+                        if response_seen:
+                            raise AgentProtocolError("A model adapter emitted events after its final response")
+                        match model_event.type:
+                            case ModelEventType.TEXT_DELTA:
+                                run.text_buffer += model_event.delta
+                                yield self._event(
+                                    context,
+                                    AgentEventType.TEXT_DELTA,
+                                    delta=model_event.delta,
+                                )
+                            case ModelEventType.REASONING_DELTA:
+                                run.reasoning_buffer += model_event.delta
+                                yield self._event(
+                                    context,
+                                    AgentEventType.REASONING_DELTA,
+                                    delta=model_event.delta,
+                                )
+                            case ModelEventType.TOOL_CALL_DELTA:
+                                if model_event.tool_call_delta is None:
+                                    raise AgentProtocolError("A tool-call event must contain a ToolCallDelta")
+                                run.append_tool_call_delta(model_event.tool_call_delta)
+                                yield self._event(
+                                    context,
+                                    AgentEventType.TOOL_CALL_DELTA,
+                                    tool_call_delta=model_event.tool_call_delta,
+                                )
+                            case ModelEventType.RESPONSE:
+                                if model_event.response is None:
+                                    raise AgentProtocolError("A response event must contain a ModelResponse")
+                                response = model_event.response
+                                response_seen = True
+                            case _:
+                                raise AgentProtocolError(f"Unsupported model event: {model_event.type}")
+                if response is None:
+                    raise AgentProtocolError("A model stream must finish with exactly one response event")
+
+                run.status = RunStatus.RUNNING
+                run.usage.add_model_usage(response.usage)
+                await self._call_extensions("after_model", context, response)
+                await self._append_message(context, response.message)
+                yield self._event(context, AgentEventType.MODEL_COMPLETED, response=response)
+                await self._checkpoint(context, CheckpointReason.MODEL_COMPLETED)
+
+                match response.message.data.tool_calls:
+                    case ():
+                        run.status = RunStatus.COMPLETED
+                        run.finished_at = utc_now()
+                        await self._call_extensions("on_run_end", context, response.message)
+                        await self._checkpoint(context, CheckpointReason.RUN_COMPLETED)
+                        result = AgentResult(response.message, context.state, dict(context.metadata))
+                        yield self._event(context, AgentEventType.RUN_COMPLETED, result=result)
+                        return
+                    case calls:
+                        for call in calls:
+                            self._ensure_not_cancelled(run, token)
+                            async for event in self._execute_tool(context, tools, call):
+                                yield event
+
+            raise AgentIterationLimitError(f"Agent exceeded {self._max_iterations} model iterations")
+        except AgentCancelledError as error:
+            run.status = RunStatus.CANCELLED
+            run.finished_at = utc_now()
+            run.error = AgentErrorInfo.from_exception(error)
+            await self._terminal_checkpoint(context, CheckpointReason.RUN_CANCELLED)
+            yield self._event(context, AgentEventType.RUN_CANCELLED, error=error)
+        except (asyncio.CancelledError, GeneratorExit) as error:
+            if run.status not in (RunStatus.COMPLETED, RunStatus.FAILED, RunStatus.CANCELLED):
+                run.status = RunStatus.CANCELLED
+                run.finished_at = utc_now()
+                run.error = AgentErrorInfo.from_exception(error)
+                await self._terminal_checkpoint(context, CheckpointReason.RUN_CANCELLED)
+            raise
+        except Exception as error:
+            run.status = RunStatus.FAILED
+            run.finished_at = utc_now()
+            run.error = AgentErrorInfo.from_exception(error)
+            await self._notify_error(context, error)
+            await self._terminal_checkpoint(context, CheckpointReason.RUN_FAILED)
+            yield self._event(context, AgentEventType.RUN_FAILED, error=error)
+            raise
+        finally:
+            await self._release_state(context)
+
+    def _normalize_messages(self, messages: AnyMessage | Sequence[AnyMessage]) -> list[AnyMessage]:
+        match messages:
+            case Message():
+                return [messages]
+            case _:
+                return list(messages)
+
+    def _create_session(self) -> SessionState[SessionDataT]:
+        data = self._state_factory() if self._state_factory is not None else cast(SessionDataT, {})
+        return SessionState(data=data)
+
+    async def _collect_tools(self, context: AgentContext[SessionDataT]) -> dict[str, AgentTool]:
+        result: dict[str, AgentTool] = {}
+        for extension in self._extensions:
+            candidates = await self._invoke_extension(extension, "tools", context)
+            for candidate in candidates or ():
+                registered = candidate if isinstance(candidate, AgentTool) else parse_tool(candidate)
+                if registered.name in result:
+                    raise ToolDefinitionError(f"Duplicate tool name: {registered.name}")
+                result[registered.name] = registered
+        return result
+
+    async def _collect_context(self, context: AgentContext[SessionDataT]) -> list[AnyMessage]:
+        messages: list[AnyMessage] = []
+        for extension in self._extensions:
+            injected = await self._invoke_extension(extension, "context_messages", context)
+            messages.extend(injected or ())
+        return messages
+
+    async def _execute_tool(
+        self,
+        context: AgentContext[SessionDataT],
+        tools: dict[str, AgentTool],
+        call: ToolCall,
+    ) -> AsyncIterator[AgentEvent]:
+        run = context.state.run
+        run.status = RunStatus.WAITING_FOR_TOOL
+        run.current_tool_call = call
+        run.usage.tool_calls += 1
+        yield self._event(context, AgentEventType.TOOL_STARTED, call=call)
+        try:
+            match tools.get(call.name):
+                case AgentTool() as registered:
+                    pass
+                case None:
+                    raise ToolDefinitionError(f"Unknown tool requested by model: {call.name}")
+                case invalid:
+                    raise ToolDefinitionError(f"Invalid tool registration for {call.name}: {invalid!r}")
+            await self._call_extensions("before_tool", context, call)
+            value = await registered(call.arguments)
+            tool_message = ToolMessage(
+                data=ToolMessageData(
+                    content=registered.serialize_result(value),
+                    tool_call_id=call.id,
+                    name=call.name,
+                ),
+                metadata=ToolMessageMetadata(success=True),
+            )
+        except Exception as error:
+            await self._notify_error(context, error)
+            tool_message = ToolMessage(
+                data=ToolMessageData(
+                    content=json.dumps({"error": str(error)}, ensure_ascii=False, separators=(",", ":")),
+                    tool_call_id=call.id,
+                    name=call.name,
+                ),
+                metadata=ToolMessageMetadata(success=False),
+            )
+            await self._append_message(context, tool_message)
+            run.status = RunStatus.RUNNING
+            run.current_tool_call = None
+            yield self._event(
+                context,
+                AgentEventType.TOOL_FAILED,
+                call=call,
+                message=tool_message,
+                error=error,
+            )
+            await self._checkpoint(context, CheckpointReason.TOOL_FAILED)
+            return
+
+        await self._append_message(context, tool_message)
+        await self._call_extensions("after_tool", context, call, tool_message)
+        run.status = RunStatus.RUNNING
+        run.current_tool_call = None
+        yield self._event(
+            context,
+            AgentEventType.TOOL_COMPLETED,
+            call=call,
+            message=tool_message,
+        )
+        await self._checkpoint(context, CheckpointReason.TOOL_COMPLETED)
+
+    async def _append_message(self, context: AgentContext[SessionDataT], message: AnyMessage) -> None:
+        context.state.run.messages.append(message)
+        await self._call_extensions("on_message", context, message)
+
+    async def _checkpoint(self, context: AgentContext[SessionDataT], reason: CheckpointReason) -> None:
+        await self._call_extensions("on_checkpoint", context, reason)
+
+    async def _terminal_checkpoint(
+        self,
+        context: AgentContext[SessionDataT],
+        reason: CheckpointReason,
+    ) -> None:
+        """Attempt terminal persistence without replacing the original terminal outcome."""
+        try:
+            await self._checkpoint(context, reason)
+        except Exception:
+            pass
+
+    async def _release_state(self, context: AgentContext[SessionDataT]) -> None:
+        """Best-effort cleanup that cannot invalidate an already emitted terminal event."""
+        for extension in self._extensions:
+            try:
+                await self._invoke_extension(extension, "release_state", context)
+            except Exception:
+                continue
+
+    async def _call_extensions(
+        self,
+        hook: str,
+        context: AgentContext[SessionDataT],
+        *arguments: Any,
+    ) -> None:
+        for extension in self._extensions:
+            await self._invoke_extension(extension, hook, context, *arguments)
+
+    async def _invoke_extension(
+        self,
+        extension: AgentExtension[SessionDataT],
+        hook: str,
+        context: AgentContext[SessionDataT],
+        *arguments: Any,
+    ) -> Any:
+        callback = getattr(extension, hook)
+        attempts = max(extension.max_retries, 0) + 1 if extension.failure_policy == ExtensionFailurePolicy.RETRY else 1
+        for attempt in range(attempts):
+            try:
+                return await callback(context, *arguments)
+            except Exception:
+                if extension.failure_policy == ExtensionFailurePolicy.CONTINUE:
+                    return None
+                if extension.failure_policy != ExtensionFailurePolicy.RETRY or attempt == attempts - 1:
+                    raise
+        return None
+
+    async def _notify_error(self, context: AgentContext[SessionDataT], error: Exception) -> None:
+        for extension in self._extensions:
+            try:
+                await self._invoke_extension(extension, "on_error", context, error)
+            except Exception:
+                continue
+
+    def _ensure_not_cancelled(self, run: RunState, token: CancellationToken) -> None:
+        if not token.cancelled:
+            return
+        run.cancellation_requested = True
+        run.status = RunStatus.CANCELLING
+        raise AgentCancelledError("Agent run was cancelled")
+
+    def _event(
+        self,
+        context: AgentContext[SessionDataT],
+        event_type: AgentEventType,
+        **values: Any,
+    ) -> AgentEvent:
+        run = context.state.run
+        return AgentEvent(
+            type=event_type,
+            run_id=run.run_id,
+            sequence=run.next_event_sequence(),
+            timestamp=utc_now(),
+            status=run.status,
+            iteration=run.iteration,
+            **values,
+        )

kcs_agent/events.py

@@ -0,0 +1,79 @@
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import StrEnum
+from typing import Any
+
+from .messages import AssistantMessage, ToolCall, ToolMessage
+from .model import ModelResponse, ToolCallDelta
+from .state import AgentState, RunStatus
+
+
+class AgentEventType(StrEnum):
+    """Observable lifecycle events emitted by the agent loop."""
+
+    RUN_STARTED = "run_started"
+    MODEL_STARTED = "model_started"
+    TEXT_DELTA = "text_delta"
+    REASONING_DELTA = "reasoning_delta"
+    TOOL_CALL_DELTA = "tool_call_delta"
+    MODEL_COMPLETED = "model_completed"
+    TOOL_STARTED = "tool_started"
+    TOOL_COMPLETED = "tool_completed"
+    TOOL_FAILED = "tool_failed"
+    RUN_COMPLETED = "run_completed"
+    RUN_FAILED = "run_failed"
+    RUN_CANCELLED = "run_cancelled"
+
+
+@dataclass(frozen=True, slots=True)
+class AgentResult[SessionDataT]:
+    """Final state returned after an agent run completes."""
+
+    message: AssistantMessage
+    state: AgentState[SessionDataT]
+    metadata: Mapping[str, Any] = field(default_factory=dict)
+
+
+@dataclass(frozen=True, slots=True)
+class AgentEvent:
+    """One ordered and observable event emitted by the agent loop."""
+
+    # Lifecycle discriminator that determines which optional payload fields are populated.
+    type: AgentEventType
+
+    # Stable identifier shared by every event produced by one agent invocation.
+    run_id: str
+
+    # One-based event number within a run, used to order events, detect gaps, and discard duplicates.
+    sequence: int
+
+    # Timezone-aware UTC timestamp captured when the core creates the event.
+    timestamp: datetime
+
+    # Run lifecycle state at emission time, such as waiting for a model, completed, or cancelled.
+    status: RunStatus
+
+    # One-based model-loop iteration, or zero before the first model step starts.
+    iteration: int = 0
+
+    # Incremental text for TEXT_DELTA or REASONING_DELTA, appended by consumers in sequence order.
+    delta: str = ""
+
+    # Partial model-generated tool ID, name, or JSON arguments; this fragment is not safe to execute.
+    tool_call_delta: ToolCallDelta | None = None
+
+    # Complete normalized invocation used by TOOL_STARTED, TOOL_COMPLETED, and TOOL_FAILED events.
+    call: ToolCall | None = None
+
+    # Authoritative normalized model response and usage, populated by MODEL_COMPLETED.
+    response: ModelResponse | None = None
+
+    # Tool result populated by TOOL_COMPLETED or TOOL_FAILED; model output is available through response.
+    message: AssistantMessage | ToolMessage | None = None
+
+    # Successful terminal output and typed state, populated only by RUN_COMPLETED.
+    result: AgentResult[Any] | None = None
+
+    # In-process failure or cancellation cause; transport adapters must serialize it safely.
+    error: Exception | None = None

kcs_agent/exceptions.py

@@ -0,0 +1,26 @@
+class AgentError(Exception):
+    """Base exception for agent runtime failures."""
+
+
+class AgentProtocolError(AgentError):
+    """Raised when a model adapter violates the streaming protocol."""
+
+
+class AgentIterationLimitError(AgentError):
+    """Raised when the model does not finish within the configured loop bound."""
+
+
+class AgentCancelledError(AgentError):
+    """Raised when a caller cooperatively cancels an active run."""
+
+
+class SessionVersionConflictError(AgentError):
+    """Raised when durable session state changed since it was loaded."""
+
+
+class ToolDefinitionError(AgentError):
+    """Raised when a decorated function cannot become a model tool."""
+
+
+class ToolInvocationError(AgentError):
+    """Raised when tool arguments are invalid or execution fails."""