Add support for LLM-generated filter parameters in VectorSearchRetrieverTool

[smurching]

Add Dynamic Filter Support to VectorSearchRetrieverTool

Summary

This PR adds opt-in support for LLM-generated filter parameters in VectorSearchRetrieverTool, enabling LLMs to dynamically construct filters based on natural language queries. This feature is controlled by a new dynamic_filter parameter (default: False) that exposes filter parameters in the tool schema when enabled.

Key Feature: The filter parameter description includes guidance for LLMs to use a fallback strategy - searching WITHOUT filters first to get broad results, then optionally adding filters to narrow down. This approach helps avoid zero results due to incorrect filter values while maintaining filtering flexibility.

Changes

Core Changes

1. New dynamic_filter Parameter

• Added dynamic_filter: bool field to VectorSearchRetrieverToolMixin (default: False)
• When True, exposes filter parameters in the tool schema for LLM-generated filters
• When False (default), maintains backward-compatible behavior with no filter parameters exposed

2. Mutual Exclusivity Validation

• Added @model_validator to ensure dynamic_filter=True and predefined filters cannot be used together
• Prevents ambiguous filter configuration by enforcing one approach or the other
• Clear error message guides users to the correct usage pattern

3. Enhanced Filter Parameter Descriptions with Fallback Strategy

• Extracts column metadata from Unity Catalog (workspace_client.tables.get())
• Includes available columns with types in the filter parameter description
• Example: "Available columns for filtering: product_category (STRING), product_sub_category (STRING)..."
• NEW: Includes guidance to search WITHOUT filters first: "IMPORTANT: If unsure about filter values, try searching WITHOUT filters first to get broad results, then optionally add filters to narrow down if needed. This ensures you don't miss relevant results due to incorrect filter values."
• Provides comprehensive operator documentation and examples

Integration Updates

OpenAI Integration (integrations/openai/src/databricks_openai/vector_search_retriever_tool.py)

• Conditionally creates EnhancedVectorSearchRetrieverToolInput (with optional filters) or BasicVectorSearchRetrieverToolInput (without filters) based on dynamic_filter setting
• Filter parameter is marked as Optional[List[FilterItem]] with default=None
• Inlines column metadata extraction during tool schema generation
• Fixed bug: Originally tried to use index.describe()["columns"] which doesn't exist; now uses Unity Catalog tables API

LangChain Integration (integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py)

• Similar conditional args_schema creation based on dynamic_filter setting
• Filter parameter is marked as optional (Optional[List[FilterItem]] with default=None)
• Maintains compatibility with LangChain's tool invocation patterns

Tests

New Test Coverage:

• test_cannot_use_both_dynamic_filter_and_predefined_filters - Validates mutual exclusivity
• test_predefined_filters_work_without_dynamic_filter - Ensures predefined filters work without dynamic mode
• test_enhanced_filter_description_with_column_metadata - Verifies column info is included
• test_enhanced_filter_description_without_column_metadata - Handles missing column info gracefully
• test_filter_item_serialization - Tests FilterItem schema

Test Results:

• ✅ OpenAI Integration: 48 tests passing
• ✅ LangChain Integration: 37 tests passing

Usage

Basic Usage (OpenAI)

from databricks_openai import VectorSearchRetrieverTool

# Enable dynamic filters
tool = VectorSearchRetrieverTool(
    index_name="catalog.schema.my_index",
    dynamic_filter=True  # Exposes optional filter parameters to LLM
)

# LLM receives guidance to try without filters first
# Then can optionally generate filters like:
# {"query": "wireless headphones", "filters": [{"key": "price <", "value": 100}]}
result = tool.execute(
    query="wireless headphones",
    filters=[{"key": "price <", "value": 100}]  # Optional!
)

Basic Usage (LangChain)

from databricks_langchain import VectorSearchRetrieverTool

tool = VectorSearchRetrieverTool(
    index_name="catalog.schema.my_index",
    dynamic_filter=True
)

# Use with LangChain agents - filter parameter is optional
result = tool.invoke({
    "query": "wireless headphones",
    "filters": [{"key": "price <", "value": 100}]  # Optional!
})

🎯 Recommendations

When to use dynamic_filter=True:

• Column values are discoverable from context (in retrieved documents)
• Filter requirements are simple and commonly understood (e.g., date ranges, numeric comparisons)
• Acceptable to have some queries return zero results due to filter mismatches
• Users can iteratively refine queries based on results
• LLM can follow the fallback strategy (search without filters first)

When to use predefined filters:

• Column values are constrained enums (product categories, status values, etc.)
• Filter logic is deterministic and known in advance
• Zero tolerance for LLM hallucination in filter values
• Consistent, predictable behavior is required

Implementation Details

Fallback Strategy Mechanism

The fallback strategy is implemented through tool description guidance rather than execution-time logic:

1. Filter Parameter Description includes: "IMPORTANT: If unsure about filter values, try searching WITHOUT filters first..."
2. Filters are Optional: Marked as Optional[List[FilterItem]] with default=None
3. LLM Follows Guidance: When the LLM sees this description, it learns to:
   • First invoke the tool without filters to get broad results
   • Examine the results to understand available filter values
   • Optionally invoke again with accurate filters to narrow results

This approach:

• ✅ Leverages LLM's ability to follow instructions in tool descriptions
• ✅ Doesn't require shipping complex merge/fallback logic
• ✅ Simple to implement (text-based guidance)
• ✅ Backward compatible
• ✅ Educates LLMs on best practices without changing execution

---

Validation

Testing in Practice

When tested with LangChain AgentExecutor, we observe that the LLM intelligently decides when to generate filters based on the user prompt and the guidance in the tool description. The fallback strategy works as intended - LLMs learn from the IMPORTANT guidance to search without filters first when unsure about filter values, then retry with filters if appropriate.

Observed LLM Behavior

When tested with LangChain AgentExecutor, the LLM demonstrates intelligent filter usage:

Example Query: "Find documentation about Data Engineering products"

LLM Actions:

1. First attempt: Tries with a filter based on the query:

   {'query': 'Data Engineering products',
    'filters': [{'key': 'product_category', 'value': 'Data Engineering'}]}

   Result: Empty (0 results) - the category doesn't exist in the index

2. Second attempt: Following the IMPORTANT guidance, automatically retries WITHOUT filters:

   {'query': 'Data Engineering'}

   Result: Success! Returns relevant results from actual categories (Software, Computers, etc.)

Key Observation: The LLM learns from the guidance to:

• Try with filters when the user query suggests specific filter criteria
• Automatically fall back to searching without filters when the first attempt fails
• Get broader, more relevant results instead of returning zero results

[smurching]

We should probably at least log a warning here

[jennsun]

This looks great! only significant comment was handling errors during the table information retrieval logic. let's be sure to update documentation afterwards to do education around the dynamic_filter parameter including these points off the top of my head:

1. When are appropriate cases to use dynamic_filter vs when predefined filters work just fine
2. Example usage

[jennsun]

not a review but wanted to comment: interesting insight here to search broadly then narrow down

[jennsun]

Have you done testing on your own to confirm this retry filtering works as intended?

[jennsun]

since the dynamic filter requires column metadata to be useful, can we raise a value error if the column_info is still empty at the end of this process?

may also be safe to add error handling/wrap the table information retrieval logic in a try catch if there are failures re permission checks/table doesn't exist or does not have entries

[jennsun]

if we add error handling per my comment in src/databricks_ai_bridge/vector_search_retriever_tool.py - can we add a test that simulates failure during WorkspaceClient.tables.get() process