pidevkit-agent 0.1.0__tar.gz

pidevkit_agent-0.1.0/.gitignore

@@ -0,0 +1,15 @@
+__pycache__/
+.pytest_cache/
+.ruff_cache/
+.venv/
+*.pyc
+*.pyo
+*.pyd
+.coverage
+.coverage.*
+htmlcov/
+dist/
+build/
+*.egg-info/
+.mypy_cache/
+.DS_Store

pidevkit_agent-0.1.0/PKG-INFO

@@ -0,0 +1,22 @@
+Metadata-Version: 2.4
+Name: pidevkit-agent
+Version: 0.1.0
+Summary: Agent runtime primitives for pidevkit
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+
+# pidevkit-agent
+
+Agent orchestration/runtime package for `pidevkit`.
+
+## Install
+
+```bash
+pip install pidevkit-agent
+```
+
+## Import
+
+```python
+from pidevkit.agent import Agent
+```

pidevkit_agent-0.1.0/README.md

@@ -0,0 +1,15 @@
+# pidevkit-agent
+
+Agent orchestration/runtime package for `pidevkit`.
+
+## Install
+
+```bash
+pip install pidevkit-agent
+```
+
+## Import
+
+```python
+from pidevkit.agent import Agent
+```

pidevkit_agent-0.1.0/pyproject.toml

@@ -0,0 +1,14 @@
+[build-system]
+requires = ["hatchling>=1.24.0"]
+build-backend = "hatchling.build"
+
+[project]
+name = "pidevkit-agent"
+version = "0.1.0"
+description = "Agent runtime primitives for pidevkit"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = []
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/pidevkit/agent"]

pidevkit_agent-0.1.0/src/pidevkit/agent/__init__.py

@@ -0,0 +1,19 @@
+from .agent import Agent
+from .agent_loop import (
+    agent_loop,
+    agent_loop_continue,
+    agentLoop,
+    agentLoopContinue,
+    validate_tool_arguments,
+)
+from .types import AgentState
+
+__all__ = [
+    "Agent",
+    "AgentState",
+    "agent_loop",
+    "agent_loop_continue",
+    "agentLoop",
+    "agentLoopContinue",
+    "validate_tool_arguments",
+]

pidevkit_agent-0.1.0/src/pidevkit/agent/agent.py

@@ -0,0 +1,418 @@
+from __future__ import annotations
+
+import asyncio
+import copy
+from dataclasses import dataclass
+from typing import Any, Callable, Mapping
+
+from .agent_loop import agent_loop, agent_loop_continue
+from .types import (
+    AgentMessage,
+    AgentOptions,
+    AgentState,
+    DEFAULT_MODEL,
+    Message,
+    Model,
+    ThinkingBudgets,
+    ThinkingLevel,
+    Transport,
+    clone_message,
+    clone_messages,
+    get_value,
+    message_content,
+    message_role,
+    now_ms,
+)
+
+
+@dataclass(slots=True)
+class AbortSignal:
+    aborted: bool = False
+
+
+class AbortController:
+    def __init__(self) -> None:
+        self.signal = AbortSignal()
+
+    def abort(self) -> None:
+        self.signal.aborted = True
+
+
+class Agent:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        options: dict[str, Any] = {}
+        if args:
+            if len(args) > 1:
+                raise TypeError("Agent accepts at most one positional options argument")
+            if not isinstance(args[0], Mapping):
+                raise TypeError("Agent positional argument must be a mapping")
+            options.update(dict(args[0]))
+        options.update(kwargs)
+
+        self._listeners: list[Callable[[dict[str, Any]], None]] = []
+        self._steering_queue: list[AgentMessage] = []
+        self._follow_up_queue: list[AgentMessage] = []
+        self._abort_controller: AbortController | None = None
+
+        self._stream_fn = get_value(options, "stream_fn", "streamFn", default=None)
+        self._convert_to_llm = get_value(options, "convert_to_llm", "convertToLlm", default=self._default_convert_to_llm)
+        self._transform_context = get_value(options, "transform_context", "transformContext", default=None)
+        self._get_api_key = get_value(options, "get_api_key", "getApiKey", default=None)
+
+        self._steering_mode = get_value(options, "steering_mode", "steeringMode", default="one-at-a-time")
+        self._follow_up_mode = get_value(options, "follow_up_mode", "followUpMode", default="one-at-a-time")
+
+        self.thinking_budgets: ThinkingBudgets = get_value(
+            options,
+            "thinking_budgets",
+            "thinkingBudgets",
+            default={"minimal": 0, "low": 0, "medium": 0, "high": 0},
+        )
+        self.transport: Transport = get_value(options, "transport", default="sse")
+        self.max_retry_delay_ms: int = int(get_value(options, "max_retry_delay_ms", "maxRetryDelayMs", default=1000) or 1000)
+
+        self._session_id: str | None = get_value(options, "session_id", "sessionId", default=None)
+
+        initial_state = get_value(options, "initial_state", "initialState", default=None)
+        self._state = self._build_initial_state(initial_state)
+
+    @property
+    def state(self) -> AgentState:
+        return self._state
+
+    def subscribe(self, listener: Callable[[dict[str, Any]], None]) -> Callable[[], None]:
+        self._listeners.append(listener)
+
+        def unsubscribe() -> None:
+            if listener in self._listeners:
+                self._listeners.remove(listener)
+
+        return unsubscribe
+
+    def _emit(self, event: dict[str, Any]) -> None:
+        for listener in list(self._listeners):
+            try:
+                listener(event)
+            except Exception:
+                continue
+
+    def set_system_prompt(self, system_prompt: str) -> None:
+        self._state.system_prompt = system_prompt
+
+    def setSystemPrompt(self, system_prompt: str) -> None:
+        self.set_system_prompt(system_prompt)
+
+    def set_model(self, model: Model) -> None:
+        self._state.model = copy.deepcopy(model)
+
+    def setModel(self, model: Model) -> None:
+        self.set_model(model)
+
+    def set_thinking_level(self, level: ThinkingLevel) -> None:
+        self._state.thinking_level = level
+
+    def setThinkingLevel(self, level: ThinkingLevel) -> None:
+        self.set_thinking_level(level)
+
+    def set_tools(self, tools: list[Any]) -> None:
+        self._state.tools = list(tools)
+
+    def setTools(self, tools: list[Any]) -> None:
+        self.set_tools(tools)
+
+    def replace_messages(self, messages: list[AgentMessage]) -> None:
+        self._state.messages = clone_messages(messages)
+
+    def replaceMessages(self, messages: list[AgentMessage]) -> None:
+        self.replace_messages(messages)
+
+    def append_message(self, message: AgentMessage) -> None:
+        self._state.messages.append(clone_message(message))
+
+    def appendMessage(self, message: AgentMessage) -> None:
+        self.append_message(message)
+
+    def clear_messages(self) -> None:
+        self._state.messages = []
+
+    def clearMessages(self) -> None:
+        self.clear_messages()
+
+    def steer(self, message: AgentMessage) -> None:
+        self._steering_queue.append(clone_message(message))
+
+    def follow_up(self, message: AgentMessage) -> None:
+        self._follow_up_queue.append(clone_message(message))
+
+    def followUp(self, message: AgentMessage) -> None:
+        self.follow_up(message)
+
+    def abort(self) -> None:
+        if self._abort_controller is not None:
+            self._abort_controller.abort()
+
+    @property
+    def session_id(self) -> str | None:
+        return self._session_id
+
+    @session_id.setter
+    def session_id(self, value: str | None) -> None:
+        self._session_id = value
+
+    @property
+    def sessionId(self) -> str | None:
+        return self._session_id
+
+    @sessionId.setter
+    def sessionId(self, value: str | None) -> None:
+        self._session_id = value
+
+    async def prompt(self, input_value: str | AgentMessage, images: list[dict[str, Any]] | None = None) -> list[AgentMessage]:
+        self._ensure_not_streaming()
+        user_message = self._make_user_message(input_value, images=images)
+        return await self._run_loop(prompts=[user_message], is_continue=False)
+
+    async def continue_(self) -> list[AgentMessage]:
+        self._ensure_not_streaming()
+
+        queued_steering = self._dequeue(self._steering_queue, self._steering_mode)
+        if queued_steering:
+            return await self._run_loop(prompts=queued_steering, is_continue=False)
+
+        queued_follow_up = self._dequeue(self._follow_up_queue, self._follow_up_mode)
+        if queued_follow_up:
+            return await self._run_loop(prompts=queued_follow_up, is_continue=False)
+
+        return await self._run_loop(prompts=[], is_continue=True)
+
+    async def resume(self) -> list[AgentMessage]:
+        return await self.continue_()
+
+    async def run_continue(self) -> list[AgentMessage]:
+        return await self.continue_()
+
+    async def wait_for_idle(self) -> None:
+        while self._state.is_streaming:
+            await asyncio.sleep(0.005)
+
+    async def _run_loop(self, prompts: list[AgentMessage], *, is_continue: bool) -> list[AgentMessage]:
+        self._state.is_streaming = True
+        self._state.error = None
+        self._state.stream_message = None
+        self._state.pending_tool_calls.clear()
+
+        self._abort_controller = AbortController()
+        skip_initial_steering_poll = len(prompts) > 0
+        steering_polled = False
+
+        async def get_steering_messages() -> list[AgentMessage]:
+            nonlocal steering_polled
+            if skip_initial_steering_poll and not steering_polled:
+                steering_polled = True
+                return []
+            steering_polled = True
+            return self._dequeue(self._steering_queue, self._steering_mode)
+
+        async def get_follow_up_messages() -> list[AgentMessage]:
+            return self._dequeue(self._follow_up_queue, self._follow_up_mode)
+
+        context: dict[str, Any] = {
+            "systemPrompt": self._state.system_prompt,
+            "messages": clone_messages(self._state.messages),
+            "tools": list(self._state.tools),
+        }
+        config = self._build_loop_config(
+            get_steering_messages=get_steering_messages,
+            get_follow_up_messages=get_follow_up_messages,
+        )
+
+        try:
+            if is_continue:
+                stream = agent_loop_continue(
+                    context,
+                    config,
+                    signal=self._abort_controller.signal,
+                    stream_fn=self._stream_fn,
+                )
+            else:
+                stream = agent_loop(
+                    prompts,
+                    context,
+                    config,
+                    signal=self._abort_controller.signal,
+                    stream_fn=self._stream_fn,
+                )
+
+            async for event in stream:
+                self._handle_stream_event(event)
+                self._emit(event)
+
+            new_messages = await stream.result()
+            if not isinstance(new_messages, list):
+                new_messages = list(new_messages)
+
+            self._state.messages.extend(clone_messages(new_messages))
+            return new_messages
+        except Exception as exc:
+            self._state.error = str(exc)
+            raise
+        finally:
+            self._state.is_streaming = False
+            self._state.stream_message = None
+            self._state.pending_tool_calls.clear()
+            self._abort_controller = None
+
+    def _handle_stream_event(self, event: Mapping[str, Any]) -> None:
+        event_type = str(get_value(event, "type", default=""))
+
+        if event_type in {"message_start", "message_update"}:
+            message = get_value(event, "message", default=None)
+            if message_role(message) == "assistant":
+                self._state.stream_message = clone_message(message)
+                self._sync_pending_tool_calls(message)
+
+        if event_type == "message_end":
+            message = get_value(event, "message", default=None)
+            if message_role(message) == "assistant":
+                self._state.stream_message = None
+
+        if event_type == "tool_execution_start":
+            tool_call_id = get_value(event, "toolCallId", "tool_call_id", default=None)
+            if isinstance(tool_call_id, str):
+                self._state.pending_tool_calls.add(tool_call_id)
+
+        if event_type == "tool_execution_end":
+            tool_call_id = get_value(event, "toolCallId", "tool_call_id", default=None)
+            if isinstance(tool_call_id, str):
+                self._state.pending_tool_calls.discard(tool_call_id)
+
+    def _sync_pending_tool_calls(self, message: AgentMessage) -> None:
+        content = message_content(message)
+        if not isinstance(content, list):
+            return
+
+        for block in content:
+            if isinstance(block, Mapping) and block.get("type") == "toolCall":
+                tool_call_id = block.get("id")
+                if isinstance(tool_call_id, str):
+                    self._state.pending_tool_calls.add(tool_call_id)
+
+    def _build_loop_config(
+        self,
+        *,
+        get_steering_messages: Callable[[], Any] | None = None,
+        get_follow_up_messages: Callable[[], Any] | None = None,
+    ) -> dict[str, Any]:
+        config: dict[str, Any] = {
+            "model": self._state.model,
+            "convertToLlm": self._convert_to_llm,
+            "getSteeringMessages": get_steering_messages or self._get_steering_messages,
+            "getFollowUpMessages": get_follow_up_messages or self._get_follow_up_messages,
+            "reasoning": self._state.thinking_level,
+            "transport": self.transport,
+            "thinkingBudgets": self.thinking_budgets,
+            "maxRetryDelayMs": self.max_retry_delay_ms,
+        }
+
+        if callable(self._transform_context):
+            config["transformContext"] = self._transform_context
+
+        if callable(self._get_api_key):
+            config["getApiKey"] = self._get_api_key
+
+        if self._session_id is not None:
+            config["sessionId"] = self._session_id
+            config["session_id"] = self._session_id
+
+        return config
+
+    async def _get_steering_messages(self) -> list[AgentMessage]:
+        return self._dequeue(self._steering_queue, self._steering_mode)
+
+    async def _get_follow_up_messages(self) -> list[AgentMessage]:
+        return self._dequeue(self._follow_up_queue, self._follow_up_mode)
+
+    @staticmethod
+    def _dequeue(queue: list[AgentMessage], mode: str) -> list[AgentMessage]:
+        if not queue:
+            return []
+
+        if mode == "all":
+            values = queue[:]
+            queue.clear()
+            return values
+
+        return [queue.pop(0)]
+
+    @staticmethod
+    def _default_convert_to_llm(messages: list[AgentMessage]) -> list[Message]:
+        converted: list[Message] = []
+        for message in messages:
+            role = message_role(message)
+            if role in {"user", "assistant", "toolResult"}:
+                converted.append(copy.deepcopy(message))
+        return converted
+
+    def _build_initial_state(self, initial_state: Mapping[str, Any] | AgentState | None) -> AgentState:
+        if isinstance(initial_state, AgentState):
+            return copy.deepcopy(initial_state)
+
+        state = AgentState()
+        if initial_state is None:
+            return state
+
+        state.system_prompt = str(get_value(initial_state, "system_prompt", "systemPrompt", default="") or "")
+
+        model_value = get_value(initial_state, "model", default=None)
+        if isinstance(model_value, Mapping):
+            state.model = copy.deepcopy(dict(model_value))
+        else:
+            state.model = copy.deepcopy(DEFAULT_MODEL)
+
+        state.thinking_level = get_value(initial_state, "thinking_level", "thinkingLevel", default="off")
+
+        tools_value = get_value(initial_state, "tools", default=[])
+        state.tools = list(tools_value) if isinstance(tools_value, list) else []
+
+        messages_value = get_value(initial_state, "messages", default=[])
+        state.messages = clone_messages(messages_value if isinstance(messages_value, list) else [])
+
+        state.is_streaming = bool(get_value(initial_state, "is_streaming", "isStreaming", default=False))
+        state.stream_message = clone_message(get_value(initial_state, "stream_message", "streamMessage", default=None))
+        pending_value = get_value(initial_state, "pending_tool_calls", "pendingToolCalls", default=set())
+        if isinstance(pending_value, set):
+            state.pending_tool_calls = set(str(item) for item in pending_value)
+        elif isinstance(pending_value, list):
+            state.pending_tool_calls = set(str(item) for item in pending_value)
+        else:
+            state.pending_tool_calls = set()
+        state.error = get_value(initial_state, "error", default=None)
+
+        return state
+
+    def _ensure_not_streaming(self) -> None:
+        if self._state.is_streaming:
+            raise RuntimeError("Agent is already processing a request")
+
+    @staticmethod
+    def _make_user_message(input_value: str | AgentMessage, images: list[dict[str, Any]] | None = None) -> AgentMessage:
+        if isinstance(input_value, Mapping) and message_role(input_value):
+            message = clone_message(input_value)
+            if get_value(message, "timestamp", default=None) is None:
+                message["timestamp"] = now_ms()
+            return message
+
+        if isinstance(input_value, str):
+            if images:
+                content: list[dict[str, Any]] = [{"type": "text", "text": input_value}, *images]
+            else:
+                content = input_value
+            return {"role": "user", "content": content, "timestamp": now_ms()}
+
+        raise TypeError("prompt() expects a string or user message mapping")
+
+
+def _resolve_options_map(options: AgentOptions | Mapping[str, Any]) -> Mapping[str, Any]:
+    if isinstance(options, Mapping):
+        return options
+    return dict(options)