diff --git a/tests/test_litellm/llms/heroku/test_heroku_chat_transformation.py b/tests/test_litellm/llms/heroku/test_heroku_chat_transformation.py index 4359bce3f70..f70392db040 100644 --- a/tests/test_litellm/llms/heroku/test_heroku_chat_transformation.py +++ b/tests/test_litellm/llms/heroku/test_heroku_chat_transformation.py @@ -80,4 +80,87 @@ class TestHerokuChatConfig: assert request.headers["Authorization"] == f"Bearer {os.environ['HEROKU_API_KEY']}" assert request.headers["Content-Type"] == "application/json" - assert response.choices[0].message.content == "It's me, Mia! How are you?" \ No newline at end of file + assert response.choices[0].message.content == "It's me, Mia! How are you?" + + @pytest.mark.respx() + def test_heroku_tool_calling(self, respx_mock): + """Test that the Heroku tool calling API is called correctly""" + config = HerokuChatConfig() + headers = {} + api_key = "fake-heroku-key" + + litellm.disable_aiohttp_transport = True + + model = "heroku/claude-4-sonnet" + + respx_mock.post("https://us.inference.heroku.com/v1/chat/completions").respond( + json={ + "id": "chatcmpl-1859428879fc791b17d73", + "object": "chat.completion", + "created": 1754506683, + "model": "claude-4-sonnet", + "system_fingerprint": "heroku-inf-cp42st", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "refusal": None, + "tool_calls": [ + { + "id": "tooluse_dV3Vtnb-S9-Z_YFicSv2Gw", + "type": "function", + "function": { + "name": "get_current_weather", + "arguments": "{\"location\":\"Portland, OR\"}" + } + } + ], + "content": "Let me check the current weather in Portland for you." + }, + "finish_reason": "tool_calls" + } + ], + "usage": { + "prompt_tokens": 354, + "completion_tokens": 69, + "total_tokens": 423 + } + }, + status_code=200, + ) + + response = completion( + model=model, + messages=[{"role": "user", "content": "What's the weather in Portland?"}], + 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. Portland, OR" + } + }, + "required": [ + "location" + ] + } + } + }], + tool_choice="auto", + ) + print(response) + assert response.choices[0].message.content == "Let me check the current weather in Portland for you." + assert response.choices[0].message.tool_calls[0].id == "tooluse_dV3Vtnb-S9-Z_YFicSv2Gw" + assert response.choices[0].message.tool_calls[0].type == "function" + assert response.choices[0].message.tool_calls[0].function.name == "get_current_weather" + assert response.choices[0].message.tool_calls[0].function.arguments == "{\"location\":\"Portland, OR\"}" + + assert response.usage.prompt_tokens == 354 + assert response.usage.completion_tokens == 69 + assert response.usage.total_tokens == 423 \ No newline at end of file