[WiP] feat(mini-chat): exa.ai web search via custom function tools + robustness improvements#4148
Conversation
…ness improvements Implement a generic custom FunctionTool trait framework enabling mini-chat models to call external APIs (exa.ai web search, future tools) without hardcoding each tool in the agentic loop. Includes vLLM adapter fixes, graceful tool error handling, and improved response quality. - Add `domain/ports/function_tool.rs` — generic FunctionTool trait with lifecycle hooks (name, definition, system prompt guard, max calls, execute, result handling). - `StreamService::resolve_function_tools()` filters and assembles tool registry per-request. - Generic agentic-loop dispatch in `provider_task.rs` mirrors `search_knowledge` pattern: replay function_call, execute, inject function_call_output, continue loop. - Per-tool call limits + graceful degradation (return "unavailable" message instead of failing). - Add `infra/llm/providers/exa_search_tool.rs` — implement exa_search function tool. Makes HTTP POST to exa.ai via OAGW, returns web results (title, URL, highlights). - Model opt-in via `enabled_function_tools: ["exa_search"]` in model catalog. - Config: `exa_search` section with host, search_type, max results, per-message call limit. - Enhanced system-prompt guard: explicit instruction to cite sources as Markdown links, no harmony citation markers【1†…】; keep reasoning out of final answer. Signed-off-by: MikeFalcon77 <3686170+MikeFalcon77@users.noreply.github.com>
📝 WalkthroughWalkthroughThe PR adds pluggable function-tool support to mini-chat, introduces Exa search configuration and provisioning, switches selected model catalog entries to local H200 vLLM, and updates the vLLM Responses adapter for function-tool serialization and reasoning SSE events. ChangesMini-chat function tools and model routing
Sequence Diagram(s)sequenceDiagram
participant StreamService
participant ContextAssembly
participant ProviderTask
participant ExaSearchTool
participant RagHttpClient
StreamService->>ContextAssembly: extra_function_tools and extra_function_guards
ContextAssembly-->>StreamService: tools and system_instructions
StreamService->>ProviderTask: function_tools registry
ProviderTask->>ExaSearchTool: execute(query)
ExaSearchTool->>RagHttpClient: POST /search
RagHttpClient-->>ExaSearchTool: search results
ExaSearchTool-->>ProviderTask: formatted output
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
config/mini-chat.yaml (1)
247-294: 🗄️ Data Integrity & Integration | 🟡 Minor
gpt-4.1advertises provider-hosted tools that the vLLM adapter silently drops, leading to incorrect surcharge billing.The entry routes to
provider_id: vllm_h200but retainstool_support.web_search: true,file_search: true, andcode_interpreter: true. The vLLM Responses adapter filters out these provider-hosted tools, transmitting onlyLlmTool::Function. Meanwhile, the token estimator applies surcharges (tool_surcharge_tokens,web_search_surcharge_tokens) based on these configuration flags.This creates a mismatch where:
- Enabled tools are silently dropped at the provider layer.
- Quota/billing estimates include surcharges for tools that never execute.
Set
web_search,file_search, andcode_interpretertofalsefor thevllm_h200entry to align with actual capabilities (as done forgpt-4.1-mini).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@config/mini-chat.yaml` around lines 247 - 294, The gpt-4.1 model configuration for the vllm_h200 provider advertises provider-hosted tools that the adapter does not actually pass through, causing tool-related surcharge estimates to be applied incorrectly. Update the gpt-4.1 entry’s tool_support settings to match the real capabilities of vLLM by setting web_search, file_search, and code_interpreter to false, consistent with the gpt-4.1-mini configuration, so the estimator and provider behavior stay aligned.
🧹 Nitpick comments (1)
gears/mini-chat/mini-chat/src/domain/service/stream_service/provider_task.rs (1)
1280-1339: 🧹 Nitpick | 🔵 TrivialCustom function-tool dispatch looks correct. Soft per-tool limit, replay of
function_call/function_call_output, and graceful error fallback mirror thesearch_knowledgepath, and the borrow ofcountdoesn't conflict with theraw_input_itemspushes.One observability note: unlike
web_search,code_interpreter, andsearch_knowledge, custom function-tool executions are neither persisted viaincrement_tool_callsnor recorded through any metric. This is acceptable for P1 settlement (the ADR reuses the flattool_surcharge_tokens), but it leaves no signal for monitoring custom-tool usage, success/failure rates, or latency. Consider adding a metric (e.g. arecord_function_toolcounter + latency histogram onMiniChatMetricsPort) so exa_search adoption and failures are observable in production.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@gears/mini-chat/mini-chat/src/domain/service/stream_service/provider_task.rs` around lines 1280 - 1339, The custom function-tool path in provider_task’s agentic dispatch is missing observability, unlike web_search, code_interpreter, and search_knowledge. Add a MiniChatMetricsPort hook for function tools (for example, a record_function_tool counter plus latency/success-failure tracking) around the existing tool.execute call and the soft-limit branch so custom tool usage, errors, and duration are visible in production.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@gears/mini-chat/mini-chat/src/domain/service/stream_service/mod.rs`:
- Around line 113-116: The `function_tools` field in `StreamService` is
formatted in a way that fails rustfmt’s check, so update that struct field
declaration to match `cargo fmt` output. Use the `StreamService` definition and
its `function_tools` field as the anchor, and apply the canonical
wrapping/indentation produced by running `cargo fmt --all` so the declaration is
either flattened or re-indented consistently with rustfmt.
- Around line 42-77: The dispatch map in resolve_function_tools_from is keyed by
the catalog name from enabled, but the LLM will call tools using
definition().name, so a name mismatch would break lookup in function_tools.get.
Update the map insertion logic to key by the tool definition name returned from
tool.definition(), while still preserving the existing tool descriptors and
system-prompt guards. Make this change in resolve_function_tools_from so the
dispatch map always aligns with the model-emitted function name.
In `@gears/mini-chat/mini-chat/src/infra/llm/providers/exa_search_tool.rs`:
- Line 156: The `query` assignment in `exa_search_tool.rs` needs to be
reformatted to satisfy rustfmt, as `cargo fmt --check` is failing on the long
chained expression. Update the `let query = input.get("query").and_then(|v|
v.as_str()).unwrap_or_default();` line in the `exa_search_tool` code path to the
multiline style rustfmt expects, then run `cargo fmt` to verify the formatting
is clean.
- Around line 93-98: The final truncation in ExaSearchTool::build_query_output
uses String::truncate with a byte-based limit and can panic on UTF-8 non-char
boundaries. Update the truncation logic in exa_search_tool.rs to ensure the
result is cut on a valid char boundary before returning, and keep the existing
accumulation logic in build_query_output unchanged. Also expand the tests around
ExaSearchTool to cover multi-byte UTF-8 input, not just ASCII, so the boundary
case is exercised.
In `@gears/mini-chat/mini-chat/src/infra/oagw_provisioning.rs`:
- Around line 114-140: The upstream registration in create_upstream is not
restart-safe because duplicate alias conflicts still fail the whole
initialization. Update the oagw_provisioning flow around
ServiceGatewayClientV1::create_upstream to handle the existing-resource conflict
case the same way create_route handles duplicates: detect the conflict/“already
exists” error, log a warning or info, and either reuse or fetch the existing
upstream so provisioning can continue. Keep the normal error path for
non-conflict failures and preserve the later alias/id usage from the upstream
value.
---
Outside diff comments:
In `@config/mini-chat.yaml`:
- Around line 247-294: The gpt-4.1 model configuration for the vllm_h200
provider advertises provider-hosted tools that the adapter does not actually
pass through, causing tool-related surcharge estimates to be applied
incorrectly. Update the gpt-4.1 entry’s tool_support settings to match the real
capabilities of vLLM by setting web_search, file_search, and code_interpreter to
false, consistent with the gpt-4.1-mini configuration, so the estimator and
provider behavior stay aligned.
---
Nitpick comments:
In
`@gears/mini-chat/mini-chat/src/domain/service/stream_service/provider_task.rs`:
- Around line 1280-1339: The custom function-tool path in provider_task’s
agentic dispatch is missing observability, unlike web_search, code_interpreter,
and search_knowledge. Add a MiniChatMetricsPort hook for function tools (for
example, a record_function_tool counter plus latency/success-failure tracking)
around the existing tool.execute call and the soft-limit branch so custom tool
usage, errors, and duration are visible in production.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 000fd74d-67e8-4516-8d7e-01379c876dab
📒 Files selected for processing (21)
config/mini-chat.yamlgears/mini-chat/docs/ADR/0004-cpt-cf-mini-chat-adr-pluggable-function-tools.mdgears/mini-chat/docs/DESIGN.mdgears/mini-chat/mini-chat-sdk/src/models.rsgears/mini-chat/mini-chat/src/config.rsgears/mini-chat/mini-chat/src/domain/llm.rsgears/mini-chat/mini-chat/src/domain/models.rsgears/mini-chat/mini-chat/src/domain/ports/function_tool.rsgears/mini-chat/mini-chat/src/domain/ports/mod.rsgears/mini-chat/mini-chat/src/domain/service/context_assembly.rsgears/mini-chat/mini-chat/src/domain/service/mod.rsgears/mini-chat/mini-chat/src/domain/service/stream_service/mod.rsgears/mini-chat/mini-chat/src/domain/service/stream_service/provider_task.rsgears/mini-chat/mini-chat/src/domain/service/test_helpers.rsgears/mini-chat/mini-chat/src/gear.rsgears/mini-chat/mini-chat/src/infra/llm/providers/exa_search_tool.rsgears/mini-chat/mini-chat/src/infra/llm/providers/mod.rsgears/mini-chat/mini-chat/src/infra/llm/providers/vllm_responses.rsgears/mini-chat/mini-chat/src/infra/llm/providers/vllm_responses_tests.rsgears/mini-chat/mini-chat/src/infra/oagw_provisioning.rsgears/mini-chat/mini-chat/src/infra/plugins/static_model_policy/tests.rs
| /// Resolved custom function tools for a request: `(dispatch map keyed by name, | ||
| /// tool descriptors for the LLM request, system-prompt guards)`. | ||
| type ResolvedFunctionTools = ( | ||
| std::collections::HashMap<String, Arc<dyn crate::domain::ports::FunctionTool>>, | ||
| Vec<crate::domain::llm::LlmTool>, | ||
| Vec<String>, | ||
| ); | ||
|
|
||
| /// Pure resolution of enabled custom function tools against a registry. | ||
| /// | ||
| /// Unknown tool names are silently ignored. Extracted as a free function so it | ||
| /// can be unit-tested without constructing a `StreamService`. | ||
| fn resolve_function_tools_from( | ||
| registry: &std::collections::HashMap<String, Arc<dyn crate::domain::ports::FunctionTool>>, | ||
| enabled: &[String], | ||
| ) -> ResolvedFunctionTools { | ||
| let mut map = std::collections::HashMap::new(); | ||
| let mut tools = Vec::new(); | ||
| let mut guards = Vec::new(); | ||
| for name in enabled { | ||
| if let Some(tool) = registry.get(name) { | ||
| let def = tool.definition(); | ||
| tools.push(crate::domain::llm::LlmTool::Function { | ||
| name: def.name, | ||
| description: def.description, | ||
| parameters: def.parameters, | ||
| }); | ||
| if let Some(guard) = tool.system_prompt_guard() { | ||
| guards.push(guard); | ||
| } | ||
| map.insert(name.clone(), Arc::clone(tool)); | ||
| } | ||
| } | ||
| (map, tools, guards) | ||
| } | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm definition().name vs name() for the exa_search tool and any other FunctionTool impls.
fd -e rs | xargs rg -nP '\bimpl\s+FunctionTool\b' -l | while read -r f; do
echo "== $f =="
rg -nP -A3 'fn (name|definition)\s*\(' "$f"
doneRepository: constructorfabric/gears-rust
Length of output: 879
Key the dispatch map by definition().name to ensure robustness against catalog name divergence.
resolve_function_tools_from currently keys the dispatch map by the enabled catalog name (name), but the model emits the name from definition().name. Although existing tools like exa_search have matching names, a future divergence would cause the dispatch lookup (function_tools.get(&name)) to fail silently. Use def.name as the map key to guarantee alignment with the model's output.
if let Some(guard) = tool.system_prompt_guard() {
guards.push(guard);
}
- map.insert(name.clone(), Arc::clone(tool));
+ map.insert(def.name.clone(), Arc::clone(tool));🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@gears/mini-chat/mini-chat/src/domain/service/stream_service/mod.rs` around
lines 42 - 77, The dispatch map in resolve_function_tools_from is keyed by the
catalog name from enabled, but the LLM will call tools using definition().name,
so a name mismatch would break lookup in function_tools.get. Update the map
insertion logic to key by the tool definition name returned from
tool.definition(), while still preserving the existing tool descriptors and
system-prompt guards. Make this change in resolve_function_tools_from so the
dispatch map always aligns with the model-emitted function name.
| /// Custom function-tool registry, keyed by tool name. Resolved per-request | ||
| /// against the model's `enabled_function_tools`. | ||
| function_tools: | ||
| std::collections::HashMap<String, Arc<dyn crate::domain::ports::FunctionTool>>, |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix cargo fmt formatting on the function_tools field.
The CI Fmt / cargo fmt (check) job fails here due to the line wrapping/indentation of the function_tools field declaration. Run cargo fmt --all to apply rustfmt's canonical layout (it will likely collapse this to a single line or re-indent the wrapped type).
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@gears/mini-chat/mini-chat/src/domain/service/stream_service/mod.rs` around
lines 113 - 116, The `function_tools` field in `StreamService` is formatted in a
way that fails rustfmt’s check, so update that struct field declaration to match
`cargo fmt` output. Use the `StreamService` definition and its `function_tools`
field as the anchor, and apply the canonical wrapping/indentation produced by
running `cargo fmt --all` so the declaration is either flattened or re-indented
consistently with rustfmt.
Source: Pipeline failures
| if out.len() >= self.max_chars { | ||
| break; | ||
| } | ||
| } | ||
| out.truncate(self.max_chars); | ||
| out |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
String::truncate can panic on a non-char-boundary.
max_chars is a byte count, but Exa results (title, url, highlights) are arbitrary web text that frequently contains multi-byte UTF-8. String::truncate(new_len) panics if new_len does not lie on a char boundary, so truncating mid-codepoint will crash the tool execution task. The existing test only exercises ASCII ("x".repeat(5000)), so it never hits this path.
🛡️ Proposed fix: truncate on a char boundary
- out.truncate(self.max_chars);
- out
+ let mut end = self.max_chars.min(out.len());
+ while end > 0 && !out.is_char_boundary(end) {
+ end -= 1;
+ }
+ out.truncate(end);
+ out📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if out.len() >= self.max_chars { | |
| break; | |
| } | |
| } | |
| out.truncate(self.max_chars); | |
| out | |
| if out.len() >= self.max_chars { | |
| break; | |
| } | |
| } | |
| let mut end = self.max_chars.min(out.len()); | |
| while end > 0 && !out.is_char_boundary(end) { | |
| end -= 1; | |
| } | |
| out.truncate(end); | |
| out |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@gears/mini-chat/mini-chat/src/infra/llm/providers/exa_search_tool.rs` around
lines 93 - 98, The final truncation in ExaSearchTool::build_query_output uses
String::truncate with a byte-based limit and can panic on UTF-8 non-char
boundaries. Update the truncation logic in exa_search_tool.rs to ensure the
result is cut on a valid char boundary before returning, and keep the existing
accumulation logic in build_query_output unchanged. Also expand the tests around
ExaSearchTool to cover multi-byte UTF-8 input, not just ASCII, so the boundary
case is exercised.
| ctx: SecurityContext, | ||
| input: serde_json::Value, | ||
| ) -> Result<String, FunctionToolError> { | ||
| let query = input.get("query").and_then(|v| v.as_str()).unwrap_or_default(); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix the cargo fmt failure.
CI (Fmt / cargo fmt --check) fails on this let query = ... assignment — rustfmt wants it wrapped across multiple lines. Run cargo fmt to unblock the pipeline.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@gears/mini-chat/mini-chat/src/infra/llm/providers/exa_search_tool.rs` at line
156, The `query` assignment in `exa_search_tool.rs` needs to be reformatted to
satisfy rustfmt, as `cargo fmt --check` is failing on the long chained
expression. Update the `let query = input.get("query").and_then(|v|
v.as_str()).unwrap_or_default();` line in the `exa_search_tool` code path to the
multiline style rustfmt expects, then run `cargo fmt` to verify the formatting
is clean.
Source: Pipeline failures
| let upstream = gateway | ||
| .create_upstream(ctx.clone(), builder.build()) | ||
| .await | ||
| .map_err(|e| anyhow::anyhow!("OAGW exa upstream registration failed: {e}"))?; | ||
| info!(alias = %upstream.alias, upstream_id = %upstream.id, "OAGW exa upstream registered"); | ||
|
|
||
| let match_rules = MatchRules { | ||
| http: Some(HttpMatch { | ||
| methods: vec![HttpMethod::Post], | ||
| path: "/search".to_owned(), | ||
| query_allowlist: vec![], | ||
| path_suffix_mode: PathSuffixMode::Disabled, | ||
| }), | ||
| grpc: None, | ||
| }; | ||
| match gateway | ||
| .create_route( | ||
| ctx.clone(), | ||
| CreateRouteRequest::builder(upstream.id, match_rules) | ||
| .enabled(true) | ||
| .build(), | ||
| ) | ||
| .await | ||
| { | ||
| Ok(route) => info!(route_id = %route.id, "OAGW exa /search route registered"), | ||
| Err(e) => warn!(error = %e, "OAGW exa route registration failed (may already exist)"), | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find create_upstream implementation and any "already exists" / conflict handling.
fd -t f -e rs . | xargs rg -nP -C4 'fn\s+create_upstream' 2>/dev/null
rg -nP -C3 'already.?exist|AlreadyExists|Conflict|duplicate' --type=rsRepository: constructorfabric/gears-rust
Length of output: 16564
Handle create_upstream conflicts to ensure restart idempotency.
Verification confirms ServiceGatewayClientV1::create_upstream is not idempotent. Unlike create_route, which tolerates existing resources by logging a warning, create_upstream returns a hard error if the upstream alias already exists. On pod restart, this will abort gear initialization.
Update gears/mini-chat/mini-chat/src/infra/oagw_provisioning.rs to catch the "already exists" error (typically DomainError::Conflict or similar) and treat it as a success by logging a warning or fetching the existing upstream, matching the pattern used for route registration.
[extract_itings: gears/system/oagw/oagw/src/domain/services/management/mod.rs lines 84-90, gears/mini-chat/mini-chat/src/infra/oagw_provisioning.rs lines 114-118]
// Current: Hard fails on duplicate
let upstream = gateway
.create_upstream(ctx.clone(), builder.build())
.await
.map_err(|e| anyhow::anyhow!("OAGW exa upstream registration failed: {e}"))?;
// Recommended: Handle conflict as "already exists"
let upstream = match gateway
.create_upstream(ctx.clone(), builder.build())
.await
{
Ok(u) => u,
Err(e) if is_conflict_error(&e) => {
// Fetch existing upstream or log warning and skip
info!("OAGW exa upstream already exists, skipping registration");
// Implement fetch logic if upstream ID is strictly required
return Ok(()); // Or appropriate flow
}
Err(e) => return Err(anyhow::anyhow!("OAGW exa upstream registration failed: {e}")),
};🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@gears/mini-chat/mini-chat/src/infra/oagw_provisioning.rs` around lines 114 -
140, The upstream registration in create_upstream is not restart-safe because
duplicate alias conflicts still fail the whole initialization. Update the
oagw_provisioning flow around ServiceGatewayClientV1::create_upstream to handle
the existing-resource conflict case the same way create_route handles
duplicates: detect the conflict/“already exists” error, log a warning or info,
and either reuse or fetch the existing upstream so provisioning can continue.
Keep the normal error path for non-conflict failures and preserve the later
alias/id usage from the upstream value.
Implement a generic custom FunctionTool trait framework enabling mini-chat models to call external APIs (exa.ai web search, future tools) without hardcoding each tool in the agentic loop. Includes vLLM adapter fixes, graceful tool error handling, and improved response quality.
Add
domain/ports/function_tool.rs— generic FunctionTool trait with lifecycle hooks (name, definition, system prompt guard, max calls, execute, result handling).StreamService::resolve_function_tools()filters and assembles tool registry per-request.Generic agentic-loop dispatch in
provider_task.rsmirrorssearch_knowledgepattern: replay function_call, execute, inject function_call_output, continue loop.Per-tool call limits + graceful degradation (return "unavailable" message instead of failing).
Add
infra/llm/providers/exa_search_tool.rs— implement exa_search function tool. Makes HTTP POST to exa.ai via OAGW, returns web results (title, URL, highlights).Model opt-in via
enabled_function_tools: ["exa_search"]in model catalog.Config:
exa_searchsection with host, search_type, max results, per-message call limit.Enhanced system-prompt guard: explicit instruction to cite sources as Markdown links, no harmony citation markers【1†…】; keep reasoning out of final answer.
Summary by CodeRabbit
New Features
Bug Fixes