agent-runtime-sdk 0.1.0__py3-none-any.whl

agent_runtime/server/management.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from fastapi import APIRouter, HTTPException
+from pydantic import BaseModel
+
+from ..config.definitions import ToolPolicy
+
+if TYPE_CHECKING:
+    from ..runtime import ManagedAgentRuntime
+
+
+def create_management_router(runtime: ManagedAgentRuntime) -> APIRouter:
+    """Create a FastAPI router for runtime tool/policy management."""
+
+    router = APIRouter(prefix="/management", tags=["management"])
+
+    def _tool_payload(tool, policy: ToolPolicy | None = None):
+        return {
+            "name": tool.name,
+            "source_mcp": tool.source_mcp,
+            "description": tool.description,
+            "inputs": tool.inputs,
+            "policy": (policy or ToolPolicy()).model_dump(),
+        }
+
+    @router.get("/tools")
+    def list_tools():
+        tools = runtime.list_tools()
+        policies = runtime.list_tool_policies()
+        return [
+            _tool_payload(t, policies.get(t.name))
+            for t in tools
+        ]
+
+    @router.get("/mcps")
+    def list_mcps():
+        return [
+            {
+                "name": mcp.name,
+                "enabled": mcp.enabled,
+                "url": mcp.url,
+                "transport": mcp.transport,
+                "tool_count": len(runtime.list_tools(mcp.name)),
+                "policy_count": len(runtime.list_tool_policies(mcp.name)),
+            }
+            for mcp in runtime.definition.mcps
+        ]
+
+    @router.get("/mcps/{mcp_name}/tools")
+    def list_mcp_tools(mcp_name: str):
+        tools = runtime.list_tools(mcp_name)
+        if not any(mcp.name == mcp_name for mcp in runtime.definition.mcps):
+            raise HTTPException(404, f"mcp '{mcp_name}' not found")
+        policies = runtime.list_tool_policies(mcp_name)
+        return [_tool_payload(tool, policies.get(tool.name)) for tool in tools]
+
+    @router.get("/mcps/{mcp_name}/tools/{tool_name}")
+    def get_mcp_tool(mcp_name: str, tool_name: str):
+        tool = runtime.get_tool(tool_name, mcp_name)
+        if not tool:
+            raise HTTPException(404, f"tool '{tool_name}' not found in mcp '{mcp_name}'")
+        policy = runtime.get_tool_policy(tool_name, mcp_name)
+        return _tool_payload(tool, policy)
+
+    @router.put("/mcps/{mcp_name}/tools/{tool_name}/policy")
+    def update_tool_policy(mcp_name: str, tool_name: str, body: ToolPolicy):
+        try:
+            runtime.set_tool_policy(tool_name, body, mcp_name=mcp_name)
+        except ValueError as exc:
+            raise HTTPException(404, str(exc)) from exc
+        return {"ok": True, "mcp_name": mcp_name, "tool_name": tool_name, "policy": body.model_dump()}
+
+    @router.delete("/mcps/{mcp_name}/tools/{tool_name}/policy")
+    def delete_tool_policy(mcp_name: str, tool_name: str):
+        removed = runtime.remove_tool_policy(tool_name, mcp_name=mcp_name)
+        if not removed:
+            raise HTTPException(404, f"no policy for '{tool_name}' in mcp '{mcp_name}'")
+        return {"ok": True}
+
+    @router.post("/reload")
+    def reload_tools():
+        runtime.reload(discover=True, skip_plugin_load=True)
+        return {"tool_count": len(runtime.discovered_tools)}
+
+    # ── Instructions ──
+
+    class InstructionsBody(BaseModel):
+        instructions: str
+
+    @router.get("/instructions")
+    def get_instructions():
+        return {"instructions": runtime.get_instructions()}
+
+    @router.put("/instructions")
+    def set_instructions(body: InstructionsBody):
+        runtime.set_instructions(body.instructions)
+        return {"ok": True, "instructions": body.instructions}
+
+    @router.get("/config")
+    def get_config():
+        dump = runtime.definition.model_dump()
+        dump.get("runtime", {}).get("model", {}).pop("api_key_env", None)
+        return dump
+
+    return router

agent_runtime/smol/__init__.py

@@ -0,0 +1,4 @@
+from .agents import ToolCallingCheckAgent
+from .llm_models import DeepSeekModel
+
+__all__ = ["DeepSeekModel", "ToolCallingCheckAgent"]

agent_runtime/smol/agents.py

@@ -0,0 +1,431 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+from collections.abc import Callable, Generator
+from typing import TYPE_CHECKING, Any
+
+from rich.live import Live
+from rich.markdown import Markdown
+from rich.text import Text
+from smolagents import (
+    ActionOutput,
+    AgentError,
+    AgentGenerationError,
+    AgentParsingError,
+    AgentToolExecutionError,
+    ChatMessage,
+    ChatMessageStreamDelta,
+    FinalAnswerStep,
+    LogLevel,
+    PlanningStep,
+    Timing,
+    ToolCall,
+    ToolCallingAgent,
+    ToolOutput,
+    YELLOW_HEX,
+    agglomerate_stream_deltas,
+    handle_agent_output_types,
+    parse_json_if_needed,
+)
+
+from ..runtime_engine.context import (
+    ActiveToolCall,
+    current_task_id,
+    current_task_pool,
+)
+from .memory import ActionStep, ToolCallbackRegistry
+
+if TYPE_CHECKING:
+    import PIL.Image
+
+
+class ToolCallingCheckAgent(ToolCallingAgent):
+    def __init__(
+        self,
+        tool_callbacks: dict[str, Callable | list[Callable]] | None = None,
+        after_tool_hook: Callable[[ToolOutput], None] | None = None,
+        **kwargs,
+    ):
+        if tool_callbacks:
+            self.interrupt_before_tool_name = list(tool_callbacks.keys())
+        self.after_tool_hook = after_tool_hook
+        self._setup_tool_callbacks(tool_callbacks)
+        super().__init__(**kwargs)
+
+    def _setup_tool_callbacks(self, tool_callbacks):
+        self.tool_callbacks = ToolCallbackRegistry()
+        if not tool_callbacks:
+            return
+        if not isinstance(tool_callbacks, dict):
+            raise ValueError("tool_callbacks must be a dict")
+        for tool_name, callbacks in tool_callbacks.items():
+            callback_list = callbacks if isinstance(callbacks, list) else [callbacks]
+            for callback in callback_list:
+                self.tool_callbacks.register(tool_name, callback)
+
+    @staticmethod
+    def _current_task_context():
+        try:
+            task_id = current_task_id.get()
+            pool = current_task_pool.get()
+        except LookupError:
+            return None, None
+        return task_id, pool.get(task_id)
+
+    def execute_tool_call(self, tool_name: str, arguments: dict[str, str] | str) -> Any:
+        task_id, task_info = self._current_task_context()
+        if task_info is not None and task_info.control.stop_requested:
+            raise task_info.control.stop_error or RuntimeError(
+                "task cancelled before tool execution"
+            )
+
+        available_tools = {**self.tools, **self.managed_agents}
+        tool = available_tools.get(tool_name)
+        active_tool: ActiveToolCall | None = None
+        if task_info is not None and tool is not None:
+            active_tool = ActiveToolCall(
+                tool_name=tool_name,
+                mcp_name=getattr(tool, "_runtime_mcp_name", None),
+                cancel_hook=getattr(tool, "_runtime_cancel_hook", None),
+            )
+            task_info.set_active_tool(active_tool)
+
+        try:
+            return super().execute_tool_call(tool_name, arguments)
+        finally:
+            if task_info is not None and task_info.active_tool is active_tool:
+                task_info.clear_active_tool()
+
+    def _run_stream(
+        self,
+        task: str,
+        max_steps: int,
+        images: list["PIL.Image.Image"] | None = None,
+    ) -> Generator[
+        ActionStep | PlanningStep | FinalAnswerStep | ChatMessageStreamDelta
+    ]:
+        self.step_number = 1
+        returned_final_answer = False
+        while not returned_final_answer and self.step_number <= max_steps:
+            if self.interrupt_switch:
+                raise AgentError("Agent interrupted.", self.logger)
+
+            if self.planning_interval is not None and (
+                self.step_number == 1
+                or (self.step_number - 1) % self.planning_interval == 0
+            ):
+                planning_start_time = time.time()
+                planning_step = None
+                for element in self._generate_planning_step(
+                    task,
+                    is_first_step=len(self.memory.steps) == 1,
+                    step=self.step_number,
+                ):
+                    yield element
+                    planning_step = element
+                assert isinstance(planning_step, PlanningStep)
+                planning_end_time = time.time()
+                planning_step.timing = Timing(
+                    start_time=planning_start_time, end_time=planning_end_time
+                )
+                self._finalize_step(planning_step)
+                self.memory.steps.append(planning_step)
+
+            action_step_start_time = time.time()
+            action_step = ActionStep(
+                step_number=self.step_number,
+                timing=Timing(start_time=action_step_start_time),
+                observations_images=images,
+            )
+            self.logger.log_rule(f"Step {self.step_number}", level=LogLevel.INFO)
+            try:
+                for output in self._step_stream(action_step):
+                    yield output
+                    if isinstance(output, ActionOutput) and output.is_final_answer:
+                        final_answer = output.output
+                        self.logger.log(
+                            Text(
+                                f"Final answer: {final_answer}",
+                                style=f"bold {YELLOW_HEX}",
+                            ),
+                            level=LogLevel.INFO,
+                        )
+                        if self.final_answer_checks:
+                            self._validate_final_answer(final_answer)
+                        returned_final_answer = True
+                        action_step.is_final_answer = True
+            except AgentGenerationError as exc:
+                raise exc
+            except AgentError as exc:
+                action_step.error = exc
+            finally:
+                self._finalize_step(action_step)
+                self._emit_step_artifacts(action_step)
+                self.memory.steps.append(action_step)
+                yield action_step
+                self.step_number += 1
+
+        if not returned_final_answer and self.step_number == max_steps + 1:
+            final_answer = self._handle_max_steps_reached(task)
+            yield action_step
+
+        yield FinalAnswerStep(handle_agent_output_types(final_answer))
+
+    def _emit_step_artifacts(self, action_step: ActionStep) -> None:
+        try:
+            task_id = current_task_id.get()
+            pool = current_task_pool.get()
+        except LookupError:
+            return
+
+        task_info = pool.get(task_id)
+        if not task_info:
+            return
+
+        updater = task_info.updater
+        loop = task_info.loop
+        if not updater or not loop or loop.is_closed():
+            return
+
+        from a2a.types import Part, TaskState
+
+        def _emit_artifact(
+            parts: list[Part], metadata: dict[str, Any], name: str
+        ) -> bool:
+            try:
+                future = asyncio.run_coroutine_threadsafe(
+                    updater.add_artifact(parts=parts, name=name, metadata=metadata),
+                    loop,
+                )
+                future.result(timeout=5)
+                return True
+            except Exception:
+                return False
+
+        step_no = action_step.step_number
+        tool_observations: dict[str, dict[str, Any]] = (
+            getattr(action_step, "tool_observations", {}) or {}
+        )
+        tool_calls = action_step.tool_calls or []
+        has_final_answer = False
+        emitted_observation_ids: set[str] = set()
+
+        for tool_call in tool_calls:
+            tool_name = ""
+            tool_args = None
+            tool_id = None
+
+            if hasattr(tool_call, "function"):
+                func = tool_call.function
+                tool_name = getattr(func, "name", "") or tool_name
+                tool_args = getattr(func, "arguments", None)
+                tool_id = getattr(tool_call, "id", None)
+            else:
+                tool_name = getattr(tool_call, "name", "") or tool_name
+                tool_args = getattr(tool_call, "arguments", None)
+                tool_id = getattr(tool_call, "id", None)
+
+            if not tool_name and hasattr(tool_call, "dict"):
+                tool_call_dict = tool_call.dict()
+                tool_name = (tool_call_dict.get("function", {}) or {}).get(
+                    "name", ""
+                ) or tool_name
+                tool_args = (tool_call_dict.get("function", {}) or {}).get(
+                    "arguments", tool_args
+                )
+                tool_id = tool_call_dict.get("id", tool_id)
+
+            if tool_name == "final_answer":
+                has_final_answer = True
+                continue
+            if not tool_name:
+                continue
+
+            normalized_args = tool_args
+            if isinstance(normalized_args, str):
+                try:
+                    normalized_args = json.loads(normalized_args)
+                except Exception:
+                    pass
+
+            _emit_artifact(
+                parts=[
+                    Part.model_validate(
+                        {
+                            "kind": "data",
+                            "data": {
+                                "tool": {
+                                    "name": tool_name,
+                                    "args": normalized_args,
+                                    "id": tool_id,
+                                }
+                            },
+                        }
+                    )
+                ],
+                metadata={
+                    "event": "tool_call",
+                    "state": TaskState.working.value,
+                    "step": step_no,
+                },
+                name="tool_call",
+            )
+
+            obs_payload = tool_observations.get(tool_id or "")
+            if obs_payload:
+                obs_text = str(obs_payload.get("observation") or "").strip()
+                if obs_text:
+                    _emit_artifact(
+                        parts=[
+                            Part.model_validate(
+                                {
+                                    "kind": "data",
+                                    "data": {
+                                        "tool": {"name": tool_name, "id": tool_id},
+                                        "observation": obs_text,
+                                    },
+                                }
+                            )
+                        ],
+                        metadata={
+                            "event": "observation",
+                            "state": TaskState.working.value,
+                            "step": step_no,
+                        },
+                        name="observation",
+                    )
+                    if tool_id:
+                        emitted_observation_ids.add(tool_id)
+
+        if not has_final_answer:
+            for obs_id, obs_payload in tool_observations.items():
+                if obs_id in emitted_observation_ids:
+                    continue
+                obs_text = str(obs_payload.get("observation") or "").strip()
+                if not obs_text:
+                    continue
+                _emit_artifact(
+                    parts=[
+                        Part.model_validate(
+                            {
+                                "kind": "data",
+                                "data": {
+                                    "tool": {
+                                        "name": obs_payload.get("tool_name"),
+                                        "id": obs_id,
+                                    },
+                                    "observation": obs_text,
+                                },
+                            }
+                        )
+                    ],
+                    metadata={
+                        "event": "observation",
+                        "state": TaskState.working.value,
+                        "step": step_no,
+                    },
+                    name="observation",
+                )
+
+    def _step_stream(
+        self,
+        memory_step: ActionStep,
+    ) -> Generator[ChatMessageStreamDelta | ToolCall | ToolOutput | ActionOutput]:
+        input_messages = self.write_memory_to_messages().copy()
+        memory_step.model_input_messages = input_messages
+        try:
+            if self.stream_outputs and hasattr(self.model, "generate_stream"):
+                output_stream = self.model.generate_stream(
+                    input_messages,
+                    stop_sequences=["Observation:", "Calling tools:"],
+                    tools_to_call_from=self.tools_and_managed_agents,
+                )
+                chat_message_stream_deltas: list[ChatMessageStreamDelta] = []
+                with Live(
+                    "", console=self.logger.console, vertical_overflow="visible"
+                ) as live:
+                    for event in output_stream:
+                        chat_message_stream_deltas.append(event)
+                        live.update(
+                            Markdown(
+                                agglomerate_stream_deltas(
+                                    chat_message_stream_deltas
+                                ).render_as_markdown()
+                            )
+                        )
+                        yield event
+                chat_message = agglomerate_stream_deltas(chat_message_stream_deltas)
+            else:
+                chat_message = self.model.generate(
+                    input_messages,
+                    stop_sequences=["Observation:", "Calling tools:"],
+                    tools_to_call_from=self.tools_and_managed_agents,
+                )
+                log_content = (
+                    str(chat_message.raw)
+                    if chat_message.content is None and chat_message.raw is not None
+                    else str(chat_message.content) or ""
+                )
+                self.logger.log_markdown(
+                    content=log_content,
+                    title="Output message of the LLM:",
+                    level=LogLevel.DEBUG,
+                )
+            memory_step.model_output_message = chat_message
+            memory_step.model_output = chat_message.content
+            memory_step.token_usage = chat_message.token_usage
+        except Exception as exc:
+            raise AgentGenerationError(
+                f"Error while generating output:\n{exc}", self.logger
+            ) from exc
+
+        if chat_message.tool_calls is None or len(chat_message.tool_calls) == 0:
+            try:
+                chat_message = self.model.parse_tool_calls(chat_message)
+            except Exception as exc:
+                raise AgentParsingError(
+                    f"Error while parsing tool call from model output: {exc}",
+                    self.logger,
+                )
+        else:
+            for tool_call in chat_message.tool_calls:
+                tool_call.function.arguments = parse_json_if_needed(
+                    tool_call.function.arguments
+                )
+
+        final_answer, got_final_answer = None, False
+        tool_observations: dict[str, dict[str, Any]] = {}
+        self.tool_callbacks.callback(chat_message, agent=self)
+
+        for output in self.process_tool_calls(chat_message, memory_step):
+            yield output
+            if isinstance(output, ToolOutput):
+                if output.id:
+                    tool_name = (
+                        getattr(output.tool_call, "name", None)
+                        if output.tool_call is not None
+                        else None
+                    )
+                    tool_observations[output.id] = {
+                        "tool_name": tool_name,
+                        "observation": output.observation,
+                    }
+                if self.after_tool_hook is not None:
+                    self.after_tool_hook(output)
+                if output.is_final_answer:
+                    if got_final_answer:
+                        raise AgentToolExecutionError(
+                            "You returned multiple final answers. Please return only one single final answer!",
+                            self.logger,
+                        )
+                    final_answer = output.output
+                    got_final_answer = True
+                    if (
+                        isinstance(final_answer, str)
+                        and final_answer in self.state.keys()
+                    ):
+                        final_answer = self.state[final_answer]
+        memory_step.tool_observations = tool_observations
+        yield ActionOutput(output=final_answer, is_final_answer=got_final_answer)