diff --git a/docs/my-website/docs/integrations/websearch_interception.md b/docs/my-website/docs/integrations/websearch_interception.md new file mode 100644 index 00000000000..0c5d8927013 --- /dev/null +++ b/docs/my-website/docs/integrations/websearch_interception.md @@ -0,0 +1,411 @@ +# Web Search Integration + +Enable transparent server-side web search execution for any LLM provider. LiteLLM automatically intercepts web search tool calls and executes them using your configured search provider (Perplexity, Tavily, etc.). + +## Quick Start + +### 1. Configure Web Search Interception + +Add to your `config.yaml`: + +```yaml +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + +litellm_settings: + callbacks: + - websearch_interception: + enabled_providers: + - openai + - minimax + - anthropic + search_tool_name: perplexity-search # Optional + +search_tools: + - search_tool_name: perplexity-search + litellm_params: + search_provider: perplexity + api_key: os.environ/PERPLEXITY_API_KEY +``` + +### 2. Use with Any Provider + +```python +import litellm + +response = await litellm.acompletion( + model="gpt-4o", + messages=[ + {"role": "user", "content": "What's the weather in San Francisco today?"} + ], + tools=[ + { + "type": "function", + "function": { + "name": "litellm_web_search", + "description": "Search the web for information", + "parameters": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query"} + }, + "required": ["query"] + } + } + } + ] +) + +# Response includes search results automatically! +print(response.choices[0].message.content) +``` + +## How It Works + +When a model makes a web search tool call, LiteLLM: + +1. **Detects** the `litellm_web_search` tool call in the response +2. **Executes** the search using your configured search provider +3. **Makes a follow-up request** with the search results +4. **Returns** the final answer to the user + +```mermaid +sequenceDiagram + participant User + participant LiteLLM + participant LLM as LLM Provider + participant Search as Search Provider + + User->>LiteLLM: Request with web_search tool + LiteLLM->>LLM: Forward request + LLM-->>LiteLLM: Response with tool_call + Note over LiteLLM: Detect web search
tool call + LiteLLM->>Search: Execute search + Search-->>LiteLLM: Search results + LiteLLM->>LLM: Follow-up with results + LLM-->>LiteLLM: Final answer + LiteLLM-->>User: Final answer with search results +``` + +**Result**: One API call from user → Complete answer with search results + +## Supported Providers + +Web search integration works with **all providers** that use: +- ✅ **Base HTTP Handler** (`BaseLLMHTTPHandler`) +- ✅ **OpenAI Completion Handler** (`OpenAIChatCompletion`) + +### Providers Using Base HTTP Handler + +| Provider | Status | Notes | +|----------|--------|-------| +| **OpenAI** | ✅ Supported | GPT-4, GPT-3.5, etc. | +| **Anthropic** | ✅ Supported | Claude models via HTTP handler | +| **MiniMax** | ✅ Supported | All MiniMax models | +| **Mistral** | ✅ Supported | Mistral AI models | +| **Cohere** | ✅ Supported | Command models | +| **Fireworks AI** | ✅ Supported | All Fireworks models | +| **Together AI** | ✅ Supported | All Together AI models | +| **Groq** | ✅ Supported | All Groq models | +| **Perplexity** | ✅ Supported | Perplexity models | +| **DeepSeek** | ✅ Supported | DeepSeek models | +| **xAI** | ✅ Supported | Grok models | +| **Hugging Face** | ✅ Supported | Inference API models | +| **OCI** | ✅ Supported | Oracle Cloud models | +| **Vertex AI** | ✅ Supported | Google Vertex AI models | +| **Bedrock** | ✅ Supported | AWS Bedrock models (converse_like route) | +| **Azure OpenAI** | ✅ Supported | Azure-hosted OpenAI models | +| **Sagemaker** | ✅ Supported | AWS Sagemaker models | +| **Databricks** | ✅ Supported | Databricks models | +| **DataRobot** | ✅ Supported | DataRobot models | +| **Hosted VLLM** | ✅ Supported | Self-hosted VLLM | +| **Heroku** | ✅ Supported | Heroku-hosted models | +| **RAGFlow** | ✅ Supported | RAGFlow models | +| **Compactif** | ✅ Supported | Compactif models | +| **Cometapi** | ✅ Supported | Comet API models | +| **A2A** | ✅ Supported | Agent-to-Agent models | +| **Bytez** | ✅ Supported | Bytez models | + +### Providers Using OpenAI Handler + +| Provider | Status | Notes | +|----------|--------|-------| +| **OpenAI** | ✅ Supported | Native OpenAI API | +| **Azure OpenAI** | ✅ Supported | Azure-hosted OpenAI | +| **OpenAI-Compatible** | ✅ Supported | Any OpenAI-compatible API | + +## Configuration + +### WebSearch Interception Parameters + +| Parameter | Type | Required | Description | Example | +|-----------|------|----------|-------------|---------| +| `enabled_providers` | List[String] | Yes | List of providers to enable web search for | `[openai, minimax, anthropic]` | +| `search_tool_name` | String | No | Specific search tool from `search_tools` config. If not set, uses first available. | `perplexity-search` | + +### Provider Values + +Use these values in `enabled_providers`: + +| Provider | Value | Provider | Value | +|----------|-------|----------|-------| +| OpenAI | `openai` | Anthropic | `anthropic` | +| MiniMax | `minimax` | Mistral | `mistral` | +| Cohere | `cohere` | Fireworks AI | `fireworks_ai` | +| Together AI | `together_ai` | Groq | `groq` | +| Perplexity | `perplexity` | DeepSeek | `deepseek` | +| xAI | `xai` | Hugging Face | `huggingface` | +| OCI | `oci` | Vertex AI | `vertex_ai` | +| Bedrock | `bedrock` | Azure | `azure` | +| Sagemaker | `sagemaker_chat` | Databricks | `databricks` | +| DataRobot | `datarobot` | VLLM | `hosted_vllm` | +| Heroku | `heroku` | RAGFlow | `ragflow` | +| Compactif | `compactif` | Cometapi | `cometapi` | +| A2A | `a2a` | Bytez | `bytez` | + +## Search Providers + +Configure which search provider to use. LiteLLM supports multiple search providers: + +| Provider | `search_provider` Value | Environment Variable | +|----------|------------------------|----------------------| +| **Perplexity AI** | `perplexity` | `PERPLEXITYAI_API_KEY` | +| **Tavily** | `tavily` | `TAVILY_API_KEY` | +| **Exa AI** | `exa_ai` | `EXA_API_KEY` | +| **Parallel AI** | `parallel_ai` | `PARALLEL_AI_API_KEY` | +| **Google PSE** | `google_pse` | `GOOGLE_PSE_API_KEY`, `GOOGLE_PSE_ENGINE_ID` | +| **DataForSEO** | `dataforseo` | `DATAFORSEO_LOGIN`, `DATAFORSEO_PASSWORD` | +| **Firecrawl** | `firecrawl` | `FIRECRAWL_API_KEY` | +| **SearXNG** | `searxng` | `SEARXNG_API_BASE` (required) | +| **Linkup** | `linkup` | `LINKUP_API_KEY` | + +See [Search Providers Documentation](../search/index.md) for detailed setup instructions. + +## Complete Configuration Example + +```yaml +model_list: + # OpenAI + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + + # MiniMax + - model_name: minimax + litellm_params: + model: minimax/MiniMax-M2.1 + api_key: os.environ/MINIMAX_API_KEY + + # Anthropic + - model_name: claude + litellm_params: + model: anthropic/claude-sonnet-4-5 + api_key: os.environ/ANTHROPIC_API_KEY + + # Azure OpenAI + - model_name: azure-gpt4 + litellm_params: + model: azure/gpt-4 + api_base: https://my-azure.openai.azure.com + api_key: os.environ/AZURE_API_KEY + +litellm_settings: + callbacks: + - websearch_interception: + enabled_providers: + - openai + - minimax + - anthropic + - azure + search_tool_name: perplexity-search + +search_tools: + - search_tool_name: perplexity-search + litellm_params: + search_provider: perplexity + api_key: os.environ/PERPLEXITY_API_KEY + + - search_tool_name: tavily-search + litellm_params: + search_provider: tavily + api_key: os.environ/TAVILY_API_KEY +``` + +## Usage Examples + +### Python SDK + +```python +import litellm + +# Configure callbacks +litellm.callbacks = ["websearch_interception"] + +# Make completion with web search tool +response = await litellm.acompletion( + model="gpt-4o", + messages=[ + {"role": "user", "content": "What are the latest AI news?"} + ], + tools=[ + { + "type": "function", + "function": { + "name": "litellm_web_search", + "description": "Search the web for current information", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query" + } + }, + "required": ["query"] + } + } + } + ] +) + +print(response.choices[0].message.content) +``` + +### Proxy Server + +```bash +# Start proxy with config +litellm --config config.yaml + +# Make request +curl http://localhost:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -d '{ + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "What is the weather in San Francisco?"} + ], + "tools": [ + { + "type": "function", + "function": { + "name": "litellm_web_search", + "description": "Search the web", + "parameters": { + "type": "object", + "properties": { + "query": {"type": "string"} + }, + "required": ["query"] + } + } + } + ] + }' +``` + +## How Search Tool Selection Works + +1. **If `search_tool_name` is specified** → Uses that specific search tool +2. **If `search_tool_name` is not specified** → Uses first search tool in `search_tools` list + +```yaml +search_tools: + - search_tool_name: perplexity-search # ← This will be used if no search_tool_name specified + litellm_params: + search_provider: perplexity + api_key: os.environ/PERPLEXITY_API_KEY + + - search_tool_name: tavily-search + litellm_params: + search_provider: tavily + api_key: os.environ/TAVILY_API_KEY +``` + +## Troubleshooting + +### Web Search Not Working + +1. **Check provider is enabled**: + ```yaml + enabled_providers: + - openai # Make sure your provider is in this list + ``` + +2. **Verify search tool is configured**: + ```yaml + search_tools: + - search_tool_name: perplexity-search + litellm_params: + search_provider: perplexity + api_key: os.environ/PERPLEXITY_API_KEY + ``` + +3. **Check API keys are set**: + ```bash + export PERPLEXITY_API_KEY=your-key + ``` + +4. **Enable debug logging**: + ```python + litellm.set_verbose = True + ``` + +### Common Issues + +**Issue**: Model returns tool_calls instead of final answer +- **Cause**: Provider not in `enabled_providers` list +- **Solution**: Add provider to `enabled_providers` + +**Issue**: "No search tool configured" error +- **Cause**: No search tools in `search_tools` config +- **Solution**: Add at least one search tool configuration + +**Issue**: "Invalid function arguments json string" error (MiniMax) +- **Cause**: Fixed in latest version - arguments weren't properly JSON serialized +- **Solution**: Update to latest LiteLLM version + +## Related Documentation + +- [Search Providers](../search/index.md) - Detailed search provider setup +- [Claude Code WebSearch](../tutorials/claude_code_websearch.md) - Using with Claude Code +- [Tool Calling](../completion/function_call.md) - General tool calling documentation +- [Callbacks](./custom_callback.md) - Custom callback documentation + +## Technical Details + +### Architecture + +Web search integration is implemented as a custom callback (`WebSearchInterceptionLogger`) that: + +1. **Pre-request Hook**: Converts native web search tools to LiteLLM standard format +2. **Post-response Hook**: Detects web search tool calls in responses +3. **Agentic Loop**: Executes searches and makes follow-up requests automatically + +### Supported APIs + +- ✅ **Chat Completions API** (OpenAI format) +- ✅ **Anthropic Messages API** (Anthropic format) +- ✅ **Streaming** (automatically converted) +- ✅ **Non-streaming** + +### Response Format Detection + +The handler automatically detects response format: +- **OpenAI format**: `tool_calls` in assistant message +- **Anthropic format**: `tool_use` blocks in content + +### Performance + +- **Latency**: Adds one additional LLM call (follow-up request with search results) +- **Caching**: Search results can be cached (depends on search provider) +- **Parallel Searches**: Multiple search queries executed in parallel + +## Contributing + +Found a bug or want to add support for a new provider? See our [Contributing Guide](https://github.com/BerriAI/litellm/blob/main/CONTRIBUTING.md). diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 343860cb158..37392ff5c18 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -96,6 +96,11 @@ const sidebars = { "proxy/prometheus" ] }, + { + type: "doc", + id: "integrations/websearch_interception", + label: "Web Search Integration" + }, { type: "category", label: "[Beta] Prompt Management", diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 07d237c4758..4a341863d4b 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -664,6 +664,37 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return final_response """ pass + + async def async_should_run_chat_completion_agentic_loop( + self, + response: Any, + model: str, + messages: List[Dict], + tools: Optional[List[Dict]], + stream: bool, + custom_llm_provider: str, + kwargs: Dict, + ) -> Tuple[bool, Dict]: + """ + Hook to determine if chat completion agentic loop should be executed. + """ + return False, {} + + async def async_run_chat_completion_agentic_loop( + self, + tools: Dict, + model: str, + messages: List[Dict], + response: Any, + optional_params: Dict, + logging_obj: "LiteLLMLoggingObj", + stream: bool, + kwargs: Dict, + ) -> Any: + """ + Hook to execute chat completion agentic loop based on context from should_run hook. + """ + pass # Useful helpers for custom logger classes diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 5d36b760afb..1277cac51d7 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -17,6 +17,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.tools import ( get_litellm_web_search_tool, is_web_search_tool, + is_web_search_tool_chat_completion, ) from litellm.integrations.websearch_interception.transformation import ( WebSearchTransformation, @@ -48,7 +49,8 @@ class WebSearchInterceptionLogger(CustomLogger): Args: enabled_providers: List of LLM providers to enable interception for. Use LlmProviders enum values (e.g., [LlmProviders.BEDROCK]) - Default: [LlmProviders.BEDROCK] + If None or empty list, enables for ALL providers. + Default: None (all providers enabled) search_tool_name: Name of search tool configured in router's search_tools. If None, will attempt to use first available search tool. """ @@ -183,10 +185,10 @@ class WebSearchInterceptionLogger(CustomLogger): verbose_logger.debug( f"WebSearchInterception: Pre-request hook called" f" - custom_llm_provider={custom_llm_provider}" - f" - enabled_providers={self.enabled_providers}" + f" - enabled_providers={self.enabled_providers or 'ALL'}" ) - if custom_llm_provider not in self.enabled_providers: + if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: verbose_logger.debug( f"WebSearchInterception: Skipping - provider {custom_llm_provider} not in {self.enabled_providers}" ) @@ -245,7 +247,12 @@ class WebSearchInterceptionLogger(CustomLogger): custom_llm_provider: str, kwargs: Dict, ) -> Tuple[bool, Dict]: - """Check if WebSearch tool interception is needed""" + """ + Check if WebSearch tool interception is needed for Anthropic Messages API. + + This is the legacy method for Anthropic-style responses. + For chat completions, use async_should_run_chat_completion_agentic_loop instead. + """ verbose_logger.debug(f"WebSearchInterception: Hook called! provider={custom_llm_provider}, stream={stream}") verbose_logger.debug(f"WebSearchInterception: Response type: {type(response)}") @@ -253,7 +260,7 @@ class WebSearchInterceptionLogger(CustomLogger): # Check if provider should be intercepted # Note: custom_llm_provider is already normalized by get_llm_provider() # (e.g., "bedrock/invoke/..." -> "bedrock") - if custom_llm_provider not in self.enabled_providers: + if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: verbose_logger.debug( f"WebSearchInterception: Skipping provider {custom_llm_provider} (not in enabled list: {self.enabled_providers})" ) @@ -267,10 +274,11 @@ class WebSearchInterceptionLogger(CustomLogger): ) return False, {} - # Detect WebSearch tool_use in response + # Detect WebSearch tool_use in response (Anthropic format) should_intercept, tool_calls = WebSearchTransformation.transform_request( response=response, stream=stream, + response_format="anthropic", ) if not should_intercept: @@ -288,6 +296,67 @@ class WebSearchInterceptionLogger(CustomLogger): "tool_calls": tool_calls, "tool_type": "websearch", "provider": custom_llm_provider, + "response_format": "anthropic", + } + return True, tools_dict + + async def async_should_run_chat_completion_agentic_loop( + self, + response: Any, + model: str, + messages: List[Dict], + tools: Optional[List[Dict]], + stream: bool, + custom_llm_provider: str, + kwargs: Dict, + ) -> Tuple[bool, Dict]: + """ + Check if WebSearch tool interception is needed for Chat Completions API. + + Similar to async_should_run_agentic_loop but for OpenAI-style chat completions. + """ + + verbose_logger.debug(f"WebSearchInterception: Chat completion hook called! provider={custom_llm_provider}, stream={stream}") + verbose_logger.debug(f"WebSearchInterception: Response type: {type(response)}") + + # Check if provider should be intercepted + if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: + verbose_logger.debug( + f"WebSearchInterception: Skipping provider {custom_llm_provider} (not in enabled list: {self.enabled_providers})" + ) + return False, {} + + # Check if tools include any web search tool (strict check for chat completions) + has_websearch_tool = any(is_web_search_tool_chat_completion(t) for t in (tools or [])) + if not has_websearch_tool: + verbose_logger.debug( + "WebSearchInterception: No litellm_web_search tool in request" + ) + return False, {} + + # Detect WebSearch tool_calls in response (OpenAI format) + should_intercept, tool_calls = WebSearchTransformation.transform_request( + response=response, + stream=stream, + response_format="openai", + ) + + if not should_intercept: + verbose_logger.debug( + "WebSearchInterception: No WebSearch tool_calls detected in response" + ) + return False, {} + + verbose_logger.debug( + f"WebSearchInterception: Detected {len(tool_calls)} WebSearch tool call(s), executing agentic loop" + ) + + # Return tools dict with tool calls + tools_dict = { + "tool_calls": tool_calls, + "tool_type": "websearch", + "provider": custom_llm_provider, + "response_format": "openai", } return True, tools_dict @@ -303,7 +372,11 @@ class WebSearchInterceptionLogger(CustomLogger): stream: bool, kwargs: Dict, ) -> Any: - """Execute agentic loop with WebSearch execution""" + """ + Execute agentic loop with WebSearch execution for Anthropic Messages API. + + This is the legacy method for Anthropic-style responses. + """ tool_calls = tools["tool_calls"] @@ -321,6 +394,41 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs=kwargs, ) + async def async_run_chat_completion_agentic_loop( + self, + tools: Dict, + model: str, + messages: List[Dict], + response: Any, + optional_params: Dict, + logging_obj: Any, + stream: bool, + kwargs: Dict, + ) -> Any: + """ + Execute agentic loop with WebSearch execution for Chat Completions API. + + Similar to async_run_agentic_loop but for OpenAI-style chat completions. + """ + + tool_calls = tools["tool_calls"] + response_format = tools.get("response_format", "openai") + + verbose_logger.debug( + f"WebSearchInterception: Executing chat completion agentic loop for {len(tool_calls)} search(es)" + ) + + return await self._execute_chat_completion_agentic_loop( + model=model, + messages=messages, + tool_calls=tool_calls, + optional_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=kwargs, + response_format=response_format, + ) + async def _execute_agentic_loop( self, model: str, @@ -382,7 +490,8 @@ class WebSearchInterceptionLogger(CustomLogger): ) # Make follow-up request with search results - follow_up_messages = messages + [assistant_message, user_message] + # Type cast: user_message is a Dict for Anthropic format (default response_format) + follow_up_messages = messages + [assistant_message, cast(Dict, user_message)] verbose_logger.debug( "WebSearchInterception: Making follow-up request with search results" @@ -521,6 +630,150 @@ class WebSearchInterceptionLogger(CustomLogger): ) raise + async def _execute_chat_completion_agentic_loop( # noqa: PLR0915 + self, + model: str, + messages: List[Dict], + tool_calls: List[Dict], + optional_params: Dict, + logging_obj: Any, + stream: bool, + kwargs: Dict, + response_format: str = "openai", + ) -> Any: + """Execute litellm.search() and make follow-up chat completion request""" + + # Extract search queries from tool_calls + search_tasks = [] + for tool_call in tool_calls: + # Handle both Anthropic-style input and OpenAI-style function.arguments + query = None + if "input" in tool_call and isinstance(tool_call["input"], dict): + query = tool_call["input"].get("query") + elif "function" in tool_call: + func = tool_call["function"] + if isinstance(func, dict): + args = func.get("arguments", {}) + if isinstance(args, dict): + query = args.get("query") + + if query: + verbose_logger.debug( + f"WebSearchInterception: Queuing search for query='{query}'" + ) + search_tasks.append(self._execute_search(query)) + else: + verbose_logger.warning( + f"WebSearchInterception: Tool call {tool_call.get('id')} has no query" + ) + # Add empty result for tools without query + search_tasks.append(self._create_empty_search_result()) + + # Execute searches in parallel + verbose_logger.debug( + f"WebSearchInterception: Executing {len(search_tasks)} search(es) in parallel" + ) + search_results = await asyncio.gather(*search_tasks, return_exceptions=True) + + # Handle any exceptions in search results + final_search_results: List[str] = [] + for i, result in enumerate(search_results): + if isinstance(result, Exception): + verbose_logger.error( + f"WebSearchInterception: Search {i} failed with error: {str(result)}" + ) + final_search_results.append( + f"Search failed: {str(result)}" + ) + elif isinstance(result, str): + final_search_results.append(cast(str, result)) + else: + verbose_logger.warning( + f"WebSearchInterception: Unexpected result type {type(result)} at index {i}" + ) + final_search_results.append(str(result)) + + # Build assistant and tool messages using transformation + assistant_message, tool_messages_or_user = WebSearchTransformation.transform_response( + tool_calls=tool_calls, + search_results=final_search_results, + response_format=response_format, + ) + + # Make follow-up request with search results + # For OpenAI format, tool_messages_or_user is a list of tool messages + if response_format == "openai": + follow_up_messages = messages + [assistant_message] + cast(List[Dict], tool_messages_or_user) + else: + # For Anthropic format (shouldn't happen in this method, but handle it) + follow_up_messages = messages + [assistant_message, cast(Dict, tool_messages_or_user)] + + verbose_logger.debug( + "WebSearchInterception: Making follow-up chat completion request with search results" + ) + verbose_logger.debug( + f"WebSearchInterception: Follow-up messages count: {len(follow_up_messages)}" + ) + + # Use litellm.acompletion for follow-up request + try: + # Remove internal parameters that shouldn't be passed to follow-up request + internal_params = { + '_websearch_interception', + 'acompletion', + 'litellm_logging_obj', + 'custom_llm_provider', + 'model_alias_map', + 'stream_response', + 'custom_prompt_dict', + } + kwargs_for_followup = { + k: v for k, v in kwargs.items() + if not k.startswith('_websearch_interception') and k not in internal_params + } + + # Get full model name from kwargs + full_model_name = model + if "custom_llm_provider" in kwargs: + custom_llm_provider = kwargs["custom_llm_provider"] + # Reconstruct full model name with provider prefix if needed + if not model.startswith(custom_llm_provider): + # Check if model already has a provider prefix + if "/" not in model: + full_model_name = f"{custom_llm_provider}/{model}" + + verbose_logger.debug( + f"WebSearchInterception: Using model name: {full_model_name}" + ) + + # Prepare tools for follow-up request (same as original) + tools_param = optional_params.get("tools") + + # Remove tools and extra_body from optional_params to avoid issues + # extra_body often contains internal LiteLLM params that shouldn't be forwarded + optional_params_clean = { + k: v for k, v in optional_params.items() + if k not in {"tools", "extra_body", "model_alias_map","stream_response", "custom_prompt_dict" } + } + + final_response = await litellm.acompletion( + model=full_model_name, + messages=follow_up_messages, + tools=tools_param, + **optional_params_clean, + **kwargs_for_followup, + ) + + verbose_logger.debug( + f"WebSearchInterception: Follow-up request completed, response type: {type(final_response)}" + ) + return final_response + except Exception as e: + verbose_logger.exception( + f"WebSearchInterception: Follow-up request failed: {str(e)}" + ) + raise + async def _create_empty_search_result(self) -> str: """Create an empty search result for tool calls without queries""" return "No search query provided" diff --git a/litellm/integrations/websearch_interception/tools.py b/litellm/integrations/websearch_interception/tools.py index 4f8b7372fe3..c39d150fb19 100644 --- a/litellm/integrations/websearch_interception/tools.py +++ b/litellm/integrations/websearch_interception/tools.py @@ -49,12 +49,57 @@ def get_litellm_web_search_tool() -> Dict[str, Any]: } +def is_web_search_tool_chat_completion(tool: Dict[str, Any]) -> bool: + """ + Check if a tool is a web search tool for Chat Completions API (strict check). + + This is a stricter version that ONLY checks for the exact LiteLLM web search tool name. + Use this for Chat Completions API to avoid false positives with user-defined tools. + + Detects ONLY: + - LiteLLM standard: name == "litellm_web_search" (Anthropic format) + - OpenAI format: type == "function" with function.name == "litellm_web_search" + + Args: + tool: Tool dictionary to check + + Returns: + True if tool is exactly the LiteLLM web search tool + + Example: + >>> is_web_search_tool_chat_completion({"name": "litellm_web_search"}) + True + >>> is_web_search_tool_chat_completion({"type": "function", "function": {"name": "litellm_web_search"}}) + True + >>> is_web_search_tool_chat_completion({"name": "web_search"}) + False + >>> is_web_search_tool_chat_completion({"name": "WebSearch"}) + False + """ + tool_name = tool.get("name", "") + tool_type = tool.get("type", "") + + # Check for OpenAI format: {"type": "function", "function": {"name": "litellm_web_search"}} + if tool_type == "function" and "function" in tool: + function_def = tool.get("function", {}) + function_name = function_def.get("name", "") + if function_name == LITELLM_WEB_SEARCH_TOOL_NAME: + return True + + # Check for LiteLLM standard tool (Anthropic format) + if tool_name == LITELLM_WEB_SEARCH_TOOL_NAME: + return True + + return False + + def is_web_search_tool(tool: Dict[str, Any]) -> bool: """ Check if a tool is a web search tool (native or LiteLLM standard). Detects: - LiteLLM standard: name == "litellm_web_search" + - OpenAI format: type == "function" with function.name == "litellm_web_search" - Anthropic native: type starts with "web_search_" (e.g., "web_search_20250305") - Claude Code: name == "web_search" with a type field - Custom: name == "WebSearch" (legacy format) @@ -68,6 +113,8 @@ def is_web_search_tool(tool: Dict[str, Any]) -> bool: Example: >>> is_web_search_tool({"name": "litellm_web_search"}) True + >>> is_web_search_tool({"type": "function", "function": {"name": "litellm_web_search"}}) + True >>> is_web_search_tool({"type": "web_search_20250305", "name": "web_search"}) True >>> is_web_search_tool({"name": "calculator"}) @@ -75,8 +122,15 @@ def is_web_search_tool(tool: Dict[str, Any]) -> bool: """ tool_name = tool.get("name", "") tool_type = tool.get("type", "") + + # Check for OpenAI format: {"type": "function", "function": {"name": "..."}} + if tool_type == "function" and "function" in tool: + function_def = tool.get("function", {}) + function_name = function_def.get("name", "") + if function_name == LITELLM_WEB_SEARCH_TOOL_NAME: + return True - # Check for LiteLLM standard tool + # Check for LiteLLM standard tool (Anthropic format) if tool_name == LITELLM_WEB_SEARCH_TOOL_NAME: return True diff --git a/litellm/integrations/websearch_interception/transformation.py b/litellm/integrations/websearch_interception/transformation.py index 313358822a5..e44ec35c3a2 100644 --- a/litellm/integrations/websearch_interception/transformation.py +++ b/litellm/integrations/websearch_interception/transformation.py @@ -1,10 +1,10 @@ """ WebSearch Tool Transformation -Transforms between Anthropic tool_use format and LiteLLM search format. +Transforms between Anthropic/OpenAI tool_use format and LiteLLM search format. """ - -from typing import Any, Dict, List, Tuple +import json +from typing import Any, Dict, List, Tuple, Union from litellm._logging import verbose_logger from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME @@ -17,28 +17,31 @@ class WebSearchTransformation: Handles transformation between: - Anthropic tool_use format → LiteLLM search requests - - LiteLLM SearchResponse → Anthropic tool_result format + - OpenAI tool_calls format → LiteLLM search requests + - LiteLLM SearchResponse → Anthropic/OpenAI tool_result format """ @staticmethod def transform_request( response: Any, stream: bool, + response_format: str = "anthropic", ) -> Tuple[bool, List[Dict]]: """ - Transform Anthropic response to extract WebSearch tool calls. + Transform model response to extract WebSearch tool calls. - Detects if response contains WebSearch tool_use blocks and extracts + Detects if response contains WebSearch tool_use/tool_calls blocks and extracts the search queries for execution. Args: - response: Model response (dict or AnthropicMessagesResponse) + response: Model response (dict, AnthropicMessagesResponse, or ModelResponse) stream: Whether response is streaming + response_format: Response format - "anthropic" or "openai" (default: "anthropic") Returns: (has_websearch, tool_calls): has_websearch: True if WebSearch tool_use found - tool_calls: List of tool_use dicts with id, name, input + tool_calls: List of tool_use/tool_calls dicts with id, name, input/function Note: Streaming requests are handled by converting stream=True to stream=False @@ -54,8 +57,11 @@ class WebSearchTransformation: ) return False, [] - # Parse non-streaming response - return WebSearchTransformation._detect_from_non_streaming_response(response) + # Parse non-streaming response based on format + if response_format == "openai": + return WebSearchTransformation._detect_from_openai_response(response) + else: + return WebSearchTransformation._detect_from_non_streaming_response(response) @staticmethod def _detect_from_non_streaming_response( @@ -114,26 +120,142 @@ class WebSearchTransformation: return len(tool_calls) > 0, tool_calls + @staticmethod + def _detect_from_openai_response( + response: Any, + ) -> Tuple[bool, List[Dict]]: + """Parse OpenAI-style response for WebSearch tool_calls""" + + # Handle both dict and ModelResponse objects + if isinstance(response, dict): + choices = response.get("choices", []) + else: + if not hasattr(response, "choices"): + verbose_logger.debug( + "WebSearchInterception: Response has no choices attribute" + ) + return False, [] + choices = response.choices or [] + + if not choices: + verbose_logger.debug( + "WebSearchInterception: Response has empty choices" + ) + return False, [] + + # Get first choice's message + first_choice = choices[0] + if isinstance(first_choice, dict): + message = first_choice.get("message", {}) + else: + message = getattr(first_choice, "message", None) + + if not message: + verbose_logger.debug( + "WebSearchInterception: First choice has no message" + ) + return False, [] + + # Get tool_calls from message + if isinstance(message, dict): + openai_tool_calls = message.get("tool_calls", []) + else: + openai_tool_calls = getattr(message, "tool_calls", None) or [] + + if not openai_tool_calls: + verbose_logger.debug( + "WebSearchInterception: Message has no tool_calls" + ) + return False, [] + + # Find all WebSearch tool calls + tool_calls = [] + for tool_call in openai_tool_calls: + # Handle both dict and object tool calls + if isinstance(tool_call, dict): + tool_id = tool_call.get("id") + tool_type = tool_call.get("type") + function = tool_call.get("function", {}) + function_name = function.get("name") if isinstance(function, dict) else getattr(function, "name", None) + function_arguments = function.get("arguments") if isinstance(function, dict) else getattr(function, "arguments", None) + else: + tool_id = getattr(tool_call, "id", None) + tool_type = getattr(tool_call, "type", None) + function = getattr(tool_call, "function", None) + function_name = getattr(function, "name", None) if function else None + function_arguments = getattr(function, "arguments", None) if function else None + + # Check for LiteLLM standard or legacy web search tools + if tool_type == "function" and function_name in ( + LITELLM_WEB_SEARCH_TOOL_NAME, "WebSearch", "web_search" + ): + # Parse arguments (might be JSON string) + if isinstance(function_arguments, str): + try: + arguments = json.loads(function_arguments) + except json.JSONDecodeError: + verbose_logger.warning( + f"WebSearchInterception: Failed to parse function arguments: {function_arguments}" + ) + arguments = {} + else: + arguments = function_arguments or {} + + # Convert to internal format (similar to Anthropic) + tool_call_dict = { + "id": tool_id, + "type": "function", + "name": function_name, + "function": { + "name": function_name, + "arguments": arguments, + }, + "input": arguments, # For compatibility with Anthropic format + } + tool_calls.append(tool_call_dict) + verbose_logger.debug( + f"WebSearchInterception: Found {function_name} tool_call with id={tool_id}" + ) + + return len(tool_calls) > 0, tool_calls + @staticmethod def transform_response( tool_calls: List[Dict], search_results: List[str], - ) -> Tuple[Dict, Dict]: + response_format: str = "anthropic", + ) -> Tuple[Dict, Union[Dict, List[Dict]]]: """ - Transform LiteLLM search results to Anthropic tool_result format. + Transform LiteLLM search results to Anthropic/OpenAI tool_result format. - Builds the assistant and user messages needed for the agentic loop + Builds the assistant and user/tool messages needed for the agentic loop follow-up request. Args: - tool_calls: List of tool_use dicts from transform_request + tool_calls: List of tool_use/tool_calls dicts from transform_request search_results: List of search result strings (one per tool_call) + response_format: Response format - "anthropic" or "openai" (default: "anthropic") Returns: - (assistant_message, user_message): - assistant_message: Message with tool_use blocks - user_message: Message with tool_result blocks + (assistant_message, user_or_tool_messages): + For Anthropic: assistant_message with tool_use blocks, user_message with tool_result blocks + For OpenAI: assistant_message with tool_calls, tool_messages list with tool results """ + if response_format == "openai": + return WebSearchTransformation._transform_response_openai( + tool_calls, search_results + ) + else: + return WebSearchTransformation._transform_response_anthropic( + tool_calls, search_results + ) + + @staticmethod + def _transform_response_anthropic( + tool_calls: List[Dict], + search_results: List[str], + ) -> Tuple[Dict, Dict]: + """Transform to Anthropic format (single user message with tool_result blocks)""" # Build assistant message with tool_use blocks assistant_message = { "role": "assistant", @@ -163,6 +285,40 @@ class WebSearchTransformation: return assistant_message, user_message + @staticmethod + def _transform_response_openai( + tool_calls: List[Dict], + search_results: List[str], + ) -> Tuple[Dict, List[Dict]]: + """Transform to OpenAI format (assistant with tool_calls, separate tool messages)""" + # Build assistant message with tool_calls + assistant_message = { + "role": "assistant", + "tool_calls": [ + { + "id": tc["id"], + "type": "function", + "function": { + "name": tc["name"], + "arguments": json.dumps(tc["input"]) if isinstance(tc["input"], dict) else str(tc["input"]), + }, + } + for tc in tool_calls + ], + } + + # Build separate tool messages (one per tool call) + tool_messages = [ + { + "role": "tool", + "tool_call_id": tool_calls[i]["id"], + "content": search_results[i], + } + for i in range(len(tool_calls)) + ] + + return assistant_message, tool_messages + @staticmethod def format_search_response(result: SearchResponse) -> str: """ diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d2ea7e872a2..95db8ec64b3 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -302,7 +302,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, signed_json_body=signed_json_body, ) - return provider_config.transform_response( + initial_response = provider_config.transform_response( model=model, raw_response=response, model_response=model_response, @@ -316,6 +316,20 @@ class BaseLLMHTTPHandler: json_mode=json_mode, ) + # Call agentic chat completion hooks + final_response = await self._call_agentic_chat_completion_hooks( + response=initial_response, + model=model, + messages=messages, + optional_params=optional_params, + logging_obj=logging_obj, + stream=False, + custom_llm_provider=custom_llm_provider, + kwargs=litellm_params, + ) + + return final_response if final_response is not None else initial_response + def completion( self, model: str, @@ -412,6 +426,11 @@ class BaseLLMHTTPHandler: }, ) + # Check if stream was converted for WebSearch interception + # This is set by the async_pre_request_hook in WebSearchInterceptionLogger + if litellm_params.get("_websearch_interception_converted_stream", False): + logging_obj.model_call_details["websearch_interception_converted_stream"] = True + if acompletion is True: if stream is True: data = self._add_stream_param_to_request_body( @@ -4361,10 +4380,10 @@ class BaseLLMHTTPHandler: kwargs: Dict, ) -> Optional[Any]: """ - Call agentic completion hooks for all custom loggers. + Call agentic completion hooks for all custom loggers (Anthropic Messages API). - 1. Call async_should_run_agentic_completion to check if agentic loop is needed - 2. If yes, call async_run_agentic_completion to execute the loop + 1. Call async_should_run_agentic_loop to check if agentic loop is needed + 2. If yes, call async_run_agentic_loop to execute the loop Returns the response from agentic loop, or None if no hook runs. """ @@ -4453,6 +4472,105 @@ class BaseLLMHTTPHandler: return None + async def _call_agentic_chat_completion_hooks( + self, + response: Any, + model: str, + messages: List[Dict], + optional_params: Dict, + logging_obj: "LiteLLMLoggingObj", + stream: bool, + custom_llm_provider: str, + kwargs: Dict, + ) -> Optional[Any]: + """ + Call agentic chat completion hooks for all custom loggers (Chat Completions API). + + 1. Call async_should_run_chat_completion_agentic_loop to check if agentic loop is needed + 2. If yes, call async_run_chat_completion_agentic_loop to execute the loop + + Returns the response from agentic loop, or None if no hook runs. + """ + from litellm._logging import verbose_logger + from litellm.integrations.custom_logger import CustomLogger + + callbacks = litellm.callbacks + ( + logging_obj.dynamic_success_callbacks or [] + ) + tools = optional_params.get("tools", []) + + for callback in callbacks: + try: + if isinstance(callback, CustomLogger): + # Check if callback has the chat completion agentic loop method + if not hasattr(callback, "async_should_run_chat_completion_agentic_loop"): + continue + + # First: Check if agentic loop should run + should_run, tool_calls = ( + await callback.async_should_run_chat_completion_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + ) + ) + + if should_run: + # Second: Execute agentic loop + # Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name + kwargs_with_provider = kwargs.copy() if kwargs else {} + kwargs_with_provider["custom_llm_provider"] = custom_llm_provider + agentic_response = await callback.async_run_chat_completion_agentic_loop( + tools=tool_calls, + model=model, + messages=messages, + response=response, + optional_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=kwargs_with_provider, + ) + # First hook that runs agentic loop wins + return agentic_response + + except Exception as e: + verbose_logger.exception( + f"LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: {str(e)}" + ) + + # Check if we need to convert response to fake stream for chat completions + # This happens when: + # 1. Stream was originally True but converted to False for WebSearch interception + # 2. No agentic loop ran (LLM didn't use the tool) + # 3. We have a non-streaming response that needs to be converted to streaming + websearch_converted_stream = ( + logging_obj.model_call_details.get("websearch_interception_converted_stream", False) + if logging_obj is not None + else False + ) + + if websearch_converted_stream: + from litellm._logging import verbose_logger + from litellm.llms.base_llm.base_model_iterator import ( + convert_model_response_to_streaming, + ) + + verbose_logger.debug( + "WebSearchInterception: No tool call made, converting non-streaming chat completion to fake stream" + ) + + # Convert the non-streaming ModelResponse to a fake stream + if hasattr(response, "choices"): + # Use the existing converter for ModelResponse + fake_stream = convert_model_response_to_streaming(response) + return fake_stream + + return None + def _handle_error( self, e: Exception, diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 8a8070240da..da87852dff5 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -501,6 +501,88 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): else: raise e + async def _call_agentic_completion_hooks_openai( + self, + response: Any, + model: str, + messages: List[Dict], + optional_params: Dict, + logging_obj: LiteLLMLoggingObj, + stream: bool, + litellm_params: Dict, + ) -> Optional[Any]: + """ + Call agentic completion hooks for all custom loggers (OpenAI Chat Completions API). + + 1. Call async_should_run_chat_completion_agentic_loop to check if agentic loop is needed + 2. If yes, call async_run_chat_completion_agentic_loop to execute the loop + + Returns the response from agentic loop, or None if no hook runs. + """ + from litellm._logging import verbose_logger + from litellm.integrations.custom_logger import CustomLogger + + callbacks = litellm.callbacks + ( + logging_obj.dynamic_success_callbacks or [] + ) + # Avoid logging full callback objects to prevent leaking sensitive data + verbose_logger.debug( + "LiteLLM.AgenticHooks: callbacks_count=%s", len(callbacks) + ) + tools = optional_params.get("tools", []) + # Avoid logging full tools payloads; they may contain sensitive parameters + verbose_logger.debug( + "LiteLLM.AgenticHooks: tools_count=%s", len(tools) if isinstance(tools, list) else 1 if tools else 0 + ) + # Get custom_llm_provider from litellm_params + custom_llm_provider = litellm_params.get("custom_llm_provider", "openai") + + for callback in callbacks: + try: + if isinstance(callback, CustomLogger): + # Check if the callback has the chat completion agentic loop methods + if not hasattr(callback, 'async_should_run_chat_completion_agentic_loop'): + continue + + # First: Check if agentic loop should run (using chat completion method) + should_run, tool_calls = ( + await callback.async_should_run_chat_completion_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=litellm_params, + ) + ) + + if should_run: + # Second: Execute agentic loop + kwargs_with_provider = litellm_params.copy() if litellm_params else {} + kwargs_with_provider["custom_llm_provider"] = custom_llm_provider + + # For OpenAI Chat Completions, use the chat completion agentic loop method + agentic_response = await callback.async_run_chat_completion_agentic_loop( + tools=tool_calls, + model=model, + messages=messages, + response=response, + optional_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=kwargs_with_provider, + ) + # First hook that runs agentic loop wins + return agentic_response + + except Exception as e: + verbose_logger.exception( + f"LiteLLM.AgenticHookError: Exception in agentic completion hooks for OpenAI: {str(e)}" + ) + + return None + def mock_streaming( self, response: ModelResponse, @@ -844,7 +926,6 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj=logging_obj, ) stringified_response = response.model_dump() - logging_obj.post_call( input=data["messages"], api_key=api_key, @@ -859,6 +940,20 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): _response_headers=headers, ) + # Call agentic completion hooks (e.g., for websearch_interception) + agentic_response = await self._call_agentic_completion_hooks_openai( + response=final_response_obj, + model=model, + messages=messages, + optional_params=optional_params, + logging_obj=logging_obj, + stream=False, + litellm_params=litellm_params, + ) + + if agentic_response is not None: + final_response_obj = agentic_response + if fake_stream is True: return self.mock_streaming( response=cast(ModelResponse, final_response_obj), diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py new file mode 100644 index 00000000000..1b53633484d --- /dev/null +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py @@ -0,0 +1,398 @@ +""" +Integration tests for WebSearch interception with chat completions API. + +Tests the end-to-end flow of websearch_interception callback with +litellm.acompletion() for transparent server-side web search execution. +""" +import os +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import litellm +from litellm.integrations.websearch_interception.handler import ( + WebSearchInterceptionLogger, +) +from litellm.types.utils import LlmProviders, ModelResponse + + +@pytest.fixture +def mock_search_response(): + """Mock search response from litellm.asearch()""" + mock_response = MagicMock() + mock_response.results = [ + MagicMock( + title="Weather in San Francisco", + url="https://weather.com/sf", + snippet="Current weather: 65°F, partly cloudy", + ) + ] + return mock_response + + +@pytest.fixture +def websearch_logger(): + """Create a WebSearchInterceptionLogger instance""" + return WebSearchInterceptionLogger( + enabled_providers=[LlmProviders.OPENAI, LlmProviders.MINIMAX] + ) + + +@pytest.mark.asyncio +@pytest.mark.skipif( + os.environ.get("OPENAI_API_KEY") is None, + reason="OPENAI_API_KEY not set", +) +async def test_websearch_chat_completion_with_openai(): + """Test websearch interception with OpenAI chat completions API. + + This test verifies that: + 1. Model calls litellm_web_search tool + 2. Server executes web search automatically + 3. Server makes follow-up request with search results + 4. User gets final answer without tool_calls + """ + # Configure WebSearch interception + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + websearch_logger = WebSearchInterceptionLogger( + enabled_providers=[LlmProviders.OPENAI] + ) + litellm.callbacks = [websearch_logger] + + try: + response = await litellm.acompletion( + model="gpt-4o-mini", # Use cheaper model for testing + messages=[ + {"role": "user", "content": "What's the weather in San Francisco today?"} + ], + tools=[ + { + "type": "function", + "function": { + "name": "litellm_web_search", + "description": "Search the web for information", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query", + } + }, + "required": ["query"], + }, + }, + } + ], + ) + + # Verify response structure + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content is not None + assert len(response.choices[0].message.content) > 0 + + # If agentic loop worked, we should NOT have tool_calls in final response + # (they should have been executed and replaced with final answer) + if hasattr(response.choices[0].message, "tool_calls"): + # If tool_calls exist, it means agentic loop didn't run + # This could happen if search tool is not configured + pytest.skip( + "Agentic loop did not execute - search tool may not be configured" + ) + + # Verify we got a meaningful response + assert response.choices[0].finish_reason in ["stop", "end_turn"] + + finally: + # Restore original callbacks + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_websearch_chat_completion_hook_detection(): + """Test that websearch hook correctly detects tool calls in response.""" + from litellm.types.utils import ( + ChatCompletionMessageToolCall, + Choices, + Function, + Message, + ) + + websearch_logger = WebSearchInterceptionLogger( + enabled_providers=[LlmProviders.OPENAI] + ) + + # Mock response with litellm_web_search tool call + mock_response = ModelResponse( + id="test-123", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_123", + type="function", + function=Function( + name="litellm_web_search", + arguments='{"query": "weather in SF"}', + ), + ) + ], + ) + ) + ], + model="gpt-4o", + object="chat.completion", + created=1234567890, + ) + + # Test should_run_chat_completion_agentic_loop + should_run, tools_dict = ( + await websearch_logger.async_should_run_chat_completion_agentic_loop( + response=mock_response, + model="gpt-4o", + messages=[{"role": "user", "content": "What's the weather?"}], + tools=[ + { + "type": "function", + "function": {"name": "litellm_web_search"}, + } + ], + stream=False, + custom_llm_provider="openai", + kwargs={}, + ) + ) + + # Verify hook detected the tool call + assert should_run is True + assert "tool_calls" in tools_dict + assert len(tools_dict["tool_calls"]) == 1 + assert tools_dict["tool_calls"][0]["name"] == "litellm_web_search" + assert tools_dict["response_format"] == "openai" + + +@pytest.mark.asyncio +async def test_websearch_not_triggered_without_tool(): + """Test that websearch hook is NOT triggered when no web search tool in request.""" + from litellm.types.utils import Choices, Message + + websearch_logger = WebSearchInterceptionLogger( + enabled_providers=[LlmProviders.OPENAI] + ) + + mock_response = ModelResponse( + id="test-123", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + role="assistant", + content="Here's the answer", + tool_calls=None, + ) + ) + ], + model="gpt-4o", + object="chat.completion", + created=1234567890, + ) + + # Test without web search tool + should_run, tools_dict = ( + await websearch_logger.async_should_run_chat_completion_agentic_loop( + response=mock_response, + model="gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + tools=[ + { + "type": "function", + "function": {"name": "some_other_tool"}, + } + ], + stream=False, + custom_llm_provider="openai", + kwargs={}, + ) + ) + + # Verify hook did NOT trigger + assert should_run is False + assert tools_dict == {} + + +@pytest.mark.asyncio +async def test_websearch_not_triggered_for_disabled_provider(): + """Test that websearch hook is NOT triggered for providers not in enabled_providers.""" + from litellm.types.utils import ( + ChatCompletionMessageToolCall, + Choices, + Function, + Message, + ) + + # Only enable bedrock + websearch_logger = WebSearchInterceptionLogger( + enabled_providers=[LlmProviders.BEDROCK] + ) + + mock_response = ModelResponse( + id="test-123", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_123", + type="function", + function=Function( + name="litellm_web_search", + arguments='{"query": "test"}', + ), + ) + ], + ) + ) + ], + model="gpt-4o", + object="chat.completion", + created=1234567890, + ) + + # Test with OpenAI provider (not enabled) + should_run, tools_dict = ( + await websearch_logger.async_should_run_chat_completion_agentic_loop( + response=mock_response, + model="gpt-4o", + messages=[{"role": "user", "content": "test"}], + tools=[ + { + "type": "function", + "function": {"name": "litellm_web_search"}, + } + ], + stream=False, + custom_llm_provider="openai", # Not in enabled_providers + kwargs={}, + ) + ) + + # Verify hook did NOT trigger + assert should_run is False + assert tools_dict == {} + + +@pytest.mark.asyncio +async def test_websearch_json_serialization_fix(): + """Test that tool call arguments are properly JSON serialized. + + Regression test for the bug where arguments were converted to Python + string representation instead of proper JSON, causing providers like + MiniMax to reject requests with 'invalid function arguments json string'. + """ + from litellm.integrations.websearch_interception.transformation import ( + WebSearchTransformation, + ) + + # Mock tool calls with dict input + tool_calls = [ + { + "id": "call_123", + "name": "litellm_web_search", + "input": {"query": "weather in SF"}, # Dict input + } + ] + + search_results = ["Weather: 65°F, partly cloudy"] + + # Transform to OpenAI format + assistant_message, tool_messages = WebSearchTransformation.transform_response( + tool_calls=tool_calls, + search_results=search_results, + response_format="openai", + ) + + # Verify arguments are properly JSON serialized + import json + + arguments_str = assistant_message["tool_calls"][0]["function"]["arguments"] + + # Should be valid JSON + parsed_args = json.loads(arguments_str) + assert parsed_args == {"query": "weather in SF"} + + # Should NOT be Python string representation like "{'query': 'weather in SF'}" + assert arguments_str == '{"query": "weather in SF"}' + assert arguments_str != "{'query': 'weather in SF'}" + + +@pytest.mark.asyncio +@pytest.mark.skipif( + os.environ.get("OPENAI_API_KEY") is None + or os.environ.get("PERPLEXITY_API_KEY") is None, + reason="OPENAI_API_KEY or PERPLEXITY_API_KEY not set", +) +async def test_websearch_streaming_conversion(): + """Test that streaming requests are converted to non-streaming for web search. + + When stream=True is passed with web search tools, the handler should: + 1. Convert stream=True to stream=False for initial request + 2. Execute web search + 3. Convert final response back to streaming + """ + websearch_logger = WebSearchInterceptionLogger( + enabled_providers=[LlmProviders.OPENAI], search_tool_name="perplexity-search" + ) + litellm.callbacks = [websearch_logger] + + try: + response = await litellm.acompletion( + model="gpt-4o-mini", + messages=[ + {"role": "user", "content": "What's the latest AI news?"} + ], + tools=[ + { + "type": "function", + "function": { + "name": "litellm_web_search", + "description": "Search the web", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + }, + } + ], + stream=True, + ) + + # Response should be a streaming iterator + chunks = [] + async for chunk in response: + chunks.append(chunk) + + # Verify we got streaming chunks + assert len(chunks) > 0 + + # Verify chunks have expected structure + for chunk in chunks: + assert hasattr(chunk, "choices") + assert len(chunk.choices) > 0 + + finally: + litellm.callbacks = [] + + +if __name__ == "__main__": + # Run with: pytest test_websearch_chat_completion.py -v -s + pytest.main([__file__, "-v", "-s"])