Add DS-Serve as an in-house search tool

[RulinShao]

Add DS Serve Search Tool Integration

Summary

This PR integrates the DS Serve API as a new search tool (ds_serve_search) for dense passage retrieval from CompactDS (2B passages) using dual ANN backends (DiskANN + IVFPQ).

Tool Details

Tool Name

ds_serve_search

Tool Parameters

• query (required): Search query string
• n_docs (default: 10, max: 1000): Number of passages to retrieve
• backend (default: "diskann"): Backend type - "diskann" or "ivfpq"
• nprobe (optional): Number of IVFPQ clusters to scan (only for IVFPQ backend)
• exact_search (default: false): Enable brute-force rerank after ANN
• diverse_search (default: false): Penalize near-duplicate passages
• lambda_param (default: 0.5): Diversity tradeoff used with diverse_search
• diskann_L (default: 500): DiskANN candidate list size (must be >= n_docs)
• diskann_W (default: 8): DiskANN beam width / I/O fan-out
• diskann_threads (optional): Override worker thread count
• min_words (default: 10): Minimum passage length filter
• base_url (optional): Custom API endpoint (default: api.ds-serve.org:30888)

Example Usage

<call_tool name="ds_serve_search">Tell me more about Albert Einstein</call_tool>
<call_tool name="ds_serve_search" n_docs="10" backend="diskann">Explain quantum physics</call_tool>
<call_tool name="ds_serve_search" n_docs="5" backend="ivfpq" nprobe="64">Who is Nikola Tesla</call_tool>

System Prompt

Added new system prompt version v20250907_with_ds_serve that includes documentation for the ds_serve_search tool. The prompt instructs the model to:

• Use ds_serve_search for dense passage retrieval from a 2B passage collection
• Understand the available backends (DiskANN and IVFPQ) and their parameters
• Use the tool for comprehensive information retrieval across diverse topics

Implementation

• API Client: agent/dr_agent/mcp_backend/apis/ds_serve_apis.py - Handles API communication and response parsing
• MCP Tool: agent/dr_agent/mcp_backend/main.py - Registers ds_serve_search as an MCP tool
• Tool Interface: agent/dr_agent/tool_interface/mcp_tools.py - DSServeSearchTool class for Python integration
• System Prompt: agent/dr_agent/shared_prompts/unified_tool_calling_v20250907_with_ds_serve.yaml
• Workflow Support: Updated auto_search_sft.py to support DS Serve as a search tool option

Configuration

The tool can be configured via YAML config files (e.g., auto_search_sft_ds_serve.yaml) with options for:

• Backend selection (diskann/ivfpq)
• Number of documents to retrieve
• Backend-specific parameters (nprobe, diskann_L, diskann_W, etc.)

[lolipopshock]

[optional] Consider removing the number of parameters in the ds_serve_search function.

If someone else uses the vanilla MCP server, all the arguments will be exposed to the llm. The model may not work relatively poorly when fed with this many options (and they are not sure how to parse it). We can simply remove some of the options from the function args?

[lolipopshock]

Perhaps we should make the name more generalizable, e.g., something like semantic search?

This is debatable, but in my view, our agent model do not need know the exact configuration of the tool -- what matters more is its intended purposes (and functions). My hunch is that if calling this with a more general name (e.g., semantic/vector search) then the LM can use some of the pre-trained knowledge to come up with search queries, compared to using a very specific name like DS_Search?

[lolipopshock]

this feels less of a test for this library, but for testing if the ds server is functioning? I guess then we can move it to the tool documentations etc (which I should have added one lmao)