115 lines
4.2 KiB
Python
115 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from app.ai_provider.config.models import ModelConfig, ParameterMapping, ProviderConfig
|
|
from app.ai_provider.llm.openai_format import OpenAIFormatAdapter
|
|
from app.ai_provider.types.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
|
|
|
|
|
class TestOpenAIMessageConversion:
|
|
@pytest.fixture
|
|
def adapter(self):
|
|
provider = ProviderConfig(
|
|
name="deepseek",
|
|
format="openai",
|
|
base_url_env="DEEPSEEK_BASE_URL",
|
|
api_key_env="DEEPSEEK_API_KEY",
|
|
parameter_mappings={
|
|
"temperature": ParameterMapping(provider_param="temperature", min=0, max=2),
|
|
"max_tokens": ParameterMapping(provider_param="max_tokens"),
|
|
},
|
|
)
|
|
return OpenAIFormatAdapter(provider_config=provider)
|
|
|
|
def test_convert_system_message(self, adapter):
|
|
msg = SystemMessage(content="You are a helpful assistant.")
|
|
result = adapter._convert_message(msg)
|
|
assert result["role"] == "system"
|
|
assert result["content"] == "You are a helpful assistant."
|
|
|
|
def test_convert_human_message(self, adapter):
|
|
msg = HumanMessage(content="Hello!")
|
|
result = adapter._convert_message(msg)
|
|
assert result["role"] == "user"
|
|
|
|
def test_convert_ai_message(self, adapter):
|
|
msg = AIMessage(
|
|
content="Sure!",
|
|
tool_calls=[{"name": "read_file", "args": {"path": "x.txt"}, "id": "call_1"}],
|
|
)
|
|
result = adapter._convert_message(msg)
|
|
assert result["role"] == "assistant"
|
|
assert "tool_calls" in result
|
|
|
|
def test_convert_tool_message(self, adapter):
|
|
msg = ToolMessage(content="File contents", tool_call_id="call_1")
|
|
result = adapter._convert_message(msg)
|
|
assert result["role"] == "tool"
|
|
assert result["tool_call_id"] == "call_1"
|
|
|
|
def test_build_request_body(self, adapter):
|
|
model = ModelConfig(
|
|
name="test", provider="deepseek", model="deepseek-v4-flash", format="openai"
|
|
)
|
|
messages = [HumanMessage(content="Hi")]
|
|
body = adapter.build_request(
|
|
model=model,
|
|
messages=messages,
|
|
params={"temperature": 0.7, "max_tokens": 4096},
|
|
)
|
|
assert body["model"] == "deepseek-v4-flash"
|
|
assert body["temperature"] == 0.7
|
|
assert body["max_tokens"] == 4096
|
|
assert len(body["messages"]) == 1
|
|
|
|
def test_parse_chunk(self, adapter):
|
|
raw = {"choices": [{"delta": {"content": "Hello"}, "index": 0}], "usage": None}
|
|
chunk = adapter.parse_chunk(raw)
|
|
assert chunk.delta_content == "Hello"
|
|
|
|
def test_parse_chunk_with_usage(self, adapter):
|
|
raw = {
|
|
"choices": [{"delta": {}, "index": 0, "finish_reason": "stop"}],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
|
}
|
|
chunk = adapter.parse_chunk(raw)
|
|
assert chunk.usage is not None
|
|
assert chunk.usage.prompt_tokens == 10
|
|
|
|
def test_parse_tool_call_chunk(self, adapter):
|
|
raw = {
|
|
"choices": [
|
|
{
|
|
"delta": {
|
|
"tool_calls": [
|
|
{
|
|
"index": 0,
|
|
"id": "call_1",
|
|
"function": {
|
|
"name": "read_file",
|
|
"arguments": '{"path": "test.txt"}',
|
|
},
|
|
}
|
|
],
|
|
},
|
|
"index": 0,
|
|
}
|
|
],
|
|
}
|
|
chunk = adapter.parse_chunk(raw)
|
|
assert len(chunk.tool_call_chunks) == 1
|
|
assert chunk.tool_call_chunks[0]["name"] == "read_file"
|
|
|
|
def test_bind_tools(self, adapter):
|
|
tools = [
|
|
{
|
|
"type": "function",
|
|
"function": {"name": "read_file", "parameters": {"type": "object"}},
|
|
}
|
|
]
|
|
body = {}
|
|
adapter.bind_tools(body, tools, tool_choice="auto")
|
|
assert body["tools"] == tools
|
|
assert body["tool_choice"] == "auto"
|