Teaching AI Agents to Use Semantic Search Tools Effectively Through Knowledge Agent Training

mixedbread_ai · Tech · 2026 · source Agentic: yes · promotional bias: med · source-backed · extracted with prompt 0.5.6 (schema v2)

Generated from extraction.json — do not edit by hand. Every claim cites verbatim source evidence; stated means the quote directly asserts the claim, implied means it supports it indirectly.

Problem

Domain: Agentic retrieval and semantic search tool use

Mixedbread AI saw a retrieval bottleneck: LLM reasoning improved quickly, but agents still generated keyword-like search queries learned from web and code-search behavior, leaving semantic search tools underused in production knowledge tasks.

Constraints:

• Agents needed to query semantic search rather than BM25-style keyword systems.
• The solution had to support general knowledge work beyond code search.
• Retrieval quality was measured with ranking metrics and trajectory quality.
• The search agent needed precision, speed, and cost-effectiveness.
• Training had to counteract pretraining behaviors that are inappropriate for domain-specific retrieval.

Architecture

Mixedbread built a custom search-agent harness with multiple semantic-search tools, corpus preview, query planning, deduplication, supervised fine-tuning from a teacher model, and on-policy reinforcement learning with retrieval and trajectory rewards.

Diagram confidence: reconstructed (reconstructed from prose).

flowchart LR
  user_query["User query"] --> corpus_preview["Corpus preview"]
  corpus_preview --> search_agent["Search agent"]
  search_agent --> semantic_search_tools["Semantic search tools"]
  semantic_search_tools --> retrieved_chunks["Retrieved chunks"]
  retrieved_chunks --> reward_models["Reward models"]
  reward_models --> search_agent

Components

Data flow

The agent receives a user query, initial semantic search results, and metadata-facet hints. It decomposes the intent into up to four semantic queries, chooses tools, retrieves and deduplicates chunks, and returns a ranked list. Training starts with teacher-model SFT and then uses RL rewards for final retrieval quality and search trajectory behavior.

Why this architecture

The architecture trains the agent to use semantic search as a tool, because prompting alone leaves agents stuck in keyword overlap behaviors that do not exploit semantic retrieval systems.

Models used: Teacher LLM · Search agent · Gemini 3.5 Flash · LLM reward judge

Scale

• The agent reached NDCG@10 of 0.4 on OBELICS Congress, compared with GPT-4's 0.18.
  “achieved an NDCG@10 of 0.4 on the OBELICS Congress benchmark (significantly outperforming GPT-4's 0.18)”summary, stated
• It reached 93.4% accuracy on Snowflake's MatchQA with Gemini 3.5 Flash.
  “93.4% accuracy on Snowflake's MatchQA benchmark when paired with Gemini 3.5 Flash.”summary, stated
• The agent can split a search into a maximum of four queries.
  “The agent splits its search intent into a maximum of four queries covering separate aspects”summary, stated

Alternatives considered (roads not taken)

Let general LLM agents generate search queries without training

Rejected because: They tend to use keyword patterns from web and code search rather than semantic query formulation.

Evaluate only final retrieved chunks

Rejected because: The training also needed trajectory rewards to shape how the agent searches.

Patterns

tool-use-orchestration (agent tools)

The search agent learns when and how to invoke semantic search tools rather than issuing naive keyword queries.

rag-retrieval (retrieval)

The system improves retrieval quality by training an agent to plan semantic search queries and rank chunks.

fine-tuning-adaptation (model adaptation)

The agent was trained with supervised fine-tuning followed by on-policy reinforcement learning.

Tradeoffs

Train a specialized search agent.

Chose SFT plus on-policy RL for semantic retrieval over Prompting a general-purpose LLM agent.

Rationale: The specialized agent directly optimized retrieval and trajectory metrics.

Provide a corpus preview before planning.

Chose Initial semantic search results and metadata hints over Blind query decomposition.

Rationale: The preview lets the agent understand corpus language and decide where to dig deeper.

Gotchas (what bit them)

Agents trained on code and web tools may misuse semantic search.

Root cause: They default to keyword overlap rather than natural-language semantic queries.

Resolution: Train a search agent specifically on semantic-search tool use and reward query formulation.

Final-result metrics alone miss poor search behavior.

Root cause: An agent can arrive at some relevant chunks through inefficient or brittle exploration.

Resolution: Use trajectory rewards that judge the search process itself.

Evaluation

Mixedbread evaluated both final retrieval rankings and agent search trajectories with NDCG, LLM relevance rubrics, and benchmark accuracy.

Metrics: NDCG@10. · OBELICS Congress benchmark. · Snowflake MatchQA accuracy. · Trajectory judge rubrics for natural query formulation and exploration.

Outcomes

• NDCG@10 reached 0.4 on OBELICS Congress versus GPT-4 at 0.18.
• MatchQA accuracy reached 93.4% with Gemini 3.5 Flash.
• The search agent learned to use semantic retrieval tools more effectively than general agents.

The case reports benchmark results rather than production traffic volume; it focuses on retrieval quality in deep search tasks.