Merge pull request #4418 from BerriAI/litellm_fireworks_ai_tool_calling

[Feat] Add Fireworks AI Tool calling support
This commit is contained in:
Ishaan Jaff 2024-06-26 08:30:06 -07:00 committed by GitHub
commit ae431eb85d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 215 additions and 0 deletions

View file

@ -817,6 +817,7 @@ from .llms.openai import (
AzureAIStudioConfig,
)
from .llms.nvidia_nim import NvidiaNimConfig
from .llms.fireworks_ai import FireworksAIConfig
from .llms.text_completion_codestral import MistralTextCompletionConfig
from .llms.azure import (
AzureOpenAIConfig,

View file

@ -0,0 +1,108 @@
import types
from typing import Literal, Optional, Union
import litellm
class FireworksAIConfig:
"""
Reference: https://docs.fireworks.ai/api-reference/post-chatcompletions
The class `FireworksAIConfig` provides configuration for the Fireworks's Chat Completions API interface. Below are the parameters:
"""
tools: Optional[list] = None
tool_choice: Optional[Union[str, dict]] = None
max_tokens: Optional[int] = None
temperature: Optional[int] = None
top_p: Optional[int] = None
top_k: Optional[int] = None
frequency_penalty: Optional[int] = None
presence_penalty: Optional[int] = None
n: Optional[int] = None
stop: Optional[Union[str, list]] = None
response_format: Optional[dict] = None
user: Optional[str] = None
# Non OpenAI parameters - Fireworks AI only params
prompt_truncate_length: Optional[int] = None
context_length_exceeded_behavior: Optional[Literal["error", "truncate"]] = None
def __init__(
self,
tools: Optional[list] = None,
tool_choice: Optional[Union[str, dict]] = None,
max_tokens: Optional[int] = None,
temperature: Optional[int] = None,
top_p: Optional[int] = None,
top_k: Optional[int] = None,
frequency_penalty: Optional[int] = None,
presence_penalty: Optional[int] = None,
n: Optional[int] = None,
stop: Optional[Union[str, list]] = None,
response_format: Optional[dict] = None,
user: Optional[str] = None,
prompt_truncate_length: Optional[int] = None,
context_length_exceeded_behavior: Optional[Literal["error", "truncate"]] = None,
) -> None:
locals_ = locals().copy()
for key, value in locals_.items():
if key != "self" and value is not None:
setattr(self.__class__, key, value)
@classmethod
def get_config(cls):
return {
k: v
for k, v in cls.__dict__.items()
if not k.startswith("__")
and not isinstance(
v,
(
types.FunctionType,
types.BuiltinFunctionType,
classmethod,
staticmethod,
),
)
and v is not None
}
def get_supported_openai_params(self):
return [
"stream",
"tools",
"tool_choice",
"max_tokens",
"temperature",
"top_p",
"top_k",
"frequency_penalty",
"presence_penalty",
"n",
"stop",
"response_format",
"user",
"prompt_truncate_length",
"context_length_exceeded_behavior",
]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
) -> dict:
supported_openai_params = self.get_supported_openai_params()
for param, value in non_default_params.items():
if param == "tool_choice":
if value == "required":
# relevant issue: https://github.com/BerriAI/litellm/issues/4416
optional_params["tool_choice"] = "any"
else:
# pass through the value of tool choice
optional_params["tool_choice"] = value
elif param in supported_openai_params:
if value is not None:
optional_params[param] = value
return optional_params

View file

@ -2073,6 +2073,30 @@
"supports_function_calling": true,
"supports_vision": true
},
"openrouter/anthropic/claude-3-haiku-20240307": {
"max_tokens": 4096,
"max_input_tokens": 200000,
"max_output_tokens": 4096,
"input_cost_per_token": 0.00000025,
"output_cost_per_token": 0.00000125,
"litellm_provider": "openrouter",
"mode": "chat",
"supports_function_calling": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 264
},
"openrouter/anthropic/claude-3.5-sonnet": {
"max_tokens": 4096,
"max_input_tokens": 200000,
"max_output_tokens": 4096,
"input_cost_per_token": 0.000003,
"output_cost_per_token": 0.000015,
"litellm_provider": "openrouter",
"mode": "chat",
"supports_function_calling": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
},
"openrouter/anthropic/claude-3-sonnet": {
"max_tokens": 200000,
"input_cost_per_token": 0.000003,

View file

@ -1222,6 +1222,44 @@ def test_completion_fireworks_ai():
pytest.fail(f"Error occurred: {e}")
def test_fireworks_ai_tool_calling():
litellm.set_verbose = True
model_name = "fireworks_ai/accounts/fireworks/models/firefunction-v2"
tools = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
},
"required": ["location"],
},
},
}
]
messages = [
{
"role": "user",
"content": "What's the weather like in Boston today in Fahrenheit?",
}
]
response = completion(
model=model_name,
messages=messages,
tools=tools,
tool_choice="required",
)
print(response)
@pytest.mark.skip(reason="this test is flaky")
def test_completion_perplexity_api():
try:

View file

@ -0,0 +1,32 @@
import os
import sys
import pytest
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
from litellm.llms.fireworks_ai import FireworksAIConfig
fireworks = FireworksAIConfig()
def test_map_openai_params_tool_choice():
# Test case 1: tool_choice is "required"
result = fireworks.map_openai_params({"tool_choice": "required"}, {}, "some_model")
assert result == {"tool_choice": "any"}
# Test case 2: tool_choice is "auto"
result = fireworks.map_openai_params({"tool_choice": "auto"}, {}, "some_model")
assert result == {"tool_choice": "auto"}
# Test case 3: tool_choice is not present
result = fireworks.map_openai_params(
{"some_other_param": "value"}, {}, "some_model"
)
assert result == {}
# Test case 4: tool_choice is None
result = fireworks.map_openai_params({"tool_choice": None}, {}, "some_model")
assert result == {"tool_choice": None}

View file

@ -3079,6 +3079,16 @@ def get_optional_params(
optional_params = litellm.NvidiaNimConfig().map_openai_params(
non_default_params=non_default_params, optional_params=optional_params
)
elif custom_llm_provider == "fireworks_ai":
supported_params = get_supported_openai_params(
model=model, custom_llm_provider=custom_llm_provider
)
_check_valid_arg(supported_params=supported_params)
optional_params = litellm.FireworksAIConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
)
elif custom_llm_provider == "groq":
supported_params = get_supported_openai_params(
model=model, custom_llm_provider=custom_llm_provider
@ -3645,6 +3655,8 @@ def get_supported_openai_params(
return litellm.OllamaChatConfig().get_supported_openai_params()
elif custom_llm_provider == "anthropic":
return litellm.AnthropicConfig().get_supported_openai_params()
elif custom_llm_provider == "fireworks_ai":
return litellm.FireworksAIConfig().get_supported_openai_params()
elif custom_llm_provider == "nvidia_nim":
return litellm.NvidiaNimConfig().get_supported_openai_params()
elif custom_llm_provider == "groq":