truegrain.tools

  1"""Tool definitions for agent frameworks.
  2
  3Most frameworks, whatever else they disagree about, accept a list of function
  4schemas in the shape OpenAI popularised, and call back with a name and a
  5dictionary of arguments. This module produces that list and dispatches those
  6calls, which is enough to wire the engine into LangChain, CrewAI, the OpenAI
  7Agents SDK, Pydantic AI or something hand-rolled, without this library growing
  8an adapter per framework. Adapters age badly; a schema does not.
  9
 10If your client speaks the Model Context Protocol, use the engine's MCP server
 11directly instead. It is the same four tools with the same guarantees.
 12
 13Example:
 14    >>> from truegrain import Client, tools
 15    >>> client = Client.from_env()
 16    >>> specs = tools.tool_specs()                     # hand to the framework
 17    >>> tools.dispatch(client, "query", {"metrics": ["sales.order_revenue"]})
 18"""
 19
 20from __future__ import annotations
 21
 22from typing import Any, Mapping
 23
 24from .client import Client
 25from .errors import Refused
 26
 27__all__ = ["tool_specs", "dispatch", "TOOL_NAMES"]
 28
 29TOOL_NAMES = ("list_metrics", "describe_metric", "list_dimensions", "query")
 30
 31_FILTER_SCHEMA: dict[str, Any] = {
 32    "type": "object",
 33    "description": (
 34        "A structured predicate. This is not SQL and no SQL fragment is accepted."
 35    ),
 36    "required": ["dimension", "op"],
 37    "properties": {
 38        "dimension": {"type": "string", "description": "dataset.field or namespace.dataset.field"},
 39        "op": {
 40            "type": "string",
 41            "enum": ["eq", "ne", "in", "not_in", "gt", "gte", "lt", "lte",
 42                     "between", "is_null", "is_not_null"],
 43        },
 44        "values": {
 45            "type": "array",
 46            "description": (
 47                "One value for the comparisons, two for between, one or more for "
 48                "in and not_in, none for the null checks."
 49            ),
 50            "items": {},
 51        },
 52    },
 53}
 54
 55
 56def tool_specs() -> list[dict[str, Any]]:
 57    """Return the engine's tools as OpenAI-style function schemas.
 58
 59    The descriptions are written for a model rather than a human: they say when
 60    to use a tool and when not to, because that is what decides whether an agent
 61    picks the right metric or invents one.
 62    """
 63    return [
 64        {
 65            "type": "function",
 66            "function": {
 67                "name": "list_metrics",
 68                "description": (
 69                    "List the metrics this semantic layer can answer, with a "
 70                    "plain-language description of each. Call this first when you "
 71                    "do not already know the exact metric name. Do not guess a "
 72                    "metric name: a name not in this list does not exist and the "
 73                    "query will be refused."
 74                ),
 75                "parameters": {
 76                    "type": "object",
 77                    "properties": {
 78                        "search": {
 79                            "type": "string",
 80                            "description": "Narrow by name, description or synonym.",
 81                        }
 82                    },
 83                },
 84            },
 85        },
 86        {
 87            "type": "function",
 88            "function": {
 89                "name": "describe_metric",
 90                "description": (
 91                    "Return one metric's full definition and exactly which "
 92                    "dimensions it can be grouped by. Call this before query "
 93                    "whenever you are unsure a metric answers the question asked, "
 94                    "or which dimensions are legal for it."
 95                ),
 96                "parameters": {
 97                    "type": "object",
 98                    "required": ["name"],
 99                    "properties": {
100                        "name": {
101                            "type": "string",
102                            "description": "Qualified as namespace.metric, or bare when unambiguous.",
103                        }
104                    },
105                },
106            },
107        },
108        {
109            "type": "function",
110            "function": {
111                "name": "list_dimensions",
112                "description": (
113                    "List the dimensions available for grouping and filtering. "
114                    "Pass a metric to get only the dimensions valid for it, which "
115                    "is almost always what you want."
116                ),
117                "parameters": {
118                    "type": "object",
119                    "properties": {
120                        "metric": {"type": "string", "description": "Restrict to this metric."}
121                    },
122                },
123            },
124        },
125        {
126            "type": "function",
127            "function": {
128                "name": "query",
129                "description": (
130                    "Answer a question by naming metrics, dimensions and filters. "
131                    "This is the only way to get numbers, and it is not a SQL "
132                    "interface. The response carries the rows, the SQL that was "
133                    "compiled and the model version, so any number can be traced "
134                    "back. If a request is refused, read the refusal: it says "
135                    "whether to change the request, wait, or stop."
136                ),
137                "parameters": {
138                    "type": "object",
139                    "required": ["metrics"],
140                    "properties": {
141                        "metrics": {
142                            "type": "array",
143                            "items": {"type": "string"},
144                            "description": "Metric names from list_metrics.",
145                        },
146                        "dimensions": {
147                            "type": "array",
148                            "items": {"type": "string"},
149                            "description": "Dimension names from list_dimensions.",
150                        },
151                        "filters": {"type": "array", "items": _FILTER_SCHEMA},
152                        "grain": {
153                            "type": "string",
154                            "enum": ["second", "minute", "hour", "day", "week",
155                                     "month", "quarter", "year"],
156                            "description": "Time bucket for the selected time dimension.",
157                        },
158                        "limit": {"type": "integer", "description": "Maximum rows."},
159                        "order_by": {
160                            "type": "array",
161                            "items": {
162                                "type": "object",
163                                "required": ["field"],
164                                "properties": {
165                                    "field": {"type": "string"},
166                                    "desc": {"type": "boolean"},
167                                },
168                            },
169                        },
170                    },
171                },
172            },
173        },
174    ]
175
176
177def dispatch(client: Client, name: str, arguments: Mapping[str, Any]) -> dict[str, Any]:
178    """Run one tool call and return a JSON-serializable result.
179
180    A refusal is returned rather than raised, because a framework feeding this
181    back to a model wants the text, not a traceback. The shape carries the
182    ``retry`` class so the model learns whether to change the request, wait, or
183    stop, instead of looping on a denial.
184    """
185    args = dict(arguments)
186    try:
187        if name == "list_metrics":
188            metrics = client.metrics(search=args.get("search"))
189            return {"metrics": [m.raw for m in metrics]}
190        if name == "describe_metric":
191            if "name" not in args:
192                raise ValueError("describe_metric needs a `name`")
193            return client.metric(args["name"]).raw
194        if name == "list_dimensions":
195            dims = client.dimensions(metric=args.get("metric"))
196            return {"dimensions": [d.raw for d in dims]}
197        if name == "query":
198            result = client.query(
199                metrics=args.get("metrics", []),
200                dimensions=args.get("dimensions"),
201                filters=args.get("filters"),
202                grain=args.get("grain"),
203                limit=args.get("limit"),
204                order_by=args.get("order_by"),
205            )
206            return result.raw
207    except Refused as refusal:
208        return {
209            "refused": True,
210            "code": refusal.code,
211            "reason": refusal.reason,
212            "hint": refusal.hint,
213            "retry": refusal.retry,
214            "note": (
215                "This is a refusal, not an empty result. Do not report a number "
216                "for this question."
217            ),
218        }
219    except ValueError as bad_argument:
220        return {"refused": True, "code": "invalid_arguments",
221                "reason": str(bad_argument), "retry": "modify"}
222
223    raise ValueError(f"unknown tool {name!r}; expected one of {', '.join(TOOL_NAMES)}")