diff --git a/src/llm/core/types.py b/src/llm/core/types.py index 84eb7cf8c..8ba80a258 100644 --- a/src/llm/core/types.py +++ b/src/llm/core/types.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json from dataclasses import dataclass, field from enum import Enum from typing import Any @@ -44,6 +45,22 @@ class Message: ] return result + def to_openai_dict(self) -> dict[str, Any]: + result = self.to_dict() + if self.tool_calls: + result["tool_calls"] = [ + { + "id": tc.id, + "type": "function", + "function": { + "name": tc.name, + "arguments": json.dumps(tc.arguments), + }, + } + for tc in self.tool_calls + ] + return result + @dataclass(frozen=True) class ToolDefinition: diff --git a/src/llm/providers/astraflow.py b/src/llm/providers/astraflow.py index 1f7cceab9..229e373d5 100644 --- a/src/llm/providers/astraflow.py +++ b/src/llm/providers/astraflow.py @@ -70,7 +70,7 @@ class _AstraflowBaseProvider(LLMProvider): try: params: dict[str, Any] = { "model": llm_input.model or self.default_model, - "messages": [msg.to_dict() for msg in llm_input.messages], + "messages": [msg.to_openai_dict() for msg in llm_input.messages], } if llm_input.temperature != 1.0: params["temperature"] = llm_input.temperature diff --git a/src/llm/providers/atlas.py b/src/llm/providers/atlas.py index b266bb614..6b34c74f6 100644 --- a/src/llm/providers/atlas.py +++ b/src/llm/providers/atlas.py @@ -83,7 +83,7 @@ class AtlasProvider(LLMProvider): try: params: dict[str, Any] = { "model": llm_input.model or self.default_model, - "messages": [msg.to_dict() for msg in llm_input.messages], + "messages": [msg.to_openai_dict() for msg in llm_input.messages], } if llm_input.temperature != 1.0: params["temperature"] = llm_input.temperature diff --git a/src/llm/providers/openai.py b/src/llm/providers/openai.py index 0bf84a33d..4759df84b 100644 --- a/src/llm/providers/openai.py +++ b/src/llm/providers/openai.py @@ -72,7 +72,7 @@ class OpenAIProvider(LLMProvider): try: params: dict[str, Any] = { "model": input.model or "gpt-4o-mini", - "messages": [msg.to_dict() for msg in input.messages], + "messages": [msg.to_openai_dict() for msg in input.messages], "temperature": input.temperature, } if input.max_tokens: diff --git a/tests/test_provider_tools.py b/tests/test_provider_tools.py index 4c9c76f92..e315bb35a 100644 --- a/tests/test_provider_tools.py +++ b/tests/test_provider_tools.py @@ -2,7 +2,9 @@ from types import SimpleNamespace import pytest -from llm.core.types import LLMInput, Message, Role, ToolDefinition +from llm.core.types import LLMInput, Message, Role, ToolCall, ToolDefinition +from llm.providers.astraflow import AstraflowProvider +from llm.providers.atlas import AtlasProvider from llm.providers.claude import ClaudeProvider from llm.providers.constants import EMPTY_FILTERED_RESPONSE_ERROR from llm.providers.openai import OpenAIProvider @@ -84,6 +86,55 @@ def test_openai_provider_serializes_tools_for_chat_completions(): ] +@pytest.mark.parametrize( + "provider_type", + [OpenAIProvider, AstraflowProvider, AtlasProvider], +) +def test_openai_compatible_provider_serializes_assistant_tool_calls(provider_type): + provider = provider_type(api_key="test") + client = _OpenAIClient() + provider.client = client + + provider.generate( + LLMInput( + messages=[ + Message(role=Role.USER, content="Find ECC"), + Message( + role=Role.ASSISTANT, + content="", + tool_calls=[ + ToolCall( + id="call_1", + name="search", + arguments={"query": "ecc"}, + ) + ], + ), + Message( + role=Role.TOOL, + content="result", + tool_call_id="call_1", + ), + ] + ) + ) + + assert client.completions.params["messages"][1] == { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "search", + "arguments": '{"query": "ecc"}', + }, + } + ], + } + + def test_openai_provider_can_be_constructed_without_credentials(monkeypatch): monkeypatch.delenv("OPENAI_API_KEY", raising=False)