From 4a9bdff5188a127f7672707f188c29040f4a4aa8 Mon Sep 17 00:00:00 2001 From: Iams4kura <126048986+Iams4kura@users.noreply.github.com> Date: Fri, 4 Sep 2026 11:16:14 +0800 Subject: [PATCH] refactor(llm): isolate OpenAI message mapping --- src/llm/core/types.py | 17 ------- src/llm/providers/astraflow.py | 3 +- src/llm/providers/atlas.py | 3 +- src/llm/providers/openai.py | 3 +- src/llm/providers/openai_compat.py | 79 ++++++++++++++++++++++++++++++ tests/test_provider_tools.py | 48 ++++++++++++------ 6 files changed, 118 insertions(+), 35 deletions(-) create mode 100644 src/llm/providers/openai_compat.py diff --git a/src/llm/core/types.py b/src/llm/core/types.py index 8ba80a258..84eb7cf8c 100644 --- a/src/llm/core/types.py +++ b/src/llm/core/types.py @@ -2,7 +2,6 @@ from __future__ import annotations -import json from dataclasses import dataclass, field from enum import Enum from typing import Any @@ -45,22 +44,6 @@ 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 229e373d5..2bcc6d6b2 100644 --- a/src/llm/providers/astraflow.py +++ b/src/llm/providers/astraflow.py @@ -16,6 +16,7 @@ from llm.core.interface import ( ) from llm.core.types import LLMInput, LLMOutput, ModelInfo, ProviderType, ToolCall from llm.providers.constants import EMPTY_FILTERED_RESPONSE_ERROR +from llm.providers.openai_compat import to_openai_message_param ASTRAFLOW_BASE_URL = "https://api.umodelverse.ai/v1" ASTRAFLOW_CN_BASE_URL = "https://api.modelverse.cn/v1" @@ -70,7 +71,7 @@ class _AstraflowBaseProvider(LLMProvider): try: params: dict[str, Any] = { "model": llm_input.model or self.default_model, - "messages": [msg.to_openai_dict() for msg in llm_input.messages], + "messages": [to_openai_message_param(msg) 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 6b34c74f6..d626509cc 100644 --- a/src/llm/providers/atlas.py +++ b/src/llm/providers/atlas.py @@ -16,6 +16,7 @@ from llm.core.interface import ( ) from llm.core.types import LLMInput, LLMOutput, ModelInfo, ProviderType, ToolCall from llm.providers.constants import EMPTY_FILTERED_RESPONSE_ERROR +from llm.providers.openai_compat import to_openai_message_param ATLAS_BASE_URL = "https://api.atlascloud.ai/v1" DEFAULT_ATLAS_MODEL = "deepseek-ai/deepseek-v4-pro" @@ -83,7 +84,7 @@ class AtlasProvider(LLMProvider): try: params: dict[str, Any] = { "model": llm_input.model or self.default_model, - "messages": [msg.to_openai_dict() for msg in llm_input.messages], + "messages": [to_openai_message_param(msg) 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 4759df84b..7d00d333c 100644 --- a/src/llm/providers/openai.py +++ b/src/llm/providers/openai.py @@ -22,6 +22,7 @@ from llm.core.types import ( ToolCall, ) from llm.providers.constants import EMPTY_FILTERED_RESPONSE_ERROR +from llm.providers.openai_compat import to_openai_message_param class OpenAIProvider(LLMProvider): @@ -72,7 +73,7 @@ class OpenAIProvider(LLMProvider): try: params: dict[str, Any] = { "model": input.model or "gpt-4o-mini", - "messages": [msg.to_openai_dict() for msg in input.messages], + "messages": [to_openai_message_param(msg) for msg in input.messages], "temperature": input.temperature, } if input.max_tokens: diff --git a/src/llm/providers/openai_compat.py b/src/llm/providers/openai_compat.py new file mode 100644 index 000000000..a497eea2f --- /dev/null +++ b/src/llm/providers/openai_compat.py @@ -0,0 +1,79 @@ +"""Message adapters shared by OpenAI-compatible providers.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING + +from llm.core.types import Message, Role, ToolCall + +if TYPE_CHECKING: + from openai.types.chat import ( + ChatCompletionAssistantMessageParam, + ChatCompletionMessageParam, + ChatCompletionMessageToolCallParam, + ChatCompletionSystemMessageParam, + ChatCompletionToolMessageParam, + ChatCompletionUserMessageParam, + ) + + +def _to_openai_tool_call_param( + tool_call: ToolCall, +) -> ChatCompletionMessageToolCallParam: + """Convert a domain tool call to the OpenAI request schema.""" + return { + "id": tool_call.id, + "type": "function", + "function": { + "name": tool_call.name, + "arguments": json.dumps(tool_call.arguments), + }, + } + + +def to_openai_message_param(message: Message) -> ChatCompletionMessageParam: + """Convert a domain message to the OpenAI chat-completions schema.""" + if message.role == Role.SYSTEM: + system_message: ChatCompletionSystemMessageParam = { + "role": "system", + "content": message.content, + } + if message.name: + system_message["name"] = message.name + return system_message + + if message.role == Role.USER: + user_message: ChatCompletionUserMessageParam = { + "role": "user", + "content": message.content, + } + if message.name: + user_message["name"] = message.name + return user_message + + if message.role == Role.ASSISTANT: + assistant_message: ChatCompletionAssistantMessageParam = { + "role": "assistant", + "content": message.content, + } + if message.name: + assistant_message["name"] = message.name + if message.tool_calls: + assistant_message["tool_calls"] = [ + _to_openai_tool_call_param(tool_call) + for tool_call in message.tool_calls + ] + return assistant_message + + if message.role == Role.TOOL: + if not message.tool_call_id: + raise ValueError("OpenAI tool messages require a tool_call_id") + tool_message: ChatCompletionToolMessageParam = { + "role": "tool", + "content": message.content, + "tool_call_id": message.tool_call_id, + } + return tool_message + + raise ValueError(f"Unsupported OpenAI message role: {message.role}") diff --git a/tests/test_provider_tools.py b/tests/test_provider_tools.py index e315bb35a..e0bea308b 100644 --- a/tests/test_provider_tools.py +++ b/tests/test_provider_tools.py @@ -1,7 +1,10 @@ +from collections.abc import Callable from types import SimpleNamespace +from typing import Any, cast import pytest +from llm.core.interface import LLMProvider from llm.core.types import LLMInput, Message, Role, ToolCall, ToolDefinition from llm.providers.astraflow import AstraflowProvider from llm.providers.atlas import AtlasProvider @@ -86,12 +89,15 @@ def test_openai_provider_serializes_tools_for_chat_completions(): ] +@pytest.mark.unit @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") +def test_openai_compatible_provider_serializes_assistant_tool_calls( + provider_type: Callable[..., LLMProvider], +) -> None: + provider = cast(Any, provider_type(api_key="test")) client = _OpenAIClient() provider.client = client @@ -119,20 +125,32 @@ def test_openai_compatible_provider_serializes_assistant_tool_calls(provider_typ ) ) - assert client.completions.params["messages"][1] == { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": { - "name": "search", - "arguments": '{"query": "ecc"}', + assert client.completions.params is not None + assert client.completions.params["messages"] == [ + {"role": "user", "content": "Find ECC"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "search", + "arguments": '{"query": "ecc"}', + }, }, - } - ], - } + ], + }, + {"role": "tool", "content": "result", "tool_call_id": "call_1"}, + ] + + client.completions.params = None + with pytest.raises(ValueError, match="tool messages require a tool_call_id"): + provider.generate( + LLMInput(messages=[Message(role=Role.TOOL, content="unlinked result")]) + ) + assert client.completions.params is None def test_openai_provider_can_be_constructed_without_credentials(monkeypatch):