"""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]