feat(snowflake): add function calling support for Snowflake Cortex REST API
Add support for function calling (tools) with Snowflake Cortex models that support it (e.g., Claude 3.5 Sonnet).
Changes:
- Add 'tools' and 'tool_choice' to supported OpenAI parameters
- Implement request transformation: OpenAI function format → Snowflake tool_spec format
- Implement response transformation: Snowflake content_list with tool_use → OpenAI tool_calls
- Add tool_choice transformation: OpenAI nested format → Snowflake array format
Request transformation:
- Transform tools from nested {"type": "function", "function": {...}} to Snowflake's {"tool_spec": {"type": "generic", "name": "...", "input_schema": {...}}}
- Transform tool_choice from {"type": "function", "function": {"name": "..."}} to {"type": "tool", "name": ["..."]}
Response transformation:
- Parse Snowflake's content_list array containing tool_use objects
- Extract tool calls with tool_use_id, name, and input
- Convert to OpenAI's tool_calls format with proper JSON serialization
Testing:
- Add 7 unit tests covering request/response transformations
- Add integration test for Responses API with tool calling
- All tests passing
Fixes issue #15218
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Javier de la Torre
parent
b348a26bdc
commit
df232a71f1
@ -1,14 +1,15 @@
|
||||
"""
|
||||
Support for Snowflake REST API
|
||||
Support for Snowflake REST API
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, List, Optional, Tuple
|
||||
import json
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import ModelResponse
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall, Function, ModelResponse
|
||||
|
||||
from ...openai_like.chat.transformation import OpenAIGPTConfig
|
||||
|
||||
@ -22,15 +23,25 @@ else:
|
||||
|
||||
class SnowflakeConfig(OpenAIGPTConfig):
|
||||
"""
|
||||
source: https://docs.snowflake.com/en/sql-reference/functions/complete-snowflake-cortex
|
||||
Reference: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api
|
||||
|
||||
Snowflake Cortex LLM REST API supports function calling with specific models (e.g., Claude 3.5 Sonnet).
|
||||
This config handles transformation between OpenAI format and Snowflake's tool_spec format.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def get_config(cls):
|
||||
return super().get_config()
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> List:
|
||||
return ["temperature", "max_tokens", "top_p", "response_format"]
|
||||
def get_supported_openai_params(self, model: str) -> List[str]:
|
||||
return [
|
||||
"temperature",
|
||||
"max_tokens",
|
||||
"top_p",
|
||||
"response_format",
|
||||
"tools",
|
||||
"tool_choice",
|
||||
]
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
@ -56,6 +67,57 @@ class SnowflakeConfig(OpenAIGPTConfig):
|
||||
optional_params[param] = value
|
||||
return optional_params
|
||||
|
||||
def _transform_tool_calls_from_snowflake_to_openai(
|
||||
self, content_list: List[Dict[str, Any]]
|
||||
) -> Tuple[str, Optional[List[ChatCompletionMessageToolCall]]]:
|
||||
"""
|
||||
Transform Snowflake tool calls to OpenAI format.
|
||||
|
||||
Args:
|
||||
content_list: Snowflake's content_list array containing text and tool_use items
|
||||
|
||||
Returns:
|
||||
Tuple of (text_content, tool_calls)
|
||||
|
||||
Snowflake format in content_list:
|
||||
{
|
||||
"type": "tool_use",
|
||||
"tool_use": {
|
||||
"tool_use_id": "tooluse_...",
|
||||
"name": "get_weather",
|
||||
"input": {"location": "Paris"}
|
||||
}
|
||||
}
|
||||
|
||||
OpenAI format (returned tool_calls):
|
||||
ChatCompletionMessageToolCall(
|
||||
id="tooluse_...",
|
||||
type="function",
|
||||
function=Function(name="get_weather", arguments='{"location": "Paris"}')
|
||||
)
|
||||
"""
|
||||
text_content = ""
|
||||
tool_calls: List[ChatCompletionMessageToolCall] = []
|
||||
|
||||
for idx, content_item in enumerate(content_list):
|
||||
if content_item.get("type") == "text":
|
||||
text_content += content_item.get("text", "")
|
||||
|
||||
## TOOL CALLING
|
||||
elif content_item.get("type") == "tool_use":
|
||||
tool_use_data = content_item.get("tool_use", {})
|
||||
tool_call = ChatCompletionMessageToolCall(
|
||||
id=tool_use_data.get("tool_use_id", ""),
|
||||
type="function",
|
||||
function=Function(
|
||||
name=tool_use_data.get("name", ""),
|
||||
arguments=json.dumps(tool_use_data.get("input", {})),
|
||||
),
|
||||
)
|
||||
tool_calls.append(tool_call)
|
||||
|
||||
return text_content, tool_calls if tool_calls else None
|
||||
|
||||
def transform_response(
|
||||
self,
|
||||
model: str,
|
||||
@ -71,6 +133,7 @@ class SnowflakeConfig(OpenAIGPTConfig):
|
||||
json_mode: Optional[bool] = None,
|
||||
) -> ModelResponse:
|
||||
response_json = raw_response.json()
|
||||
|
||||
logging_obj.post_call(
|
||||
input=messages,
|
||||
api_key="",
|
||||
@ -78,6 +141,26 @@ class SnowflakeConfig(OpenAIGPTConfig):
|
||||
additional_args={"complete_input_dict": request_data},
|
||||
)
|
||||
|
||||
## RESPONSE TRANSFORMATION
|
||||
# Snowflake returns content_list (not content) with tool_use objects
|
||||
# We need to transform this to OpenAI's format with content + tool_calls
|
||||
if "choices" in response_json and len(response_json["choices"]) > 0:
|
||||
choice = response_json["choices"][0]
|
||||
if "message" in choice and "content_list" in choice["message"]:
|
||||
content_list = choice["message"]["content_list"]
|
||||
(
|
||||
text_content,
|
||||
tool_calls,
|
||||
) = self._transform_tool_calls_from_snowflake_to_openai(content_list)
|
||||
|
||||
# Update the choice message with OpenAI format
|
||||
choice["message"]["content"] = text_content
|
||||
if tool_calls:
|
||||
choice["message"]["tool_calls"] = tool_calls
|
||||
|
||||
# Remove Snowflake-specific content_list
|
||||
del choice["message"]["content_list"]
|
||||
|
||||
returned_response = ModelResponse(**response_json)
|
||||
|
||||
returned_response.model = "snowflake/" + (returned_response.model or "")
|
||||
@ -150,6 +233,95 @@ class SnowflakeConfig(OpenAIGPTConfig):
|
||||
|
||||
return api_base
|
||||
|
||||
def _transform_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Transform OpenAI tool format to Snowflake tool format.
|
||||
|
||||
Args:
|
||||
tools: List of tools in OpenAI format
|
||||
|
||||
Returns:
|
||||
List of tools in Snowflake format
|
||||
|
||||
OpenAI format:
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "...",
|
||||
"parameters": {...}
|
||||
}
|
||||
}
|
||||
|
||||
Snowflake format:
|
||||
{
|
||||
"tool_spec": {
|
||||
"type": "generic",
|
||||
"name": "get_weather",
|
||||
"description": "...",
|
||||
"input_schema": {...}
|
||||
}
|
||||
}
|
||||
"""
|
||||
snowflake_tools: List[Dict[str, Any]] = []
|
||||
for tool in tools:
|
||||
if tool.get("type") == "function":
|
||||
function = tool.get("function", {})
|
||||
snowflake_tool: Dict[str, Any] = {
|
||||
"tool_spec": {
|
||||
"type": "generic",
|
||||
"name": function.get("name"),
|
||||
"input_schema": function.get(
|
||||
"parameters",
|
||||
{"type": "object", "properties": {}},
|
||||
),
|
||||
}
|
||||
}
|
||||
# Add description if present
|
||||
if "description" in function:
|
||||
snowflake_tool["tool_spec"]["description"] = function[
|
||||
"description"
|
||||
]
|
||||
|
||||
snowflake_tools.append(snowflake_tool)
|
||||
|
||||
return snowflake_tools
|
||||
|
||||
def _transform_tool_choice(
|
||||
self, tool_choice: Union[str, Dict[str, Any]]
|
||||
) -> Union[str, Dict[str, Any]]:
|
||||
"""
|
||||
Transform OpenAI tool_choice format to Snowflake format.
|
||||
|
||||
Args:
|
||||
tool_choice: Tool choice in OpenAI format (str or dict)
|
||||
|
||||
Returns:
|
||||
Tool choice in Snowflake format
|
||||
|
||||
OpenAI format:
|
||||
{"type": "function", "function": {"name": "get_weather"}}
|
||||
|
||||
Snowflake format:
|
||||
{"type": "tool", "name": ["get_weather"]}
|
||||
|
||||
Note: String values ("auto", "required", "none") pass through unchanged.
|
||||
"""
|
||||
if isinstance(tool_choice, str):
|
||||
# "auto", "required", "none" pass through as-is
|
||||
return tool_choice
|
||||
|
||||
if isinstance(tool_choice, dict):
|
||||
if tool_choice.get("type") == "function":
|
||||
function_name = tool_choice.get("function", {}).get("name")
|
||||
if function_name:
|
||||
return {
|
||||
"type": "tool",
|
||||
"name": [function_name], # Snowflake expects array
|
||||
}
|
||||
|
||||
return tool_choice
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
model: str,
|
||||
@ -160,6 +332,18 @@ class SnowflakeConfig(OpenAIGPTConfig):
|
||||
) -> dict:
|
||||
stream: bool = optional_params.pop("stream", None) or False
|
||||
extra_body = optional_params.pop("extra_body", {})
|
||||
|
||||
## TOOL CALLING
|
||||
# Transform tools from OpenAI format to Snowflake's tool_spec format
|
||||
tools = optional_params.pop("tools", None)
|
||||
if tools:
|
||||
optional_params["tools"] = self._transform_tools(tools)
|
||||
|
||||
# Transform tool_choice from OpenAI format to Snowflake's tool name array format
|
||||
tool_choice = optional_params.pop("tool_choice", None)
|
||||
if tool_choice:
|
||||
optional_params["tool_choice"] = self._transform_tool_choice(tool_choice)
|
||||
|
||||
return {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
|
||||
@ -6,7 +6,7 @@ from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
import pytest
|
||||
|
||||
from litellm import completion, acompletion
|
||||
from litellm import completion, acompletion, responses
|
||||
from litellm.exceptions import APIConnectionError
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@ -87,3 +87,70 @@ async def test_chat_completion_snowflake_stream(sync_mode):
|
||||
raise # Re-raise if it's a different APIConnectionError
|
||||
except Exception as e:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Requires Snowflake credentials - run manually when needed")
|
||||
def test_snowflake_tool_calling_responses_api():
|
||||
"""
|
||||
Test Snowflake tool calling with Responses API.
|
||||
Requires SNOWFLAKE_JWT and SNOWFLAKE_ACCOUNT_ID environment variables.
|
||||
"""
|
||||
import litellm
|
||||
|
||||
# Skip if credentials not available
|
||||
if not os.getenv("SNOWFLAKE_JWT") or not os.getenv("SNOWFLAKE_ACCOUNT_ID"):
|
||||
pytest.skip("Snowflake credentials not available")
|
||||
|
||||
litellm.drop_params = False # We now support tools!
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "get_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",
|
||||
}
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
try:
|
||||
# Test with tool_choice to force tool use
|
||||
response = responses(
|
||||
model="snowflake/claude-3-5-sonnet",
|
||||
input="What's the weather in Paris?",
|
||||
tools=tools,
|
||||
tool_choice={"type": "function", "function": {"name": "get_weather"}},
|
||||
max_output_tokens=200,
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert hasattr(response, "output")
|
||||
assert len(response.output) > 0
|
||||
|
||||
# Verify tool call was made
|
||||
tool_call_found = False
|
||||
for item in response.output:
|
||||
if hasattr(item, "type") and item.type == "function_call":
|
||||
tool_call_found = True
|
||||
assert item.name == "get_weather"
|
||||
assert hasattr(item, "arguments")
|
||||
print(f"✅ Tool call detected: {item.name}({item.arguments})")
|
||||
break
|
||||
|
||||
assert tool_call_found, "Expected tool call but none was found"
|
||||
|
||||
except APIConnectionError as e:
|
||||
if "JWT token is invalid" in str(e):
|
||||
pytest.skip("Invalid Snowflake JWT token")
|
||||
elif "Application failed to respond" in str(e) or "502" in str(e):
|
||||
pytest.skip(f"Snowflake API unavailable: {e}")
|
||||
else:
|
||||
raise
|
||||
|
||||
@ -0,0 +1,315 @@
|
||||
"""
|
||||
Unit tests for Snowflake chat transformation
|
||||
Tests tool calling request/response transformations
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.snowflake.chat.transformation import SnowflakeConfig
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
||||
class TestSnowflakeToolTransformation:
|
||||
"""Test suite for Snowflake tool calling transformations"""
|
||||
|
||||
def test_transform_request_with_tools(self):
|
||||
"""
|
||||
Test that OpenAI tool format is correctly transformed to Snowflake's tool_spec format.
|
||||
"""
|
||||
config = SnowflakeConfig()
|
||||
|
||||
# OpenAI format tools
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_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"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
optional_params = {"tools": tools}
|
||||
|
||||
transformed_request = config.transform_request(
|
||||
model="claude-3-5-sonnet",
|
||||
messages=[{"role": "user", "content": "What's the weather?"}],
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
# Verify tools were transformed to Snowflake format
|
||||
assert "tools" in transformed_request
|
||||
assert len(transformed_request["tools"]) == 1
|
||||
|
||||
snowflake_tool = transformed_request["tools"][0]
|
||||
assert "tool_spec" in snowflake_tool
|
||||
assert snowflake_tool["tool_spec"]["type"] == "generic"
|
||||
assert snowflake_tool["tool_spec"]["name"] == "get_weather"
|
||||
assert snowflake_tool["tool_spec"]["description"] == "Get the current weather in a given location"
|
||||
assert "input_schema" in snowflake_tool["tool_spec"]
|
||||
assert snowflake_tool["tool_spec"]["input_schema"]["type"] == "object"
|
||||
assert "location" in snowflake_tool["tool_spec"]["input_schema"]["properties"]
|
||||
|
||||
def test_transform_request_with_tool_choice(self):
|
||||
"""
|
||||
Test that OpenAI tool_choice format is correctly transformed to Snowflake format.
|
||||
"""
|
||||
config = SnowflakeConfig()
|
||||
|
||||
# OpenAI format tool_choice
|
||||
tool_choice = {"type": "function", "function": {"name": "get_weather"}}
|
||||
|
||||
optional_params = {"tool_choice": tool_choice}
|
||||
|
||||
transformed_request = config.transform_request(
|
||||
model="claude-3-5-sonnet",
|
||||
messages=[{"role": "user", "content": "What's the weather?"}],
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
# Verify tool_choice was transformed to Snowflake format
|
||||
assert "tool_choice" in transformed_request
|
||||
assert transformed_request["tool_choice"]["type"] == "tool"
|
||||
assert transformed_request["tool_choice"]["name"] == ["get_weather"] # Array format
|
||||
|
||||
def test_transform_request_with_string_tool_choice(self):
|
||||
"""
|
||||
Test that string tool_choice values pass through unchanged.
|
||||
"""
|
||||
config = SnowflakeConfig()
|
||||
|
||||
for value in ["auto", "required", "none"]:
|
||||
optional_params = {"tool_choice": value}
|
||||
|
||||
transformed_request = config.transform_request(
|
||||
model="claude-3-5-sonnet",
|
||||
messages=[{"role": "user", "content": "Test"}],
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert transformed_request["tool_choice"] == value
|
||||
|
||||
def test_transform_response_with_tool_calls(self):
|
||||
"""
|
||||
Test that Snowflake's content_list with tool_use is transformed to OpenAI format.
|
||||
"""
|
||||
config = SnowflakeConfig()
|
||||
|
||||
# Mock Snowflake response with tool call
|
||||
mock_snowflake_response = {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content_list": [
|
||||
{"type": "text", "text": ""},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"tool_use": {
|
||||
"tool_use_id": "tooluse_abc123",
|
||||
"name": "get_weather",
|
||||
"input": {"location": "Paris, France", "unit": "celsius"},
|
||||
},
|
||||
},
|
||||
]
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
|
||||
}
|
||||
|
||||
response = httpx.Response(
|
||||
status_code=200,
|
||||
json=mock_snowflake_response,
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
|
||||
model_response = ModelResponse(
|
||||
choices=[litellm.Choices(index=0, message=litellm.Message())]
|
||||
)
|
||||
|
||||
logging_obj = MagicMock()
|
||||
|
||||
result = config.transform_response(
|
||||
model="claude-3-5-sonnet",
|
||||
raw_response=response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
request_data={},
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding={},
|
||||
)
|
||||
|
||||
# General assertions
|
||||
assert isinstance(result, ModelResponse)
|
||||
assert len(result.choices) == 1
|
||||
|
||||
choice = result.choices[0]
|
||||
assert isinstance(choice, litellm.Choices)
|
||||
|
||||
# Message and tool_calls assertions
|
||||
message = choice.message
|
||||
assert isinstance(message, litellm.Message)
|
||||
assert hasattr(message, "tool_calls")
|
||||
assert isinstance(message.tool_calls, list)
|
||||
assert len(message.tool_calls) == 1
|
||||
|
||||
# Specific tool_call assertions
|
||||
tool_call = message.tool_calls[0]
|
||||
assert isinstance(tool_call, litellm.utils.ChatCompletionMessageToolCall)
|
||||
assert tool_call.id == "tooluse_abc123"
|
||||
assert tool_call.type == "function"
|
||||
assert tool_call.function.name == "get_weather"
|
||||
|
||||
# Verify arguments are properly JSON serialized
|
||||
arguments = json.loads(tool_call.function.arguments)
|
||||
assert arguments["location"] == "Paris, France"
|
||||
assert arguments["unit"] == "celsius"
|
||||
|
||||
# Verify content_list was removed and content was set
|
||||
assert message.content == ""
|
||||
|
||||
def test_transform_response_with_mixed_content(self):
|
||||
"""
|
||||
Test that responses with both text and tool calls are handled correctly.
|
||||
"""
|
||||
config = SnowflakeConfig()
|
||||
|
||||
# Mock Snowflake response with text and tool call
|
||||
mock_snowflake_response = {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content_list": [
|
||||
{"type": "text", "text": "Let me check the weather for you. "},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"tool_use": {
|
||||
"tool_use_id": "tooluse_xyz789",
|
||||
"name": "get_weather",
|
||||
"input": {"location": "Tokyo, Japan"},
|
||||
},
|
||||
},
|
||||
]
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 15, "completion_tokens": 25, "total_tokens": 40},
|
||||
}
|
||||
|
||||
response = httpx.Response(
|
||||
status_code=200,
|
||||
json=mock_snowflake_response,
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
|
||||
model_response = ModelResponse(
|
||||
choices=[litellm.Choices(index=0, message=litellm.Message())]
|
||||
)
|
||||
|
||||
logging_obj = MagicMock()
|
||||
|
||||
result = config.transform_response(
|
||||
model="claude-3-5-sonnet",
|
||||
raw_response=response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
request_data={},
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding={},
|
||||
)
|
||||
|
||||
# Verify text content was extracted
|
||||
message = result.choices[0].message
|
||||
assert message.content == "Let me check the weather for you. "
|
||||
|
||||
# Verify tool call was also extracted
|
||||
assert len(message.tool_calls) == 1
|
||||
assert message.tool_calls[0].function.name == "get_weather"
|
||||
|
||||
def test_transform_response_without_tool_calls(self):
|
||||
"""
|
||||
Test that regular text responses (without tools) work correctly.
|
||||
"""
|
||||
config = SnowflakeConfig()
|
||||
|
||||
# Mock Snowflake response without tool calls (standard response)
|
||||
mock_snowflake_response = {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": "Hello! I'm doing well, thank you for asking.",
|
||||
"role": "assistant",
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 15, "total_tokens": 25},
|
||||
}
|
||||
|
||||
response = httpx.Response(
|
||||
status_code=200,
|
||||
json=mock_snowflake_response,
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
|
||||
model_response = ModelResponse(
|
||||
choices=[litellm.Choices(index=0, message=litellm.Message())]
|
||||
)
|
||||
|
||||
logging_obj = MagicMock()
|
||||
|
||||
result = config.transform_response(
|
||||
model="mistral-7b",
|
||||
raw_response=response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
request_data={},
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding={},
|
||||
)
|
||||
|
||||
# Verify standard response works
|
||||
assert isinstance(result, ModelResponse)
|
||||
assert result.choices[0].message.content == "Hello! I'm doing well, thank you for asking."
|
||||
|
||||
def test_get_supported_openai_params_includes_tools(self):
|
||||
"""
|
||||
Test that tools and tool_choice are in supported params.
|
||||
"""
|
||||
config = SnowflakeConfig()
|
||||
supported_params = config.get_supported_openai_params("claude-3-5-sonnet")
|
||||
|
||||
assert "tools" in supported_params
|
||||
assert "tool_choice" in supported_params
|
||||
assert "temperature" in supported_params
|
||||
assert "max_tokens" in supported_params
|
||||
Loading…
Reference in New Issue
Block a user