| """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"] |
| 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") |
| if not isinstance(schema, dict): |
| schema = raw.get("inputSchema") |
| parameters = _parse_parameters(schema) if isinstance(schema, dict) else () |
| kw_raw = raw.get("keywords") |
| 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)] |
|
|
