simbarag/tests/unit/test_artifacts.py

"""Tests for artifact creation (blueprints/artifacts/logic.py) and the owner-scoped API."""

from datetime import timedelta
from types import SimpleNamespace

import pytest
from quart import Quart
from quart_jwt_extended import JWTManager, create_access_token
from tortoise import Tortoise

from blueprints.artifacts import artifacts_blueprint
from blueprints.artifacts.logic import ARTIFACT_MAX_BYTES, create_artifact_from_output
from blueprints.artifacts.models import Artifact
from blueprints.conversation.models import Conversation, ConversationMessage
from blueprints.users.models import User


@pytest.fixture
async def db():
    await Tortoise.init(
        db_url="sqlite://:memory:",
        modules={
            "models": [
                "blueprints.users.models",
                "blueprints.conversation.models",
                "blueprints.conversation.run_models",
                "blueprints.artifacts.models",
            ]
        },
    )
    await Tortoise.generate_schemas()
    yield
    await Tortoise.close_connections()


@pytest.fixture
async def owner(db):
    user = await User.create(username="owner", email="owner@example.com")
    conversation = await Conversation.create(name="c", user=user)
    message = await ConversationMessage.create(
        text="[mealie_get_recipe result]\nSoup",
        conversation=conversation,
        speaker="tool",
    )
    return user, conversation, message


def _output(artifact):
    return SimpleNamespace(content="text", artifact=artifact)


RECIPE = {"kind": "mealie.recipe", "title": "Soup", "data": {"slug": "soup"}}


class TestCreateArtifact:
    async def test_creates_and_links(self, owner):
        user, conversation, message = owner
        artifact = await create_artifact_from_output(
            output=_output({**RECIPE, "source_ref": {"system": "mealie"}}),
            tool="mealie_get_recipe",
            user=user,
            conversation=conversation,
            message=message,
        )
        assert artifact is not None
        saved = await Artifact.get(id=artifact.id)
        assert saved.message_id == message.id
        assert saved.kind == "mealie.recipe"
        assert saved.data == {"slug": "soup"}
        assert saved.source_ref == {"system": "mealie"}
        assert saved.version == 1

    @pytest.mark.parametrize(
        "artifact", [None, [{"page_content": "doc"}], {"title": "no kind"}, "str"]
    )
    async def test_ignores_non_artifacts(self, owner, artifact):
        user, conversation, message = owner
        result = await create_artifact_from_output(
            output=_output(artifact),
            tool="simba_search",
            user=user,
            conversation=conversation,
            message=message,
        )
        assert result is None
        assert await Artifact.all().count() == 0

    async def test_plain_string_output(self, owner):
        user, conversation, message = owner
        result = await create_artifact_from_output(
            output="just text", tool="t", user=user, conversation=conversation
        )
        assert result is None

    async def test_drops_oversized(self, owner):
        user, conversation, message = owner
        big = {**RECIPE, "data": {"blob": "x" * (ARTIFACT_MAX_BYTES + 1)}}
        result = await create_artifact_from_output(
            output=_output(big), tool="t", user=user, conversation=conversation
        )
        assert result is None
        assert await Artifact.all().count() == 0

    async def test_drops_unserialisable(self, owner):
        user, conversation, message = owner
        bad = {**RECIPE, "data": {"obj": object()}}
        result = await create_artifact_from_output(
            output=_output(bad), tool="t", user=user, conversation=conversation
        )
        assert result is None

    async def test_message_delete_sets_null(self, owner):
        user, conversation, message = owner
        artifact = await create_artifact_from_output(
            output=_output(RECIPE),
            tool="t",
            user=user,
            conversation=conversation,
            message=message,
        )
        await message.delete()
        assert (await Artifact.get(id=artifact.id)).message_id is None

    async def test_conversation_delete_cascades(self, owner):
        user, conversation, message = owner
        await create_artifact_from_output(
            output=_output(RECIPE), tool="t", user=user, conversation=conversation
        )
        await conversation.delete()
        assert await Artifact.all().count() == 0


@pytest.fixture
def app():
    app = Quart(__name__)
    app.config["JWT_SECRET_KEY"] = "test-secret"
    app.config["JWT_ACCESS_TOKEN_EXPIRES"] = timedelta(hours=1)
    JWTManager(app)
    app.register_blueprint(artifacts_blueprint)
    return app


async def _headers(app, user):
    async with app.app_context():
        token = create_access_token(identity=str(user.id))
    return {"Authorization": f"Bearer {token}"}


class TestArtifactRoutes:
    async def _setup(self, owner):
        user, conversation, message = owner
        artifact = await create_artifact_from_output(
            output=_output(RECIPE),
            tool="mealie_get_recipe",
            user=user,
            conversation=conversation,
            message=message,
        )
        other = await User.create(username="other", email="other@example.com")
        return user, other, conversation, message, artifact

    async def test_list_for_owner(self, app, owner):
        user, _, conversation, message, artifact = await self._setup(owner)
        resp = await app.test_client().get(
            f"/api/artifacts?conversation_id={conversation.id}",
            headers=await _headers(app, user),
        )
        assert resp.status_code == 200
        body = await resp.get_json()
        assert len(body) == 1
        assert body[0]["id"] == str(artifact.id)
        assert body[0]["kind"] == "mealie.recipe"
        assert body[0]["title"] == "Soup"
        assert body[0]["tool"] == "mealie_get_recipe"
        assert body[0]["message_id"] == str(message.id)
        assert "created_at" in body[0]
        assert "data" not in body[0]

    async def test_get_for_owner(self, app, owner):
        user, _, conversation, message, artifact = await self._setup(owner)
        resp = await app.test_client().get(
            f"/api/artifacts/{artifact.id}", headers=await _headers(app, user)
        )
        assert resp.status_code == 200
        body = await resp.get_json()
        assert body["data"] == {"slug": "soup"}
        assert body["source_ref"] is None
        assert body["conversation_id"] == str(conversation.id)

    async def test_other_user_gets_404(self, app, owner):
        _, other, conversation, _, artifact = await self._setup(owner)
        headers = await _headers(app, other)
        client = app.test_client()
        resp = await client.get(f"/api/artifacts/{artifact.id}", headers=headers)
        assert resp.status_code == 404
        resp = await client.get(
            f"/api/artifacts?conversation_id={conversation.id}", headers=headers
        )
        assert resp.status_code == 404

    async def test_invalid_id_is_404(self, app, owner):
        user, *_ = await self._setup(owner)
        resp = await app.test_client().get(
            "/api/artifacts/not-a-uuid", headers=await _headers(app, user)
        )
        assert resp.status_code == 404

    async def test_invalid_id_is_404_without_querying(self, app, owner, monkeypatch):
        user, *_ = await self._setup(owner)

        async def _boom(*args, **kwargs):
            raise AssertionError("queried with an invalid id")

        monkeypatch.setattr(Artifact, "get_or_none", _boom)
        resp = await app.test_client().get(
            "/api/artifacts/not-a-uuid", headers=await _headers(app, user)
        )
        assert resp.status_code == 404

    async def test_invalid_conversation_id_is_400_without_querying(
        self, app, owner, monkeypatch
    ):
        user, *_ = await self._setup(owner)

        async def _boom(*args, **kwargs):
            raise AssertionError("queried with an invalid id")

        monkeypatch.setattr(Conversation, "get_or_none", _boom)
        resp = await app.test_client().get(
            "/api/artifacts?conversation_id=garbage", headers=await _headers(app, user)
        )
        assert resp.status_code == 400

    async def test_requires_auth(self, app, owner):
        await self._setup(owner)
        resp = await app.test_client().get("/api/artifacts/whatever")
        assert resp.status_code == 401

    async def test_missing_conversation_id_is_400(self, app, owner):
        user, *_ = await self._setup(owner)
        resp = await app.test_client().get(
            "/api/artifacts", headers=await _headers(app, user)
        )
        assert resp.status_code == 400


class TestConversationPayload:
    async def test_tool_messages_gain_tool_label_artifact(self, owner):
        from blueprints.conversation import conversation_blueprint

        user, conversation, message = owner
        await ConversationMessage.create(
            text="[web_search result]\nresults",
            conversation=conversation,
            speaker="tool",
        )
        await ConversationMessage.create(
            text="hello", conversation=conversation, speaker="user"
        )
        artifact = await create_artifact_from_output(
            output=_output(RECIPE),
            tool="mealie_get_recipe",
            user=user,
            conversation=conversation,
            message=message,
        )

        app = Quart(__name__)
        app.config["JWT_SECRET_KEY"] = "test-secret"
        JWTManager(app)
        app.register_blueprint(conversation_blueprint)
        resp = await app.test_client().get(
            f"/api/conversation/{conversation.id}", headers=await _headers(app, user)
        )
        assert resp.status_code == 200
        by_text = {m["text"]: m for m in (await resp.get_json())["messages"]}

        recipe = by_text["[mealie_get_recipe result]"]
        assert recipe["tool"] == "mealie_get_recipe"
        assert recipe["label"] == "Fetching recipe…"
        assert recipe["artifact"] == {
            "id": str(artifact.id),
            "kind": "mealie.recipe",
            "title": "Soup",
        }
        search = by_text["[web_search result]"]
        assert search["tool"] == "web_search"
        assert search["artifact"] is None
        assert "tool" not in by_text["hello"]


class _FakeAgent:
    def __init__(self, events):
        self._events = events

    async def astream_events(self, payload, version, config):
        for event in self._events:
            yield event


class TestStreamQuery:
    async def test_emits_v2_tool_events_and_saves_artifact(self, owner, monkeypatch):
        import json

        import blueprints.conversation as conv
        from langchain_core.messages import AIMessage, ToolMessage

        user, conversation, _ = owner
        events = [
            {
                "event": "on_tool_start",
                "name": "mealie_get_recipe",
                "run_id": "r1",
                "data": {"input": {"recipe_slug": "soup", "config": {}}},
            },
            {
                "event": "on_tool_start",
                "name": "web_search",
                "run_id": "r2",
                "data": {"input": {"query": "cats"}},
            },
            {
                "event": "on_tool_end",
                "name": "mealie_get_recipe",
                "run_id": "r1",
                "data": {
                    "output": ToolMessage(
                        content="Soup recipe",
                        artifact=RECIPE,
                        tool_call_id="c1",
                        name="mealie_get_recipe",
                    )
                },
            },
            {
                "event": "on_tool_error",
                "name": "web_search",
                "run_id": "r2",
                "data": {"error": RuntimeError("boom")},
            },
            {
                "event": "on_chain_end",
                "name": "agent",
                "data": {"output": {"messages": [AIMessage(content="Done")]}},
            },
        ]
        monkeypatch.setattr(conv.runtime, "get_agent_for", lambda u: _FakeAgent(events))

        async def _prompt(user_id, query=""):
            return "system"

        monkeypatch.setattr(conv.runtime, "build_system_prompt_with_memories", _prompt)

        app = Quart(__name__)
        app.config["JWT_SECRET_KEY"] = "test-secret"
        JWTManager(app)
        app.register_blueprint(conv.conversation_blueprint)
        resp = await app.test_client().post(
            "/api/conversation/stream-query",
            json={"query": "soup?", "conversation_id": str(conversation.id)},
            headers=await _headers(app, user),
        )
        body = (await resp.get_data()).decode()
        frames = [
            json.loads(line[len("data: ") :])
            for line in body.splitlines()
            if line.startswith("data: {")
        ]
        by_type = {}
        for f in frames:
            by_type.setdefault(f["type"], []).append(f)

        starts = by_type["tool_start"]
        assert starts[0]["input"] == {"recipe_slug": "soup"}
        assert starts[0]["label"] == "Fetching recipe…"

        ends = {f["run_id"]: f for f in by_type["tool_end"]}
        artifact = await Artifact.get(tool="mealie_get_recipe")
        assert ends["r1"]["status"] == "ok"
        assert ends["r1"]["artifact"] == {
            "id": str(artifact.id),
            "kind": "mealie.recipe",
            "title": "Soup",
        }
        assert ends["r2"]["status"] == "error"
        assert ends["r2"]["artifact"] is None
        assert by_type["response"][0]["message"] == "Done"

        tool_message = await ConversationMessage.get(id=artifact.message_id)
        assert tool_message.text.startswith("[mealie_get_recipe result]")