wrolpi/modules/ai/test/test_chat_api.py

"""End-to-end tests for /api/chat: SSE framing, the agent loop, and its guardrails.

llama-server is scripted (llama_chat_stream is patched); tool calls execute against the real
/api/ai blueprint through the test client."""
import json
from http import HTTPStatus
from unittest import mock
from urllib.parse import urlencode

import pytest

from modules.ai.config import get_ai_config


@pytest.fixture
async def ai_enabled(async_client, test_ai_config, test_directory):
    """AI enabled with a (fake) downloaded model, llama-server 'healthy'."""
    models_dir = test_directory / 'ai/models'
    models_dir.mkdir(parents=True)
    (models_dir / 'test.gguf').write_bytes(b'GGUF')
    config = get_ai_config()
    config.active_model = 'test.gguf'
    config.enabled = True

    with mock.patch('modules.ai.service.llama_healthy', return_value=True), \
            mock.patch('modules.ai.chat.service.ensure_llama_running', return_value=True):
        yield config


def _llama_script(*turns):
    """Patch llama_chat_stream with scripted turns.

    Each turn is a list of events ({'type': 'token'|'stop', ...}); turn N answers call N."""
    calls = dict(count=0, messages=[])

    def factory(messages, tools):
        calls['messages'].append(messages)
        turn = turns[min(calls['count'], len(turns) - 1)]
        calls['count'] += 1

        async def generator():
            for event in turn:
                yield event

        return generator()

    return mock.patch('modules.ai.chat.llama_chat_stream', side_effect=factory), calls


def _endpoint_shim(async_client):
    async def shim(method, path, query, body):
        url = f'{path}?{urlencode(query)}' if query else path
        if method == 'post':
            request, response = await async_client.post(url, content=json.dumps(body or {}))
        else:
            request, response = await async_client.get(url)
        return response.status_code, response.json

    return mock.patch('modules.ai.agent.call_ai_endpoint', side_effect=shim)


def _parse_sse(text):
    """[(event, data_dict), ...] from an SSE body."""
    events = []
    for block in text.split('\n\n'):
        event = data = None
        for line in block.splitlines():
            if line.startswith('event:'):
                event = line[6:].strip()
            elif line.startswith('data:'):
                data = json.loads(line[5:].strip())
        if event:
            events.append((event, data))
    return events


@pytest.mark.asyncio
async def test_chat_tool_call_then_answer(async_client, ai_enabled, video_factory, test_session):
    """The loop executes a tool call against the blueprint, then streams the final answer."""
    video_factory(title='pressure canning')
    test_session.commit()

    patch, calls = _llama_script(
        [dict(type='stop', tool_calls=[
            dict(id='call_1', name='search_videos', arguments='{"search_str": "canning"}')])],
        [dict(type='token', content='Found '), dict(type='token', content='one video.'),
         dict(type='stop', tool_calls=[])],
    )
    body = dict(mode='research', messages=[dict(role='user', content='Find videos about canning')])
    with patch, _endpoint_shim(async_client):
        request, response = await async_client.post('/api/chat/', content=json.dumps(body))
    assert response.status_code == HTTPStatus.OK

    events = _parse_sse(response.text)
    names = [e for e, _ in events]
    assert names == ['tool_call', 'tool_result', 'token', 'token', 'done']
    tool_call = dict(events)['tool_call']
    assert tool_call == dict(tool='search_videos', args=dict(search_str='canning'))
    assert dict(events)['tool_result'] == dict(tool='search_videos', success=True)
    assert dict(events)['done'] == dict(content='Found one video.')

    # The second llama call saw the system prompt, the user message, and the tool result.
    second_call_messages = calls['messages'][1]
    assert second_call_messages[0]['role'] == 'system'
    assert 'librarian' in second_call_messages[0]['content']
    assert second_call_messages[-1]['role'] == 'tool'
    assert 'pressure canning' in second_call_messages[-1]['content']


@pytest.mark.asyncio
async def test_chat_plain_answer(async_client, ai_enabled):
    """A toolless answer streams tokens then done."""
    patch, _ = _llama_script(
        [dict(type='token', content='Hello!'), dict(type='stop', tool_calls=[])],
    )
    body = dict(mode='help', messages=[dict(role='user', content='hi')])
    with patch:
        request, response = await async_client.post('/api/chat/', content=json.dumps(body))
    events = _parse_sse(response.text)
    assert events == [('token', dict(content='Hello!')), ('done', dict(content='Hello!'))]


@pytest.mark.asyncio
async def test_chat_tool_loop_bounded(async_client, ai_enabled):
    """A model that never stops calling tools hits max_tool_calls and gets the fallback."""
    get_ai_config().max_tool_calls = 2
    patch, calls = _llama_script(
        [dict(type='stop', tool_calls=[dict(id='x', name='search_help', arguments='{"search_str": "a"}')])],
    )
    body = dict(mode='help', messages=[dict(role='user', content='loop forever')])
    with patch, _endpoint_shim(async_client):
        request, response = await async_client.post('/api/chat/', content=json.dumps(body))
    events = _parse_sse(response.text)
    assert events[-1][0] == 'error'
    assert calls['count'] == 3  # initial + max_tool_calls


@pytest.mark.asyncio
async def test_chat_invalid_tool_args(async_client, ai_enabled):
    """Bad tool JSON and disallowed tools come back to the model as errors, not crashes."""
    patch, calls = _llama_script(
        [dict(type='stop', tool_calls=[dict(id='x', name='search_help', arguments='{not json')])],
        [dict(type='stop', tool_calls=[dict(id='y', name='list_disks', arguments='{}')])],
        [dict(type='token', content='Sorry.'), dict(type='stop', tool_calls=[])],
    )
    body = dict(mode='help', messages=[dict(role='user', content='hi')])
    with patch, _endpoint_shim(async_client):
        request, response = await async_client.post('/api/chat/', content=json.dumps(body))
    events = _parse_sse(response.text)
    results = [d for e, d in events if e == 'tool_result']
    assert results[0]['success'] is False  # invalid JSON arguments
    assert results[1]['success'] is False  # list_disks is not a help-mode tool
    assert events[-1] == ('done', dict(content='Sorry.'))
    # The model saw instructive errors.
    tool_messages = [m for m in calls['messages'][-1] if m['role'] == 'tool']
    assert 'not valid JSON' in tool_messages[0]['content']
    assert 'Unknown tool' in tool_messages[1]['content']


@pytest.mark.asyncio
async def test_chat_validation(async_client, ai_enabled):
    """Bad requests are refused before the stream starts."""
    request, response = await async_client.post(
        '/api/chat/', content=json.dumps(dict(mode='nope', messages=[dict(role='user', content='x')])))
    assert response.status_code == HTTPStatus.BAD_REQUEST

    request, response = await async_client.post(
        '/api/chat/', content=json.dumps(dict(mode='help', messages=[])))
    assert response.status_code == HTTPStatus.BAD_REQUEST


@pytest.mark.asyncio
async def test_chat_disabled(async_client, test_ai_config):
    """Chat refuses when AI is disabled or no model is selected."""
    body = dict(mode='help', messages=[dict(role='user', content='x')])
    request, response = await async_client.post('/api/chat/', content=json.dumps(body))
    assert response.status_code == HTTPStatus.CONFLICT
    assert response.json['code'] == 'AI_DISABLED'


@pytest.mark.asyncio
async def test_chat_modes_endpoint(async_client):
    """The Chat tab can list modes and their suggestion chips."""
    request, response = await async_client.get('/api/chat/modes')
    assert response.status_code == HTTPStatus.OK
    modes = {m['name']: m for m in response.json['modes']}
    assert set(modes) == {'help', 'research', 'system'}
    assert modes['research']['suggestions']