mirror of
https://github.com/affaan-m/ECC.git
synced 2026-10-14 12:15:29 +02:00
523 lines
19 KiB
Python
523 lines
19 KiB
Python
import asyncio
|
|
import json
|
|
import urllib.request
|
|
from collections.abc import Callable
|
|
from io import BytesIO
|
|
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
|
|
from llm.providers.claude import ClaudeProvider
|
|
from llm.providers.constants import EMPTY_FILTERED_RESPONSE_ERROR
|
|
from llm.providers.ollama import OllamaProvider
|
|
from llm.providers.openai import OpenAIProvider
|
|
from llm.tools import ReActAgent, ToolExecutor, ToolRegistry
|
|
|
|
|
|
def _tool() -> ToolDefinition:
|
|
return ToolDefinition(
|
|
name="search",
|
|
description="Search",
|
|
parameters={"type": "object", "properties": {"query": {"type": "string"}}},
|
|
)
|
|
|
|
|
|
class _OpenAICompletions:
|
|
def __init__(self, response: SimpleNamespace | None = None) -> None:
|
|
self.params = None
|
|
self.response = response
|
|
|
|
def create(self, **params):
|
|
self.params = params
|
|
if self.response:
|
|
return self.response
|
|
return _openai_response(model=params["model"])
|
|
|
|
|
|
class _OpenAIClient:
|
|
def __init__(self, response: SimpleNamespace | None = None) -> None:
|
|
self.completions = _OpenAICompletions(response=response)
|
|
self.chat = SimpleNamespace(completions=self.completions)
|
|
|
|
|
|
class _AnthropicMessages:
|
|
def __init__(self) -> None:
|
|
self.params = None
|
|
|
|
def create(self, **params):
|
|
self.params = params
|
|
return SimpleNamespace(
|
|
content=[SimpleNamespace(text="ok", type="text")],
|
|
model=params["model"],
|
|
usage=SimpleNamespace(input_tokens=1, output_tokens=1),
|
|
stop_reason="end_turn",
|
|
)
|
|
|
|
|
|
class _AnthropicClient:
|
|
def __init__(self) -> None:
|
|
self.messages = _AnthropicMessages()
|
|
self.api_key = "test"
|
|
|
|
|
|
def _openai_response(**overrides) -> SimpleNamespace:
|
|
defaults = {
|
|
"choices": [SimpleNamespace(message=SimpleNamespace(content="ok", tool_calls=None), finish_reason="stop")],
|
|
"model": "gpt-4o-mini",
|
|
"usage": SimpleNamespace(prompt_tokens=1, completion_tokens=1, total_tokens=2),
|
|
}
|
|
defaults.update(overrides)
|
|
return SimpleNamespace(**defaults)
|
|
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize(
|
|
("default_model", "request_model", "expected_model"),
|
|
[
|
|
(None, None, "gpt-4o-mini"),
|
|
(None, "gpt-4o", "gpt-4o"),
|
|
("gpt-4o", None, "gpt-4o"),
|
|
("gpt-4o", "gpt-4-turbo", "gpt-4-turbo"),
|
|
],
|
|
)
|
|
def test_openai_generate_respects_configured_default_and_request_model(
|
|
default_model: str | None, request_model: str | None, expected_model: str
|
|
) -> None:
|
|
provider_options = {"api_key": "test"}
|
|
if default_model is not None:
|
|
provider_options["default_model"] = default_model
|
|
provider = OpenAIProvider(**provider_options)
|
|
client = _OpenAIClient()
|
|
provider.client = client
|
|
expected_default = default_model or "gpt-4o-mini"
|
|
|
|
assert provider.get_default_model() == expected_default
|
|
|
|
provider.generate(
|
|
LLMInput(messages=[Message(role=Role.USER, content="hi")], model=request_model)
|
|
)
|
|
|
|
assert client.completions.params is not None
|
|
assert client.completions.params["model"] == expected_model
|
|
assert provider.get_default_model() == expected_default
|
|
|
|
|
|
def test_openai_provider_serializes_tools_for_chat_completions():
|
|
provider = OpenAIProvider(api_key="test")
|
|
client = _OpenAIClient()
|
|
provider.client = client
|
|
|
|
provider.generate(LLMInput(messages=[Message(role=Role.USER, content="hi")], tools=[_tool()]))
|
|
|
|
assert client.completions.params["tools"] == [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "search",
|
|
"description": "Search",
|
|
"parameters": {"type": "object", "properties": {"query": {"type": "string"}}},
|
|
"strict": True,
|
|
},
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize(
|
|
"provider_type",
|
|
[OpenAIProvider, AstraflowProvider, AtlasProvider],
|
|
)
|
|
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
|
|
|
|
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 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"},
|
|
]
|
|
|
|
submitted_params = client.completions.params
|
|
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 submitted_params
|
|
|
|
for invalid_arguments in ({"score": float("nan")}, {"value": object()}):
|
|
with pytest.raises(
|
|
ValueError,
|
|
match="tool call arguments must be valid JSON",
|
|
):
|
|
provider.generate(
|
|
LLMInput(
|
|
messages=[
|
|
Message(
|
|
role=Role.ASSISTANT,
|
|
content="",
|
|
tool_calls=[
|
|
ToolCall(
|
|
id="call_invalid",
|
|
name="score",
|
|
arguments=invalid_arguments,
|
|
)
|
|
],
|
|
)
|
|
]
|
|
)
|
|
)
|
|
assert client.completions.params is submitted_params
|
|
|
|
|
|
def test_openai_provider_can_be_constructed_without_credentials(monkeypatch):
|
|
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
|
|
|
provider = OpenAIProvider()
|
|
|
|
assert provider.validate_config() is False
|
|
|
|
|
|
def test_openai_provider_rejects_empty_or_filtered_responses():
|
|
provider = OpenAIProvider(api_key="test")
|
|
|
|
for response in [
|
|
_openai_response(choices=[]),
|
|
_openai_response(choices=[SimpleNamespace(message=None, finish_reason="content_filter")]),
|
|
]:
|
|
provider.client = _OpenAIClient(response=response)
|
|
with pytest.raises(ValueError, match=EMPTY_FILTERED_RESPONSE_ERROR):
|
|
provider.generate(LLMInput(messages=[Message(role=Role.USER, content="hi")]))
|
|
|
|
|
|
def test_openai_provider_allows_missing_usage():
|
|
provider = OpenAIProvider(api_key="test")
|
|
provider.client = _OpenAIClient(response=_openai_response(usage=None))
|
|
|
|
output = provider.generate(LLMInput(messages=[Message(role=Role.USER, content="hi")]))
|
|
|
|
assert output.content == "ok"
|
|
assert output.usage is None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("max_tokens", "temperature", "expected_options"),
|
|
[
|
|
(128, 1.0, {"num_predict": 128}),
|
|
(128, 0.2, {"temperature": 0.2, "num_predict": 128}),
|
|
(128, 0.0, {"temperature": 0.0, "num_predict": 128}),
|
|
(0, 1.0, {"num_predict": 0}),
|
|
(None, 1.0, {}),
|
|
(None, 0.2, {"temperature": 0.2}),
|
|
],
|
|
)
|
|
def test_ollama_provider_serializes_generation_options(
|
|
monkeypatch, max_tokens, temperature, expected_options
|
|
):
|
|
requests = []
|
|
|
|
def fake_urlopen(request, timeout):
|
|
requests.append((request, timeout))
|
|
return BytesIO(b'{"message": {"content": "ok"}, "done_reason": "stop"}')
|
|
|
|
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
|
|
provider = OllamaProvider(base_url="http://localhost:11434", default_model="llama3.2")
|
|
|
|
output = provider.generate(
|
|
LLMInput(
|
|
messages=[Message(role=Role.USER, content="hi")],
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
)
|
|
)
|
|
|
|
assert len(requests) == 1
|
|
request, timeout = requests[0]
|
|
expected_payload = {
|
|
"model": "llama3.2",
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"stream": False,
|
|
}
|
|
if expected_options:
|
|
expected_payload["options"] = expected_options
|
|
assert json.loads(request.data) == expected_payload
|
|
assert request.full_url == "http://localhost:11434/api/chat"
|
|
assert request.get_method() == "POST"
|
|
assert request.get_header("Content-type") == "application/json"
|
|
assert timeout == 60
|
|
assert output.content == "ok"
|
|
assert output.model == "llama3.2"
|
|
assert output.stop_reason == "stop"
|
|
|
|
|
|
class _FakeOllama:
|
|
"""Serves /api/show and replays queued /api/chat replies, recording requests."""
|
|
|
|
def __init__(
|
|
self,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
replies: list[dict[str, Any]],
|
|
capabilities: list[str] | None = None,
|
|
) -> None:
|
|
self.replies = list(replies)
|
|
self.capabilities = capabilities
|
|
self.chats: list[dict[str, Any]] = []
|
|
self.shows: list[str] = []
|
|
monkeypatch.setattr(urllib.request, "urlopen", self.urlopen)
|
|
|
|
def urlopen(self, request: urllib.request.Request, timeout: float) -> BytesIO:
|
|
body = json.loads(request.data)
|
|
if request.full_url.endswith("/api/show"):
|
|
self.shows.append(body["model"])
|
|
if self.capabilities is None:
|
|
raise OSError("connection refused")
|
|
return BytesIO(json.dumps({"capabilities": self.capabilities}).encode())
|
|
self.chats.append(body)
|
|
return BytesIO(json.dumps(self.replies.pop(0)).encode())
|
|
|
|
|
|
def _ollama_tool_calls(*calls: tuple[str, dict[str, Any]]) -> dict[str, Any]:
|
|
return {
|
|
"message": {
|
|
"content": "",
|
|
"tool_calls": [{"function": {"name": name, "arguments": args}} for name, args in calls],
|
|
},
|
|
"done_reason": "stop",
|
|
}
|
|
|
|
|
|
def _ollama_answer(text: str) -> dict[str, Any]:
|
|
return {"message": {"content": text}, "done_reason": "stop"}
|
|
|
|
|
|
def _lookup_tool() -> ToolDefinition:
|
|
return ToolDefinition(
|
|
name="lookup",
|
|
description="Look up a city",
|
|
parameters={"type": "object", "properties": {"city": {"type": "string"}}},
|
|
)
|
|
|
|
|
|
def test_ollama_provider_serializes_tools_and_parses_tool_calls(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
server = _FakeOllama(
|
|
monkeypatch, [_ollama_tool_calls(("search", {"query": "ollama"}))], capabilities=["completion", "tools"]
|
|
)
|
|
provider = OllamaProvider(base_url="http://localhost:11434", default_model="qwen3")
|
|
|
|
output = provider.generate(LLMInput(messages=[Message(role=Role.USER, content="hi")], tools=[_tool()]))
|
|
|
|
assert server.shows == ["qwen3"]
|
|
assert server.chats[0]["tools"] == [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "search",
|
|
"description": "Search",
|
|
"parameters": {"type": "object", "properties": {"query": {"type": "string"}}},
|
|
"strict": True,
|
|
},
|
|
}
|
|
]
|
|
assert output.has_tool_calls
|
|
assert output.tool_calls[0].name == "search"
|
|
assert output.tool_calls[0].arguments == {"query": "ollama"}
|
|
assert output.tool_calls[0].id
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("model", "capabilities", "expect_tools", "expect_shows"),
|
|
[
|
|
("llama3.2", None, True, []),
|
|
("codellama", None, False, []),
|
|
("qwen3", ["completion", "tools"], True, ["qwen3"]),
|
|
("gemma3", ["completion", "vision"], False, ["gemma3"]),
|
|
("unknown", None, False, ["unknown", "unknown"]),
|
|
],
|
|
)
|
|
def test_ollama_provider_sends_tools_only_to_models_that_declare_them(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
model: str,
|
|
capabilities: list[str] | None,
|
|
expect_tools: bool,
|
|
expect_shows: list[str],
|
|
) -> None:
|
|
server = _FakeOllama(monkeypatch, [_ollama_answer("ok"), _ollama_answer("ok")], capabilities=capabilities)
|
|
provider = OllamaProvider(default_model=model)
|
|
request = LLMInput(messages=[Message(role=Role.USER, content="hi")], tools=[_tool()])
|
|
|
|
provider.generate(request)
|
|
provider.generate(request)
|
|
|
|
assert all(("tools" in chat) is expect_tools for chat in server.chats)
|
|
# An answered capability lookup is cached per model; a failed one is retried.
|
|
assert server.shows == expect_shows
|
|
|
|
|
|
def test_ollama_provider_retries_capability_lookup_after_a_failure(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
server = _FakeOllama(monkeypatch, [_ollama_answer("ok"), _ollama_answer("ok"), _ollama_answer("ok")])
|
|
provider = OllamaProvider(default_model="qwen3")
|
|
request = LLMInput(messages=[Message(role=Role.USER, content="hi")], tools=[_tool()])
|
|
|
|
provider.generate(request)
|
|
server.capabilities = ["completion", "tools"]
|
|
provider.generate(request)
|
|
provider.generate(request)
|
|
|
|
assert ["tools" in chat for chat in server.chats] == [False, True, True]
|
|
assert server.shows == ["qwen3", "qwen3"]
|
|
|
|
|
|
def test_ollama_react_agent_correlates_multiple_calls_without_ids(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
server = _FakeOllama(
|
|
monkeypatch,
|
|
[
|
|
_ollama_tool_calls(("lookup", {"city": "Hefei"}), ("search", {"query": "weather"})),
|
|
_ollama_answer("Hefei is sunny."),
|
|
],
|
|
capabilities=["completion", "tools"],
|
|
)
|
|
registry = ToolRegistry()
|
|
registry.register(_lookup_tool(), lambda city: f"city={city}")
|
|
registry.register(_tool(), lambda query: f"query={query}")
|
|
agent = ReActAgent(OllamaProvider(default_model="qwen3"), ToolExecutor(registry))
|
|
|
|
output = asyncio.run(
|
|
agent.run(
|
|
LLMInput(
|
|
messages=[Message(role=Role.USER, content="Weather in Hefei?")],
|
|
tools=[_lookup_tool(), _tool()],
|
|
)
|
|
)
|
|
)
|
|
|
|
assert output.content == "Hefei is sunny."
|
|
follow_up = server.chats[1]
|
|
assert "tools" in follow_up
|
|
assistant, first, second = follow_up["messages"][1:]
|
|
ids = [call["id"] for call in assistant["tool_calls"]]
|
|
assert all(ids) and len(set(ids)) == 2
|
|
assert [call["function"]["name"] for call in assistant["tool_calls"]] == ["lookup", "search"]
|
|
assert first == {"role": "tool", "content": "city=Hefei", "tool_call_id": ids[0], "tool_name": "lookup"}
|
|
assert second == {"role": "tool", "content": "query=weather", "tool_call_id": ids[1], "tool_name": "search"}
|
|
|
|
|
|
def test_ollama_react_agent_reports_unknown_tool_under_its_own_name(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
server = _FakeOllama(
|
|
monkeypatch,
|
|
[_ollama_tool_calls(("missing", {}), ("search", {"query": "x"})), _ollama_answer("done")],
|
|
capabilities=["completion", "tools"],
|
|
)
|
|
registry = ToolRegistry()
|
|
registry.register(_tool(), lambda query: f"query={query}")
|
|
agent = ReActAgent(OllamaProvider(default_model="qwen3"), ToolExecutor(registry))
|
|
|
|
asyncio.run(agent.run(LLMInput(messages=[Message(role=Role.USER, content="go")], tools=[_tool()])))
|
|
|
|
results = server.chats[1]["messages"][2:]
|
|
assert [(r["tool_name"], r["content"]) for r in results] == [
|
|
("missing", "Error: Tool 'missing' not found"),
|
|
("search", "query=x"),
|
|
]
|
|
|
|
|
|
def test_ollama_react_agent_without_tool_support_answers_directly(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
server = _FakeOllama(monkeypatch, [_ollama_answer("I cannot call tools.")])
|
|
agent = ReActAgent(OllamaProvider(default_model="codellama"), ToolExecutor())
|
|
|
|
output = asyncio.run(agent.run(LLMInput(messages=[Message(role=Role.USER, content="hi")], tools=[_tool()])))
|
|
|
|
assert output.content == "I cannot call tools."
|
|
assert "tools" not in server.chats[0]
|
|
assert server.shows == []
|
|
|
|
|
|
def test_claude_provider_serializes_tools_for_messages_api():
|
|
provider = ClaudeProvider(api_key="test")
|
|
client = _AnthropicClient()
|
|
provider.client = client
|
|
|
|
provider.generate(LLMInput(messages=[Message(role=Role.USER, content="hi")], tools=[_tool()]))
|
|
|
|
assert client.messages.params["tools"] == [
|
|
{
|
|
"name": "search",
|
|
"description": "Search",
|
|
"input_schema": {"type": "object", "properties": {"query": {"type": "string"}}},
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize('capabilities', ['tools', {'tools': True}, ['tools', 1]])
|
|
def test_ollama_malformed_capability_declaration_is_not_authority(monkeypatch, capabilities):
|
|
server = _FakeOllama(monkeypatch, [_ollama_answer('ok'), _ollama_answer('ok')], capabilities=capabilities)
|
|
provider = OllamaProvider(default_model='custom')
|
|
request = LLMInput(messages=[Message(role=Role.USER, content='hi')], tools=[_tool()])
|
|
provider.generate(request)
|
|
server.capabilities = ['completion', 'tools']
|
|
provider.generate(request)
|
|
assert ['tools' in chat for chat in server.chats] == [False, True]
|
|
assert server.shows == ['custom', 'custom']
|
|
|
|
|
|
def test_ollama_mixed_history_does_not_infer_tool_names_from_empty_ids(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
server = _FakeOllama(monkeypatch, [_ollama_answer("ok")])
|
|
messages = [
|
|
Message(role=Role.ASSISTANT, content="", tool_calls=[ToolCall(id="", name="first", arguments={}), ToolCall(id="", name="last", arguments={}), ToolCall(id="real-id", name="identified", arguments={})]),
|
|
Message(role=Role.TOOL, content="unnamed missing"),
|
|
Message(role=Role.TOOL, content="unnamed empty", tool_call_id=""),
|
|
Message(role=Role.TOOL, content="explicit empty", tool_call_id="", name="explicit"),
|
|
Message(role=Role.TOOL, content="identified result", tool_call_id="real-id"),
|
|
Message(role=Role.TOOL, content="explicit override", tool_call_id="real-id", name="override"),
|
|
]
|
|
OllamaProvider().generate(LLMInput(messages=messages))
|
|
results = server.chats[0]["messages"][1:]
|
|
assert "tool_name" not in results[0]
|
|
assert "tool_name" not in results[1]
|
|
assert [item["tool_name"] for item in results[2:]] == ["explicit", "identified", "override"]
|
|
assert results[3]["tool_call_id"] == "real-id"
|