From 380eb31103466469ea41b44fa277bb25e02e2a5d Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 5 Jul 2025 10:38:37 -0700 Subject: [PATCH] feat(vertex_ai/): add new deepseek-ai api service (#12312) * feat(vertex_ai/): add new deepseek-ai api service Closes https://github.com/BerriAI/litellm/issues/12192 * test: cleanup test --- .../vertex_ai_partner_models/main.py | 5 +- litellm/llms/vertex_ai/vertex_llm_base.py | 16 +- .../test_amazing_vertex_completion.py | 380 ++++++++---------- 3 files changed, 181 insertions(+), 220 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index 78b281494e..7303ab0786 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -31,7 +31,7 @@ class VertexAIError(Exception): class VertexAIPartnerModels(VertexBase): def __init__(self) -> None: pass - + @staticmethod def is_vertex_partner_model(model: str): """ @@ -43,6 +43,7 @@ class VertexAIPartnerModels(VertexBase): """ if ( model.startswith("meta/") + or model.startswith("deepseek-ai") or model.startswith("mistral") or model.startswith("codestral") or model.startswith("jamba") @@ -114,7 +115,7 @@ class VertexAIPartnerModels(VertexBase): optional_params["stream"] = stream - if "llama" in model: + if "llama" in model or "deepseek-ai" in model: partner = VertexPartnerProvider.llama elif "mistral" in model or "codestral" in model: partner = VertexPartnerProvider.mistralai diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index f45549368a..22f119dab2 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -83,7 +83,11 @@ class VertexBase: if "type" in json_obj and json_obj["type"] == "external_account": # If environment_id key contains "aws" value it corresponds to an AWS config file credential_source = json_obj.get("credential_source", {}) - environment_id = credential_source.get("environment_id", "") if isinstance(credential_source, dict) else "" + environment_id = ( + credential_source.get("environment_id", "") + if isinstance(credential_source, dict) + else "" + ) if isinstance(environment_id, str) and "aws" in environment_id: creds = self._credentials_from_identity_pool_with_aws(json_obj) else: @@ -130,7 +134,7 @@ class VertexBase: from google.auth import identity_pool return identity_pool.Credentials.from_info(json_obj) - + def _credentials_from_identity_pool_with_aws(self, json_obj): from google.auth import aws @@ -183,7 +187,7 @@ class VertexBase: api_base = api_base or f"https://{vertex_location}-aiplatform.googleapis.com" if partner == VertexPartnerProvider.llama: - return f"{api_base}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/endpoints/openapi/chat/completions" + return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/endpoints/openapi/chat/completions" elif partner == VertexPartnerProvider.mistralai: if stream: return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/mistralai/models/{model}:streamRawPredict" @@ -490,7 +494,7 @@ class VertexBase: headers.update(extra_headers) return headers - + @staticmethod def get_vertex_ai_project(litellm_params: dict) -> Optional[str]: return ( @@ -499,7 +503,7 @@ class VertexBase: or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") ) - + @staticmethod def get_vertex_ai_credentials(litellm_params: dict) -> Optional[str]: return ( @@ -507,7 +511,7 @@ class VertexBase: or litellm_params.pop("vertex_ai_credentials", None) or get_secret_str("VERTEXAI_CREDENTIALS") ) - + @staticmethod def get_vertex_ai_location(litellm_params: dict) -> Optional[str]: return ( diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index f74307ae28..7ec88b5413 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -489,6 +489,7 @@ async def test_async_vertexai_response(): @pytest.mark.asyncio async def test_async_vertexai_streaming_response(): import random + litellm._turn_on_debug() load_vertex_ai_credentials() @@ -1534,6 +1535,7 @@ async def test_gemini_pro_json_schema_args_sent_httpx( assert resp.model == model.split("/")[1] + @pytest.mark.asyncio async def test_anthropic_message_via_anthropic_messages(): from litellm.llms.custom_httpx.llm_http_handler import AsyncHTTPHandler @@ -1552,20 +1554,30 @@ async def test_anthropic_message_via_anthropic_messages(): call_2_kwargs = {} with patch.object(client, "post", new=httpx_response) as mock_call: messages = [{"role": "user", "content": "List 5 cookie recipes"}] - response = await litellm.anthropic_messages(model="vertex_ai/claude-3-5-sonnet@20240620", messages=messages, max_tokens=100, client=client) + response = await litellm.anthropic_messages( + model="vertex_ai/claude-3-5-sonnet@20240620", + messages=messages, + max_tokens=100, + client=client, + ) print(f"response: {response}") assert mock_call.call_count == 1 call_1_kwargs = mock_call.call_args.kwargs with patch.object(client, "post", new=httpx_response) as mock_call: - response_2 = await litellm.acompletion(model="vertex_ai/claude-3-5-sonnet@20240620", messages=messages, max_tokens=100, client=client) + response_2 = await litellm.acompletion( + model="vertex_ai/claude-3-5-sonnet@20240620", + messages=messages, + max_tokens=100, + client=client, + ) print(f"response_2: {response_2}") call_args = mock_call.call_args print(f"call_args: {call_args}") call_2_kwargs = mock_call.call_args.kwargs call_2_kwargs["url"] = call_args[0][0] - + """ Compare Call 1 and Call 2 @@ -1578,9 +1590,15 @@ async def test_anthropic_message_via_anthropic_messages(): """ print(f"call_1_kwargs: {call_1_kwargs}") print(f"call_2_kwargs: {call_2_kwargs}") - assert call_1_kwargs["url"] == call_2_kwargs["url"], f"Expected url to be the same, but got {call_1_kwargs['url']} and Expected {call_2_kwargs['url']}" - assert "Authorization".lower() in [k.lower() for k in call_1_kwargs["headers"].keys()], f"Expected Authorization header to be present in call_1_kwargs, but got {call_1_kwargs['headers'].keys()}" - assert "content-type".lower() in [k.lower() for k in call_1_kwargs["headers"].keys()], f"Expected Content-Type header to be present in call_1_kwargs, but got {call_1_kwargs['headers'].keys()}" + assert ( + call_1_kwargs["url"] == call_2_kwargs["url"] + ), f"Expected url to be the same, but got {call_1_kwargs['url']} and Expected {call_2_kwargs['url']}" + assert "Authorization".lower() in [ + k.lower() for k in call_1_kwargs["headers"].keys() + ], f"Expected Authorization header to be present in call_1_kwargs, but got {call_1_kwargs['headers'].keys()}" + assert "content-type".lower() in [ + k.lower() for k in call_1_kwargs["headers"].keys() + ], f"Expected Content-Type header to be present in call_1_kwargs, but got {call_1_kwargs['headers'].keys()}" ## validate request body print(f"call 1 kwargs keys: {call_1_kwargs.keys()}") @@ -1588,7 +1606,10 @@ async def test_anthropic_message_via_anthropic_messages(): print(f"call_1_kwargs['data']: {type(call_1_kwargs['data'])}") call_1_kwargs_data = json.loads(call_1_kwargs["data"]) for k, v in call_2_kwargs["json"].items(): - assert k in call_1_kwargs_data, f"Expected {k} to be present in call_1_kwargs['data'], but got {call_1_kwargs_data.keys()}" + assert ( + k in call_1_kwargs_data + ), f"Expected {k} to be present in call_1_kwargs['data'], but got {call_1_kwargs_data.keys()}" + @pytest.mark.parametrize( "model, vertex_location, supports_response_schema", @@ -1903,7 +1924,6 @@ async def test_gemini_pro_function_calling_streaming(sync_mode): pass - # asyncio.run(gemini_pro_async_function_calling()) @@ -3132,7 +3152,9 @@ async def test_vertexai_embedding_finetuned(respx_mock: MockRouter): """ load_vertex_ai_credentials() litellm.set_verbose = True - litellm.disable_aiohttp_transport = True # since this uses respx, we need to set use_aiohttp_transport to False + litellm.disable_aiohttp_transport = ( + True # since this uses respx, we need to set use_aiohttp_transport to False + ) # Test input input_text = ["good morning from litellm", "this is another item"] @@ -3201,7 +3223,9 @@ async def test_vertexai_model_garden_model_completion( Using OpenAI compatible models from Vertex Model Garden """ - litellm.disable_aiohttp_transport = True # since this uses respx, we need to set use_aiohttp_transport to False + litellm.disable_aiohttp_transport = ( + True # since this uses respx, we need to set use_aiohttp_transport to False + ) litellm.module_level_aclient = httpx.AsyncClient() load_vertex_ai_credentials() litellm.set_verbose = True @@ -3278,7 +3302,6 @@ async def test_vertexai_model_garden_model_completion( assert response.usage.total_tokens == 172 - def vertex_ai_anthropic_thinking_mock_response(*args, **kwargs): mock_response = MagicMock() mock_response.status_code = 200 @@ -3537,12 +3560,9 @@ def test_gemini_tool_calling_working_demo(): "messages": [ { "content": "\n You are a helpful assistant who can help with questions on customers business or personal finances.\n Use the results from the available tools to answer the question.\n ", - "role": "system" + "role": "system", }, - { - "content": "Hello", - "role": "user" - } + {"content": "Hello", "role": "user"}, ], "max_completion_tokens": 1000, "temperature": 0.0, @@ -3557,91 +3577,57 @@ def test_gemini_tool_calling_working_demo(): "state": { "properties": { "messages": { - "items": { - "type": "object" - }, - "type": "array" + "items": {"type": "object"}, + "type": "array", }, - "conversation_id": { - "type": "string" - } + "conversation_id": {"type": "string"}, }, - "required": [ - "messages", - "conversation_id" - ], - "type": "object" + "required": ["messages", "conversation_id"], + "type": "object", }, "config": { "description": "Configuration for a Runnable.", "properties": { "tags": { - "items": { - "type": "string" - }, - "type": "array" - }, - "metadata": { - "type": "object" + "items": {"type": "string"}, + "type": "array", }, + "metadata": {"type": "object"}, "callbacks": { "anyOf": [ {"type": "array"}, {"type": "object"}, - {"type": "null"} + {"type": "null"}, ], }, - "run_name": { - "type": "string" - }, + "run_name": {"type": "string"}, "max_concurrency": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ] - }, - "recursion_limit": { - "type": "integer" - }, - "configurable": { - "type": "object" + "anyOf": [{"type": "integer"}, {"type": "null"}] }, + "recursion_limit": {"type": "integer"}, + "configurable": {"type": "object"}, "run_id": { "anyOf": [ - { - "format": "uuid", - "type": "string" - }, - { - "type": "null" - } + {"format": "uuid", "type": "string"}, + {"type": "null"}, ] - } + }, }, - "type": "object" + "type": "object", }, - "kwargs": { - "default": None, - "type": "object" - } + "kwargs": {"default": None, "type": "object"}, }, - "required": [ - "state", - "config" - ], - "type": "object" - } - } + "required": ["state", "config"], + "type": "object", + }, + }, } - ] + ], } response = completion(model="vertex_ai/gemini-2.0-flash", **args) print(response) + def test_gemini_tool_calling_not_working(): load_vertex_ai_credentials() litellm._turn_on_debug() @@ -3649,152 +3635,104 @@ def test_gemini_tool_calling_not_working(): "messages": [ { "content": "\n You are a helpful assistant who can help with questions on customers business or personal finances.\n Use the results from the available tools to answer the question.\n ", - "role": "system" + "role": "system", }, - { - "content": "Hello", - "role": "user" - } + {"content": "Hello", "role": "user"}, ], "max_completion_tokens": 1000, "temperature": 0.0, "tools": [ - { - "type": "function", - "function": { - "name": "test_agent", - "description": "This tool helps find relevant help content", - "parameters": { - "properties": { - "state": { - "properties": { - "messages": { - "items": {}, - "type": "array" + { + "type": "function", + "function": { + "name": "test_agent", + "description": "This tool helps find relevant help content", + "parameters": { + "properties": { + "state": { + "properties": { + "messages": {"items": {}, "type": "array"}, + "conversation_id": {"type": "string"}, }, - "conversation_id": { - "type": "string" - } + "required": ["messages", "conversation_id"], + "type": "object", }, - "required": [ - "messages", - "conversation_id" - ], - "type": "object" - }, - "config": { - "description": "Configuration for a Runnable.", - "properties": { - "tags": { - "items": { - "type": "string" + "config": { + "description": "Configuration for a Runnable.", + "properties": { + "tags": { + "items": {"type": "string"}, + "type": "array", + }, + "metadata": {"type": "object"}, + "callbacks": { + "anyOf": [ + {"items": {}, "type": "array"}, + {}, + {"type": "null"}, + ] + }, + "run_name": {"type": "string"}, + "max_concurrency": { + "anyOf": [{"type": "integer"}, {"type": "null"}] + }, + "recursion_limit": {"type": "integer"}, + "configurable": {"type": "object"}, + "run_id": { + "anyOf": [ + {"format": "uuid", "type": "string"}, + {"type": "null"}, + ] }, - "type": "array" }, - "metadata": { - "type": "object" - }, - "callbacks": { - "anyOf": [ - { - "items": {}, - "type": "array" - }, - {}, - { - "type": "null" - } - ] - }, - "run_name": { - "type": "string" - }, - "max_concurrency": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ] - }, - "recursion_limit": { - "type": "integer" - }, - "configurable": { - "type": "object" - }, - "run_id": { - "anyOf": [ - { - "format": "uuid", - "type": "string" - }, - { - "type": "null" - } - ] - } + "type": "object", }, - "type": "object" + "kwargs": {"default": None, "type": "object"}, }, - "kwargs": { - "default": None, - "type": "object" - } + "required": ["state", "config"], + "type": "object", }, - "required": [ - "state", - "config" - ], - "type": "object" - } + }, } - } - ] + ], } response = completion(model="vertex_ai/gemini-2.0-flash", **args) print(response) + def test_vertex_ai_llama_tool_calling(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") load_vertex_ai_credentials() litellm._turn_on_debug() args = { - "model": "vertex_ai/meta/llama-4-maverick-17b-128e-instruct-maas", - "messages": [ - { - "role": "user", - "content": "What is the weather in Boston, MA today?" - } - ], - "tools": [ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get current temperature for a given location.", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "City and country e.g. Bogotá, Colombia" - } - }, - "required": [ - "location" - ], - "additionalProperties": False - } - } - } - ], - "vertex_location": "us-east5" + "model": "vertex_ai/meta/llama-4-maverick-17b-128e-instruct-maas", + "messages": [ + {"role": "user", "content": "What is the weather in Boston, MA today?"} + ], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current temperature for a given location.", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "City and country e.g. Bogotá, Colombia", + } + }, + "required": ["location"], + "additionalProperties": False, + }, + }, + } + ], + "vertex_location": "us-east5", } - try: + try: response = completion(**args) except litellm.RateLimitError: pytest.skip("Rate limit error") @@ -3812,7 +3750,6 @@ def test_vertex_schema_test(): def tool_call(text: str | None) -> str: return text or "No text provided" - tool = { "type": "function", "function": { @@ -3823,15 +3760,18 @@ def test_vertex_schema_test(): "properties": { "repo_path": {"title": "Repo Path", "type": "string"}, "branch_name": {"title": "Branch Name", "type": "string"}, - "base_branch": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": None, "title": "Base Branch"}, + "base_branch": { + "anyOf": [{"type": "string"}, {"type": "null"}], + "default": None, + "title": "Base Branch", + }, }, "required": ["repo_path", "branch_name"], "title": "GitCreateBranch", - } - } + }, + }, } - response = litellm.completion( model="vertex_ai/gemini-2.5-flash-preview-05-20", messages=[{"role": "user", "content": "call the tool"}], @@ -3842,11 +3782,10 @@ def test_vertex_schema_test(): print(response) - - def test_vertex_ai_response_id(): """Test that litellm preserves the response ID from Vertex AI's API for non-streaming responses""" from litellm.llms.custom_httpx.http_handler import HTTPHandler + load_vertex_ai_credentials() client = HTTPHandler() @@ -3896,6 +3835,7 @@ def test_vertex_ai_streaming_response_id(): from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( make_sync_call, ) + load_vertex_ai_credentials() client = HTTPHandler() @@ -3956,7 +3896,10 @@ def test_vertex_ai_gemini_2_5_pro_streaming(): has_real_content = False for chunk in response: print(chunk) - if chunk.choices[0].delta.content is not None and len(chunk.choices[0].delta.content) > 0: + if ( + chunk.choices[0].delta.content is not None + and len(chunk.choices[0].delta.content) > 0 + ): has_real_content = True assert has_real_content @@ -3969,12 +3912,9 @@ def test_vertex_ai_gemini_audio_ogg(): messages=[ { "content": [ - { - "text": "generate a transcript of the speech.", - "type": "text" - } + {"text": "generate a transcript of the speech.", "type": "text"} ], - "role": "user" + "role": "user", }, { "content": [ @@ -3982,16 +3922,32 @@ def test_vertex_ai_gemini_audio_ogg(): "file": { "file_id": "https://upload.wikimedia.org/wikipedia/commons/5/5f/En-us-public.ogg" }, - "type": "file" + "type": "file", } ], - "role": "user" - } + "role": "user", + }, ], ) print(response) +@pytest.mark.asyncio +async def test_vertex_ai_deepseek(): + load_vertex_ai_credentials() + litellm._turn_on_debug() + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + client = AsyncHTTPHandler() + with patch.object(client, "post", return_value=MagicMock()) as mock_post: + response = await acompletion( + model="vertex_ai/deepseek-ai/deepseek-r1-0528-maas", + messages=[{"role": "user", "content": "Hi!"}], + client=client, + ) + mock_post.assert_called_once() + print(f"mock_post.call_args: {mock_post.call_args[0][0]}") + assert "v1beta1" not in mock_post.call_args[0][0] + assert "v1" in mock_post.call_args[0][0]