sadpandajoe commented on code in PR #43237: URL: https://github.com/apache/superset/pull/43237#discussion_r3818556400
########## superset/ai/tools/base.py: ########## @@ -0,0 +1,624 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +The tool contract and the registry that dispatches to it. + +A tool is a small, self-authorizing unit of work. "Self-authorizing" is the +important half: the registry does not check permissions on a tool's behalf, and +neither does the policy chain in :mod:`superset.ai.policy`, which answers the +coarser question of whether a *shape* of call should be attempted at all. Any +tool that returns or mutates a data-bearing object performs its own +``security_manager`` check, because that is the only place with enough context to +know which object is being touched. + +A tool returns two things. :attr:`ToolOutput.content` is what the model reads. +:attr:`ToolOutput.display` is a summary for the UI, so a user can expand what the +assistant did and see the SQL it ran and the rows it got back. Both are +size-bounded here rather than in each tool: ``display`` is persisted on the +message and shipped to the browser, so an unbounded one would be a second way to +blow up a response. +""" + +from __future__ import annotations + +import logging +import time +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any, ClassVar + +from superset.ai.llm.base import ToolCall, ToolDefinition, ToolResult +from superset.mcp_service.utils.sanitization import ( + LLM_CONTEXT_CLOSE_DELIMITER, + LLM_CONTEXT_ESCAPED_CLOSE_DELIMITER, + LLM_CONTEXT_ESCAPED_OPEN_DELIMITER, + LLM_CONTEXT_OPEN_DELIMITER, +) +from superset.utils import json + +logger = logging.getLogger(__name__) + +#: Keys added to a payload whose budget was exceeded. Phrased for the model: it +#: says what was lost and what to do differently, because a model told only +#: "truncated" reissues the identical call. +TRUNCATION_KEY = "_truncated" +TRUNCATION_NOTE_KEY = "_truncation_note" + +#: Share of the response budget the UI summary may use. The model's copy is the +#: one that has to be complete enough to reason over; ``display`` only has to be +#: enough to render, and it is persisted, so it gets the smaller share. +_DISPLAY_BUDGET_FRACTION = 0.5 + +#: Fallback response budget for use outside an application context, matching the +#: shipped ``AI_AGENT_MAX_RESULT_BYTES`` default. +_DEFAULT_MAX_BYTES = 256 * 1024 + + +class ToolError(Exception): + """ + A failure that should be shown to the model rather than raised at the user. + + Tools raise this for conditions the model can act on — a database it may not + read, a column that does not resolve, SQL that will not parse. The registry + turns it into a :class:`~superset.ai.llm.base.ToolResult` with + ``is_error=True`` so the turn continues and the model can correct itself. + + The message is model-visible, so it must never carry a driver exception, a + connection string, or a stack trace. + """ + + +@dataclass(frozen=True) +class ToolOutput: + """ + What a tool returns. + + Build one with :meth:`of` rather than by hand, so that ``content`` and + ``payload`` cannot disagree. + """ + + #: Model-facing text. Becomes :attr:`ToolResult.content`. + content: str + + #: JSON-serialisable summary for the UI, or ``None`` when there is nothing + #: worth rendering. Must never carry credentials, a connection string, or a + #: full result set. + display: dict[str, Any] | None = None + + #: The structure ``content`` was serialised from. Retained so the registry + #: can shrink an oversized result by dropping rows rather than cutting JSON + #: mid-token. ``None`` when the tool supplied text directly. + payload: Any = None + + @classmethod + def of( + cls, + payload: Any, + display: dict[str, Any] | None = None, + ) -> ToolOutput: + """Serialise ``payload`` as the model-facing content.""" + return cls( + content=json.dumps(payload, default=str), + display=display, + payload=payload, + ) + + +@dataclass +class ToolInvocation: + """ + One completed dispatch, with everything a caller might need. + + Exists because two consumers want different things from the same call: the + provider needs a :class:`~superset.ai.llm.base.ToolResult`, while the event + stream and message persistence want the UI summary and the timing. + """ + + call_id: str + tool_name: str + result: ToolResult + display: dict[str, Any] | None = None + duration_ms: int = 0 + truncated: bool = False + arguments: dict[str, Any] = field(default_factory=dict) + + @property + def is_error(self) -> bool: + """Whether the call failed.""" + return self.result.is_error + + def to_tool_result(self) -> ToolResult: + """The provider-neutral result to feed back to the model.""" + return self.result + + +class AITool(ABC): + """ + One capability offered to the model. + + Subclasses set :attr:`name`, :attr:`description` and :attr:`input_schema`, + and implement :meth:`run`. Everything else — size capping, error + translation, timing — is the registry's job. + """ + + #: Stable identifier the model calls, and the key an operator types when + #: configuring which tools an agent profile may use. Renaming one is a + #: breaking change: it appears in stored conversation history and in + #: deployment configuration. + name: ClassVar[str] = "" + + #: Shown to the model verbatim. This is the tool's entire user manual, so it + #: should say when to reach for the tool and what it returns, not merely + #: what it is called. + description: ClassVar[str] = "" + + #: JSON Schema for the arguments object. Providers translate it into + #: whatever their API expects. + input_schema: ClassVar[dict[str, Any]] = {"type": "object", "properties": {}} + + @abstractmethod + def run(self, **kwargs: Any) -> ToolOutput: + """ + Perform the work. + + Raise :class:`ToolError` for anything the model should see and be able + to recover from. Any other exception is treated as a defect: it is + logged with a traceback and reported to the model as a generic failure, + so that an unexpected driver error cannot leak its message. + """ + + def definition(self) -> ToolDefinition: + """Provider-neutral description of this tool.""" + return ToolDefinition( + name=self.name, + description=self.description, + input_schema=self.input_schema, + ) + + +def truncate_payload(payload: Any, max_bytes: int) -> tuple[str, bool]: + """ + Serialise ``payload`` and bound it to ``max_bytes``. + + Returns valid JSON text and whether anything was dropped. Mapping fields are + shortened largest-first so useful smaller fields survive without cutting JSON + mid-token. + """ + text = json.dumps(payload, default=str) + if len(text.encode("utf-8")) <= max_bytes: + return text, False + + if isinstance(payload, dict): + candidate = dict(payload) + list_counts = [ + f"{len(value)} {key}" + for key, value in payload.items() + if isinstance(value, list) and value + ] + candidate[TRUNCATION_KEY] = True + candidate[TRUNCATION_NOTE_KEY] = ( + f"Payload exceeded the {max_bytes} byte response budget" + + (f" ({', '.join(list_counts)})" if list_counts else "") + + ". Large fields were shortened; narrow the request to see the rest." + ) + + while True: + text = json.dumps(candidate, default=str) + if len(text.encode("utf-8")) <= max_bytes: + return text, True + values = { + key: value + for key, value in candidate.items() + if key not in {TRUNCATION_KEY, TRUNCATION_NOTE_KEY} + } + if not values: + break + key = max( + values, + key=lambda item: len( + json.dumps(values[item], default=str).encode("utf-8") + ), + ) + value = values[key] + if isinstance(value, str) and len(value) > 1: Review Comment: With a small configured result budget, a value that reaches the one-character-plus-ellipsis form is replaced with the same value forever while the truncation metadata still exceeds the budget, so this can spin the tool request instead of returning the fallback. Could this drop the field when a shrink step makes no progress? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
