fix(ollama_chat.py): fix key error + remove redundant code

This commit is contained in:
Krrish Dholakia
2025-05-12 16:03:27 -07:00
parent a4fb1da2d9
commit 2eb4aae26f
4 changed files with 424 additions and 457 deletions
+14 -13
View File
@@ -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
View File
@@ -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}