"""Tools a model may call while it answers. A registry of named callables with a JSON schema each: offered to the endpoint, executed here when it asks. MCP servers and agentic execution plug in at the same place, which is why the registry is keyed and grouped rather than being a handful of if-statements. Three things gate whether a tool is offered: * the instance is configured for it (web search has a provider, and so on), * the reader has the permission, and * the chat's model is marked as having that tool. The last is not optional politeness. Sending a ``tools`` array to an endpoint that does not implement tool calling fails the entire request, exactly the way sending image parts to a model without vision does. Tools that *write* -- notes, memories, skills -- need a database session and a user, and they run inside a background generation that outlives the request. So they are handed a `ToolContext` carrying an owner id rather than a live session, and open their own scope, the same way `services.generation` does. """ from __future__ import annotations import json import logging from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from typing import Any 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 search as search_service from lembas.services import settings_store from lembas.services.library import documents as documents_service from lembas.services.library import memories as memories_service from lembas.services.library import notes as notes_service from lembas.services.library import skills as skills_service from lembas.services.search.base import SearchError log = logging.getLogger(__name__) # How many times a model may call tools before it has to answer with words. # Not a safety limit so much as a termination one: a small model that has # decided searching is the answer will otherwise search until the context runs # out, and each round costs a full request. MAX_ROUNDS = 3 # Tool families, matching the per-model capability flags and the permission # keys. The three names differ by prefix only, which is deliberate: adding a # family means adding one entry here and one permission. FAMILY_SEARCH = "web_search" FAMILY_KNOWLEDGE = "knowledge" FAMILY_NOTES = "notes" FAMILY_MEMORY = "memory" FAMILY_SKILLS = "skills" FAMILIES = (FAMILY_SEARCH, FAMILY_KNOWLEDGE, FAMILY_NOTES, FAMILY_MEMORY, FAMILY_SKILLS) @dataclass class ToolContext: """What a tool needs to do its work, without holding a session open. `owner_id` rather than a User for the same reason `Endpoint` is a frozen snapshot rather than a Connection: a generation outlives the request that started it, and a detached SQLAlchemy instance is a trap. """ 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) @dataclass class ToolOutcome: """What running a tool produced, for the model and for the reader. The two are deliberately different. `content` is the flat text the model reads back; `event` is what the transcript shows, and keeps results structured so they can be rendered as links rather than as a wall of URLs. """ content: str event: dict[str, Any] = field(default_factory=dict) Runner = Callable[[ToolContext, dict[str, Any]], Awaitable[ToolOutcome]] @dataclass(frozen=True) class ToolDef: name: str family: str description: str parameters: dict[str, Any] run: Runner @property def schema(self) -> dict[str, Any]: return { "type": "function", "function": { "name": self.name, "description": self.description, "parameters": self.parameters, }, } def _object(properties: dict[str, Any], required: list[str]) -> dict[str, Any]: return {"type": "object", "properties": properties, "required": required} _STRING = {"type": "string"} # --- Web search -------------------------------------------------------------- async def _run_web_search(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: query = str(args.get("query") or "").strip() if not query: return ToolOutcome( "No search query was given.", {"name": "web_search", "status": "error", "error": "No query was given."}, ) limit = args.get("max_results") try: limit = int(limit) if limit is not None else None except (TypeError, ValueError): limit = None try: results = await search_service.run(context.search_config, query, limit=limit) except SearchError as exc: log.info("web search failed for %r: %s", query[:60], exc.message) return ToolOutcome( f"The search failed: {exc.message}", {"name": "web_search", "query": query, "status": "error", "error": exc.message}, ) event = { "name": "web_search", "query": query, "status": "ok", "results": [ {"title": r.title, "url": r.url, "snippet": r.snippet, "host": r.host} for r in results ], } if not results: return ToolOutcome(f"No results were found for {query!r}.", event) lines = [f"Search results for {query!r}:"] for index, result in enumerate(results, start=1): lines.append(f"\n[{index}] {result.title}\n{result.url}\n{result.snippet}") return ToolOutcome("\n".join(lines), event) # --- Knowledge --------------------------------------------------------------- async def _run_knowledge_search(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: query = str(args.get("query") or "").strip() if not query: return ToolOutcome( "No search terms were given.", {"name": "knowledge_search", "status": "error", "error": "No query."}, ) with session_scope() as db: user = db.get(User, context.owner_id) found = documents_service.search( db, user, query, limit=6, base_ids=context.base_ids ) event = { "name": "knowledge_search", "query": query, "status": "ok", "results": [ {"title": d.title, "id": d.id, "kind": d.kind, "host": d.source_url} for d in found ], } if not found: return ToolOutcome( f"Nothing in the knowledge library matches {query!r}.", event ) lines = [f"Knowledge library matches for {query!r}:"] for document in found: lines.append( f"\n[{document.id}] {document.title}\n" f"{documents_service.snippet(document)}" ) lines.append( "\nUse knowledge_get with an id in brackets to read a document in full." ) return ToolOutcome("\n".join(lines), event) async def _run_knowledge_get(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: document_id = str(args.get("id") or "").strip() with session_scope() as db: user = db.get(User, context.owner_id) document = documents_service.get(db, document_id, user) if document is None: return ToolOutcome( "There is no such document, or it is not available to you.", {"name": "knowledge_get", "status": "error", "error": "Not found."}, ) event = { "name": "knowledge_get", "query": document.title, "status": "ok", "results": [{"title": document.title, "id": document.id}], } body = document.extracted_text or document.extraction_error or "(no text)" return ToolOutcome(f"{document.title}\n\n{body}", event) # --- Notes ------------------------------------------------------------------- async def _run_notes_search(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: query = str(args.get("query") or "").strip() with session_scope() as db: user = db.get(User, context.owner_id) found = ( notes_service.search(db, user, query, limit=8) if query else notes_service.recent(db, user, limit=8) ) event = { "name": "notes_search", "query": query, "status": "ok", "results": [{"title": n.title, "id": n.id} for n in found], } if not found: return ToolOutcome("There are no notes matching that.", event) lines = ["Notes:"] for note in found: lines.append(f"\n[{note.id}] {note.title}\n{notes_service.snippet(note)}") lines.append("\nUse notes_get with an id to read one in full.") return ToolOutcome("\n".join(lines), event) async def _run_notes_get(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: with session_scope() as db: user = db.get(User, context.owner_id) note = notes_service.get(db, str(args.get("id") or ""), user) if note is None: return ToolOutcome( "There is no such note, or it is not available to you.", {"name": "notes_get", "status": "error", "error": "Not found."}, ) return ToolOutcome( f"{note.title}\n\n{note.body}", { "name": "notes_get", "query": note.title, "status": "ok", "results": [{"title": note.title, "id": note.id}], }, ) async def _run_notes_create(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: title = str(args.get("title") or "").strip() body = str(args.get("body") or "").strip() if not body: return ToolOutcome( "A note needs a body.", {"name": "notes_create", "status": "error", "error": "Empty body."}, ) with session_scope() as db: user = db.get(User, context.owner_id) note = notes_service.create( db, owner=user, title=title, body=body, author=AUTHOR_MODEL ) return ToolOutcome( f"Saved note {note.id} — {note.title!r}.", { "name": "notes_create", "query": note.title, "status": "ok", "results": [{"title": note.title, "id": note.id}], }, ) async def _run_notes_edit(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: with session_scope() as db: user = db.get(User, context.owner_id) note = notes_service.get(db, str(args.get("id") or ""), user) if note is None or note.owner_id != context.owner_id: return ToolOutcome( "There is no such note, or it belongs to someone else. A note " "shared with you can be read but not changed.", {"name": "notes_edit", "status": "error", "error": "Not writable."}, ) notes_service.update( db, note, title=args.get("title"), body=args.get("body"), ) return ToolOutcome( f"Updated note {note.id}.", { "name": "notes_edit", "query": note.title, "status": "ok", "results": [{"title": note.title, "id": note.id}], }, ) async def _run_notes_delete(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: with session_scope() as db: user = db.get(User, context.owner_id) note = notes_service.get(db, str(args.get("id") or ""), user) if note is None or note.owner_id != context.owner_id: return ToolOutcome( "There is no such note, or it belongs to someone else.", {"name": "notes_delete", "status": "error", "error": "Not writable."}, ) title = note.title notes_service.delete(db, note) return ToolOutcome( f"Deleted note {title!r}.", {"name": "notes_delete", "query": title, "status": "ok", "results": []}, ) # --- Memory ------------------------------------------------------------------ async def _run_memory_add(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: content = str(args.get("content") or "").strip() with session_scope() as db: user = db.get(User, context.owner_id) try: memory = memories_service.add( db, owner=user, content=content, author=AUTHOR_MODEL ) except ValueError as exc: return ToolOutcome( str(exc), {"name": "memory_add", "status": "error", "error": str(exc)} ) note = "" if len(content) > memories_service.MAX_MEMORY_CHARS: # Trimmed rather than refused, with the model told so -- it can then # decide to put the long version in a note. note = ( f" It was shortened to {memories_service.MAX_MEMORY_CHARS} characters; " f"use notes for anything longer." ) return ToolOutcome( f"Remembered: {memory.content}{note}", { "name": "memory_add", "query": memory.content, "status": "ok", "results": [], }, ) async def _run_memory_forget(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: wanted = str(args.get("content") or "").strip().lower() with session_scope() as db: user = db.get(User, context.owner_id) records = memories_service.all_for(db, user) match = next((m for m in records if wanted and wanted in m.content.lower()), None) if match is None: return ToolOutcome( "No memory matches that. The full list is in the prompt already.", {"name": "memory_forget", "status": "error", "error": "No match."}, ) content = match.content memories_service.delete(db, match) return ToolOutcome( f"Forgotten: {content}", {"name": "memory_forget", "query": content, "status": "ok", "results": []}, ) # --- Skills ------------------------------------------------------------------ async def _run_skill_get(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: name = str(args.get("name") or "").strip() with session_scope() as db: user = db.get(User, context.owner_id) skill = skills_service.by_name(db, name, user) if skill is None: return ToolOutcome( f"There is no skill called {name!r}.", {"name": "skill_get", "status": "error", "error": "Not found."}, ) return ToolOutcome( f"Skill {skill.name}: {skill.description}\n\n{skill.body}", { "name": "skill_get", "query": skill.name, "status": "ok", "results": [{"title": skill.name, "id": skill.id}], }, ) async def _run_skill_create(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: with session_scope() as db: user = db.get(User, context.owner_id) try: skill = skills_service.create( db, owner=user, name=str(args.get("name") or ""), description=str(args.get("description") or ""), body=str(args.get("body") or ""), author=AUTHOR_MODEL, ) except skills_service.SkillError as exc: return ToolOutcome( str(exc), {"name": "skill_create", "status": "error", "error": str(exc)} ) return ToolOutcome( f"Created skill {skill.name!r}.", { "name": "skill_create", "query": skill.name, "status": "ok", "results": [{"title": skill.name, "id": skill.id}], }, ) async def _run_skill_edit(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: with session_scope() as db: user = db.get(User, context.owner_id) skill = skills_service.by_name(db, str(args.get("name") or ""), user) if skill is None or skill.owner_id != context.owner_id: return ToolOutcome( "There is no such skill, or it belongs to someone else.", {"name": "skill_edit", "status": "error", "error": "Not writable."}, ) skills_service.update( db, skill, description=args.get("description"), body=args.get("body"), author=AUTHOR_MODEL, note=str(args.get("reason") or "")[:200], ) return ToolOutcome( f"Updated skill {skill.name!r}. The previous version was kept and can " f"be restored.", { "name": "skill_edit", "query": skill.name, "status": "ok", "results": [{"title": skill.name, "id": skill.id}], }, ) # --- The registry ------------------------------------------------------------ REGISTRY: dict[str, ToolDef] = { tool.name: tool for tool in ( ToolDef( name="web_search", family=FAMILY_SEARCH, description=( "Search the web for current information. Use this when the answer " "depends on recent events, on facts you are unsure of, or on " "anything that may have changed since your training data. Returns " "a numbered list of results with titles, URLs and short extracts." ), parameters=_object( { "query": { "type": "string", "description": "The search terms. Keep them short and specific.", }, "max_results": { "type": "integer", "description": "How many results to return.", }, }, ["query"], ), run=_run_web_search, ), ToolDef( name="knowledge_search", family=FAMILY_KNOWLEDGE, description=( "Search the user's own collected documents, files and saved web " "pages. Use this before searching the web when the question is " "about their material rather than about the world." ), parameters=_object( {"query": {**_STRING, "description": "Words likely to appear in the document."}}, ["query"], ), run=_run_knowledge_search, ), ToolDef( name="knowledge_get", family=FAMILY_KNOWLEDGE, description="Read one knowledge document in full, by the id a search returned.", parameters=_object({"id": _STRING}, ["id"]), run=_run_knowledge_get, ), ToolDef( name="notes_search", family=FAMILY_NOTES, description=( "Search your notes. These are things you or the user wrote down in " "earlier conversations. With no query, returns the most recent." ), parameters=_object({"query": _STRING}, []), run=_run_notes_search, ), ToolDef( name="notes_get", family=FAMILY_NOTES, description="Read one note in full, by the id a search returned.", parameters=_object({"id": _STRING}, ["id"]), run=_run_notes_get, ), ToolDef( name="notes_create", family=FAMILY_NOTES, description=( "Write a note. Use this for something worth having in a later " "conversation that is too long or too detailed for a memory: a " "procedure, a summary, a set of preferences with reasons." ), parameters=_object( {"title": _STRING, "body": {**_STRING, "description": "Markdown."}}, ["title", "body"], ), run=_run_notes_create, ), ToolDef( name="notes_edit", family=FAMILY_NOTES, description="Change a note you can write to. Omit a field to leave it alone.", parameters=_object({"id": _STRING, "title": _STRING, "body": _STRING}, ["id"]), run=_run_notes_edit, ), ToolDef( name="notes_delete", family=FAMILY_NOTES, description="Delete a note that is no longer true or useful.", parameters=_object({"id": _STRING}, ["id"]), run=_run_notes_delete, ), ToolDef( name="memory_add", family=FAMILY_MEMORY, description=( "Remember one short, durable fact about the user — a preference, a " "constraint, how they like to be addressed. You are shown every " "memory on every turn, so keep them few and short, and never store " "passwords, keys or anything else secret." ), parameters=_object( {"content": {**_STRING, "description": "One fact, in one sentence."}}, ["content"], ), run=_run_memory_add, ), ToolDef( name="memory_forget", family=FAMILY_MEMORY, description=( "Remove a memory that has become wrong. Give enough of its text to " "identify it." ), parameters=_object({"content": _STRING}, ["content"]), run=_run_memory_forget, ), ToolDef( name="skill_get", family=FAMILY_SKILLS, description=( "Read the full instructions for one of the skills listed in your " "prompt. Do this before following a skill — the list gives only its " "name and what it is for." ), parameters=_object({"name": _STRING}, ["name"]), run=_run_skill_get, ), ToolDef( name="skill_create", family=FAMILY_SKILLS, description=( "Write a new skill: a reusable procedure for a task you expect to be " "asked again. The description must say when to use it, since that is " "all you will see next time." ), parameters=_object( { "name": {**_STRING, "description": "Short slug, e.g. 'weekly-report'."}, "description": {**_STRING, "description": "When to use this skill."}, "body": {**_STRING, "description": "The instructions, in Markdown."}, }, ["name", "description", "body"], ), run=_run_skill_create, ), ToolDef( name="skill_edit", family=FAMILY_SKILLS, description=( "Improve one of your skills. The previous version is kept and can be " "restored, so say why you changed it." ), parameters=_object( { "name": _STRING, "description": _STRING, "body": _STRING, "reason": {**_STRING, "description": "Why the change was made."}, }, ["name"], ), run=_run_skill_edit, ), ) } def _family_allowed( family: str, *, config: dict, capabilities: dict, allowed: dict ) -> bool: """Whether one family is on for this chat. A model configured before the per-tool flags existed has no `tool_*` keys. 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. """ default = bool(capabilities.get("tools")) if not capabilities.get(f"tool_{family}", default): return False if family == 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")) def enabled_tools(db: DBSession, chat: Chat, user: User | None) -> list[dict[str, Any]]: """The tool schemas to offer for this chat.""" from lembas.security import permissions from lembas.services import chat as chat_service capabilities = {} model = chat_service.model_for(db, chat) if model is not None: capabilities = model.capabilities_json or {} if not capabilities.get("tools"): return [] 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] def context_for(db: DBSession, user: User | None, chat: Chat | 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 [], ) async def run_tool(context: ToolContext, name: str, arguments: str) -> ToolOutcome: """Execute one tool call. 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. """ tool = REGISTRY.get(name) if tool is None: return ToolOutcome( f"There is no tool called {name!r}.", {"name": name, "status": "error", "error": "Unknown tool."}, ) try: parsed = json.loads(arguments) if arguments.strip() else {} 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()} if not isinstance(parsed, dict): parsed = {"query": str(parsed)} try: return await tool.run(context, parsed) except Exception as exc: # noqa: BLE001 - a tool must never kill the reply log.exception("tool %s failed", name) return ToolOutcome( f"The {name} tool failed: {exc}", {"name": name, "status": "error", "error": str(exc)[:200]}, ) class ToolCallAccumulator: """Reassembles tool calls arriving as streamed fragments. An endpoint sends ``delta.tool_calls`` as a list of partial objects: the id and the function name arrive once, and ``arguments`` arrives as a string split across however many chunks the tokeniser produced. Entries are keyed by ``index`` because that is the only field guaranteed on every fragment -- the id is absent from continuations, and matching on name breaks the moment a model calls the same tool twice in one turn. """ def __init__(self) -> None: self._calls: dict[int, dict[str, Any]] = {} def feed(self, fragments: list[dict[str, Any]]) -> None: for fragment in fragments: if not isinstance(fragment, dict): continue index = fragment.get("index") if not isinstance(index, int): # Some servers omit index entirely when there is only one call. index = 0 call = self._calls.setdefault(index, {"id": "", "name": "", "arguments": ""}) if fragment.get("id"): call["id"] = str(fragment["id"]) function = fragment.get("function") or {} if isinstance(function, dict): if function.get("name"): call["name"] = str(function["name"]) arguments = function.get("arguments") if isinstance(arguments, str): call["arguments"] += arguments @property def calls(self) -> list[dict[str, Any]]: """Completed calls, in the order the endpoint indexed them.""" return [ { # An id is required when the results are sent back, and not # every server supplies one. "id": call["id"] or f"call_{index}", "name": call["name"], "arguments": call["arguments"], } for index, call in sorted(self._calls.items()) if call["name"] ] def __bool__(self) -> bool: return bool(self.calls) def assistant_turn(calls: list[dict[str, Any]], content: str) -> dict[str, Any]: """The assistant message to send back with the tool results. The endpoint needs its own tool_calls echoed before the tool replies, or it has nothing to match the tool_call_ids against. """ return { "role": "assistant", "content": content or None, "tool_calls": [ { "id": call["id"], "type": "function", "function": {"name": call["name"], "arguments": call["arguments"]}, } for call in calls ], } def tool_turn(call: dict[str, Any], content: str) -> dict[str, Any]: return { "role": "tool", "tool_call_id": call["id"], "name": call["name"], "content": content, } __all__ = [ "FAMILIES", "MAX_ROUNDS", "REGISTRY", "ToolCallAccumulator", "ToolContext", "ToolDef", "ToolOutcome", "assistant_turn", "context_for", "enabled_tools", "run_tool", "tool_turn", ]