toolset.py

"""Tool-catalog loading for the embedding tool selector.

A *pack* is JSON in the standard function-calling shape shared by OpenAI function
calling, plain JSON-Schema tool definitions, and MCP ``tools/list``::

    {"name": "ecommerce", "tools": [
        {"name": "search_products",
         "description": "Search the catalog by keyword and filters.",
         "parameters": {"type": "object",
             "properties": {
                 "query": {"type": "string"},
                 "category": {"type": "string", "enum": ["lighting", "rugs"]}},
             "required": ["query"]}}]}

Seven curated packs ship in ``packs/`` (ecommerce, devops, travel, support,
finance, healthcare, workplace — 151 tools total). They are merged into one
catalog and embedded once at startup; typing a request retrieves the few tools
that matter. The loader is liberal in what it accepts (a ``{"tools": [...]}``
wrapper, an MCP result with params under ``inputSchema``, a bare OpenAI tools
array, or a single tool), so any catalog that speaks the schema works unchanged.

Enum-valued parameters become part of a tool's indexed text: the discriminating
word in a request is often a *value* ("aws", "business class", "urgent"), so
surfacing the allowed vocabulary lets the retriever match it.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

PACKS_DIR = Path(__file__).resolve().parent / "packs"


@dataclass(frozen=True)
class Parameter:
    """One tool parameter parsed from a JSON-Schema property."""

    name: str
    type: str
    description: str
    enum: tuple[str, ...] | None
    required: bool

    @property
    def is_enum(self) -> bool:
        return self.enum is not None and len(self.enum) > 0


@dataclass(frozen=True)
class Tool:
    """A single callable tool, its parameters, and the pack it came from."""

    name: str
    description: str
    domain: str = ""
    parameters: tuple[Parameter, ...] = ()
    keywords: tuple[str, ...] = ()

    @property
    def routing_text(self) -> str:
        """Text indexed for retrieval: name + description + parameter names + enum values + keywords.

        Including the allowed enum values markedly improves routing on mixed catalogs
        (combined-catalog R@1 0.82 -> 0.95 on the original eval): the discriminating
        word is often a value ("aws", "business class", "urgent"), so surfacing the
        vocabulary lets the retriever match it.

        ``keywords`` is *document expansion* for implicit queries: a 350M retriever does
        not reliably infer "table" -> "furniture" -> product catalog, so we add the
        concrete words users say ("table", "lamp", "desk") to the indexed text. The words
        are index-only — they never appear in the card or the JSON tool definition.
        """
        desc = self.description.strip()
        base = f"{self.name}: {desc}" if desc else self.name
        extras: list[str] = []
        if self.parameters:
            extras.append("parameters: " + ", ".join(p.name.replace("_", " ") for p in self.parameters))
            values = [v.replace("_", " ") for p in self.parameters if p.is_enum for v in (p.enum or ())]
            if values:
                extras.append("options: " + ", ".join(values))
        if self.keywords:
            extras.append("examples: " + ", ".join(self.keywords))
        return f"{base} ({'; '.join(extras)})" if extras else base

    def to_schema(self) -> dict[str, Any]:
        """The standard function-calling schema — exactly what a downstream LLM receives."""
        properties: dict[str, Any] = {}
        required: list[str] = []
        for p in self.parameters:
            spec: dict[str, Any] = {"type": p.type}
            if p.description:
                spec["description"] = p.description
            if p.enum:
                spec["enum"] = list(p.enum)
            properties[p.name] = spec
            if p.required:
                required.append(p.name)
        params: dict[str, Any] = {"type": "object", "properties": properties}
        if required:
            params["required"] = required
        return {"name": self.name, "description": self.description, "parameters": params}

    def to_public(self) -> dict[str, Any]:
        """JSON-able view sent to the frontend: card data + the full tool definition."""
        return {
            "name": self.name,
            "description": self.description,
            "domain": self.domain,
            "params": [
                {"name": p.name, "type": p.type, "description": p.description,
                 "enum": list(p.enum) if p.enum else None, "required": p.required}
                for p in self.parameters
            ],
            "schema": self.to_schema(),
        }


class ToolSetError(ValueError):
    """Raised when a payload cannot be parsed as a pack."""


def _parse_parameters(schema: dict[str, Any]) -> tuple[Parameter, ...]:
    properties = schema.get("properties")
    if not isinstance(properties, dict):
        return ()
    required_raw = schema.get("required", [])
    required = {str(r) for r in required_raw} if isinstance(required_raw, list) else set()
    params: list[Parameter] = []
    for raw_name, raw_spec in properties.items():
        spec = raw_spec if isinstance(raw_spec, dict) else {}
        enum_raw = spec.get("enum")
        enum = tuple(str(v) for v in enum_raw) if isinstance(enum_raw, list) and enum_raw else None
        params.append(
            Parameter(
                name=str(raw_name),
                type=str(spec.get("type", "string")),
                description=str(spec.get("description", "")),
                enum=enum,
                required=str(raw_name) in required,
            )
        )
    return tuple(params)


def tool_from_dict(raw: dict[str, Any], domain: str = "") -> Tool:
    """Parse one tool dict (OpenAI/JSON-Schema/MCP), tolerant of the common shapes."""
    if raw.get("type") == "function" and isinstance(raw.get("function"), dict):
        raw = raw["function"]  # OpenAI wraps each tool as {"type": "function", "function": {...}}
    name = raw.get("name")
    if not isinstance(name, str) or not name:
        raise ToolSetError("each tool needs a non-empty 'name'")
    schema = raw.get("parameters")  # JSON-Schema/OpenAI use 'parameters'; MCP uses 'inputSchema'
    if not isinstance(schema, dict):
        schema = raw.get("inputSchema")
    parameters = _parse_parameters(schema) if isinstance(schema, dict) else ()
    kw_raw = raw.get("keywords")  # optional index-only document expansion for implicit queries
    keywords = tuple(str(k) for k in kw_raw) if isinstance(kw_raw, list) else ()
    return Tool(name=name, description=str(raw.get("description", "")), domain=domain,
                parameters=parameters, keywords=keywords)


def _extract(data: Any) -> tuple[str, list[dict[str, Any]]]:
    """Pull (pack_name, [tool_dict, ...]) out of any accepted top-level shape."""
    if isinstance(data, list):
        return "tools", [d for d in data if isinstance(d, dict)]
    if isinstance(data, dict):
        tools = data.get("tools")
        if isinstance(tools, list):
            return str(data.get("name", "tools")), [d for d in tools if isinstance(d, dict)]
        if data.get("name") and (data.get("parameters") or data.get("inputSchema")):
            return str(data["name"]), [data]
    raise ToolSetError("expected a tools array, a {'tools': [...]} object, or one tool")


def load_pack(path: Path) -> list[Tool]:
    """Load one curated pack file; the pack's declared name is each tool's domain."""
    data = json.loads(path.read_text(encoding="utf-8"))
    domain, dicts = _extract(data)
    return [tool_from_dict(d, domain=domain) for d in dicts]


def available_packs() -> list[Path]:
    """All curated packs shipped with the demo, sorted by name (deterministic order)."""
    return sorted(PACKS_DIR.glob("*.json"))


def load_catalog() -> list[Tool]:
    """Every curated pack merged into one catalog, in a deterministic order."""
    return [tool for path in available_packs() for tool in load_pack(path)]