tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-15 04:32:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #9465 from BerriAI/litellm_add_openai_web_search

Browse Source 
[Feat] Add OpenAI Web Search Tool Call Support - Initial support
This commit is contained in:
Ishaan Jaff
2025-03-22 13:45:42 -07:00
committed by GitHub
parent 6bc6859224 cae75fe46b
commit d468a2fa34
10 changed files with 298 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py
+1
View File
@@ -494,6 +494,7 @@ def convert_to_model_response_object( # noqa: PLR0915
provider_specific_fields=provider_specific_fields,
reasoning_content=reasoning_content,
thinking_blocks=thinking_blocks,
annotations=choice["message"].get("annotations", None),
)
finish_reason = choice.get("finish_reason", None)
if finish_reason is None:
litellm/litellm_core_utils/streaming_handler.py
+4 -1
View File
@@ -799,6 +799,10 @@ class CustomStreamWrapper:
"provider_specific_fields" in response_obj
and response_obj["provider_specific_fields"] is not None
)
or (
"annotations" in model_response.choices[0].delta
and model_response.choices[0].delta.annotations is not None
)
):
return True
else:
@@ -939,7 +943,6 @@ class CustomStreamWrapper:
and model_response.choices[0].delta.audio is not None
):
return model_response
else:
if hasattr(model_response, "usage"):
self.chunks.append(model_response)
litellm/model_prices_and_context_window_backup.json
+62 -1
View File
@@ -15,6 +15,12 @@
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.0000,
"search_context_size_medium": 0.0000,
"search_context_size_high": 0.0000
},
"deprecation_date": "date when the model becomes deprecated in the format YYYY-MM-DD"
},
"omni-moderation-latest": {
@@ -59,6 +65,31 @@
},
"gpt-4o": {
"max_tokens": 16384,
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"input_cost_per_token": 0.0000025,
"output_cost_per_token": 0.000010,
"input_cost_per_token_batches": 0.00000125,
"output_cost_per_token_batches": 0.00000500,
"cache_read_input_token_cost": 0.00000125,
"litellm_provider": "openai",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.030,
"search_context_size_medium": 0.035,
"search_context_size_high": 0.050
}
},
"gpt-4o-search-preview": {
"max_tokens": 16384,
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"input_cost_per_token": 0.0000025,
@@ -201,6 +232,31 @@
"supports_system_messages": true,
"supports_tool_choice": true
},
"gpt-4o-mini-search-preview": {
"max_tokens": 16384,
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"input_cost_per_token": 0.00000015,
"output_cost_per_token": 0.00000060,
"input_cost_per_token_batches": 0.000000075,
"output_cost_per_token_batches": 0.00000030,
"cache_read_input_token_cost": 0.000000075,
"litellm_provider": "openai",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.025,
"search_context_size_medium": 0.0275,
"search_context_size_high": 0.030
}
},
"gpt-4o-mini-2024-07-18": {
"max_tokens": 16384,
"max_input_tokens": 128000,
@@ -218,7 +274,12 @@
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"search_context_cost_per_query": {
"search_context_size_low": 30.00,
"search_context_size_medium": 35.00,
"search_context_size_high": 50.00
}
},
"o1-pro": {
"max_tokens": 100000,
litellm/types/llms/openai.py
+22
View File
@@ -382,6 +382,28 @@ class ChatCompletionThinkingBlock(TypedDict, total=False):
cache_control: Optional[Union[dict, ChatCompletionCachedContent]]
class ChatCompletionAnnotationURLCitation(TypedDict, total=False):
end_index: int
"""The index of the last character of the URL citation in the message."""
start_index: int
"""The index of the first character of the URL citation in the message."""
title: str
"""The title of the web resource."""
url: str
"""The URL of the web resource."""
class ChatCompletionAnnotation(TypedDict, total=False):
type: Literal["url_citation"]
"""The type of the URL citation. Always `url_citation`."""
url_citation: ChatCompletionAnnotationURLCitation
"""A URL citation when using web search."""
class OpenAIChatCompletionTextObject(TypedDict):
type: Literal["text"]
text: str
litellm/types/utils.py
+15
View File
@@ -7,6 +7,7 @@ from typing import Any, Dict, List, Literal, Optional, Tuple, Union
from aiohttp import FormData
from openai._models import BaseModel as OpenAIObject
from openai.types.audio.transcription_create_params import FileTypes # type: ignore
from openai.types.chat.chat_completion import ChatCompletion
from openai.types.completion_usage import (
CompletionTokensDetails,
CompletionUsage,
@@ -27,6 +28,7 @@ from ..litellm_core_utils.core_helpers import map_finish_reason
from .guardrails import GuardrailEventHooks
from .llms.openai import (
Batch,
ChatCompletionAnnotation,
ChatCompletionThinkingBlock,
ChatCompletionToolCallChunk,
ChatCompletionUsageBlock,
@@ -527,6 +529,7 @@ class Message(OpenAIObject):
provider_specific_fields: Optional[Dict[str, Any]] = Field(
default=None, exclude=True
)
annotations: Optional[List[ChatCompletionAnnotation]] = None
def __init__(
self,
@@ -538,6 +541,7 @@ class Message(OpenAIObject):
provider_specific_fields: Optional[Dict[str, Any]] = None,
reasoning_content: Optional[str] = None,
thinking_blocks: Optional[List[ChatCompletionThinkingBlock]] = None,
annotations: Optional[List[ChatCompletionAnnotation]] = None,
**params,
):
init_values: Dict[str, Any] = {
@@ -566,6 +570,9 @@ class Message(OpenAIObject):
if thinking_blocks is not None:
init_values["thinking_blocks"] = thinking_blocks
if annotations is not None:
init_values["annotations"] = annotations
if reasoning_content is not None:
init_values["reasoning_content"] = reasoning_content
@@ -623,6 +630,7 @@ class Delta(OpenAIObject):
audio: Optional[ChatCompletionAudioResponse] = None,
reasoning_content: Optional[str] = None,
thinking_blocks: Optional[List[ChatCompletionThinkingBlock]] = None,
annotations: Optional[List[ChatCompletionAnnotation]] = None,
**params,
):
super(Delta, self).__init__(**params)
@@ -633,6 +641,7 @@ class Delta(OpenAIObject):
self.function_call: Optional[Union[FunctionCall, Any]] = None
self.tool_calls: Optional[List[Union[ChatCompletionDeltaToolCall, Any]]] = None
self.audio: Optional[ChatCompletionAudioResponse] = None
self.annotations: Optional[List[ChatCompletionAnnotation]] = None
if reasoning_content is not None:
self.reasoning_content = reasoning_content
@@ -646,6 +655,12 @@ class Delta(OpenAIObject):
# ensure default response matches OpenAI spec
del self.thinking_blocks
# Add annotations to the delta, ensure they are only on Delta if they exist (Match OpenAI spec)
if annotations is not None:
self.annotations = annotations
else:
del self.annotations
if function_call is not None and isinstance(function_call, dict):
self.function_call = FunctionCall(**function_call)
else:
litellm/utils.py
+21
View File
@@ -1975,6 +1975,27 @@ def supports_system_messages(model: str, custom_llm_provider: Optional[str]) ->
)
def supports_web_search(model: str, custom_llm_provider: Optional[str]) -> bool:
"""
Check if the given model supports web search and return a boolean value.
Parameters:
model (str): The model name to be checked.
custom_llm_provider (str): The provider to be checked.
Returns:
bool: True if the model supports web search, False otherwise.
Raises:
Exception: If the given model is not found in model_prices_and_context_window.json.
"""
return _supports_factory(
model=model,
custom_llm_provider=custom_llm_provider,
key="supports_web_search",
)
def supports_native_streaming(model: str, custom_llm_provider: Optional[str]) -> bool:
"""
Check if the given model supports native streaming and return a boolean value.
model_prices_and_context_window.json
+62 -1
View File
@@ -15,6 +15,12 @@
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.0000,
"search_context_size_medium": 0.0000,
"search_context_size_high": 0.0000
},
"deprecation_date": "date when the model becomes deprecated in the format YYYY-MM-DD"
},
"omni-moderation-latest": {
@@ -59,6 +65,31 @@
},
"gpt-4o": {
"max_tokens": 16384,
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"input_cost_per_token": 0.0000025,
"output_cost_per_token": 0.000010,
"input_cost_per_token_batches": 0.00000125,
"output_cost_per_token_batches": 0.00000500,
"cache_read_input_token_cost": 0.00000125,
"litellm_provider": "openai",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.030,
"search_context_size_medium": 0.035,
"search_context_size_high": 0.050
}
},
"gpt-4o-search-preview": {
"max_tokens": 16384,
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"input_cost_per_token": 0.0000025,
@@ -201,6 +232,31 @@
"supports_system_messages": true,
"supports_tool_choice": true
},
"gpt-4o-mini-search-preview": {
"max_tokens": 16384,
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"input_cost_per_token": 0.00000015,
"output_cost_per_token": 0.00000060,
"input_cost_per_token_batches": 0.000000075,
"output_cost_per_token_batches": 0.00000030,
"cache_read_input_token_cost": 0.000000075,
"litellm_provider": "openai",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.025,
"search_context_size_medium": 0.0275,
"search_context_size_high": 0.030
}
},
"gpt-4o-mini-2024-07-18": {
"max_tokens": 16384,
"max_input_tokens": 128000,
@@ -218,7 +274,12 @@
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"search_context_cost_per_query": {
"search_context_size_low": 30.00,
"search_context_size_medium": 35.00,
"search_context_size_high": 50.00
}
},
"o1-pro": {
"max_tokens": 100000,
tests/litellm/litellm_core_utils/test_streaming_handler.py
+34
View File
@@ -136,6 +136,40 @@ def test_is_chunk_non_empty(initialized_custom_stream_wrapper: CustomStreamWrapp
)
def test_is_chunk_non_empty_with_annotations(
initialized_custom_stream_wrapper: CustomStreamWrapper,
):
"""Unit test if non-empty when annotations are present"""
chunk = {
"id": "e89b6501-8ac2-464c-9550-7cd3daf94350",
"object": "chat.completion.chunk",
"created": 1741037890,
"model": "deepseek-reasoner",
"system_fingerprint": "fp_5417b77867_prod0225",
"choices": [
{
"index": 0,
"delta": {
"content": None,
"annotations": [
{"type": "url_citation", "url": "https://www.google.com"}
],
},
"logprobs": None,
"finish_reason": None,
}
],
}
assert (
initialized_custom_stream_wrapper.is_chunk_non_empty(
completion_obj=MagicMock(),
model_response=ModelResponseStream(**chunk),
response_obj=MagicMock(),
)
is True
)
def test_optional_combine_thinking_block_in_choices(
initialized_custom_stream_wrapper: CustomStreamWrapper,
):
tests/llm_translation/test_openai.py
+67
View File
@@ -3,6 +3,7 @@ import os
import sys
from datetime import datetime
from unittest.mock import AsyncMock, patch
from typing import Optional
sys.path.insert(
0, os.path.abspath("../..")
@@ -17,6 +18,10 @@ import litellm
from litellm import Choices, Message, ModelResponse
from base_llm_unit_tests import BaseLLMChatTest
import asyncio
from litellm.types.llms.openai import (
ChatCompletionAnnotation,
ChatCompletionAnnotationURLCitation,
)
def test_openai_prediction_param():
@@ -391,3 +396,65 @@ def test_openai_chat_completion_streaming_handler_reasoning_content():
)
assert response.choices[0].delta.reasoning_content == "."
def validate_response_url_citation(url_citation: ChatCompletionAnnotationURLCitation):
assert "end_index" in url_citation
assert "start_index" in url_citation
assert "url" in url_citation
def validate_web_search_annotations(annotations: ChatCompletionAnnotation):
"""validates litellm response contains web search annotations"""
print("annotations: ", annotations)
assert annotations is not None
assert isinstance(annotations, list)
for annotation in annotations:
assert annotation["type"] == "url_citation"
url_citation: ChatCompletionAnnotationURLCitation = annotation["url_citation"]
validate_response_url_citation(url_citation)
def test_openai_web_search():
"""Makes a simple web search request and validates the response contains web search annotations and all expected fields are present"""
litellm._turn_on_debug()
response = litellm.completion(
model="openai/gpt-4o-search-preview",
messages=[
{
"role": "user",
"content": "What was a positive news story from today?",
}
],
)
print("litellm response: ", response.model_dump_json(indent=4))
message = response.choices[0].message
annotations: ChatCompletionAnnotation = message.annotations
validate_web_search_annotations(annotations)
def test_openai_web_search_streaming():
"""Makes a simple web search request and validates the response contains web search annotations and all expected fields are present"""
# litellm._turn_on_debug()
test_openai_web_search: Optional[ChatCompletionAnnotation] = None
response = litellm.completion(
model="openai/gpt-4o-search-preview",
messages=[
{
"role": "user",
"content": "What was a positive news story from today?",
}
],
stream=True,
)
for chunk in response:
print("litellm response chunk: ", chunk)
if (
hasattr(chunk.choices[0].delta, "annotations")
and chunk.choices[0].delta.annotations is not None
):
test_openai_web_search = chunk.choices[0].delta.annotations
# Assert this request has at-least one web search annotation
assert test_openai_web_search is not None
validate_web_search_annotations(test_openai_web_search)
tests/local_testing/test_get_model_info.py
+10
View File
@@ -500,6 +500,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"supports_tool_choice": {"type": "boolean"},
"supports_video_input": {"type": "boolean"},
"supports_vision": {"type": "boolean"},
"supports_web_search": {"type": "boolean"},
"tool_use_system_prompt_tokens": {"type": "number"},
"tpm": {"type": "number"},
"supported_endpoints": {
@@ -518,6 +519,15 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
],
},
},
"search_context_cost_per_query": {
"type": "object",
"properties": {
"search_context_size_low": {"type": "number"},
"search_context_size_medium": {"type": "number"},
"search_context_size_high": {"type": "number"},
},
"additionalProperties": False,
},
"supported_modalities": {
"type": "array",
"items": {