diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index ab13a51137..cdd3fce6c6 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -398,6 +398,8 @@ curl http://localhost:4000/v1/chat/completions \ +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)** diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 749b6d9428..e7d3d2b060 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -240,6 +240,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): gtool_func_declarations = [] googleSearch: Optional[dict] = None googleSearchRetrieval: Optional[dict] = None + enterpriseWebSearch: Optional[dict] = None code_execution: Optional[dict] = None # remove 'additionalProperties' from tools value = _remove_additional_properties(value) @@ -273,6 +274,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): googleSearch = tool["googleSearch"] elif tool.get("googleSearchRetrieval", None) is not None: googleSearchRetrieval = tool["googleSearchRetrieval"] + elif tool.get("enterpriseWebSearch", None) is not None: + enterpriseWebSearch = tool["enterpriseWebSearch"] elif tool.get("code_execution", None) is not None: code_execution = tool["code_execution"] elif openai_function_object is not None: @@ -299,6 +302,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): _tools["googleSearch"] = googleSearch if googleSearchRetrieval is not None: _tools["googleSearchRetrieval"] = googleSearchRetrieval + if enterpriseWebSearch is not None: + _tools["enterpriseWebSearch"] = enterpriseWebSearch if code_execution is not None: _tools["code_execution"] = code_execution return [_tools] diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 7fa167938f..2e25f259b0 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -187,6 +187,7 @@ class Tools(TypedDict, total=False): function_declarations: List[FunctionDeclaration] googleSearch: dict googleSearchRetrieval: dict + enterpriseWebSearch: dict code_execution: dict retrieval: Retrieval diff --git a/tests/llm_translation/test_vertex.py b/tests/llm_translation/test_vertex.py index d821fb415e..9118d94a6f 100644 --- a/tests/llm_translation/test_vertex.py +++ b/tests/llm_translation/test_vertex.py @@ -141,6 +141,7 @@ def test_build_vertex_schema(): [ ([{"googleSearch": {}}], "googleSearch"), ([{"googleSearchRetrieval": {}}], "googleSearchRetrieval"), + ([{"enterpriseWebSearch": {}}], "enterpriseWebSearch"), ([{"code_execution": {}}], "code_execution"), ], )