AI v2 — Tool Retrieval

Audience: backend engineers working on the AI v2 units agent or its tools. TL;DR: The units agent has 44 tools — ~12k tokens of JSON Schema on every model call. A semantic retriever embeds each tool's description once and ranks it against the user's message, so a turn is sent the always-on grounding tools plus the 12 tools the message actually needs (17 of 44, a ~61% cut). Ranking is by meaning, not keywords, so it works identically in Arabic, English, Turkish or anything else. Fail-open: any retrieval failure sends the whole roster.

Why

The tool block is the largest fixed cost in a turn. It is charged on every iteration of a ReAct loop budgeted at up to MAX_SCHEMA_TOOL_CALLS (25) rounds, and it does not shrink as the conversation goes on the way history does. Past ~30–50 tools a model's selection accuracy also starts to drop as look-alikes multiply, so sending fewer, better-matched tools helps twice.

It used to be lexical, and that was the bug. The previous retriever scored tools by token overlap against /[a-z0-9]+/g. An Arabic, Turkish or Persian message tokenizes to nothing under that regex, so it always fell open to the full roster. Narrowing therefore only ever applied to English traffic — the same product behaving differently by language, in an Arabic-first CRM. Longer keyword lists could never have fixed that; only multilingual embeddings could.

A note on prompt caching. An earlier version of this document argued the full roster was worth keeping because a byte-stable tool block caches better than a per-message subset. That reasoning does not hold here: AI_UNITS_PROVIDER=kimi, and neither the Kimi nor the Gemini client reads request.promptCache — the value set in cortexa-chat-model.ts is discarded, and the Anthropic client that implements cache_control is never constructed without an API key. With no cache-write premium to pay, a smaller tool block is an unoffset saving. Revisit this if an agent is ever pinned to Anthropic, where a cache write costs 1.25× and the tool block sits first in the cacheable prefix.

The pipeline

1. Metadata — @AiV2ToolDef

Authorization descriptors (verb, permission) deliberately live in the policy registry (policy/tool-policy.registry.ts), not here — this metadata is for discovery; the policy registry is the single source of truth for access.

2. Registry — domain-grouped delivery

AiV2ToolRegistry discovers every @AiV2ToolDef tool at boot and sorts each agent's roster by DOMAIN_RANK: records → objects → fields → relations → views → comments → tasks → catalog → meta. sort() is stable, so discovery order holds within a domain, and a selected subset stays domain-grouped and byte-stable.

3. Vector store — ToolVectorStore

Each tool's retrieval document is name + description + keywords + category. The description carries most of the signal and is already maintained, because the model reads it — unlike a keyword list, which someone has to remember to update.

• Embedded with RETRIEVAL_DOCUMENT, the message with RETRIEVAL_QUERY. Retrieval models are asymmetric; using the same task type for both measurably degrades ranking.
• Built on first use, not at boot, so an unreachable provider can never stop the app from starting. Concurrent first turns share one in-flight build.
• Cached in Redis under ai-v2:tool-vectors:<model>:<corpusHash>, as one packed float32 buffer. The key hashes the corpus, so editing any tool description writes a new key rather than serving vectors built from the old wording.
• The key is deliberately not tenant- or workspace-scoped, and the store never touches repoProvider. Tool descriptions are global application code, identical for every tenant — there is no tenant data here to isolate, and putting it behind RLS would miscategorise code as customer data. (Embedding records would be the opposite case: that IS workspace data and would live under RLS.)

4. Retriever — ToolRetriever.select(catalog, query)

Keeps the always-on core unconditionally, adds the top MAX_RETRIEVED_TOOLS (12) ranked tools, and returns them in catalog order so each domain still reaches the model contiguously.

There is no minimum-similarity floor, and that is measured rather than forgotten. On this corpus the weakest true match (0.604) and the strongest off-topic match (0.604) score identically, so any floor that dropped junk would also drop real matches. An off-topic message costing a few unused schemas is the cheap side of that trade.

5. Escape hatch — find_tools

Retrieval runs once, on the opening message, and the bound set is otherwise fixed for the whole ReAct loop. A plan often only reveals what it needs partway through — inspect the fields, then realise a relation is required. Without an escape hatch, a tool retrieval narrowed away is unreachable for the rest of the turn.

So the model gets find_tools: it describes the capability it needs in plain language (any language), and the best MAX_DISCOVERED_TOOLS (5) matches become callable from its next step. This is the two-phase pattern the tool-retrieval literature converges on — narrow by default, let the model widen on demand.

Two properties keep it safe and cheap:

• Visibility ≠ permission. Tools held back are still in the graph's execution map; only the bound set is narrowed. Revealing one grants nothing, because PolicyGuard authorizes every call by name at execution time regardless of what was bound.
• Discoveries reset per turn. They are read off the transcript (discoveredToolNames), not held in a state channel. A channel is checkpointed with the thread, so discoveries would accumulate and creep the tool block back to the full roster — the bug that has already shipped twice here (toolCallCount, questionRounds).

The cap is 5 rather than 12 because by the time the model asks, it has stopped guessing from an opening message and is naming a capability it knows it needs — precision beats recall, and it can always ask again.

The safety properties

1. Nothing to shrink, nothing done. A roster with no more candidates than MAX_RETRIEVED_TOOLS is passed through whole — retrieval would return all of them anyway, so the round-trip is skipped entirely.
2. Fail-open. An embeddings failure, an unconfigured key, or a blank message returns the full roster. A hot-path network call must never be able to strand a turn; the worst case is exactly the pre-retrieval behaviour.
3. Core always kept. The agent can always inspect schema and find records even when the message doesn't name those tools.
4. Resume never retrieves. A paused write resumes with the full roster (agents/schema-manager/schema-manager-agent.service.ts), so an approved write can never be missing its tool. Retrieval runs on a fresh turn only.

Adding a tool

Tag the class with @AiV2ToolDef({ agents, category, keywords }) and list it in the module's providers. The registry rosters it, the retriever ranks it. Then add a case to the eval in both languages — a tool with no case is a tool whose recall nobody is watching, and the eval fails if any non-core tool is uncovered.

Write the description for a reader who does not already know the tool exists. It is the primary retrieval signal now; a thin one makes the tool hard to reach no matter how many keywords it carries (the eval enforces a minimum length).

The eval — bilingual retrieval recall

tool-registry/__tests__/tool-retrieval.eval.spec.ts. Two halves:

• Metadata checks (always run, no network). The catalog is the real 44-tool roster built exactly as production builds it; every tool has substantial retrieval text, and every non-core tool has at least one eval case.
• Recall (calls the embeddings provider). 39 realistic requests in English and Arabic, asserting for each that the needed tool is in the selection and that the selection actually narrowed. Both assertions are required: recall alone passes vacuously when the retriever falls open to the full roster — which is precisely what the old lexical scorer did for Arabic, so an English-only recall suite would have called that bug a pass.

It needs GEMINI_API_KEY and self-skips without one, so a keyless CI stays green.

npm test -- tool-retrieval.eval

A failure prints the exact "prompt" → missing tool lines. Fix by improving that tool's description (or its keywords), not by loosening the assertion.

Current measurement — gemini-embedding-001, 44 tools, 39 cases per language:

Arabic ranks slightly lower at exact top-1 but is identical at top-12, which is what the retriever actually uses. MAX_RETRIEVED_TOOLS = 12 comes from this table — retune it here, not by feel.

Model choice

gemini-embedding-001, on the GEMINI_API_KEY already in use, overridable via AI_EMBEDDINGS_MODEL. Chosen by measurement: it is genuinely multilingual, and Arabic recall matches English. The other candidate on the same key, text-embedding-004, is not served (404 on this endpoint) and is documented English-only regardless.

Cost is a rounding error and should not drive this choice: one query embedding per turn is ~50 tokens, and the 44-document corpus (~6.6k tokens) is embedded once per deploy and then cached. Select an embedding model on multilingual recall alone.

Deliberately not built

• Splitting units into 4 specialist agents. The literature recommends 5–10 tools per agent, and 44 is well past that. It is not done because agent dispatch probes each agent's pending gate serially (ai-graph.service.ts) — those probes share the one workspace-pinned Postgres connection and so cannot be parallelised, meaning every extra specialist adds a DB round-trip to every turn. find_tools buys most of the benefit without that cost. Revisit if the pending-gate probe is ever made connection-free.
• Merging tools behind an action enum (e.g. one record_mutate instead of create/update/delete). TOOL_POLICIES keys risk tier on the tool NAME, and the roster spans 17 SAFE / 25 WRITE / 6 DANGER. Collapsing create_record (WRITE) into the same tool as delete_record (DANGER) forces one tier for both — either everything over-confirms, or deletes stop being confirmed. The second is a security regression, so cross-tier merges are off the table regardless of the token saving.
• Narrowing the resume roster. With writeGate: 'per-write', every confirmed write is start (narrowed) → interrupt → resume (full), so the tool block flips twice per write. Sending only the approved tools on resume would be cheaper, but the resume decision carries approved call ids, not tool names, and resolving them means reading pending calls back out of the checkpoint — the interrupt-vs-getState interaction that has already caused two bugs here. Not worth risking "an approved write is missing its tool".
• Hybrid keyword + embedding retrieval. Rejected on purpose: it reinstates the English-biased lexical layer as a primary gate and inherits its blind spot.
• Per-session retrieval cache. The corpus vectors are already cached; only the query embedding is per-turn, and caching that would hit almost never.
• A task planner before retrieval. The units agent is a single ReAct loop by design; long-horizon planning is the workspace builder's job.

Scaling past this

Two independent axes, and they need different fixes:

1. One agent's roster grows. Retrieval already handles it — MAX_RETRIEVED_TOOLS bounds what the model sees regardless of roster size. Keep the eval honest as tools are added.
2. The number of AGENTS grows. The supervisor router picks ONE agent per turn and select runs on that agent's catalog only, so the router's choice IS the coarse first retrieval and this is the fine second one. A large agent count means the router needs retrieval over agent descriptions — a separate layer from this file. Keep agent domains cohesive and non-overlapping, or the failure moves UP to the router as a mis-route.