forked from phoenix/litellm-mirror
Merge pull request #4418 from BerriAI/litellm_fireworks_ai_tool_calling
[Feat] Add Fireworks AI Tool calling support
This commit is contained in:
commit
ae431eb85d
6 changed files with 215 additions and 0 deletions
|
@ -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,
|
||||
|
|
108
litellm/llms/fireworks_ai.py
Normal file
108
litellm/llms/fireworks_ai.py
Normal 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
|
|
@ -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,
|
||||
|
|
|
@ -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:
|
||||
|
|
32
litellm/tests/test_fireworks_ai.py
Normal file
32
litellm/tests/test_fireworks_ai.py
Normal 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}
|
|
@ -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":
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue