tests/test_chat_tool_requests.py

import pytest
import os

os.environ.setdefault("CEREBRAS_API_KEY", "test-key")
import app.websockets.handlers as handlers


class _FakeToolResult:
    def __init__(self, payload):
        self._payload = payload

    def model_dump(self):
        return self._payload


class _FakeTutor:
    def __init__(self):
        self.flashcard_calls = []
        self.quiz_calls = []

    def generate_flashcards(self, label, chunks, intention):
        self.flashcard_calls.append((label, chunks, intention))
        return _FakeToolResult({"cards": []})

    def generate_quiz(self, label, chunks, intention):
        self.quiz_calls.append((label, chunks, intention))
        return _FakeToolResult({"questions": []})


class _FakeConnectionManager:
    def __init__(self):
        self.sent = []

    async def send(self, project_id, event_type, payload):
        self.sent.append((project_id, event_type, payload))


class _FakeCache:
    def __init__(self):
        self._store = {}

    def make_key(self, *parts):
        return "|".join(str(p) for p in parts)

    def get(self, key):
        return self._store.get(key)

    def put(self, key, value):
        self._store[key] = value


class _FakeDecision:
    def __init__(self, modality, decline_reason=""):
        self.modality = modality
        self.decline_reason = decline_reason


class _FakeRouter:
    def __init__(self, decision):
        self._decision = decision
        self.classify_calls = []

    def classify(self, *args):
        self.classify_calls.append(args)
        return self._decision


class _FakeVisuals:
    def __init__(self):
        self.generate_calls = []

    def generate(self, *args):
        self.generate_calls.append(args)
        return _FakeToolResult({"html_code": "<div></div>", "animation_type": args[1]})


class _RaisingVisuals:
    def generate(self, *args):
        raise AssertionError("generate should not be called when the router declines")


@pytest.mark.asyncio
async def test_flashcards_request_uses_topic_and_whole_project(monkeypatch):
    calls = []
    tutor = _FakeTutor()

    async def fake_get_chunks(project_id, query, n=5, chunk_type=None, document_ids=None):
        calls.append({
            "project_id": project_id,
            "query": query,
            "chunk_type": chunk_type,
            "document_ids": document_ids,
        })
        return [{"text": "attention chunk", "document_id": "paper-a"}]

    monkeypatch.setattr(handlers, "_get_chunks", fake_get_chunks)
    monkeypatch.setattr(handlers, "_get_tutor", lambda: tutor)
    monkeypatch.setattr(handlers, "_cm", _FakeConnectionManager())
    monkeypatch.setattr(handlers, "_cache", _FakeCache())

    await handlers.handle_event("project-1", "FLASHCARDS_REQUEST", {
        "node_id": "old-node",
        "node_label": "old node label",
        "topic": "attention mechanisms",
        "intention": "graduate",
    })

    assert calls[0]["query"] == "attention mechanisms"
    assert calls[0]["chunk_type"] == "question"
    assert calls[0]["document_ids"] is None
    assert tutor.flashcard_calls[0][0] == "attention mechanisms"


@pytest.mark.asyncio
async def test_quiz_request_uses_topic_and_whole_project(monkeypatch):
    calls = []
    tutor = _FakeTutor()

    async def fake_get_chunks(project_id, query, n=5, chunk_type=None, document_ids=None):
        calls.append({
            "project_id": project_id,
            "query": query,
            "chunk_type": chunk_type,
            "document_ids": document_ids,
        })
        return [{"text": "transformer chunk", "document_id": "paper-b"}]

    monkeypatch.setattr(handlers, "_get_chunks", fake_get_chunks)
    monkeypatch.setattr(handlers, "_get_tutor", lambda: tutor)
    monkeypatch.setattr(handlers, "_cm", _FakeConnectionManager())
    monkeypatch.setattr(handlers, "_cache", _FakeCache())

    await handlers.handle_event("project-1", "QUIZ_REQUEST", {
        "node_id": "old-node",
        "node_label": "old node label",
        "topic": "transformers",
        "intention": "graduate",
    })

    assert calls[0]["query"] == "transformers"
    assert calls[0]["chunk_type"] == "question"
    assert calls[0]["document_ids"] is None
    assert tutor.quiz_calls[0][0] == "transformers"


@pytest.mark.asyncio
async def test_tool_requests_require_explicit_topic(monkeypatch):
    tutor = _FakeTutor()
    cm = _FakeConnectionManager()

    async def fake_get_chunks(*args, **kwargs):
        raise AssertionError("retrieval should not run without an explicit topic")

    monkeypatch.setattr(handlers, "_get_chunks", fake_get_chunks)
    monkeypatch.setattr(handlers, "_get_tutor", lambda: tutor)
    monkeypatch.setattr(handlers, "_cm", cm)
    monkeypatch.setattr(handlers, "_cache", _FakeCache())

    await handlers.handle_event("project-1", "FLASHCARDS_REQUEST", {
        "node_id": "old-node",
        "node_label": "old node label",
        "intention": "graduate",
    })
    await handlers.handle_event("project-1", "QUIZ_REQUEST", {
        "node_id": "old-node",
        "node_label": "old node label",
        "intention": "graduate",
    })

    assert [event for _, event, _ in cm.sent] == ["ERROR", "ERROR"]
    assert tutor.flashcard_calls == []
    assert tutor.quiz_calls == []


@pytest.mark.asyncio
async def test_visualize_request_requires_explicit_topic(monkeypatch):
    cm = _FakeConnectionManager()
    visuals = _RaisingVisuals()

    async def fake_get_chunks(*args, **kwargs):
        raise AssertionError("retrieval should not run without an explicit topic")

    monkeypatch.setattr(handlers, "_get_chunks", fake_get_chunks)
    monkeypatch.setattr(handlers, "_get_visuals", lambda: visuals)
    monkeypatch.setattr(handlers, "_cm", cm)
    monkeypatch.setattr(handlers, "_cache", _FakeCache())

    await handlers.handle_event("project-1", "VISUALIZE_REQUEST", {
        "intention": "graduate",
    })

    assert [event for _, event, _ in cm.sent] == ["ERROR"]
    assert cm.sent[0][2]["event_type"] == "VISUALIZE_REQUEST"


@pytest.mark.asyncio
async def test_visualize_request_decline_never_calls_generate(monkeypatch):
    cm = _FakeConnectionManager()
    router = _FakeRouter(_FakeDecision("decline", decline_reason="Nothing to plot here."))
    visuals = _RaisingVisuals()

    async def fake_get_chunks(project_id, query, n=8, chunk_type=None, document_ids=None):
        return [{"text": "some chunk", "document_id": "paper-a"}]

    monkeypatch.setattr(handlers, "_get_chunks", fake_get_chunks)
    monkeypatch.setattr(handlers, "_get_router", lambda: router)
    monkeypatch.setattr(handlers, "_get_visuals", lambda: visuals)
    monkeypatch.setattr(handlers, "_cm", cm)
    monkeypatch.setattr(handlers, "_cache", _FakeCache())

    await handlers.handle_event("project-1", "VISUALIZE_REQUEST", {
        "topic": "entropy",
        "intention": "graduate",
    })

    assert [event for _, event, _ in cm.sent] == ["VISUALIZE_READY"]
    payload = cm.sent[0][2]
    assert payload["visual"] is None
    assert payload["decline_reason"] == "Nothing to plot here."


@pytest.mark.asyncio
async def test_visualize_request_generates_visual_for_real_modality(monkeypatch):
    cm = _FakeConnectionManager()
    router = _FakeRouter(_FakeDecision("graph"))
    visuals = _FakeVisuals()
    chunks = [{"text": "some chunk", "document_id": "paper-a"}]

    async def fake_get_chunks(project_id, query, n=8, chunk_type=None, document_ids=None):
        return chunks

    monkeypatch.setattr(handlers, "_get_chunks", fake_get_chunks)
    monkeypatch.setattr(handlers, "_get_router", lambda: router)
    monkeypatch.setattr(handlers, "_get_visuals", lambda: visuals)
    monkeypatch.setattr(handlers, "_cm", cm)
    monkeypatch.setattr(handlers, "_cache", _FakeCache())

    await handlers.handle_event("project-1", "VISUALIZE_REQUEST", {
        "topic": "entropy",
        "intention": "graduate",
    })

    assert visuals.generate_calls == [("entropy", "graph", "graduate", chunks)]
    assert [event for _, event, _ in cm.sent] == ["VISUALIZE_READY"]
    payload = cm.sent[0][2]
    assert payload["topic"] == "entropy"
    assert payload["visual"] == {"html_code": "<div></div>", "animation_type": "graph"}


@pytest.mark.asyncio
async def test_visualize_request_cache_hit_skips_router_and_generate(monkeypatch):
    cm = _FakeConnectionManager()
    router = _FakeRouter(_FakeDecision("graph"))
    visuals = _RaisingVisuals()
    cache = _FakeCache()
    chunks = [{"text": "some chunk", "document_id": "paper-a"}]

    async def fake_get_chunks(project_id, query, n=8, chunk_type=None, document_ids=None):
        return chunks

    monkeypatch.setattr(handlers, "_get_chunks", fake_get_chunks)
    monkeypatch.setattr(handlers, "_get_router", lambda: router)
    monkeypatch.setattr(handlers, "_get_visuals", lambda: visuals)
    monkeypatch.setattr(handlers, "_cm", cm)
    monkeypatch.setattr(handlers, "_cache", cache)

    anchor_id = "visualize:entropy"
    cache_key = cache.make_key(
        "VISUALIZE_REQUEST", "graduate", anchor_id, [c["text"] for c in chunks], "entropy"
    )
    cache.put(cache_key, {"html_code": "<div>cached</div>", "animation_type": "graph"})

    await handlers.handle_event("project-1", "VISUALIZE_REQUEST", {
        "topic": "entropy",
        "intention": "graduate",
    })

    assert router.classify_calls == []
    assert [event for _, event, _ in cm.sent] == ["VISUALIZE_READY"]
    payload = cm.sent[0][2]
    assert payload["visual"] == {"html_code": "<div>cached</div>", "animation_type": "graph"}