Custom HTTP tools an administrator defines
A row in custom_tools becomes a ToolDef like any built-in, offered beside the thirteen. The registry had to stop being an import-time constant for that: `resolve_tools` now returns the schemas *and* the runners together, carried to the loop on the ToolContext. That closes a hole on the way. `run_tool` looked names up in the global REGISTRY with no reference to what had been offered, so a model naming a tool its chat was gated out of -- a family switched off, a permission the reader lacks -- had it run anyway. The resolved set is now authoritative. Arguments come from a model, so an argument may fill a hole but never move the target: the scheme and host of a URL template are literal, values are escaped for where they land, and the origin is pinned afterwards. Every redirect hop is checked the way services/fetch.py checks one, and the secret is dropped if a hop leaves the origin it was issued for. Also fixes the tool-activity block claiming every library tool had "searched the web", which it has done since the second family landed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Jaroslav Beneš
@@ -0,0 +1,431 @@
|
||||
"""Running the HTTP tools an administrator defined.
|
||||
|
||||
A row in `custom_tools` becomes a `ToolDef` like any built-in: same schema in
|
||||
the same array, same `ToolOutcome` back. What is different is that the arguments
|
||||
come from a model and the destination comes from a template, so two things have
|
||||
to hold.
|
||||
|
||||
**An argument may fill a hole; it may not move the target.** The scheme and host
|
||||
of the template are literal, checked when the row is saved and again here in
|
||||
case a row predates the check, and every value is escaped for the position it
|
||||
lands in -- percent-encoded with nothing safe in a URL, JSON-escaped in a body,
|
||||
stripped of line breaks in a header. `quote(value, safe="")` is what stops a
|
||||
value adding a path segment, a query parameter or a fragment; pinning the origin
|
||||
afterwards is what catches anything that got past it.
|
||||
|
||||
**Every hop is checked.** This is the same request-forgery problem
|
||||
`services/fetch.py` exists to solve, and the same answer: resolve and check the
|
||||
address, follow redirects by hand, refuse private ranges unless this particular
|
||||
row was allowed them. `fetch.fetch` itself cannot be reused -- it is GET-only,
|
||||
has no body, and refuses any content type that is not HTML or text, which is
|
||||
every JSON API there is -- so its redirect loop is deliberately copied rather
|
||||
than the module bent into a general HTTP client.
|
||||
|
||||
The secret is decrypted into the snapshot and goes nowhere else: not into the
|
||||
event, not into a log line, and not across a redirect that leaves the origin it
|
||||
was issued for.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
from urllib.parse import quote, urlparse
|
||||
|
||||
import httpx
|
||||
from sqlalchemy.orm import Session as DBSession
|
||||
|
||||
from lembas.db.models import (
|
||||
RESPONSE_JSON,
|
||||
RESPONSE_RAW,
|
||||
RESPONSE_TEXT,
|
||||
SECRET_BEARER,
|
||||
SECRET_HEADER,
|
||||
SECRET_QUERY,
|
||||
CustomTool,
|
||||
User,
|
||||
)
|
||||
from lembas.services import fetch as fetch_service
|
||||
from lembas.services import tool_access
|
||||
from lembas.services.crypto import decrypt
|
||||
from lembas.services.prompts import VARIABLE_PATTERN
|
||||
from lembas.services.tools import ToolContext, ToolDef, ToolOutcome
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# What a response may weigh before it is cut off. Well below the page fetcher's
|
||||
# ceiling, because this is text that will be sent back to a model rather than
|
||||
# stored for a person to read.
|
||||
MAX_RESPONSE_BYTES = 2 * 1024 * 1024
|
||||
|
||||
# How much of the response is kept on the message row for the transcript. Capped
|
||||
# separately from `max_chars`: what the model reads is spent once, what the event
|
||||
# holds is stored on every message forever.
|
||||
MAX_EVENT_CHARS = 2000
|
||||
|
||||
# How much of the arguments the transcript summarises.
|
||||
MAX_SUMMARY_CHARS = 200
|
||||
|
||||
ALLOWED_METHODS = ("GET", "POST", "PUT", "PATCH", "DELETE")
|
||||
|
||||
# Bounds an administrator's number is clamped into. A tool that may return
|
||||
# 400 000 characters is a tool that can fill the context window in one call.
|
||||
MIN_CHARS, MAX_CHARS = 200, 40_000
|
||||
MIN_TIMEOUT, MAX_TIMEOUT = 1, 120
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HttpSpec:
|
||||
"""Everything one custom tool needs, read while the session was open.
|
||||
|
||||
A frozen snapshot rather than the row, for the reason `Endpoint` is one: a
|
||||
generation outlives the request that started it, and a detached SQLAlchemy
|
||||
instance is a trap. The decrypted secret lives here and nowhere else.
|
||||
"""
|
||||
|
||||
slug: str
|
||||
label: str
|
||||
method: str
|
||||
url_template: str
|
||||
body_template: str = ""
|
||||
headers: dict[str, str] = field(default_factory=dict)
|
||||
secret: str = ""
|
||||
secret_placement: str = SECRET_BEARER
|
||||
secret_name: str = "Authorization"
|
||||
response_mode: str = RESPONSE_TEXT
|
||||
response_path: str = ""
|
||||
max_chars: int = 8000
|
||||
timeout: int = 20
|
||||
allow_private: bool = False
|
||||
parameters: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def secret_header(self) -> str:
|
||||
"""The header the secret rides in, if it rides in one."""
|
||||
if not self.secret or self.secret_placement not in (SECRET_BEARER, SECRET_HEADER):
|
||||
return ""
|
||||
return self.secret_name or "Authorization"
|
||||
|
||||
|
||||
def spec_from(row: CustomTool) -> HttpSpec:
|
||||
"""Snapshot a row, decrypting its secret. Call this with a session open."""
|
||||
return HttpSpec(
|
||||
slug=row.slug,
|
||||
label=row.name or row.slug,
|
||||
method=(row.method or "GET").upper(),
|
||||
url_template=row.url_template or "",
|
||||
body_template=row.body_template or "",
|
||||
headers=dict(row.headers_json or {}),
|
||||
secret=decrypt(row.secret_encrypted),
|
||||
secret_placement=row.secret_placement,
|
||||
secret_name=row.secret_name or "Authorization",
|
||||
response_mode=row.response_mode,
|
||||
response_path=row.response_path or "",
|
||||
max_chars=min(max(int(row.max_chars or 0), MIN_CHARS), MAX_CHARS),
|
||||
timeout=min(max(int(row.timeout or 0), MIN_TIMEOUT), MAX_TIMEOUT),
|
||||
allow_private=bool(row.allow_private),
|
||||
parameters=dict(row.parameters_json or {}),
|
||||
)
|
||||
|
||||
|
||||
def tool_defs(
|
||||
db: DBSession, user: User | None, *, everything: bool = False
|
||||
) -> list[ToolDef]:
|
||||
"""One `ToolDef` per custom tool this user may be offered."""
|
||||
return [
|
||||
ToolDef(
|
||||
name=row.slug,
|
||||
family=f"custom:{row.slug}",
|
||||
description=row.description or f"Call the {row.name} tool.",
|
||||
parameters=_schema_of(row),
|
||||
run=_runner(spec_from(row)),
|
||||
)
|
||||
for row in tool_access.visible_custom_tools(db, user, everything=everything)
|
||||
]
|
||||
|
||||
|
||||
def _schema_of(row: CustomTool) -> dict[str, Any]:
|
||||
schema = dict(row.parameters_json or {})
|
||||
if schema.get("type") != "object":
|
||||
# An endpoint expects an object here; anything else it will reject
|
||||
# outright, which fails the whole request rather than the one tool.
|
||||
return {"type": "object", "properties": {}}
|
||||
return schema
|
||||
|
||||
|
||||
def _runner(spec: HttpSpec):
|
||||
async def run(context: ToolContext, args: dict[str, Any]) -> ToolOutcome:
|
||||
return await call(spec, args)
|
||||
|
||||
return run
|
||||
|
||||
|
||||
# --- Filling the template ----------------------------------------------------
|
||||
def _scalar(value: Any) -> str:
|
||||
"""One argument as text, before it is escaped for wherever it is going."""
|
||||
if value is None:
|
||||
return ""
|
||||
if isinstance(value, bool):
|
||||
return "true" if value else "false"
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, int | float):
|
||||
return str(value)
|
||||
return json.dumps(value, ensure_ascii=False)
|
||||
|
||||
|
||||
def _for_url(value: str) -> str:
|
||||
# safe="" is the whole point: an argument must not be able to introduce a
|
||||
# path segment, a query separator or a fragment.
|
||||
return quote(value, safe="")
|
||||
|
||||
|
||||
def _for_body(value: str) -> str:
|
||||
# The inside of a JSON string, so a quote or a backslash in an argument
|
||||
# cannot end it early and add a field of its own.
|
||||
return json.dumps(value, ensure_ascii=False)[1:-1]
|
||||
|
||||
|
||||
def _for_header(value: str) -> str:
|
||||
# A newline in a header value is header injection. Other control characters
|
||||
# go with it; none of them mean anything in a header.
|
||||
return "".join(character for character in value if character.isprintable())
|
||||
|
||||
|
||||
def _substitute(template: str, spec: HttpSpec, args: dict[str, Any], escape) -> str:
|
||||
"""Fill `{{name}}` from the call's arguments.
|
||||
|
||||
Not `prompts.substitute`, though the grammar is shared. The rules differ,
|
||||
and the differences are the point: a name the tool does not declare never
|
||||
substitutes, an unrecognised one becomes nothing rather than passing through
|
||||
verbatim -- a literal `{{x}}` in a URL is not a feature -- and every value
|
||||
is escaped for where it lands.
|
||||
"""
|
||||
declared = set(spec.parameters.get("properties") or {})
|
||||
|
||||
def swap(match) -> str:
|
||||
name = match.group(1)
|
||||
if name not in declared:
|
||||
return ""
|
||||
return escape(_scalar(args.get(name)))
|
||||
|
||||
return VARIABLE_PATTERN.sub(swap, template)
|
||||
|
||||
|
||||
def _origin(url: str) -> tuple[str, str]:
|
||||
parsed = urlparse(url)
|
||||
if parsed.scheme not in ("http", "https"):
|
||||
raise fetch_service.FetchError("A tool's URL must start with http:// or https://")
|
||||
if not parsed.netloc:
|
||||
raise fetch_service.FetchError("A tool's URL has no host.")
|
||||
return parsed.scheme, parsed.netloc
|
||||
|
||||
|
||||
def fill_url(spec: HttpSpec, args: dict[str, Any]) -> str:
|
||||
"""Fill the URL template, refusing anything that moved the host.
|
||||
|
||||
Checked twice over: the template's own scheme and authority must be literal,
|
||||
and the filled URL must still point at them. The first check is what stops
|
||||
`https://{{host}}/x` from ever being saved; the second is what catches a row
|
||||
that predates it, or an escaping mistake.
|
||||
"""
|
||||
template = spec.url_template.strip()
|
||||
scheme, netloc = _origin(template)
|
||||
if VARIABLE_PATTERN.search(f"{scheme}://{netloc}"):
|
||||
raise fetch_service.FetchError(
|
||||
"A tool's scheme and host must be literal, not filled from an argument."
|
||||
)
|
||||
|
||||
filled = _substitute(template, spec, args, _for_url)
|
||||
if _origin(filled) != (scheme, netloc):
|
||||
raise fetch_service.FetchError("That call would have pointed somewhere else.")
|
||||
return filled
|
||||
|
||||
|
||||
def _prepare(spec: HttpSpec, args: dict[str, Any]) -> tuple[str, dict[str, str], bytes | None]:
|
||||
"""The URL, headers and body for one call, secret included."""
|
||||
url = fill_url(spec, args)
|
||||
headers = {
|
||||
"User-Agent": fetch_service.USER_AGENT,
|
||||
"Accept": "application/json, text/*;q=0.9, */*;q=0.5",
|
||||
}
|
||||
for name, value in spec.headers.items():
|
||||
clean = _for_header(str(name)).strip()
|
||||
if clean:
|
||||
headers[clean] = _substitute(str(value), spec, args, _for_header)
|
||||
|
||||
body: bytes | None = None
|
||||
if spec.body_template.strip() and spec.method != "GET":
|
||||
body = _substitute(spec.body_template, spec, args, _for_body).encode("utf-8")
|
||||
headers.setdefault("Content-Type", "application/json")
|
||||
|
||||
if spec.secret:
|
||||
if spec.secret_placement == SECRET_BEARER:
|
||||
headers[spec.secret_name or "Authorization"] = f"Bearer {spec.secret}"
|
||||
elif spec.secret_placement == SECRET_HEADER:
|
||||
headers[spec.secret_name or "Authorization"] = spec.secret
|
||||
elif spec.secret_placement == SECRET_QUERY:
|
||||
# Only on the URL this call starts at. A redirect's Location
|
||||
# replaces the query, so the credential does not travel on by
|
||||
# itself -- which is the behaviour wanted anyway.
|
||||
joiner = "&" if urlparse(url).query else "?"
|
||||
url = f"{url}{joiner}{quote(spec.secret_name)}={quote(spec.secret, safe='')}"
|
||||
|
||||
return url, headers, body
|
||||
|
||||
|
||||
# --- Reading the response ----------------------------------------------------
|
||||
def _narrow(payload: Any, path: str) -> Any:
|
||||
"""Walk a dotted path into a decoded JSON document.
|
||||
|
||||
Integer segments index a list, so "data.0.title" works. A path that does not
|
||||
lead anywhere yields the whole document rather than nothing: an unhelpful
|
||||
answer beats a silent empty one when the model has to explain itself.
|
||||
"""
|
||||
current = payload
|
||||
for segment in [part for part in path.split(".") if part]:
|
||||
if isinstance(current, dict) and segment in current:
|
||||
current = current[segment]
|
||||
elif isinstance(current, list) and segment.lstrip("-").isdigit():
|
||||
try:
|
||||
current = current[int(segment)]
|
||||
except IndexError:
|
||||
return payload
|
||||
else:
|
||||
return payload
|
||||
return current
|
||||
|
||||
|
||||
def _decode(payload: bytes, response: httpx.Response) -> str:
|
||||
return payload.decode(response.encoding or "utf-8", "replace")
|
||||
|
||||
|
||||
def _as_text(spec: HttpSpec, payload: bytes, response: httpx.Response) -> str:
|
||||
content_type = response.headers.get("content-type", "")
|
||||
|
||||
if spec.response_mode == RESPONSE_JSON:
|
||||
try:
|
||||
document = json.loads(_decode(payload, response))
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
# Falling back rather than failing: a JSON API answering with an
|
||||
# HTML error page is a thing the model can report usefully.
|
||||
return _decode(payload, response)
|
||||
value = _narrow(document, spec.response_path)
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
return json.dumps(value, indent=2, ensure_ascii=False)
|
||||
|
||||
if spec.response_mode == RESPONSE_RAW:
|
||||
return _decode(payload, response)
|
||||
|
||||
text = _decode(payload, response)
|
||||
if "html" in content_type or text.lstrip()[:1] == "<":
|
||||
_, text = fetch_service.html_to_text(text)
|
||||
return text
|
||||
|
||||
|
||||
def _clip(text: str, limit: int) -> str:
|
||||
if len(text) <= limit:
|
||||
return text
|
||||
return text[:limit].rstrip() + "\n… (truncated)"
|
||||
|
||||
|
||||
def _summary(args: dict[str, Any]) -> str:
|
||||
"""What the transcript shows the tool was asked for."""
|
||||
parts = [f"{name}={_scalar(value)!r}" for name, value in args.items()]
|
||||
return _clip(", ".join(parts), MAX_SUMMARY_CHARS)
|
||||
|
||||
|
||||
def _event(spec: HttpSpec, args: dict[str, Any], *, status: str, **extra: Any) -> dict[str, Any]:
|
||||
return {
|
||||
"name": spec.slug,
|
||||
"kind": "custom",
|
||||
"label": spec.label,
|
||||
"query": _summary(args),
|
||||
# The host, never the filled URL: a path or query segment can carry an
|
||||
# argument, and the event is rendered and stored.
|
||||
"detail": f"{spec.method} {urlparse(spec.url_template).netloc}",
|
||||
"status": status,
|
||||
"results": [],
|
||||
**extra,
|
||||
}
|
||||
|
||||
|
||||
# --- Making the call ---------------------------------------------------------
|
||||
async def call(spec: HttpSpec, args: dict[str, Any]) -> ToolOutcome:
|
||||
"""Run one custom tool. Reports its own failures rather than raising."""
|
||||
try:
|
||||
url, headers, body = _prepare(spec, args)
|
||||
current = fetch_service.check_url(url, allow_private=spec.allow_private)
|
||||
origin = _origin(current)
|
||||
response = await _send(spec, current, headers, body, origin)
|
||||
except fetch_service.FetchError as exc:
|
||||
return ToolOutcome(
|
||||
f"The {spec.label} tool could not be called: {exc.message}",
|
||||
_event(spec, args, status="error", error=exc.message),
|
||||
)
|
||||
except httpx.RequestError as exc:
|
||||
message = f"Could not reach the {spec.label} tool: {exc}"
|
||||
return ToolOutcome(message, _event(spec, args, status="error", error=str(exc)[:200]))
|
||||
|
||||
payload = response.content[:MAX_RESPONSE_BYTES]
|
||||
text = _clip(_as_text(spec, payload, response).strip(), spec.max_chars)
|
||||
|
||||
if response.status_code >= 400:
|
||||
note = f"{spec.label} returned HTTP {response.status_code}."
|
||||
return ToolOutcome(
|
||||
f"{note}\n\n{text}" if text else note,
|
||||
_event(
|
||||
spec,
|
||||
args,
|
||||
status="error",
|
||||
error=f"HTTP {response.status_code}",
|
||||
text=text[:MAX_EVENT_CHARS],
|
||||
),
|
||||
)
|
||||
|
||||
if not text:
|
||||
return ToolOutcome(
|
||||
f"{spec.label} returned nothing.",
|
||||
_event(spec, args, status="ok", text=""),
|
||||
)
|
||||
|
||||
return ToolOutcome(text, _event(spec, args, status="ok", text=text[:MAX_EVENT_CHARS]))
|
||||
|
||||
|
||||
async def _send(
|
||||
spec: HttpSpec,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
body: bytes | None,
|
||||
origin: tuple[str, str],
|
||||
) -> httpx.Response:
|
||||
"""Send the request, following redirects by hand so each hop is checked."""
|
||||
current = url
|
||||
async with httpx.AsyncClient(timeout=spec.timeout, follow_redirects=False) as client:
|
||||
for _ in range(fetch_service.MAX_REDIRECTS + 1):
|
||||
response = await client.request(
|
||||
spec.method, current, headers=headers, content=body
|
||||
)
|
||||
if not response.is_redirect:
|
||||
return response
|
||||
|
||||
location = response.headers.get("location", "")
|
||||
if not location:
|
||||
raise fetch_service.FetchError("That tool redirected to nowhere.")
|
||||
current = fetch_service.check_url(
|
||||
str(response.url.join(location)), allow_private=spec.allow_private
|
||||
)
|
||||
if _origin(current) != origin:
|
||||
# A server that can redirect us anywhere must not be able to
|
||||
# redirect us at somebody else carrying the key.
|
||||
if spec.secret_header:
|
||||
headers.pop(spec.secret_header, None)
|
||||
origin = _origin(current)
|
||||
|
||||
raise fetch_service.FetchError("That tool redirected too many times.")
|
||||
|
||||
|
||||
__all__ = ["HttpSpec", "call", "fill_url", "spec_from", "tool_defs"]
|
||||
@@ -239,7 +239,11 @@ async def _run(generation: Generation) -> None:
|
||||
owner = db.get(User, chat.user_id)
|
||||
|
||||
# Read while the session is open: everything below outlives it.
|
||||
offered = tools_service.enabled_tools(db, chat, owner)
|
||||
# Resolved once, so that what the loop is allowed to *run* is the
|
||||
# same set the endpoint was *offered* -- not whatever happens to
|
||||
# exist by the time a call comes back.
|
||||
toolset = tools_service.resolve_tools(db, chat, owner)
|
||||
offered = toolset.schemas
|
||||
payload = chat_service.build_request(
|
||||
db, chat, upto=message, tools=offered, user=owner
|
||||
)
|
||||
@@ -248,7 +252,7 @@ async def _run(generation: Generation) -> None:
|
||||
# Read here, with the rest, because titling happens after this
|
||||
# session has closed and must not open another one.
|
||||
title_prompt = prompts_service.resolve(db, "task.title")
|
||||
tool_context = tools_service.context_for(db, owner, chat)
|
||||
tool_context = tools_service.context_for(db, owner, chat, tools=toolset)
|
||||
|
||||
model = chat_service.model_for(db, chat)
|
||||
generation.context_limit = model.context_length if model is not None else 0
|
||||
|
||||
@@ -57,16 +57,22 @@ MAX_HARNESS_CHARS = 8000
|
||||
MAX_NAMED_DOCUMENTS = 5
|
||||
|
||||
|
||||
def _families(tools: list[dict[str, Any]]) -> list[str]:
|
||||
"""Which families are represented in an offered tool list, in a fixed order."""
|
||||
from lembas.services.tools import FAMILIES, REGISTRY
|
||||
def _families(db: DBSession, tools: list[dict[str, Any]]) -> list[str]:
|
||||
"""Which families are represented in an offered tool list, in a fixed order.
|
||||
|
||||
Resolved against the database rather than the import-time registry, because
|
||||
an administrator-defined tool is a row and would otherwise contribute no
|
||||
family at all -- which is to say its guidance would never be admitted.
|
||||
"""
|
||||
from lembas.services import tools as tools_service
|
||||
|
||||
book = tools_service.registry(db)
|
||||
offered = {
|
||||
REGISTRY[name].family
|
||||
book[name].family
|
||||
for tool in tools
|
||||
if (name := (tool.get("function") or {}).get("name")) in REGISTRY
|
||||
if (name := (tool.get("function") or {}).get("name")) in book
|
||||
}
|
||||
return [family for family in FAMILIES if family in offered]
|
||||
return [family for family in tools_service.families(db) if family in offered]
|
||||
|
||||
|
||||
def _tool_names(tools: list[dict[str, Any]]) -> str:
|
||||
@@ -109,7 +115,7 @@ def context_variables(
|
||||
from lembas.services import tools as tools_service
|
||||
|
||||
offered = tools or []
|
||||
families = _families(offered)
|
||||
families = _families(db, offered)
|
||||
stamp = datetime.now().astimezone()
|
||||
|
||||
values: dict[str, str] = {
|
||||
@@ -185,7 +191,7 @@ def compose(
|
||||
return compose_from(
|
||||
db,
|
||||
variables=context_variables(db, user, offered, chat),
|
||||
families=_families(offered),
|
||||
families=_families(db, offered),
|
||||
has_tools=bool(offered),
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Who may be offered which administrator-defined tool.
|
||||
|
||||
One definition, used by the resolver that builds a chat's tool set and by the
|
||||
admin screens that say what a row will reach. It mirrors
|
||||
`permissions.models_visible_to` rather than `services/sharing.py`, and the
|
||||
difference matters: sharing has no admin branch because reading someone's
|
||||
private notes is not something a role should grant, while a tool is instance
|
||||
configuration an administrator could give themselves in one click anyway.
|
||||
|
||||
`everything=True` is for callers that need to map a tool *name* back to what it
|
||||
belongs to -- the harness, the prompt preview -- rather than to decide what to
|
||||
offer. Gating on visibility there would mean an administrator's own preview
|
||||
disagreeing with what a reader actually gets.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session as DBSession
|
||||
|
||||
from lembas.db.models import CustomTool, McpServer, User
|
||||
|
||||
|
||||
def _permitted(rows: list, user: User | None, *, everything: bool) -> list:
|
||||
if everything or (user is not None and user.is_admin):
|
||||
return rows
|
||||
if user is None:
|
||||
return []
|
||||
member_of = {group.id for group in user.groups}
|
||||
return [
|
||||
row for row in rows if row.public or member_of.intersection({g.id for g in row.groups})
|
||||
]
|
||||
|
||||
|
||||
def visible_custom_tools(
|
||||
db: DBSession, user: User | None, *, everything: bool = False
|
||||
) -> list[CustomTool]:
|
||||
"""Enabled custom tools this user may be offered, in position order."""
|
||||
rows = list(
|
||||
db.scalars(
|
||||
select(CustomTool)
|
||||
.where(CustomTool.enabled.is_(True))
|
||||
.order_by(CustomTool.position, CustomTool.slug)
|
||||
)
|
||||
)
|
||||
return _permitted(rows, user, everything=everything)
|
||||
|
||||
|
||||
def visible_mcp_servers(
|
||||
db: DBSession, user: User | None, *, everything: bool = False
|
||||
) -> list[McpServer]:
|
||||
"""Enabled MCP servers this user may be offered, in position order."""
|
||||
rows = list(
|
||||
db.scalars(
|
||||
select(McpServer)
|
||||
.where(McpServer.enabled.is_(True))
|
||||
.order_by(McpServer.position, McpServer.slug)
|
||||
)
|
||||
)
|
||||
return _permitted(rows, user, everything=everything)
|
||||
|
||||
|
||||
__all__ = ["visible_custom_tools", "visible_mcp_servers"]
|
||||
+180
-18
@@ -29,10 +29,12 @@ from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session as DBSession
|
||||
|
||||
from lembas.db.models import AUTHOR_MODEL, Chat, User
|
||||
from lembas.db.session import session_scope
|
||||
from lembas.services import prompts as prompts_service
|
||||
from lembas.services import search as search_service
|
||||
from lembas.services import settings_store
|
||||
from lembas.services.library import documents as documents_service
|
||||
@@ -58,8 +60,24 @@ FAMILY_NOTES = "notes"
|
||||
FAMILY_MEMORY = "memory"
|
||||
FAMILY_SKILLS = "skills"
|
||||
|
||||
# A tool that is a database row gets a family of its own, so that it can carry
|
||||
# its own guidance -- "custom:weather", "mcp:github". Everything before the
|
||||
# colon is the *gate*: the capability flag and the permission are per gate, not
|
||||
# per row, because a server advertising forty tools must not mean forty
|
||||
# checkboxes on every model.
|
||||
FAMILY_CUSTOM = "custom"
|
||||
FAMILY_MCP = "mcp"
|
||||
|
||||
# The built-in families, in the order they are offered.
|
||||
FAMILIES = (FAMILY_SEARCH, FAMILY_KNOWLEDGE, FAMILY_NOTES, FAMILY_MEMORY, FAMILY_SKILLS)
|
||||
|
||||
GATES = (*FAMILIES, FAMILY_CUSTOM, FAMILY_MCP)
|
||||
|
||||
|
||||
def gate_of(family: str) -> str:
|
||||
"""The part a capability flag and a permission are named after."""
|
||||
return family.split(":", 1)[0]
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolContext:
|
||||
@@ -72,10 +90,14 @@ class ToolContext:
|
||||
|
||||
owner_id: str
|
||||
search_config: dict[str, Any] = field(default_factory=dict)
|
||||
allow_private_fetch: bool = False
|
||||
# Which knowledge bases this chat is scoped to. Empty means "everything the
|
||||
# owner can see", which is what a chat with none attached should do.
|
||||
base_ids: list[str] = field(default_factory=list)
|
||||
# Name -> definition for the tools actually offered on this request. None
|
||||
# means nobody resolved a set, and only then does `run_tool` fall back to
|
||||
# the import-time registry. A dict, *even an empty one*, is authoritative:
|
||||
# a model naming a tool it was not offered must not get it run.
|
||||
tools: dict[str, ToolDef] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -114,6 +136,31 @@ class ToolDef:
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ToolSet:
|
||||
"""What one request may call: the schemas to send, and how to run them.
|
||||
|
||||
The two halves have to travel together. `enabled_tools` used to return
|
||||
schemas alone, which worked only because every runner was reachable through
|
||||
the import-time `REGISTRY`. A tool that is a database row is not, so the
|
||||
resolution has to be carried from the session that made it to the loop that
|
||||
uses it.
|
||||
"""
|
||||
|
||||
defs: tuple[ToolDef, ...] = ()
|
||||
|
||||
@property
|
||||
def schemas(self) -> list[dict[str, Any]]:
|
||||
return [tool.schema for tool in self.defs]
|
||||
|
||||
@property
|
||||
def by_name(self) -> dict[str, ToolDef]:
|
||||
return {tool.name: tool for tool in self.defs}
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
return bool(self.defs)
|
||||
|
||||
|
||||
def _object(properties: dict[str, Any], required: list[str]) -> dict[str, Any]:
|
||||
return {"type": "object", "properties": properties, "required": required}
|
||||
|
||||
@@ -645,21 +692,69 @@ def _family_allowed(
|
||||
Absent counts as on when `tools` is on, so an upgrade does not silently take
|
||||
web search away from every model already set up for it.
|
||||
"""
|
||||
gate = gate_of(family)
|
||||
default = bool(capabilities.get("tools"))
|
||||
if not capabilities.get(f"tool_{family}", default):
|
||||
if not capabilities.get(f"tool_{gate}", default):
|
||||
return False
|
||||
|
||||
if family == FAMILY_SEARCH:
|
||||
if gate == FAMILY_SEARCH:
|
||||
return bool(
|
||||
allowed.get("tools.web_search")
|
||||
and config.get("enabled")
|
||||
and not search_service.availability(str(config.get("provider") or "ddgs"))
|
||||
)
|
||||
return bool(allowed.get(f"tools.{family}") and allowed.get("library.use"))
|
||||
if gate in (FAMILY_CUSTOM, FAMILY_MCP):
|
||||
# Deliberately without `library.use`: an HTTP endpoint an administrator
|
||||
# wrote has nothing to do with this person's own documents and notes,
|
||||
# and requiring the library permission for it would be a coincidence of
|
||||
# naming rather than a rule.
|
||||
return bool(allowed.get(f"tools.{gate}"))
|
||||
return bool(allowed.get(f"tools.{gate}") and allowed.get("library.use"))
|
||||
|
||||
|
||||
def enabled_tools(db: DBSession, chat: Chat, user: User | None) -> list[dict[str, Any]]:
|
||||
"""The tool schemas to offer for this chat."""
|
||||
def _row_defs(db: DBSession, user: User | None, *, everything: bool = False) -> list[ToolDef]:
|
||||
"""Tool definitions built from rows, in the order they claim names.
|
||||
|
||||
Imported here rather than at the top because `custom_tools` needs `ToolDef`
|
||||
from this module.
|
||||
"""
|
||||
from lembas.services import custom_tools
|
||||
|
||||
return custom_tools.tool_defs(db, user, everything=everything)
|
||||
|
||||
|
||||
def _book(defs: list[ToolDef]) -> dict[str, ToolDef]:
|
||||
"""Keyed by name, first claim winning.
|
||||
|
||||
The built-ins are laid down first, so a row can never shadow one -- a tool
|
||||
called `notes_delete` that turns out to be somebody's HTTP endpoint is the
|
||||
kind of surprise that has no good failure mode.
|
||||
"""
|
||||
book = dict(REGISTRY)
|
||||
for tool in defs:
|
||||
book.setdefault(tool.name, tool)
|
||||
return book
|
||||
|
||||
|
||||
def registry(db: DBSession) -> dict[str, ToolDef]:
|
||||
"""Every tool that exists on this instance, keyed by name, ungated.
|
||||
|
||||
`REGISTRY` holds the built-ins alone, because it is built at import time and
|
||||
an administrator-defined tool is a row. Callers that only need to map a name
|
||||
back to a family use this; callers deciding what to *offer* use
|
||||
`resolve_tools`, which applies the gates as well.
|
||||
"""
|
||||
return _book(_row_defs(db, None, everything=True))
|
||||
|
||||
|
||||
def families(db: DBSession) -> tuple[str, ...]:
|
||||
"""Every family that exists, the built-ins in their fixed order first."""
|
||||
rows = tuple(tool.family for tool in _row_defs(db, None, everything=True))
|
||||
return (*FAMILIES, *rows)
|
||||
|
||||
|
||||
def resolve_tools(db: DBSession, chat: Chat, user: User | None) -> ToolSet:
|
||||
"""Every tool this chat may call right now, with its runner attached."""
|
||||
from lembas.security import permissions
|
||||
from lembas.services import chat as chat_service
|
||||
|
||||
@@ -669,25 +764,47 @@ def enabled_tools(db: DBSession, chat: Chat, user: User | None) -> list[dict[str
|
||||
capabilities = model.capabilities_json or {}
|
||||
|
||||
if not capabilities.get("tools"):
|
||||
return []
|
||||
return ToolSet()
|
||||
|
||||
allowed = permissions.resolve(db, user)
|
||||
config = settings_store.search(db)
|
||||
|
||||
families = {
|
||||
family
|
||||
for family in FAMILIES
|
||||
if _family_allowed(family, config=config, capabilities=capabilities, allowed=allowed)
|
||||
}
|
||||
return [tool.schema for tool in REGISTRY.values() if tool.family in families]
|
||||
# Resolved against what this reader may see, not against everything that
|
||||
# exists: a tool restricted to a group is not offered outside it.
|
||||
book = _book(_row_defs(db, user))
|
||||
return ToolSet(
|
||||
tuple(
|
||||
tool
|
||||
for tool in book.values()
|
||||
if _family_allowed(
|
||||
tool.family, config=config, capabilities=capabilities, allowed=allowed
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def context_for(db: DBSession, user: User | None, chat: Chat | None = None) -> ToolContext:
|
||||
def enabled_tools(db: DBSession, chat: Chat, user: User | None) -> list[dict[str, Any]]:
|
||||
"""The tool schemas to offer for this chat.
|
||||
|
||||
The shape is unchanged on purpose: the inspector and `build_request` want
|
||||
exactly this. Anything that will also *run* a tool wants `resolve_tools`.
|
||||
"""
|
||||
return resolve_tools(db, chat, user).schemas
|
||||
|
||||
|
||||
def context_for(
|
||||
db: DBSession,
|
||||
user: User | None,
|
||||
chat: Chat | None = None,
|
||||
*,
|
||||
tools: ToolSet | None = None,
|
||||
) -> ToolContext:
|
||||
"""The snapshot a running tool needs, taken while the session is open."""
|
||||
return ToolContext(
|
||||
owner_id=user.id if user else "",
|
||||
search_config=settings_store.search(db),
|
||||
base_ids=[base.id for base in chat.knowledge_bases] if chat is not None else [],
|
||||
tools=tools.by_name if tools is not None else None,
|
||||
)
|
||||
|
||||
|
||||
@@ -697,8 +814,15 @@ async def run_tool(context: ToolContext, name: str, arguments: str) -> ToolOutco
|
||||
Never raises. A tool that fails hands the model an explanation and lets it
|
||||
carry on -- a failed lookup should produce "I could not find that" rather
|
||||
than killing the whole reply.
|
||||
|
||||
The lookup is against what was *offered*, not against everything that
|
||||
exists. Reaching for the registry directly meant a model naming a tool its
|
||||
chat was gated out of -- a family switched off for the model, a permission
|
||||
the reader does not have -- had it run anyway, because only the offer was
|
||||
ever filtered.
|
||||
"""
|
||||
tool = REGISTRY.get(name)
|
||||
book = REGISTRY if context.tools is None else context.tools
|
||||
tool = book.get(name)
|
||||
if tool is None:
|
||||
return ToolOutcome(
|
||||
f"There is no tool called {name!r}.",
|
||||
@@ -710,9 +834,12 @@ async def run_tool(context: ToolContext, name: str, arguments: str) -> ToolOutco
|
||||
except json.JSONDecodeError:
|
||||
# Small models emit malformed argument JSON often enough that this is a
|
||||
# normal path, not an exceptional one. Treat the whole string as the
|
||||
# first required argument rather than giving up.
|
||||
required = tool.parameters.get("required") or ["query"]
|
||||
parsed = {required[0]: arguments.strip()}
|
||||
# tool's first argument rather than giving up: what it says is required,
|
||||
# else the first thing it declares, and only then a guess -- a schema
|
||||
# somebody else wrote need not have either.
|
||||
properties = tool.parameters.get("properties") or {}
|
||||
names = tool.parameters.get("required") or list(properties) or ["query"]
|
||||
parsed = {str(names[0]): arguments.strip()}
|
||||
if not isinstance(parsed, dict):
|
||||
parsed = {"query": str(parsed)}
|
||||
|
||||
@@ -808,6 +935,31 @@ def tool_turn(call: dict[str, Any], content: str) -> dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def _row_source(db: DBSession):
|
||||
"""One harness fragment per administrator-defined tool.
|
||||
|
||||
The seam `prompts.register_source` exists for. The row supplies the default
|
||||
text and the admin page supplies the override, which is why a tool deleted
|
||||
and recreated under the same slug keeps whatever wording somebody chose for
|
||||
it -- the override outlives the row.
|
||||
|
||||
Gated on the tool's own family, so the guidance appears exactly when the
|
||||
tool it describes is offered and never otherwise.
|
||||
"""
|
||||
from lembas.db.models import CustomTool
|
||||
|
||||
for row in db.scalars(select(CustomTool).order_by(CustomTool.position, CustomTool.slug)):
|
||||
yield prompts_service.Fragment(
|
||||
key=f"tool.custom_{row.slug}",
|
||||
label=row.name or row.slug,
|
||||
group=prompts_service.GROUP_TOOLS,
|
||||
order=500 + row.position,
|
||||
families=(f"{FAMILY_CUSTOM}:{row.slug}",),
|
||||
hint=f"Appears when the {row.slug} tool is offered.",
|
||||
default=row.guidance or "",
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"FAMILIES",
|
||||
"MAX_ROUNDS",
|
||||
@@ -816,9 +968,19 @@ __all__ = [
|
||||
"ToolContext",
|
||||
"ToolDef",
|
||||
"ToolOutcome",
|
||||
"ToolSet",
|
||||
"assistant_turn",
|
||||
"context_for",
|
||||
"enabled_tools",
|
||||
"families",
|
||||
"registry",
|
||||
"resolve_tools",
|
||||
"run_tool",
|
||||
"tool_turn",
|
||||
]
|
||||
|
||||
|
||||
# Registered at import. `services.tools` is imported by the chat routes, the
|
||||
# generation service and the prompts admin, so the source is in place before
|
||||
# anything renders a catalogue.
|
||||
prompts_service.register_source(_row_source)
|
||||
|
||||
Reference in New Issue
Block a user