toolproxy 0.1.0__py3-none-any.whl

toolproxy/planner.py

@@ -0,0 +1,241 @@
+"""
+Planner module for toolproxy.
+
+The Planner is responsible for:
+1. Building system prompts (native or emulated mode).
+2. Calling the LLM client.
+3. Parsing the response into Action objects (emulated) or ToolCall lists (native).
+4. Retrying on malformed JSON up to parse_retries times.
+"""
+from __future__ import annotations
+
+import json
+import re
+from typing import List, Optional, Union
+
+from .exceptions import SchemaValidationError
+from .llm_client import LLMClient, ModelResponse
+from .schemas import Action, Message, ToolCall
+from .tools import ToolRegistry
+
+# ---------------------------------------------------------------------------
+# Prompt templates
+# ---------------------------------------------------------------------------
+
+EMULATED_SYSTEM_TEMPLATE = """\
+You are a helpful AI assistant with access to the following tools:
+
+{tool_descriptions}
+
+## CRITICAL INSTRUCTIONS
+
+You MUST respond with ONLY a valid JSON object. No other text, no markdown, no explanation.
+
+Your response must match EXACTLY one of these two formats:
+
+FORMAT 1 - When you need to use a tool:
+{{"type": "tool_call", "tool": {{"tool_name": "<name>", "arguments": {{<key>: <value>}}}}}}
+
+FORMAT 2 - When you have a final answer:
+{{"type": "final", "content": "<your answer here>"}}
+
+Rules:
+- Respond ONLY with raw JSON, nothing else.
+- Always choose the correct tool from the list above.
+- Arguments must match the tool's parameter types exactly.
+- When you have enough information to answer, use format 2.
+"""
+
+TOOL_RESULT_TEMPLATE = "Tool '{name}' returned: {result}"
+
+RETRY_PROMPT_TEMPLATE = """\
+Your previous response was not valid JSON or did not match the required schema.
+Error: {error}
+
+Please try again. You MUST respond with ONLY a valid JSON object matching one of these formats:
+- {{"type": "tool_call", "tool": {{"tool_name": "<name>", "arguments": {{}}}}}}
+- {{"type": "final", "content": "<answer>"}}
+"""
+
+
+# ---------------------------------------------------------------------------
+# JSON extraction helpers
+# ---------------------------------------------------------------------------
+
+def _extract_json(text: str) -> str:
+    """
+    Try to extract a JSON object from *text*.
+
+    Handles cases where the model wraps JSON in markdown code fences or
+    adds leading/trailing commentary.
+    """
+    # Remove markdown fences
+    text = re.sub(r"```(?:json)?", "", text).strip().rstrip("`").strip()
+
+    # Try to find the outermost JSON object
+    start = text.find("{")
+    if start == -1:
+        return text
+
+    depth = 0
+    in_string = False
+    escape_next = False
+    for i, ch in enumerate(text[start:], start=start):
+        if escape_next:
+            escape_next = False
+            continue
+        if ch == "\\" and in_string:
+            escape_next = True
+            continue
+        if ch == '"':
+            in_string = not in_string
+        if not in_string:
+            if ch == "{":
+                depth += 1
+            elif ch == "}":
+                depth -= 1
+                if depth == 0:
+                    return text[start : i + 1]
+    return text[start:]
+
+
+# ---------------------------------------------------------------------------
+# Planner
+# ---------------------------------------------------------------------------
+
+class Planner:
+    """
+    Coordinates model calls, prompt construction, and response parsing.
+
+    In native mode:
+        - Passes tool schemas to the LLM client.
+        - Converts provider tool-call objects to internal ToolCall list.
+
+    In emulated mode:
+        - Injects tool descriptions into the system prompt.
+        - Instructs the model to output a JSON Action.
+        - Validates and retries on schema errors.
+    """
+
+    def __init__(
+        self,
+        client: LLMClient,
+        registry: ToolRegistry,
+        mode: str = "auto",
+        parse_retries: int = 3,
+    ) -> None:
+        self._client = client
+        self._registry = registry
+        self._parse_retries = parse_retries
+
+        if mode == "auto":
+            self._use_native = client.supports_native_tools()
+        elif mode == "native_only":
+            self._use_native = True
+        else:
+            self._use_native = False
+
+    @property
+    def use_native(self) -> bool:
+        return self._use_native
+
+    # ------------------------------------------------------------------
+    # Main entry point
+    # ------------------------------------------------------------------
+
+    def plan(
+        self, messages: List[Message]
+    ) -> Union[List[ToolCall], Action]:
+        """
+        Given *messages*, ask the model what to do next.
+
+        Returns either:
+        - A list of ToolCall objects (native mode, could be multiple).
+        - An Action object (emulated mode).
+        """
+        if self._use_native:
+            return self._plan_native(messages)
+        return self._plan_emulated(messages)
+
+    # ------------------------------------------------------------------
+    # Native mode
+    # ------------------------------------------------------------------
+
+    def _plan_native(self, messages: List[Message]) -> Union[List[ToolCall], Action]:
+        tool_schemas = self._registry.to_openai_schema()
+        response: ModelResponse = self._client.generate(messages, tools=tool_schemas)
+
+        if response.has_tool_calls:
+            return response.tool_calls
+
+        # Model gave a text response with no tool calls → final answer
+        return Action(type="final", content=response.raw_text or "")
+
+    # ------------------------------------------------------------------
+    # Emulated mode
+    # ------------------------------------------------------------------
+
+    def _build_system_message(self) -> Message:
+        tool_descriptions = self._registry.to_prompt_description()
+        content = EMULATED_SYSTEM_TEMPLATE.format(tool_descriptions=tool_descriptions)
+        return Message(role="system", content=content)
+
+    def _inject_system(self, messages: List[Message]) -> List[Message]:
+        """Prepend or replace the system message with the emulated tool prompt."""
+        if messages and messages[0].role == "system":
+            # Replace existing system message content
+            existing_content = messages[0].content
+            new_content = self._build_system_message().content
+            if existing_content not in new_content:
+                new_content = new_content + "\n\n# Additional instructions:\n" + existing_content
+            updated = Message(role="system", content=new_content)
+            return [updated] + messages[1:]
+        return [self._build_system_message()] + messages
+
+    def _plan_emulated(self, messages: List[Message]) -> Action:
+        full_messages = self._inject_system(messages)
+        last_error: Optional[str] = None
+
+        for attempt in range(self._parse_retries):
+            if attempt > 0 and last_error:
+                # Append a retry nudge as assistant + user alternation
+                retry_content = RETRY_PROMPT_TEMPLATE.format(error=last_error)
+                full_messages = full_messages + [
+                    Message(role="user", content=retry_content)
+                ]
+
+            response: ModelResponse = self._client.generate(full_messages)
+            raw = response.raw_text.strip()
+
+            try:
+                action = self._parse_action(raw)
+                return action
+            except (SchemaValidationError, json.JSONDecodeError, ValueError) as exc:
+                last_error = str(exc)
+
+        raise SchemaValidationError(
+            f"Failed to parse valid Action after {self._parse_retries} attempts. "
+            f"Last error: {last_error}"
+        )
+
+    def _parse_action(self, raw_text: str) -> Action:
+        """Parse *raw_text* into an Action, raising on failure."""
+        if not raw_text:
+            raise SchemaValidationError("Model returned empty response.")
+
+        extracted = _extract_json(raw_text)
+        try:
+            data = json.loads(extracted)
+        except json.JSONDecodeError as exc:
+            raise SchemaValidationError(f"Invalid JSON: {exc}. Raw text: {raw_text!r}") from exc
+
+        try:
+            action = Action.model_validate(data)
+        except Exception as exc:
+            raise SchemaValidationError(f"Schema mismatch: {exc}") from exc
+
+        # Validate tool_call has a tool
+        if action.type == "tool_call" and action.tool is None:
+            raise SchemaValidationError("Action type='tool_call' but 'tool' field is missing.")
+
+        return action

toolproxy/schemas.py

@@ -0,0 +1,125 @@
+"""
+Pydantic schemas shared across toolproxy components.
+
+Covers:
+- Conversation messages
+- Tool calls (native and emulated)
+- Agent Action (emulated mode output schema)
+- Tool results
+- Trace and final response
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, List, Literal, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+# ---------------------------------------------------------------------------
+# Conversation message types
+# ---------------------------------------------------------------------------
+
+class Message(BaseModel):
+    """A single conversation message."""
+
+    role: Literal["system", "user", "assistant", "tool"] = "user"
+    content: str
+    # For tool role messages
+    tool_name: Optional[str] = None
+    tool_call_id: Optional[str] = None
+
+
+# ---------------------------------------------------------------------------
+# Tool-call representations
+# ---------------------------------------------------------------------------
+
+class ToolCall(BaseModel):
+    """
+    Represents a request to invoke a single tool.
+
+    Used both in emulated mode (parsed from model JSON output) and
+    native mode (converted from provider-specific format).
+    """
+
+    tool_name: str
+    arguments: Dict[str, Any] = Field(default_factory=dict)
+    # Provider-supplied ID, present only in native mode
+    call_id: Optional[str] = None
+
+
+# ---------------------------------------------------------------------------
+# Emulated-mode Action schema
+# This is the JSON schema the LLM must emit in emulated mode.
+# ---------------------------------------------------------------------------
+
+class Action(BaseModel):
+    """
+    The structured output the LLM must emit in emulated mode.
+
+    type="tool_call"  → the model wants to call a tool.
+    type="final"      → the model is ready to give a final answer.
+    """
+
+    type: Literal["tool_call", "final"]
+    tool: Optional[ToolCall] = None
+    content: Optional[str] = None
+
+    model_config = ConfigDict(extra="ignore")
+
+
+# ---------------------------------------------------------------------------
+# Tool result
+# ---------------------------------------------------------------------------
+
+class ToolResult(BaseModel):
+    """The outcome of executing a single tool call."""
+
+    tool_name: str
+    call_id: Optional[str] = None
+    output: Optional[str] = None
+    error: Optional[str] = None
+
+    @property
+    def success(self) -> bool:
+        return self.error is None
+
+    def as_message_content(self) -> str:
+        """Serialize the result for injection back into the conversation."""
+        if self.success:
+            return f"[Tool '{self.tool_name}' result]: {self.output}"
+        return f"[Tool '{self.tool_name}' error]: {self.error}"
+
+
+# ---------------------------------------------------------------------------
+# Trace (debug info)
+# ---------------------------------------------------------------------------
+
+class TraceEntry(BaseModel):
+    """One step in the agent trace."""
+
+    step: int
+    tool_call: Optional[ToolCall] = None
+    tool_result: Optional[ToolResult] = None
+    model_output: Optional[str] = None
+
+
+class AgentTrace(BaseModel):
+    """Full trace of an agent run."""
+
+    steps: List[TraceEntry] = Field(default_factory=list)
+
+    @property
+    def tool_calls(self) -> List[ToolCall]:
+        return [e.tool_call for e in self.steps if e.tool_call is not None]
+
+
+# ---------------------------------------------------------------------------
+# Final agent response
+# ---------------------------------------------------------------------------
+
+class AgentResponse(BaseModel):
+    """The final response returned by UniversalAgent.run()."""
+
+    content: str
+    trace: Optional[AgentTrace] = None
+    steps_taken: int = 0

toolproxy/tools.py

@@ -0,0 +1,221 @@
+"""
+Tool registry and @tool decorator for toolproxy.
+
+Allows users to register Python callables as tools with automatic
+Pydantic schema generation from type hints or explicit args_schema.
+"""
+from __future__ import annotations
+
+import inspect
+import json
+from dataclasses import dataclass, field
+from functools import wraps
+from typing import Any, Callable, Dict, Optional, Type, get_type_hints
+
+from pydantic import BaseModel, create_model
+
+from .exceptions import ToolNotFoundError
+
+
+# ---------------------------------------------------------------------------
+# ToolDefinition
+# ---------------------------------------------------------------------------
+
+@dataclass
+class ToolDefinition:
+    """Metadata + callable for a single registered tool."""
+
+    name: str
+    description: str
+    callable: Callable[..., Any]
+    args_schema: Type[BaseModel]
+
+    def to_openai_schema(self) -> Dict[str, Any]:
+        """
+        Return an OpenAI-compatible function/tool schema dict.
+
+        Compatible with OpenAI, OpenRouter, and most OpenAI-style providers.
+        """
+        schema = self.args_schema.model_json_schema()
+        # Remove title from root and from nested properties for cleaner prompts
+        schema.pop("title", None)
+        for prop in schema.get("properties", {}).values():
+            prop.pop("title", None)
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": schema,
+            },
+        }
+
+    def to_prompt_description(self) -> str:
+        """
+        Return a human-readable description of the tool for emulated-mode prompts.
+        """
+        schema = self.args_schema.model_json_schema()
+        props = schema.get("properties", {})
+        required = schema.get("required", [])
+        params_lines = []
+        for param_name, param_info in props.items():
+            req_marker = " (required)" if param_name in required else " (optional)"
+            param_type = param_info.get("type", "any")
+            desc = param_info.get("description", "")
+            line = f"    - {param_name}: {param_type}{req_marker}"
+            if desc:
+                line += f" — {desc}"
+            params_lines.append(line)
+        params_str = "\n".join(params_lines) if params_lines else "    (no parameters)"
+        return (
+            f"Tool: {self.name}\n"
+            f"Description: {self.description}\n"
+            f"Parameters:\n{params_str}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Schema auto-generation from function signature
+# ---------------------------------------------------------------------------
+
+def _build_args_schema(func: Callable[..., Any]) -> Type[BaseModel]:
+    """
+    Introspect *func*'s type hints and build a Pydantic model for its args.
+
+    Only non-return, non-self/cls parameters are included. Parameters without
+    type annotations default to `Any`.
+    """
+    hints = get_type_hints(func)
+    hints.pop("return", None)
+    sig = inspect.signature(func)
+    fields: Dict[str, Any] = {}
+    for param_name, param in sig.parameters.items():
+        if param_name in ("self", "cls"):
+            continue
+        annotation = hints.get(param_name, Any)
+        if param.default is inspect.Parameter.empty:
+            fields[param_name] = (annotation, ...)
+        else:
+            fields[param_name] = (annotation, param.default)
+    model_name = f"{func.__name__.capitalize()}Args"
+    return create_model(model_name, **fields)  # type: ignore[call-overload]
+
+
+# ---------------------------------------------------------------------------
+# @tool decorator
+# ---------------------------------------------------------------------------
+
+def tool(
+    func: Optional[Callable[..., Any]] = None,
+    *,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    args_schema: Optional[Type[BaseModel]] = None,
+) -> Any:
+    """
+    Decorator that marks a Python callable as a toolproxy tool.
+
+    Usage (auto-detected schema)::
+
+        @tool
+        def search_web(query: str, top_k: int = 5) -> str:
+            \"\"\"Search the web for recent information.\"\"\"
+            ...
+
+    Usage (explicit schema)::
+
+        class SearchArgs(BaseModel):
+            query: str
+            top_k: int = 5
+
+        @tool(name="search_web", args_schema=SearchArgs, description="Search the web")
+        def search_web_tool(args: SearchArgs) -> str:
+            ...
+
+    Returns the original function (unchanged), but attaches a
+    `._tool_definition` attribute that ToolRegistry can read.
+    """
+    def _decorate(fn: Callable[..., Any]) -> Callable[..., Any]:
+        tool_name = name or fn.__name__
+        tool_desc = description or (inspect.getdoc(fn) or "")
+        schema = args_schema or _build_args_schema(fn)
+
+        definition = ToolDefinition(
+            name=tool_name,
+            description=tool_desc,
+            callable=fn,
+            args_schema=schema,
+        )
+
+        @wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            return fn(*args, **kwargs)
+
+        wrapper._tool_definition = definition  # type: ignore[attr-defined]
+        return wrapper
+
+    if func is not None:
+        # Called as @tool (no parentheses)
+        return _decorate(func)
+    # Called as @tool(...) with arguments
+    return _decorate
+
+
+# ---------------------------------------------------------------------------
+# ToolRegistry
+# ---------------------------------------------------------------------------
+
+class ToolRegistry:
+    """Stores and provides access to all registered tools."""
+
+    def __init__(self) -> None:
+        self._tools: Dict[str, ToolDefinition] = {}
+
+    def register(self, tool_or_definition: Any) -> None:
+        """
+        Register a tool.
+
+        Accepts either:
+        - A callable decorated with @tool (has ._tool_definition attribute).
+        - A ToolDefinition instance directly.
+        """
+        if isinstance(tool_or_definition, ToolDefinition):
+            defn = tool_or_definition
+        elif hasattr(tool_or_definition, "_tool_definition"):
+            defn = tool_or_definition._tool_definition
+        else:
+            raise ValueError(
+                f"Cannot register '{tool_or_definition}': "
+                "must be a @tool-decorated callable or a ToolDefinition."
+            )
+        if defn.name in self._tools:
+            raise ValueError(f"A tool named '{defn.name}' is already registered.")
+        self._tools[defn.name] = defn
+
+    def get(self, name: str) -> ToolDefinition:
+        """Return the ToolDefinition for *name*, raising ToolNotFoundError if absent."""
+        try:
+            return self._tools[name]
+        except KeyError:
+            raise ToolNotFoundError(name)
+
+    def list_tools(self) -> list[ToolDefinition]:
+        """Return all registered ToolDefinitions."""
+        return list(self._tools.values())
+
+    def to_openai_schema(self) -> list[Dict[str, Any]]:
+        """Return all tools as a list of OpenAI-compatible tool schemas."""
+        return [t.to_openai_schema() for t in self._tools.values()]
+
+    def to_prompt_description(self) -> str:
+        """Return a multi-tool prompt block for emulated mode."""
+        if not self._tools:
+            return "No tools available."
+        blocks = [t.to_prompt_description() for t in self._tools.values()]
+        return "\n\n".join(blocks)
+
+    def __len__(self) -> int:
+        return len(self._tools)
+
+    def __contains__(self, name: str) -> bool:
+        return name in self._tools