app/agents/tooling.py

"""Lightweight tooling primitives for Lancer agents.

This module provides a small subset of the ideas from JadeAgent:
- declarative tool schemas from Python signatures
- a registry for reusable tools
- structured tool-call objects that can be returned by LLM backends

It intentionally stays small and synchronous so the existing agent code can
adopt it incrementally without turning Lancer into a generic framework.
"""

from __future__ import annotations

import inspect
from dataclasses import dataclass, field
from typing import Any, Callable, get_args, get_origin, get_type_hints


JSON_TYPE_MAP = {
    str: "string",
    int: "integer",
    float: "number",
    bool: "boolean",
    list: "array",
    dict: "object",
}


def _python_type_to_json_schema(py_type: type) -> dict[str, Any]:
    """Map simple Python annotations to JSON schema fragments."""
    origin = get_origin(py_type)
    if origin is not None:
        args = get_args(py_type)
        if origin is list:
            item_type = args[0] if args else str
            return {"type": "array", "items": _python_type_to_json_schema(item_type)}
        if origin is dict:
            return {"type": "object"}
        non_none = [arg for arg in args if arg is not type(None)]
        if non_none:
            return _python_type_to_json_schema(non_none[0])

    return {"type": JSON_TYPE_MAP.get(py_type, "string")}


def _extract_param_description(docstring: str | None, param_name: str) -> str | None:
    """Extract a simple parameter description from an Args: section."""
    if not docstring:
        return None

    lines = docstring.splitlines()
    in_args = False

    for raw_line in lines:
        line = raw_line.strip()
        if line.lower().startswith("args:"):
            in_args = True
            continue
        if not in_args:
            continue
        if not line:
            continue
        if line.startswith(f"{param_name}:"):
            return line.split(":", 1)[1].strip() or None
        if line.startswith(f"{param_name} "):
            parts = line.split(":", 1)
            if len(parts) > 1:
                return parts[1].strip() or None
        # Stop when we leave the Args section.
        if not raw_line.startswith((" ", "\t")):
            break

    return None


@dataclass(frozen=True)
class ToolCall:
    """Structured tool call emitted by the LLM layer."""

    id: str
    name: str
    arguments: dict[str, Any]


@dataclass(frozen=True)
class ToolSchema:
    """OpenAI-compatible tool schema."""

    name: str
    description: str
    parameters: dict[str, Any]

    def to_openai_tool(self) -> dict[str, Any]:
        """Convert schema into OpenAI-compatible tool format."""
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters,
            },
        }


@dataclass
class Tool:
    """Registered callable tool with generated schema."""

    func: Callable[..., Any]
    name: str | None = None
    description: str | None = None
    schema: ToolSchema = field(init=False)

    def __post_init__(self):
        if self.name is None:
            self.name = self.func.__name__
        if self.description is None:
            self.description = self.func.__doc__ or f"Tool: {self.name}"
        self.schema = self._build_schema()

    def _build_schema(self) -> ToolSchema:
        sig = inspect.signature(self.func)
        hints = get_type_hints(self.func)
        properties: dict[str, Any] = {}
        required: list[str] = []

        for param_name, param in sig.parameters.items():
            if param_name in {"self", "cls"}:
                continue

            py_type = hints.get(param_name, str)
            prop = _python_type_to_json_schema(py_type)
            description = _extract_param_description(self.func.__doc__, param_name)
            if description:
                prop["description"] = description
            properties[param_name] = prop

            if param.default is inspect.Parameter.empty:
                required.append(param_name)

        parameters: dict[str, Any] = {
            "type": "object",
            "properties": properties,
        }
        if required:
            parameters["required"] = required

        return ToolSchema(
            name=str(self.name),
            description=str(self.description),
            parameters=parameters,
        )

    def execute(self, arguments: dict[str, Any]) -> Any:
        """Execute the tool with validated arguments."""
        return self.func(**arguments)


def tool(
    func: Callable[..., Any] | None = None,
    *,
    name: str | None = None,
    description: str | None = None,
) -> Tool | Callable[[Callable[..., Any]], Tool]:
    """Decorator for creating Tool objects from plain Python callables."""

    def decorator(inner: Callable[..., Any]) -> Tool:
        return Tool(func=inner, name=name, description=description)

    if func is not None:
        return decorator(func)

    return decorator


class ToolRegistry:
    """Small registry of reusable tools."""

    def __init__(self, tools: list[Tool | Callable[..., Any]] | None = None):
        self._tools: dict[str, Tool] = {}
        for item in tools or []:
            self.register(item)

    def register(self, item: Tool | Callable[..., Any]):
        """Register a Tool or plain callable."""
        if isinstance(item, Tool):
            tool_obj = item
        elif callable(item):
            tool_obj = Tool(func=item)
        else:
            raise TypeError(f"Expected Tool or callable, got {type(item)!r}")

        self._tools[str(tool_obj.name)] = tool_obj

    def get(self, name: str) -> Tool | None:
        """Get a tool by name."""
        return self._tools.get(name)

    @property
    def names(self) -> list[str]:
        """Registered tool names."""
        return list(self._tools.keys())

    @property
    def schemas(self) -> list[ToolSchema]:
        """Structured schemas for all tools."""
        return [tool_obj.schema for tool_obj in self._tools.values()]

    def as_openai_tools(self) -> list[dict[str, Any]]:
        """Serialize all tools into OpenAI-compatible schema format."""
        return [schema.to_openai_tool() for schema in self.schemas]

    def execute(self, tool_call: ToolCall) -> Any:
        """Execute a structured tool call."""
        tool_obj = self.get(tool_call.name)
        if tool_obj is None:
            raise KeyError(f"Unknown tool '{tool_call.name}'. Available: {self.names}")
        return tool_obj.execute(tool_call.arguments)

    def __len__(self) -> int:
        return len(self._tools)