app/backend/tests/test_ai.py

"""The assistant: who may set it up, and that it acts as the person chatting."""

import json
from collections.abc import Iterator
from pathlib import Path
from typing import Any

import pytest
from fastapi.testclient import TestClient

from fluksio.ai import agent
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) -> dict[str, Any]:
        seen.append(kwargs["messages"])
        assert any(t["function"]["name"] == "save_flow" for t in kwargs["tools"])
        return {"choices": [{"message": next(answers)}]}

    monkeypatch.setattr(agent.client, "chat_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",
        "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() == []