simbarag/blueprints/conversation/tool_routing.py

"""Tool selection per model call.

Static mode (the default) binds the full role-appropriate set. Dynamic mode
routes per call by keyword; discovery grants live in tool messages, not globals.
"""

import json
import re

from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langchain_core.utils.function_calling import convert_to_openai_tool

GROUP_TERMS = {
    "pets": "simba cat cats pet pets vet veterinary medical weight litter drinking insulin medication dosage vaccine",
    "meals": "mealie meal meals recipe recipes ingredient ingredients dinner lunch breakfast grocery groceries shopping cook cooking eat food",
    "notes": "obsidian note notes vault journal task tasks todo markdown",
    "finance": "ynab budget budgets spending spent spend transaction transactions account accounts balance balances money finance financial cost costs",
    "restaurants": "beli restaurant restaurants dining dine cuisine cuisines cafe cafes takeout places",
    "calendar": "calendar event events appointment appointments schedule scheduled meeting meetings tomorrow availability",
    "email": "gmail email emails mail inbox message messages sender subject",
    "inventory": "homebox inventory belongings item items appliance appliances warranty warranties serial",
    "knowledge": "bookstack wiki book books chapter chapters page pages documentation",
    "web": "web internet online search news weather current latest research",
}
ADMIN_TOOL_NAMES = {
    "get_calendar_events",
    "create_calendar_event",
    "gmail_search_emails",
    "gmail_get_email",
    "ynab_sql",
}
CORE_TOOL_NAMES = {"get_current_date", "save_user_memory"}
DISCOVERY_NAME = "discover_tool_groups"
TOOL_ROUTING_MODES = ("static", "dynamic")


def words(text):
    return set(re.findall(r"[a-z0-9]+", text.casefold()))


def message_text(message):
    content = message.content
    if isinstance(content, str):
        return content
    return " ".join(
        p.get("text", "") if isinstance(p, dict) else str(p) for p in content
    )


def tool_name(value):
    """Accept LangChain tools, callables and provider schema dictionaries."""
    if hasattr(value, "name"):
        return value.name
    schema = convert_to_openai_tool(value)
    return schema.get("function", schema)["name"]


def tool_group(name):
    for prefix, group in (
        ("simbakit_", "pets"),
        ("simba_", "pets"),
        ("mealie_", "meals"),
        ("obsidian_", "notes"),
        ("journal_", "notes"),
        ("ynab_", "finance"),
        ("beli_", "restaurants"),
        ("gmail_", "email"),
        ("homebox", "inventory"),
        ("bookstack", "knowledge"),
    ):
        if name.startswith(prefix):
            return group
    if name in {"get_calendar_events", "create_calendar_event"}:
        return "calendar"
    if name == "web_search":
        return "web"
    # Unknown integration tools remain discoverable as their own group.
    return name


class ToolRouter:
    def __init__(self, tools, *, admin_tools=()):
        self.tools = list(tools)
        self.admin_names = ADMIN_TOOL_NAMES | {tool_name(t) for t in admin_tools}
        self.admin_names |= {
            tool_name(t)
            for t in self.tools
            if tool_name(t).startswith(("homebox", "bookstack"))
        }
        self.groups = {}
        for t in self.tools:
            if tool_name(t) not in CORE_TOOL_NAMES:
                self.groups.setdefault(tool_group(tool_name(t)), []).append(t)
        self.discovery = self._discovery_tool()

    def available_groups(self, is_admin):
        return {
            group: [
                t for t in tools if is_admin or tool_name(t) not in self.admin_names
            ]
            for group, tools in self.groups.items()
            if any(is_admin or tool_name(t) not in self.admin_names for t in tools)
        }

    def matching_groups(self, query, available):
        tokens = words(query)
        return {
            group
            for group in available
            if tokens & words(GROUP_TERMS.get(group, group.replace("_", " ")))
            or any(tool_name(t) in query.casefold() for t in available[group])
        }

    def _discovery_tool(self):
        @tool
        def discover_tool_groups(groups: list[str], config: RunnableConfig) -> str:
            """Enable additional tool groups on the next model call.

            Pass group names, ["all"] for recovery, or [] to list available groups.
            """
            available = self.available_groups(
                config.get("configurable", {}).get("is_admin", False)
            )
            enabled = sorted(
                available if "all" in groups else set(groups) & available.keys()
            )
            return json.dumps(
                {
                    "enabled_tool_groups": enabled,
                    "available_tool_groups": sorted(available),
                    "unknown_or_unavailable": sorted(
                        set(groups) - available.keys() - {"all"}
                    ),
                }
            )

        # Catalog is small, and tells the model what to request without binding
        # every schema. Runtime filtering also gates the discovery result.
        discover_tool_groups.description += " Group names: " + ", ".join(
            sorted(self.groups)
        )
        return discover_tool_groups

    def select_static(self, *, is_admin=False):
        """Every role-appropriate tool in registration order, without discovery.

        The bound set never varies by query or turn, so the rendered tool block
        stays byte-identical and llama-server can reuse its cached prefix.
        """
        return [
            t for t in self.tools if is_admin or tool_name(t) not in self.admin_names
        ]

    def select(self, messages, *, is_admin=False):
        available = self.available_groups(is_admin)
        human_indices = [
            i for i, m in enumerate(messages) if isinstance(m, HumanMessage)
        ]
        start = human_indices[-1] if human_indices else 0
        query = message_text(messages[start]) if human_indices else ""
        selected = self.matching_groups(query, available)
        current = messages[start:]
        previous = messages[human_indices[-2] : start] if len(human_indices) > 1 else []
        # Deictic followups inherit just the preceding turn, never all history.
        followup = bool(
            words(query)
            & {
                "it",
                "its",
                "that",
                "those",
                "them",
                "these",
                "same",
                "another",
                "again",
                "instead",
                "also",
                "more",
                "there",
                "ones",
            }
        ) or query.casefold().startswith(("what about", "and "))
        if (
            words(query)
            and len(words(query)) <= 5
            and words(query)
            <= words(
                "next last this previous tomorrow yesterday today week month year weekend please"
            )
        ):
            followup = True
        if followup:
            if previous:
                selected |= self.matching_groups(message_text(previous[0]), available)
        relevant_messages = (previous if followup else []) + current
        # IDs, paths and slugs copied from a previous result also preserve its group.
        reference_words = words(query) - words(
            "read show get open please the a an my me"
        )
        if not followup and any(
            isinstance(m, ToolMessage)
            and any(len(w) > 3 and w in words(message_text(m)) for w in reference_words)
            for m in previous
        ):
            relevant_messages = previous + current
        calls = {}
        for m in relevant_messages:
            if isinstance(m, AIMessage):
                for call in m.tool_calls:
                    calls[call["id"]] = call["name"]
                    selected.add(tool_group(call["name"]))
            if (
                isinstance(m, ToolMessage)
                and calls.get(m.tool_call_id) == DISCOVERY_NAME
                and m.status == "success"
            ):
                try:
                    grants = json.loads(message_text(m)).get("enabled_tool_groups", [])
                    if isinstance(grants, list):
                        selected.update(g for g in grants if isinstance(g, str))
                except (ValueError, AttributeError):
                    pass
        names = {
            tool_name(t)
            for group in selected & available.keys()
            for t in available[group]
        }
        return [
            t
            for t in self.tools
            if tool_name(t) in names or tool_name(t) in CORE_TOOL_NAMES
        ] + [self.discovery]