ai/client.py streams the completion and puts the chunks back together (Assembler), covering OpenAI's piecewise tool-call arguments, Ollama's whole calls, OpenRouter's comment lines, unnumbered calls and in-stream errors. The turn emits a delta event per piece and a message per round as before, and the chat grows the last answer as deltas arrive. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WgBfkm4rCmtd1T4vHWjEPx
253 lines
11 KiB
Python
253 lines
11 KiB
Python
"""The assistant: who may set it up, and that it acts as the person chatting."""
|
|
|
|
import json
|
|
from collections.abc import AsyncIterator, Iterator
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
from fluksio.ai import agent, client
|
|
from fluksio.core.config import settings
|
|
|
|
|
|
@pytest.fixture
|
|
def ai_files(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
|
|
monkeypatch.setattr(settings, "AI_CONFIG_FILE", tmp_path / "ai.json")
|
|
monkeypatch.setattr(settings, "CHATS_DIR", tmp_path / "chats")
|
|
return tmp_path
|
|
|
|
|
|
def _events(response: Any) -> list[dict[str, Any]]:
|
|
return [
|
|
json.loads(line[5:])
|
|
for line in response.iter_lines()
|
|
if line.startswith("data:")
|
|
]
|
|
|
|
|
|
def test_the_provider_is_a_superuser_setting_and_the_key_stays_sealed(
|
|
client: TestClient,
|
|
superuser_token_headers: dict[str, str],
|
|
normal_user_token_headers: dict[str, str],
|
|
ai_files: Path,
|
|
) -> None:
|
|
assert client.get(
|
|
"/api/v1/ai/status", headers=normal_user_token_headers
|
|
).json() == {
|
|
"available": False,
|
|
"model": None,
|
|
}
|
|
body = {"api_base": "http://llm.test/v1", "api_key": "sk-test-123456", "model": "m"}
|
|
refused = client.put(
|
|
"/api/v1/ai/config", json=body, headers=normal_user_token_headers
|
|
)
|
|
assert refused.status_code == 403
|
|
|
|
saved = client.put("/api/v1/ai/config", json=body, headers=superuser_token_headers)
|
|
assert saved.status_code == 200, saved.text
|
|
assert saved.json()["api_key_masked"] == "sk-…3456"
|
|
assert "sk-test" not in (ai_files / "ai.json").read_text()
|
|
assert client.get(
|
|
"/api/v1/ai/status", headers=normal_user_token_headers
|
|
).json() == {
|
|
"available": True,
|
|
"model": "m",
|
|
}
|
|
# A save that names no key keeps the one there is.
|
|
kept = client.put(
|
|
"/api/v1/ai/config",
|
|
json={"api_base": "http://llm.test/v1", "model": "m2"},
|
|
headers=superuser_token_headers,
|
|
)
|
|
assert kept.json()["api_key_set"] is True
|
|
assert kept.json()["model"] == "m2"
|
|
|
|
|
|
@pytest.mark.usefixtures("ai_files")
|
|
def test_a_turn_runs_the_tools_as_the_person_chatting(
|
|
client: TestClient,
|
|
superuser_token_headers: dict[str, str],
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
client.put(
|
|
"/api/v1/ai/config",
|
|
json={"api_base": "http://llm.test/v1", "api_key": "k", "model": "m"},
|
|
headers=superuser_token_headers,
|
|
)
|
|
seen: list[list[dict[str, Any]]] = []
|
|
answers: Iterator[dict[str, Any]] = iter(
|
|
[
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {"name": "list_flows", "arguments": "{}"},
|
|
}
|
|
],
|
|
},
|
|
{"role": "assistant", "content": "Counted them."},
|
|
]
|
|
)
|
|
|
|
async def fake_completion(**kwargs: Any) -> AsyncIterator[str | dict[str, Any]]:
|
|
seen.append(kwargs["messages"])
|
|
assert any(t["function"]["name"] == "save_flow" for t in kwargs["tools"])
|
|
answer = next(answers)
|
|
if answer["content"]:
|
|
yield "Counted "
|
|
yield "them."
|
|
yield answer
|
|
|
|
monkeypatch.setattr(agent.client, "stream_completion", fake_completion)
|
|
|
|
with client.stream(
|
|
"POST",
|
|
"/api/v1/ai/chat",
|
|
json={"message": "How many flows are there?"},
|
|
headers=superuser_token_headers,
|
|
) as response:
|
|
assert response.status_code == 200
|
|
events = _events(response)
|
|
|
|
assert [e["type"] for e in events] == [
|
|
"start",
|
|
"tool_call",
|
|
"tool_result",
|
|
"delta",
|
|
"delta",
|
|
"message",
|
|
"done",
|
|
]
|
|
# The tool went through the REST API with the caller's own token: an
|
|
# unauthenticated call would have come back as an error the model reads.
|
|
assert events[2] == {"type": "tool_result", "name": "list_flows", "ok": True}
|
|
answered = seen[-1][-1]
|
|
assert answered["role"] == "tool"
|
|
assert "count" in json.loads(answered["content"])
|
|
assert seen[-1][0]["role"] == "system"
|
|
|
|
transcript = client.get("/api/v1/ai/chat", headers=superuser_token_headers).json()
|
|
assert [m["role"] for m in transcript] == ["user", "assistant", "tool", "assistant"]
|
|
assert transcript[2] == {
|
|
"role": "tool",
|
|
"content": None,
|
|
"tool_calls": None,
|
|
"ok": True,
|
|
"error": None,
|
|
}
|
|
client.delete("/api/v1/ai/chat", headers=superuser_token_headers)
|
|
assert client.get("/api/v1/ai/chat", headers=superuser_token_headers).json() == []
|
|
|
|
|
|
# One answer per provider as it streams it. Ollama's text is recorded from a
|
|
# local server, the rest follows the shape each endpoint documents.
|
|
STREAMS = {
|
|
"ollama": """
|
|
data: {"id":"chatcmpl-888","object":"chat.completion.chunk","model":"gemma3:4b-it-qat","system_fingerprint":"fp_ollama","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}
|
|
|
|
data: {"id":"chatcmpl-888","object":"chat.completion.chunk","model":"gemma3:4b-it-qat","system_fingerprint":"fp_ollama","choices":[{"index":0,"delta":{"role":"assistant","content":" there"},"finish_reason":null}]}
|
|
|
|
data: {"id":"chatcmpl-888","object":"chat.completion.chunk","model":"gemma3:4b-it-qat","system_fingerprint":"fp_ollama","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":"stop"}]}
|
|
|
|
data: [DONE]
|
|
""",
|
|
"ollama-tools": """
|
|
data: {"id":"chatcmpl-12","object":"chat.completion.chunk","model":"qwen3","system_fingerprint":"fp_ollama","choices":[{"index":0,"delta":{"role":"assistant","content":"","reasoning":"The user wants"},"finish_reason":null}]}
|
|
|
|
data: {"id":"chatcmpl-12","object":"chat.completion.chunk","model":"qwen3","system_fingerprint":"fp_ollama","choices":[{"index":0,"delta":{"role":"assistant","content":"","tool_calls":[{"id":"call_k2","index":0,"type":"function","function":{"name":"get_flow","arguments":"{\\"name\\":\\"demo\\"}"}}]},"finish_reason":null}]}
|
|
|
|
data: {"id":"chatcmpl-12","object":"chat.completion.chunk","model":"qwen3","system_fingerprint":"fp_ollama","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":"tool_calls"}]}
|
|
|
|
data: [DONE]
|
|
""",
|
|
"openai": """
|
|
data: {"id":"chatcmpl-9","object":"chat.completion.chunk","model":"gpt-4.1","choices":[{"index":0,"delta":{"role":"assistant","content":null,"tool_calls":[{"index":0,"id":"call_a","type":"function","function":{"name":"get_flow","arguments":""}}],"refusal":null},"finish_reason":null}]}
|
|
|
|
data: {"id":"chatcmpl-9","object":"chat.completion.chunk","model":"gpt-4.1","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\\"name\\""}}]},"finish_reason":null}]}
|
|
|
|
data: {"id":"chatcmpl-9","object":"chat.completion.chunk","model":"gpt-4.1","choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"id":"call_b","type":"function","function":{"name":"list_flows","arguments":""}}]},"finish_reason":null}]}
|
|
|
|
data: {"id":"chatcmpl-9","object":"chat.completion.chunk","model":"gpt-4.1","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":":\\"demo\\"}"}}]},"finish_reason":null}]}
|
|
|
|
data: {"id":"chatcmpl-9","object":"chat.completion.chunk","model":"gpt-4.1","choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"{}"}}]},"finish_reason":null}]}
|
|
|
|
data: {"id":"chatcmpl-9","object":"chat.completion.chunk","model":"gpt-4.1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}
|
|
|
|
data: {"id":"chatcmpl-9","object":"chat.completion.chunk","model":"gpt-4.1","choices":[],"usage":{"prompt_tokens":9,"completion_tokens":12,"total_tokens":21}}
|
|
|
|
data: [DONE]
|
|
""",
|
|
"anthropic": """
|
|
data: {"id":"msg_01","object":"chat.completion.chunk","model":"claude-sonnet-4-5","choices":[{"index":0,"delta":{"role":"assistant","content":"Reading "},"finish_reason":null}]}
|
|
|
|
data: {"id":"msg_01","object":"chat.completion.chunk","model":"claude-sonnet-4-5","choices":[{"index":0,"delta":{"content":"it."},"finish_reason":null}]}
|
|
|
|
data: {"id":"msg_01","object":"chat.completion.chunk","model":"claude-sonnet-4-5","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"toolu_01","type":"function","function":{"name":"get_flow","arguments":""}}]},"finish_reason":null}]}
|
|
|
|
data: {"id":"msg_01","object":"chat.completion.chunk","model":"claude-sonnet-4-5","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\\"name\\": \\"de"}}]},"finish_reason":null}]}
|
|
|
|
data: {"id":"msg_01","object":"chat.completion.chunk","model":"claude-sonnet-4-5","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"mo\\"}"}}]},"finish_reason":null}]}
|
|
|
|
data: {"id":"msg_01","object":"chat.completion.chunk","model":"claude-sonnet-4-5","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}
|
|
|
|
data: [DONE]
|
|
""",
|
|
"openrouter": """
|
|
: OPENROUTER PROCESSING
|
|
|
|
data: {"id":"gen-1","provider":"Google","model":"google/gemini-2.5-flash","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant","content":"","reasoning":"Short answer."},"finish_reason":null}]}
|
|
|
|
: OPENROUTER PROCESSING
|
|
|
|
data: {"id":"gen-1","provider":"Google","model":"google/gemini-2.5-flash","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}
|
|
|
|
data: {"id":"gen-1","provider":"Google","model":"google/gemini-2.5-flash","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant","content":" there"},"finish_reason":"stop"}],"usage":{"prompt_tokens":4,"completion_tokens":2,"total_tokens":6}}
|
|
|
|
data: [DONE]
|
|
""",
|
|
# Gemini's compatibility endpoint sends each call whole and unnumbered.
|
|
"gemini": """
|
|
data: {"choices":[{"delta":{"role":"assistant","tool_calls":[{"id":"c1","type":"function","function":{"name":"get_flow","arguments":"{\\"name\\":\\"demo\\"}"}},{"id":"c2","type":"function","function":{"name":"list_flows","arguments":"{}"}}]},"finish_reason":"tool_calls","index":0}],"object":"chat.completion.chunk"}
|
|
|
|
data: [DONE]
|
|
""",
|
|
}
|
|
|
|
GET_DEMO = {"name": "get_flow", "arguments": '{"name":"demo"}'}
|
|
LIST = {"name": "list_flows", "arguments": "{}"}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("provider", "content", "calls"),
|
|
[
|
|
("ollama", "Hello there", []),
|
|
("ollama-tools", None, [GET_DEMO]),
|
|
("openai", None, [GET_DEMO, LIST]),
|
|
("anthropic", "Reading it.", [{**GET_DEMO, "arguments": '{"name": "demo"}'}]),
|
|
("openrouter", "Hello there", []),
|
|
("gemini", None, [GET_DEMO, LIST]),
|
|
],
|
|
)
|
|
def test_a_streamed_answer_comes_back_whole(
|
|
provider: str, content: str | None, calls: list[dict[str, str]]
|
|
) -> None:
|
|
assembler = client.Assembler()
|
|
pieces = [assembler.feed(line) for line in STREAMS[provider].splitlines()]
|
|
message = assembler.message()
|
|
assert "".join(pieces) == (content or "")
|
|
assert message["content"] == content
|
|
assert [c["function"] for c in message["tool_calls"] or []] == calls
|
|
assert all(c["id"] and c["type"] == "function" for c in message["tool_calls"] or [])
|
|
|
|
|
|
def test_an_error_inside_the_stream_stops_it() -> None:
|
|
assembler = client.Assembler()
|
|
line = 'data: {"error":{"code":502,"message":"Provider disconnected"},"choices":[{"index":0,"delta":{"content":""},"finish_reason":"error"}]}'
|
|
with pytest.raises(client.ProviderError, match="Provider disconnected"):
|
|
assembler.feed(line)
|