LLeMbas/tests/test_chat.py

"""Chat, folders, and the streaming reply path against a mocked endpoint."""

from __future__ import annotations

import httpx
import pytest
from fastapi.testclient import TestClient
from sqlalchemy import select

from lembas.db.models import Chat, Connection, Folder, Message, Model
from lembas.services import chat as chat_service
from lembas.services.crypto import encrypt
from lembas.services.llm.openai_client import Endpoint, LLMError, delta_text, list_models
from lembas.services.sse import event


# --- SSE framing -------------------------------------------------------------
def test_sse_event_framing():
    assert event("token", "hello") == "event: token\ndata: hello\n\n"


def test_sse_splits_newlines_across_data_lines():
    """A payload with a newline must become several data: lines. Sending a raw
    newline truncates the event, which is what breaks the first code block a
    model emits."""
    assert event("token", "a\nb") == "event: token\ndata: a\ndata: b\n\n"


def test_sse_round_trips_through_the_browser_rejoin_rule():
    payload = "line one\nline two\n\nline four"
    framed = event("token", payload)
    body = framed.split("\n", 1)[1]
    rejoined = "\n".join(
        line.removeprefix("data: ") for line in body.split("\n") if line.startswith("data:")
    )
    assert rejoined == payload


# --- Delta parsing -----------------------------------------------------------
def test_delta_text_reads_the_normal_shape():
    assert delta_text({"choices": [{"delta": {"content": "hi"}}]}) == "hi"


def test_delta_text_handles_typed_content_parts():
    chunk = {"choices": [{"delta": {"content": [{"type": "text", "text": "hi"}]}}]}
    assert delta_text(chunk) == "hi"


@pytest.mark.parametrize(
    "chunk", [{}, {"choices": []}, {"choices": [{}]}, {"choices": [{"delta": {}}]}]
)
def test_delta_text_tolerates_junk(chunk):
    """Providers vary; an unexpected chunk shape must not kill a reply."""
    assert delta_text(chunk) == ""


# --- Endpoint URL handling ---------------------------------------------------
@pytest.mark.parametrize(
    ("base", "expected"),
    [
        ("http://host:1234", "http://host:1234/v1/models"),
        ("http://host:1234/v1", "http://host:1234/v1/models"),
        ("http://host:1234/", "http://host:1234/v1/models"),
    ],
)
def test_base_url_with_or_without_v1(base, expected):
    """Users should not have to guess which form is expected."""
    assert Endpoint(base_url=base.rstrip("/"), api_key="", extra_headers={}).url(
        "models"
    ) == expected


def test_no_authorization_header_without_a_key():
    """Local runners often reject an empty bearer token outright."""
    assert "Authorization" not in Endpoint("http://h", "", {}).headers()
    assert Endpoint("http://h", "k", {}).headers()["Authorization"] == "Bearer k"


# --- Model discovery ---------------------------------------------------------
async def test_list_models_accepts_the_bare_list_shape():
    """The spec says {"data": [...]}, but some servers return a bare list."""

    def handler(_request):
        return httpx.Response(200, json=[{"id": "a"}, {"id": "b"}])

    original = httpx.AsyncClient

    class Patched(original):
        def __init__(self, **kwargs):
            super().__init__(transport=httpx.MockTransport(handler), **kwargs)

    httpx.AsyncClient = Patched
    try:
        assert [m["id"] for m in await list_models(Endpoint("http://h", "", {}))] == ["a", "b"]
    finally:
        httpx.AsyncClient = original


async def test_list_models_reports_a_rejected_key_readably():
    def handler(_request):
        return httpx.Response(401, json={"error": {"message": "Incorrect API key."}})

    original = httpx.AsyncClient

    class Patched(original):
        def __init__(self, **kwargs):
            super().__init__(transport=httpx.MockTransport(handler), **kwargs)

    httpx.AsyncClient = Patched
    try:
        with pytest.raises(LLMError) as caught:
            await list_models(Endpoint("http://h", "bad", {}))
        assert "rejected" in caught.value.message
        assert "Incorrect API key." in caught.value.message
    finally:
        httpx.AsyncClient = original


# --- Titles ------------------------------------------------------------------
def test_fallback_title_keeps_a_short_message_intact():
    assert chat_service.fallback_title("What is lembas?") == "What is lembas?"


def test_fallback_title_trims_on_a_word_boundary():
    title = chat_service.fallback_title("word " * 60)
    assert len(title) <= chat_service.MAX_TITLE_LENGTH + 1
    assert title.endswith("…")


def test_fallback_title_of_nothing():
    assert chat_service.fallback_title("   ") == "New chat"


# --- Chats and folders (through the API) -------------------------------------
def _add_connection(db) -> Connection:
    # Port 1 refuses connections, which is what the error-path test relies on.
    connection = Connection(
        name="Test", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt("")
    )
    db.add(connection)
    db.commit()
    db.add(Model(connection_id=connection.id, model_id="test-model"))
    db.commit()
    return connection


def test_new_chat_redirects_to_its_own_url(client: TestClient, db, registered):
    _add_connection(db)
    response = client.post("/api/chats", headers={"HX-Request": "true"})
    assert response.status_code == 204
    assert response.headers["HX-Redirect"].startswith("/chat/")


def test_new_chat_picks_up_the_default_model(client: TestClient, db, registered):
    _add_connection(db)
    client.post("/api/chats", headers={"HX-Request": "true"})
    assert db.scalar(select(Chat)).model_id == "test-model"


def test_posting_a_message_stores_both_turns(client: TestClient, db, registered):
    _add_connection(db)
    chat_id = client.post("/api/chats").headers["HX-Redirect"].rsplit("/", 1)[1]

    response = client.post(f"/api/chats/{chat_id}/messages", data={"content": "Hello there"})
    assert response.status_code == 200

    messages = db.scalars(select(Message).order_by(Message.created_at)).all()
    assert [m.role for m in messages] == ["user", "assistant"]
    assert messages[0].content == "Hello there"
    # The assistant row is created empty and incomplete; that is what carries
    # the sse-connect the browser uses to start the stream.
    assert messages[1].content == ""
    assert messages[1].complete is False
    assert "sse-connect" in response.text


def test_empty_message_is_ignored(client: TestClient, db, registered):
    _add_connection(db)
    chat_id = client.post("/api/chats").headers["HX-Redirect"].rsplit("/", 1)[1]
    assert client.post(f"/api/chats/{chat_id}/messages", data={"content": "   "}).status_code == 204
    assert db.scalar(select(Message)) is None


def test_a_chat_belonging_to_someone_else_is_not_found(client: TestClient, db, registered):
    _add_connection(db)
    chat_id = client.post("/api/chats").headers["HX-Redirect"].rsplit("/", 1)[1]

    client.post("/auth/logout", follow_redirects=False)
    client.post(
        "/auth/register",
        data={"name": "Sam", "email": "sam@shire.test", "password": "potatoes-po-ta-toes"},
        follow_redirects=False,
    )
    # 404 and not 403: whether that id exists is not this endpoint's to reveal.
    assert client.get(f"/chat/{chat_id}").status_code == 404


def test_renaming_a_chat_stops_it_being_auto_titled(client: TestClient, db, registered):
    _add_connection(db)
    chat_id = client.post("/api/chats").headers["HX-Redirect"].rsplit("/", 1)[1]
    client.patch(f"/api/chats/{chat_id}", data={"title": "My own title"})

    chat = db.get(Chat, chat_id)
    db.refresh(chat)
    assert chat.title == "My own title"
    assert chat.title_generated is True


def test_deleting_a_chat_removes_its_messages(client: TestClient, db, registered):
    _add_connection(db)
    chat_id = client.post("/api/chats").headers["HX-Redirect"].rsplit("/", 1)[1]
    client.post(f"/api/chats/{chat_id}/messages", data={"content": "Hello"})

    client.delete(f"/api/chats/{chat_id}")
    assert db.scalar(select(Chat)) is None
    assert db.scalar(select(Message)) is None


def test_deleting_a_folder_keeps_the_chats_inside_it(client: TestClient, db, registered):
    """Losing a conversation to a mis-clicked folder delete is unforgivable."""
    _add_connection(db)
    client.post("/api/folders", data={"name": "Quests"})
    folder = db.scalar(select(Folder))

    chat_id = client.post("/api/chats").headers["HX-Redirect"].rsplit("/", 1)[1]
    client.patch(f"/api/chats/{chat_id}", data={"folder_id": folder.id})

    client.delete(f"/api/folders/{folder.id}")

    chat = db.get(Chat, chat_id)
    db.refresh(chat)
    assert chat is not None
    assert chat.folder_id is None


def test_a_folder_cannot_be_moved_inside_itself(client: TestClient, db, registered):
    client.post("/api/folders", data={"name": "Outer"})
    folder = db.scalar(select(Folder))
    response = client.patch(f"/api/folders/{folder.id}", data={"parent_id": folder.id})
    assert response.status_code == 400


# --- Request building --------------------------------------------------------
def test_request_forwards_only_known_sampling_parameters(db, user_id):
    """A stray key in params_json must not become a 400 from the provider that
    looks like a LLeMbas bug."""
    connection = _add_connection(db)
    chat = Chat(
        user_id=user_id,
        model_id="test-model",
        connection_id=connection.id,
        params_json={"temperature": 0.4, "nonsense": "drop me"},
    )
    db.add(chat)
    db.commit()

    payload = chat_service.build_request(db, chat)
    assert payload["temperature"] == 0.4
    assert "nonsense" not in payload


def test_history_skips_failed_and_empty_turns(db, user_id):
    """Sending an empty assistant message upsets several providers."""
    connection = _add_connection(db)
    chat = Chat(user_id=user_id, model_id="test-model", connection_id=connection.id)
    db.add(chat)
    db.commit()

    db.add_all(
        [
            Message(chat_id=chat.id, role="user", content="one"),
            Message(chat_id=chat.id, role="assistant", content="", error="boom"),
            Message(chat_id=chat.id, role="user", content="two"),
        ]
    )
    db.commit()

    contents = [m["content"] for m in chat_service.build_request(db, chat)["messages"]]
    assert contents == ["one", "two"]


def test_system_prompt_leads_the_message_list(db, user_id):
    connection = _add_connection(db)
    chat = Chat(
        user_id=user_id,
        model_id="test-model",
        connection_id=connection.id,
        system_prompt="You are terse.",
    )
    db.add(chat)
    db.commit()

    messages = chat_service.build_request(db, chat)["messages"]
    assert messages[0] == {"role": "system", "content": "You are terse."}


def test_streaming_reports_an_unreachable_endpoint_in_the_thread(
    client: TestClient, db, registered
):
    """A failed turn must never be an unexplained blank bubble."""
    _add_connection(db)  # points at 127.0.0.1:1, which refuses connections
    chat_id = client.post("/api/chats").headers["HX-Redirect"].rsplit("/", 1)[1]
    client.post(f"/api/chats/{chat_id}/messages", data={"content": "Hello"})
    message = db.scalar(select(Message).where(Message.role == "assistant"))

    response = client.get(f"/api/chats/{chat_id}/messages/{message.id}/stream")
    assert response.status_code == 200
    assert "Could not reach" in response.text
    assert "alert--error" in response.text

    db.refresh(message)
    assert message.complete is True
    assert message.error