tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 06:22:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(vertex_ai/): add new deepseek-ai api service (#12312)

Browse Source 
* feat(vertex_ai/): add new deepseek-ai api service

Closes https://github.com/BerriAI/litellm/issues/12192

* test: cleanup test
This commit is contained in:
Krish Dholakia
2025-07-05 10:38:37 -07:00
committed by GitHub
parent e7f1fa26ab
commit 380eb31103
3 changed files with 181 additions and 220 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/vertex_ai/vertex_ai_partner_models/main.py
+3 -2
View File
@@ -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
litellm/llms/vertex_ai/vertex_llm_base.py
+10 -6
View File
@@ -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 (
tests/local_testing/test_amazing_vertex_completion.py
+168 -212
View File
@@ -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]