phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: add enterpriseWebSearch tool for vertex-ai (#9856)

Browse source
This commit is contained in:
qvalentin 2025-04-09 20:17:48 +00:00 • committed by GitHub
parent d5e362459c
commit 93532e00db
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 9 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

2
docs/my-website/docs/providers/vertex.md
View file

@ -398,6 +398,8 @@ curl http://localhost:4000/v1/chat/completions \
</TabItem> </TabItem>
</Tabs> </Tabs>
You can also use the `enterpriseWebSearch` tool for an [enterprise compliant search](https://cloud.google.com/vertex-ai/generative-ai/docs/grounding/web-grounding-enterprise).
#### **Moving from Vertex AI SDK to LiteLLM (GROUNDING)** #### **Moving from Vertex AI SDK to LiteLLM (GROUNDING)**

5
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -240,6 +240,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
gtool_func_declarations = [] gtool_func_declarations = []
googleSearch: Optional[dict] = None googleSearch: Optional[dict] = None
googleSearchRetrieval: Optional[dict] = None googleSearchRetrieval: Optional[dict] = None
enterpriseWebSearch: Optional[dict] = None
code_execution: Optional[dict] = None code_execution: Optional[dict] = None
# remove 'additionalProperties' from tools # remove 'additionalProperties' from tools
value = _remove_additional_properties(value) value = _remove_additional_properties(value)
@ -273,6 +274,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
googleSearch = tool["googleSearch"] googleSearch = tool["googleSearch"]
elif tool.get("googleSearchRetrieval", None) is not None: elif tool.get("googleSearchRetrieval", None) is not None:
googleSearchRetrieval = tool["googleSearchRetrieval"] googleSearchRetrieval = tool["googleSearchRetrieval"]
elif tool.get("enterpriseWebSearch", None) is not None:
enterpriseWebSearch = tool["enterpriseWebSearch"]
elif tool.get("code_execution", None) is not None: elif tool.get("code_execution", None) is not None:
code_execution = tool["code_execution"] code_execution = tool["code_execution"]
elif openai_function_object is not None: elif openai_function_object is not None:
@ -299,6 +302,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
_tools["googleSearch"] = googleSearch _tools["googleSearch"] = googleSearch
if googleSearchRetrieval is not None: if googleSearchRetrieval is not None:
_tools["googleSearchRetrieval"] = googleSearchRetrieval _tools["googleSearchRetrieval"] = googleSearchRetrieval
if enterpriseWebSearch is not None:
_tools["enterpriseWebSearch"] = enterpriseWebSearch
if code_execution is not None: if code_execution is not None:
_tools["code_execution"] = code_execution _tools["code_execution"] = code_execution
return [_tools] return [_tools]

1
litellm/types/llms/vertex_ai.py
View file

@ -187,6 +187,7 @@ class Tools(TypedDict, total=False):
function_declarations: List[FunctionDeclaration] function_declarations: List[FunctionDeclaration]
googleSearch: dict googleSearch: dict
googleSearchRetrieval: dict googleSearchRetrieval: dict
enterpriseWebSearch: dict
code_execution: dict code_execution: dict
retrieval: Retrieval retrieval: Retrieval

1
tests/llm_translation/test_vertex.py
View file

@ -141,6 +141,7 @@ def test_build_vertex_schema():
[ [
([{"googleSearch": {}}], "googleSearch"), ([{"googleSearch": {}}], "googleSearch"),
([{"googleSearchRetrieval": {}}], "googleSearchRetrieval"), ([{"googleSearchRetrieval": {}}], "googleSearchRetrieval"),
([{"enterpriseWebSearch": {}}], "enterpriseWebSearch"),
([{"code_execution": {}}], "code_execution"), ([{"code_execution": {}}], "code_execution"),
], ],
) )