mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-21 00:21:38 +00:00
fix(ollama_chat.py): fix key error + remove redundant code
This commit is contained in:
+14
-13
@@ -1,3 +1,4 @@
|
||||
import inspect
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
@@ -6,7 +7,6 @@ from typing import Any, List, Optional, Union
|
||||
import aiohttp
|
||||
import httpx
|
||||
from pydantic import BaseModel
|
||||
import inspect
|
||||
|
||||
import litellm
|
||||
from litellm import verbose_logger
|
||||
@@ -142,13 +142,6 @@ class OllamaChatConfig(OpenAIGPTConfig):
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
value = non_default_params["response_format"]
|
||||
if inspect.isclass(value) and issubclass(value, BaseModel):
|
||||
non_default_params["response_format"] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {"schema": value.model_json_schema()}
|
||||
}
|
||||
|
||||
for param, value in non_default_params.items():
|
||||
if param == "max_tokens" or param == "max_completion_tokens":
|
||||
optional_params["num_predict"] = value
|
||||
@@ -164,9 +157,17 @@ class OllamaChatConfig(OpenAIGPTConfig):
|
||||
optional_params["repeat_penalty"] = value
|
||||
if param == "stop":
|
||||
optional_params["stop"] = value
|
||||
if param == "response_format" and isinstance(value, dict) and value.get("type") == "json_object":
|
||||
if (
|
||||
param == "response_format"
|
||||
and isinstance(value, dict)
|
||||
and value.get("type") == "json_object"
|
||||
):
|
||||
optional_params["format"] = "json"
|
||||
if param == "response_format" and isinstance(value, dict) and value.get("type") == "json_schema":
|
||||
if (
|
||||
param == "response_format"
|
||||
and isinstance(value, dict)
|
||||
and value.get("type") == "json_schema"
|
||||
):
|
||||
if value.get("json_schema") and value["json_schema"].get("schema"):
|
||||
optional_params["format"] = value["json_schema"]["schema"]
|
||||
### FUNCTION CALLING LOGIC ###
|
||||
@@ -207,9 +208,9 @@ class OllamaChatConfig(OpenAIGPTConfig):
|
||||
litellm.add_function_to_prompt = (
|
||||
True # so that main.py adds the function call to the prompt
|
||||
)
|
||||
optional_params["functions_unsupported_model"] = non_default_params.get(
|
||||
"functions"
|
||||
)
|
||||
optional_params[
|
||||
"functions_unsupported_model"
|
||||
] = non_default_params.get("functions")
|
||||
non_default_params.pop("tool_choice", None) # causes ollama requests to hang
|
||||
non_default_params.pop("functions", None) # causes ollama requests to hang
|
||||
return optional_params
|
||||
|
||||
@@ -1,82 +1,72 @@
|
||||
import inspect
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
import inspect
|
||||
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../../..')))
|
||||
sys.path.insert(
|
||||
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))
|
||||
)
|
||||
|
||||
from litellm.llms.ollama_chat import OllamaChatConfig
|
||||
from litellm.utils import get_optional_params
|
||||
|
||||
|
||||
class TestEvent(BaseModel):
|
||||
name: str
|
||||
value: int
|
||||
|
||||
class TestOllamaChatConfigResponseFormat:
|
||||
def test_map_openai_params_with_pydantic_model(self):
|
||||
config = OllamaChatConfig()
|
||||
|
||||
non_default_params = {
|
||||
"response_format": TestEvent
|
||||
}
|
||||
optional_params = {}
|
||||
|
||||
expected_schema_structure = TestEvent.model_json_schema()
|
||||
|
||||
config.map_openai_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params=optional_params,
|
||||
class TestOllamaChatConfigResponseFormat:
|
||||
def test_get_optional_params_with_pydantic_model(self):
|
||||
optional_params = get_optional_params(
|
||||
model="ollama_chat/test-model",
|
||||
drop_params=False
|
||||
response_format=TestEvent,
|
||||
custom_llm_provider="ollama_chat",
|
||||
)
|
||||
|
||||
assert "format" in optional_params, "Transformed 'format' key not found in optional_params"
|
||||
|
||||
print(f"optional_params: {optional_params}")
|
||||
|
||||
assert "format" in optional_params
|
||||
transformed_format = optional_params["format"]
|
||||
|
||||
assert transformed_format == expected_schema_structure, \
|
||||
f"Transformed schema does not match expected. Got: {transformed_format}, Expected: {expected_schema_structure}"
|
||||
|
||||
expected_schema_structure = TestEvent.model_json_schema()
|
||||
transformed_format.pop("additionalProperties")
|
||||
|
||||
assert (
|
||||
transformed_format == expected_schema_structure
|
||||
), f"Transformed schema does not match expected. Got: {transformed_format}, Expected: {expected_schema_structure}"
|
||||
|
||||
def test_map_openai_params_with_dict_json_schema(self):
|
||||
config = OllamaChatConfig()
|
||||
|
||||
|
||||
direct_schema = TestEvent.model_json_schema()
|
||||
response_format_dict = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {"schema": direct_schema}
|
||||
"json_schema": {"schema": direct_schema},
|
||||
}
|
||||
|
||||
non_default_params = {
|
||||
"response_format": response_format_dict
|
||||
}
|
||||
optional_params = {}
|
||||
|
||||
config.map_openai_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params=optional_params,
|
||||
|
||||
non_default_params = {"response_format": response_format_dict}
|
||||
|
||||
optional_params = get_optional_params(
|
||||
model="ollama_chat/test-model",
|
||||
drop_params=False
|
||||
response_format=response_format_dict,
|
||||
custom_llm_provider="ollama_chat",
|
||||
)
|
||||
|
||||
|
||||
assert "format" in optional_params
|
||||
assert optional_params["format"] == direct_schema, \
|
||||
f"Schema from dict did not pass through correctly. Got: {optional_params['format']}, Expected: {direct_schema}"
|
||||
assert (
|
||||
optional_params["format"] == direct_schema
|
||||
), f"Schema from dict did not pass through correctly. Got: {optional_params['format']}, Expected: {direct_schema}"
|
||||
|
||||
def test_map_openai_params_with_json_object(self):
|
||||
config = OllamaChatConfig()
|
||||
|
||||
non_default_params = {
|
||||
"response_format": {"type": "json_object"}
|
||||
}
|
||||
optional_params = {}
|
||||
|
||||
config.map_openai_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params=optional_params,
|
||||
optional_params = get_optional_params(
|
||||
model="ollama_chat/test-model",
|
||||
drop_params=False
|
||||
response_format={"type": "json_object"},
|
||||
custom_llm_provider="ollama_chat",
|
||||
)
|
||||
|
||||
|
||||
assert "format" in optional_params
|
||||
assert optional_params["format"] == "json", \
|
||||
f"Expected 'json' for type 'json_object', got: {optional_params['format']}"
|
||||
assert (
|
||||
optional_params["format"] == "json"
|
||||
), f"Expected 'json' for type 'json_object', got: {optional_params['format']}"
|
||||
|
||||
+369
-2
@@ -1,17 +1,19 @@
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
import litellm
|
||||
from litellm.utils import get_optional_params_image_gen
|
||||
|
||||
# Adds the parent directory to the system path
|
||||
|
||||
|
||||
def test_get_optional_params_image_gen():
|
||||
from litellm.llms.azure.image_generation import AzureGPTImageGenerationConfig
|
||||
@@ -28,3 +30,368 @@ def test_get_optional_params_image_gen():
|
||||
assert optional_params is not None
|
||||
assert "response_format" not in optional_params
|
||||
assert optional_params["n"] == 3
|
||||
|
||||
|
||||
def return_mocked_response(model: str):
|
||||
if model == "bedrock/mistral.mistral-large-2407-v1:0":
|
||||
return {
|
||||
"metrics": {"latencyMs": 316},
|
||||
"output": {
|
||||
"message": {
|
||||
"content": [{"text": "Hello! How are you doing today? How can"}],
|
||||
"role": "assistant",
|
||||
}
|
||||
},
|
||||
"stopReason": "max_tokens",
|
||||
"usage": {"inputTokens": 5, "outputTokens": 10, "totalTokens": 15},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"bedrock/mistral.mistral-large-2407-v1:0",
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio()
|
||||
async def test_bedrock_max_completion_tokens(model: str):
|
||||
"""
|
||||
Tests that:
|
||||
- max_completion_tokens is passed as max_tokens to bedrock models
|
||||
"""
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
litellm.set_verbose = True
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
|
||||
mock_response = return_mocked_response(model)
|
||||
_model = model.split("/")[1]
|
||||
print("\n\nmock_response: ", mock_response)
|
||||
|
||||
with patch.object(client, "post") as mock_client:
|
||||
try:
|
||||
response = await litellm.acompletion(
|
||||
model=model,
|
||||
max_completion_tokens=10,
|
||||
messages=[{"role": "user", "content": "Hello!"}],
|
||||
client=client,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
|
||||
mock_client.assert_called_once()
|
||||
request_body = json.loads(mock_client.call_args.kwargs["data"])
|
||||
|
||||
print("request_body: ", request_body)
|
||||
|
||||
assert request_body == {
|
||||
"messages": [{"role": "user", "content": [{"text": "Hello!"}]}],
|
||||
"additionalModelRequestFields": {},
|
||||
"system": [],
|
||||
"inferenceConfig": {"maxTokens": 10},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
["anthropic/claude-3-sonnet-20240229", "anthropic/claude-3-opus-20240229"],
|
||||
)
|
||||
@pytest.mark.asyncio()
|
||||
async def test_anthropic_api_max_completion_tokens(model: str):
|
||||
"""
|
||||
Tests that:
|
||||
- max_completion_tokens is passed as max_tokens to anthropic models
|
||||
"""
|
||||
litellm.set_verbose = True
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
mock_response = {
|
||||
"content": [{"text": "Hi! My name is Claude.", "type": "text"}],
|
||||
"id": "msg_013Zva2CMHLNnXjNJJKqJ2EF",
|
||||
"model": "claude-3-5-sonnet-20240620",
|
||||
"role": "assistant",
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"type": "message",
|
||||
"usage": {"input_tokens": 2095, "output_tokens": 503},
|
||||
}
|
||||
|
||||
client = HTTPHandler()
|
||||
|
||||
print("\n\nmock_response: ", mock_response)
|
||||
|
||||
with patch.object(client, "post") as mock_client:
|
||||
try:
|
||||
response = await litellm.acompletion(
|
||||
model=model,
|
||||
max_completion_tokens=10,
|
||||
messages=[{"role": "user", "content": "Hello!"}],
|
||||
client=client,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
mock_client.assert_called_once()
|
||||
request_body = mock_client.call_args.kwargs["json"]
|
||||
|
||||
print("request_body: ", request_body)
|
||||
|
||||
assert request_body == {
|
||||
"messages": [
|
||||
{"role": "user", "content": [{"type": "text", "text": "Hello!"}]}
|
||||
],
|
||||
"max_tokens": 10,
|
||||
"model": model.split("/")[-1],
|
||||
}
|
||||
|
||||
|
||||
def test_all_model_configs():
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import (
|
||||
VertexAIAi21Config,
|
||||
)
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.llama3.transformation import (
|
||||
VertexAILlama3Config,
|
||||
)
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in VertexAILlama3Config().get_supported_openai_params(model="llama3")
|
||||
)
|
||||
assert VertexAILlama3Config().map_openai_params(
|
||||
{"max_completion_tokens": 10}, {}, "llama3", drop_params=False
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
assert "max_completion_tokens" in VertexAIAi21Config().get_supported_openai_params(
|
||||
model="jamba-1.5-mini@001"
|
||||
)
|
||||
assert VertexAIAi21Config().map_openai_params(
|
||||
{"max_completion_tokens": 10}, {}, "jamba-1.5-mini@001", drop_params=False
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig
|
||||
|
||||
assert "max_completion_tokens" in FireworksAIConfig().get_supported_openai_params(
|
||||
model="llama3"
|
||||
)
|
||||
assert FireworksAIConfig().map_openai_params(
|
||||
model="llama3",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.nvidia_nim.chat.transformation import NvidiaNimConfig
|
||||
|
||||
assert "max_completion_tokens" in NvidiaNimConfig().get_supported_openai_params(
|
||||
model="llama3"
|
||||
)
|
||||
assert NvidiaNimConfig().map_openai_params(
|
||||
model="llama3",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.ollama_chat import OllamaChatConfig
|
||||
|
||||
assert "max_completion_tokens" in OllamaChatConfig().get_supported_openai_params(
|
||||
model="llama3"
|
||||
)
|
||||
assert OllamaChatConfig().map_openai_params(
|
||||
model="llama3",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"num_predict": 10}
|
||||
|
||||
from litellm.llms.predibase.chat.transformation import PredibaseConfig
|
||||
|
||||
assert "max_completion_tokens" in PredibaseConfig().get_supported_openai_params(
|
||||
model="llama3"
|
||||
)
|
||||
assert PredibaseConfig().map_openai_params(
|
||||
model="llama3",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_new_tokens": 10}
|
||||
|
||||
from litellm.llms.codestral.completion.transformation import (
|
||||
CodestralTextCompletionConfig,
|
||||
)
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in CodestralTextCompletionConfig().get_supported_openai_params(model="llama3")
|
||||
)
|
||||
assert CodestralTextCompletionConfig().map_openai_params(
|
||||
model="llama3",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.volcengine import VolcEngineConfig
|
||||
|
||||
assert "max_completion_tokens" in VolcEngineConfig().get_supported_openai_params(
|
||||
model="llama3"
|
||||
)
|
||||
assert VolcEngineConfig().map_openai_params(
|
||||
model="llama3",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.ai21.chat.transformation import AI21ChatConfig
|
||||
|
||||
assert "max_completion_tokens" in AI21ChatConfig().get_supported_openai_params(
|
||||
"jamba-1.5-mini@001"
|
||||
)
|
||||
assert AI21ChatConfig().map_openai_params(
|
||||
model="jamba-1.5-mini@001",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIConfig
|
||||
|
||||
assert "max_completion_tokens" in AzureOpenAIConfig().get_supported_openai_params(
|
||||
model="gpt-3.5-turbo"
|
||||
)
|
||||
assert AzureOpenAIConfig().map_openai_params(
|
||||
model="gpt-3.5-turbo",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
api_version="2022-12-01",
|
||||
drop_params=False,
|
||||
) == {"max_completion_tokens": 10}
|
||||
|
||||
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in AmazonConverseConfig().get_supported_openai_params(
|
||||
model="anthropic.claude-3-sonnet-20240229-v1:0"
|
||||
)
|
||||
)
|
||||
assert AmazonConverseConfig().map_openai_params(
|
||||
model="anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"maxTokens": 10}
|
||||
|
||||
from litellm.llms.codestral.completion.transformation import (
|
||||
CodestralTextCompletionConfig,
|
||||
)
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in CodestralTextCompletionConfig().get_supported_openai_params(model="llama3")
|
||||
)
|
||||
assert CodestralTextCompletionConfig().map_openai_params(
|
||||
model="llama3",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm import AmazonAnthropicClaude3Config, AmazonAnthropicConfig
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in AmazonAnthropicClaude3Config().get_supported_openai_params(
|
||||
model="anthropic.claude-3-sonnet-20240229-v1:0"
|
||||
)
|
||||
)
|
||||
|
||||
assert AmazonAnthropicClaude3Config().map_openai_params(
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
model="anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in AmazonAnthropicConfig().get_supported_openai_params(model="")
|
||||
)
|
||||
|
||||
assert AmazonAnthropicConfig().map_openai_params(
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
model="",
|
||||
drop_params=False,
|
||||
) == {"max_tokens_to_sample": 10}
|
||||
|
||||
from litellm.llms.databricks.chat.transformation import DatabricksConfig
|
||||
|
||||
assert "max_completion_tokens" in DatabricksConfig().get_supported_openai_params()
|
||||
|
||||
assert DatabricksConfig().map_openai_params(
|
||||
model="databricks/llama-3-70b-instruct",
|
||||
drop_params=False,
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import (
|
||||
VertexAIAnthropicConfig,
|
||||
)
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in VertexAIAnthropicConfig().get_supported_openai_params(
|
||||
model="claude-3-5-sonnet-20240620"
|
||||
)
|
||||
)
|
||||
|
||||
assert VertexAIAnthropicConfig().map_openai_params(
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
model="claude-3-5-sonnet-20240620",
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig,
|
||||
)
|
||||
|
||||
assert "max_completion_tokens" in VertexGeminiConfig().get_supported_openai_params(
|
||||
model="gemini-1.0-pro"
|
||||
)
|
||||
|
||||
assert VertexGeminiConfig().map_openai_params(
|
||||
model="gemini-1.0-pro",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_output_tokens": 10}
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in GoogleAIStudioGeminiConfig().get_supported_openai_params(
|
||||
model="gemini-1.0-pro"
|
||||
)
|
||||
)
|
||||
|
||||
assert GoogleAIStudioGeminiConfig().map_openai_params(
|
||||
model="gemini-1.0-pro",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_output_tokens": 10}
|
||||
|
||||
assert "max_completion_tokens" in VertexGeminiConfig().get_supported_openai_params(
|
||||
model="gemini-1.0-pro"
|
||||
)
|
||||
|
||||
assert VertexGeminiConfig().map_openai_params(
|
||||
model="gemini-1.0-pro",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_output_tokens": 10}
|
||||
|
||||
@@ -1,391 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../..")
|
||||
) # Adds the parent directory to the system path
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
import httpx
|
||||
import pytest
|
||||
from respx import MockRouter
|
||||
from unittest.mock import patch, MagicMock, AsyncMock
|
||||
|
||||
import litellm
|
||||
from litellm import Choices, Message, ModelResponse
|
||||
|
||||
# Adds the parent directory to the system path
|
||||
|
||||
|
||||
def return_mocked_response(model: str):
|
||||
if model == "bedrock/mistral.mistral-large-2407-v1:0":
|
||||
return {
|
||||
"metrics": {"latencyMs": 316},
|
||||
"output": {
|
||||
"message": {
|
||||
"content": [{"text": "Hello! How are you doing today? How can"}],
|
||||
"role": "assistant",
|
||||
}
|
||||
},
|
||||
"stopReason": "max_tokens",
|
||||
"usage": {"inputTokens": 5, "outputTokens": 10, "totalTokens": 15},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"bedrock/mistral.mistral-large-2407-v1:0",
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio()
|
||||
async def test_bedrock_max_completion_tokens(model: str):
|
||||
"""
|
||||
Tests that:
|
||||
- max_completion_tokens is passed as max_tokens to bedrock models
|
||||
"""
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
litellm.set_verbose = True
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
|
||||
mock_response = return_mocked_response(model)
|
||||
_model = model.split("/")[1]
|
||||
print("\n\nmock_response: ", mock_response)
|
||||
|
||||
with patch.object(client, "post") as mock_client:
|
||||
try:
|
||||
response = await litellm.acompletion(
|
||||
model=model,
|
||||
max_completion_tokens=10,
|
||||
messages=[{"role": "user", "content": "Hello!"}],
|
||||
client=client,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
|
||||
mock_client.assert_called_once()
|
||||
request_body = json.loads(mock_client.call_args.kwargs["data"])
|
||||
|
||||
print("request_body: ", request_body)
|
||||
|
||||
assert request_body == {
|
||||
"messages": [{"role": "user", "content": [{"text": "Hello!"}]}],
|
||||
"additionalModelRequestFields": {},
|
||||
"system": [],
|
||||
"inferenceConfig": {"maxTokens": 10},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
["anthropic/claude-3-sonnet-20240229", "anthropic/claude-3-opus-20240229"],
|
||||
)
|
||||
@pytest.mark.asyncio()
|
||||
async def test_anthropic_api_max_completion_tokens(model: str):
|
||||
"""
|
||||
Tests that:
|
||||
- max_completion_tokens is passed as max_tokens to anthropic models
|
||||
"""
|
||||
litellm.set_verbose = True
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
mock_response = {
|
||||
"content": [{"text": "Hi! My name is Claude.", "type": "text"}],
|
||||
"id": "msg_013Zva2CMHLNnXjNJJKqJ2EF",
|
||||
"model": "claude-3-5-sonnet-20240620",
|
||||
"role": "assistant",
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"type": "message",
|
||||
"usage": {"input_tokens": 2095, "output_tokens": 503},
|
||||
}
|
||||
|
||||
client = HTTPHandler()
|
||||
|
||||
print("\n\nmock_response: ", mock_response)
|
||||
|
||||
with patch.object(client, "post") as mock_client:
|
||||
try:
|
||||
response = await litellm.acompletion(
|
||||
model=model,
|
||||
max_completion_tokens=10,
|
||||
messages=[{"role": "user", "content": "Hello!"}],
|
||||
client=client,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
mock_client.assert_called_once()
|
||||
request_body = mock_client.call_args.kwargs["json"]
|
||||
|
||||
print("request_body: ", request_body)
|
||||
|
||||
assert request_body == {
|
||||
"messages": [
|
||||
{"role": "user", "content": [{"type": "text", "text": "Hello!"}]}
|
||||
],
|
||||
"max_tokens": 10,
|
||||
"model": model.split("/")[-1],
|
||||
}
|
||||
|
||||
|
||||
def test_all_model_configs():
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import (
|
||||
VertexAIAi21Config,
|
||||
)
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.llama3.transformation import (
|
||||
VertexAILlama3Config,
|
||||
)
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in VertexAILlama3Config().get_supported_openai_params(model="llama3")
|
||||
)
|
||||
assert VertexAILlama3Config().map_openai_params(
|
||||
{"max_completion_tokens": 10}, {}, "llama3", drop_params=False
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
assert "max_completion_tokens" in VertexAIAi21Config().get_supported_openai_params(
|
||||
model="jamba-1.5-mini@001"
|
||||
)
|
||||
assert VertexAIAi21Config().map_openai_params(
|
||||
{"max_completion_tokens": 10}, {}, "jamba-1.5-mini@001", drop_params=False
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.fireworks_ai.chat.transformation import (
|
||||
FireworksAIConfig,
|
||||
)
|
||||
|
||||
assert "max_completion_tokens" in FireworksAIConfig().get_supported_openai_params(
|
||||
model="llama3"
|
||||
)
|
||||
assert FireworksAIConfig().map_openai_params(
|
||||
model="llama3",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.nvidia_nim.chat.transformation import NvidiaNimConfig
|
||||
|
||||
assert "max_completion_tokens" in NvidiaNimConfig().get_supported_openai_params(
|
||||
model="llama3"
|
||||
)
|
||||
assert NvidiaNimConfig().map_openai_params(
|
||||
model="llama3",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.ollama_chat import OllamaChatConfig
|
||||
|
||||
assert "max_completion_tokens" in OllamaChatConfig().get_supported_openai_params(
|
||||
model="llama3"
|
||||
)
|
||||
assert OllamaChatConfig().map_openai_params(
|
||||
model="llama3",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"num_predict": 10}
|
||||
|
||||
from litellm.llms.predibase.chat.transformation import PredibaseConfig
|
||||
|
||||
assert "max_completion_tokens" in PredibaseConfig().get_supported_openai_params(
|
||||
model="llama3"
|
||||
)
|
||||
assert PredibaseConfig().map_openai_params(
|
||||
model="llama3",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_new_tokens": 10}
|
||||
|
||||
from litellm.llms.codestral.completion.transformation import (
|
||||
CodestralTextCompletionConfig,
|
||||
)
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in CodestralTextCompletionConfig().get_supported_openai_params(model="llama3")
|
||||
)
|
||||
assert CodestralTextCompletionConfig().map_openai_params(
|
||||
model="llama3",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.volcengine import VolcEngineConfig
|
||||
|
||||
assert "max_completion_tokens" in VolcEngineConfig().get_supported_openai_params(
|
||||
model="llama3"
|
||||
)
|
||||
assert VolcEngineConfig().map_openai_params(
|
||||
model="llama3",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.ai21.chat.transformation import AI21ChatConfig
|
||||
|
||||
assert "max_completion_tokens" in AI21ChatConfig().get_supported_openai_params(
|
||||
"jamba-1.5-mini@001"
|
||||
)
|
||||
assert AI21ChatConfig().map_openai_params(
|
||||
model="jamba-1.5-mini@001",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIConfig
|
||||
|
||||
assert "max_completion_tokens" in AzureOpenAIConfig().get_supported_openai_params(
|
||||
model="gpt-3.5-turbo"
|
||||
)
|
||||
assert AzureOpenAIConfig().map_openai_params(
|
||||
model="gpt-3.5-turbo",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
api_version="2022-12-01",
|
||||
drop_params=False,
|
||||
) == {"max_completion_tokens": 10}
|
||||
|
||||
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in AmazonConverseConfig().get_supported_openai_params(
|
||||
model="anthropic.claude-3-sonnet-20240229-v1:0"
|
||||
)
|
||||
)
|
||||
assert AmazonConverseConfig().map_openai_params(
|
||||
model="anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"maxTokens": 10}
|
||||
|
||||
from litellm.llms.codestral.completion.transformation import (
|
||||
CodestralTextCompletionConfig,
|
||||
)
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in CodestralTextCompletionConfig().get_supported_openai_params(model="llama3")
|
||||
)
|
||||
assert CodestralTextCompletionConfig().map_openai_params(
|
||||
model="llama3",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm import (
|
||||
AmazonAnthropicClaude3Config,
|
||||
AmazonAnthropicConfig,
|
||||
)
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in AmazonAnthropicClaude3Config().get_supported_openai_params(
|
||||
model="anthropic.claude-3-sonnet-20240229-v1:0"
|
||||
)
|
||||
)
|
||||
|
||||
assert AmazonAnthropicClaude3Config().map_openai_params(
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
model="anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in AmazonAnthropicConfig().get_supported_openai_params(model="")
|
||||
)
|
||||
|
||||
assert AmazonAnthropicConfig().map_openai_params(
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
model="",
|
||||
drop_params=False,
|
||||
) == {"max_tokens_to_sample": 10}
|
||||
|
||||
from litellm.llms.databricks.chat.transformation import DatabricksConfig
|
||||
|
||||
assert "max_completion_tokens" in DatabricksConfig().get_supported_openai_params()
|
||||
|
||||
assert DatabricksConfig().map_openai_params(
|
||||
model="databricks/llama-3-70b-instruct",
|
||||
drop_params=False,
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import (
|
||||
VertexAIAnthropicConfig,
|
||||
)
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in VertexAIAnthropicConfig().get_supported_openai_params(
|
||||
model="claude-3-5-sonnet-20240620"
|
||||
)
|
||||
)
|
||||
|
||||
assert VertexAIAnthropicConfig().map_openai_params(
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
model="claude-3-5-sonnet-20240620",
|
||||
drop_params=False,
|
||||
) == {"max_tokens": 10}
|
||||
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig,
|
||||
)
|
||||
from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig
|
||||
|
||||
assert "max_completion_tokens" in VertexGeminiConfig().get_supported_openai_params(
|
||||
model="gemini-1.0-pro"
|
||||
)
|
||||
|
||||
assert VertexGeminiConfig().map_openai_params(
|
||||
model="gemini-1.0-pro",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_output_tokens": 10}
|
||||
|
||||
assert (
|
||||
"max_completion_tokens"
|
||||
in GoogleAIStudioGeminiConfig().get_supported_openai_params(
|
||||
model="gemini-1.0-pro"
|
||||
)
|
||||
)
|
||||
|
||||
assert GoogleAIStudioGeminiConfig().map_openai_params(
|
||||
model="gemini-1.0-pro",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_output_tokens": 10}
|
||||
|
||||
assert "max_completion_tokens" in VertexGeminiConfig().get_supported_openai_params(
|
||||
model="gemini-1.0-pro"
|
||||
)
|
||||
|
||||
assert VertexGeminiConfig().map_openai_params(
|
||||
model="gemini-1.0-pro",
|
||||
non_default_params={"max_completion_tokens": 10},
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
) == {"max_output_tokens": 10}
|
||||
Reference in New Issue
Block a user