From 84883c36f6b8ddd5cc01f77a1660037c7dba5a9e Mon Sep 17 00:00:00 2001 From: Roland Date: Tue, 25 Aug 2026 21:47:50 -0600 Subject: [PATCH] =?UTF-8?q?ai:=20the=20/api/chat=20agent=20loop=20?= =?UTF-8?q?=E2=80=94=20SSE=20streaming,=20modes,=20tool=20execution?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tool definitions are generated in-process from the /api/ai blueprint's own OpenAPI operations (operationId = tool name, description = tool description, dataclass bodies + declared parameters = one flat JSON schema) — the spec IS the catalog, nothing to drift. Manage endpoints never become tools. A mode is a server-enforced system prompt + tool allowlist (Help / Research / System per ai.md); the executor can only reach allowlisted /api/ai endpoints, path parameters are encoded and traversal-rejected so a crafted argument cannot escape an endpoint. POST /api/chat streams SSE (status/token/tool_call/tool_result/done/ error), starts llama-server on demand, validates tool args server-side with instructive errors the model can retry on, caps tool results with a paging hint, and bounds the loop by ai.yaml max_tool_calls plus a hard timeout. Caddy passes the stream unbuffered (flush_interval -1) in both deployments. Co-Authored-By: Claude Fable 5 --- docker/web/Caddyfile | 7 + etc/raspberrypios/Caddyfile | 7 + modules/ai/agent.py | 236 +++++++++++++++++++++++++++++++ modules/ai/api.py | 10 ++ modules/ai/chat.py | 200 ++++++++++++++++++++++++++ modules/ai/modes.py | 81 +++++++++++ modules/ai/schema.py | 12 +- modules/ai/test/test_agent.py | 135 ++++++++++++++++++ modules/ai/test/test_chat_api.py | 188 ++++++++++++++++++++++++ wrolpi/root_api.py | 2 + 10 files changed, 872 insertions(+), 6 deletions(-) create mode 100644 modules/ai/agent.py create mode 100644 modules/ai/chat.py create mode 100644 modules/ai/modes.py create mode 100644 modules/ai/test/test_agent.py create mode 100644 modules/ai/test/test_chat_api.py diff --git a/docker/web/Caddyfile b/docker/web/Caddyfile index 1d745032..b775e57a 100644 --- a/docker/web/Caddyfile +++ b/docker/web/Caddyfile @@ -51,6 +51,13 @@ } } + # AI chat streams Server-Sent Events; never buffer them. + handle /api/chat* { + reverse_proxy api:8081 { + flush_interval -1 + } + } + # Main API. handle /api/* { reverse_proxy api:8081 diff --git a/etc/raspberrypios/Caddyfile b/etc/raspberrypios/Caddyfile index 5a6f2484..e0a2a2fe 100644 --- a/etc/raspberrypios/Caddyfile +++ b/etc/raspberrypios/Caddyfile @@ -51,6 +51,13 @@ } } + # AI chat streams Server-Sent Events; never buffer them. + handle /api/chat* { + reverse_proxy 0.0.0.0:8081 { + flush_interval -1 + } + } + # Main API. handle /api/* { reverse_proxy 0.0.0.0:8081 diff --git a/modules/ai/agent.py b/modules/ai/agent.py new file mode 100644 index 00000000..bb3f379c --- /dev/null +++ b/modules/ai/agent.py @@ -0,0 +1,236 @@ +"""Tool definitions and tool execution for the agent loop. + +The /api/ai blueprint IS the tool catalog: tool definitions are generated from the blueprint's +own OpenAPI operations (in-process — the spec endpoint is not proxied by Caddy), so there is no +separate registry to drift. operationIds are the tool names; @openapi.description strings are +the tool descriptions; @validate dataclasses and @openapi.parameter declarations become the +JSON-schema parameters.""" +import dataclasses +import json +from typing import Dict, List, Optional, Tuple +from urllib.parse import quote, urlencode + +from sanic_ext.extensions.openapi.builders import OperationStore, SpecificationBuilder +from sanic_ext.utils.route import get_all_routes + +from wrolpi.api_utils import api_app +from wrolpi.common import aiohttp_get, aiohttp_post, logger +from wrolpi.vars import API_PORT + +from .lib import PAGE_SIZE + +logger = logger.getChild(__name__) + +AI_PATH_PREFIX = '/api/ai/' +# Manage endpoints are for the Manage tab, never the model. +EXCLUDED_PREFIX = '/api/ai/manage' + +# Tool results are capped so one result cannot blow a small context; content endpoints page. +TOOL_RESULT_MAX_CHARS = PAGE_SIZE + 1_000 + +_TYPE_MAP = {'integer': int, 'number': float, 'boolean': bool, 'string': str} + + +@dataclasses.dataclass +class ToolDefinition: + name: str + description: str + method: str # 'get' or 'post' + path: str # e.g. '/api/ai/videos/{file_group_id}/captions' + path_params: Dict[str, dict] # name -> JSON schema + query_params: Dict[str, dict] + body_params: Dict[str, dict] + + def parameters_schema(self) -> dict: + """One flat JSON-schema object; the executor routes each arg to its place.""" + properties = {**self.body_params, **self.query_params, **self.path_params} + required = sorted(self.path_params) + schema = dict(type='object', properties=properties) + if required: + schema['required'] = required + return schema + + def as_openai_tool(self) -> dict: + return dict(type='function', function=dict( + name=self.name, + description=self.description, + parameters=self.parameters_schema(), + )) + + +def _build_ai_spec() -> dict: + """Populate and serialize the OpenAPI spec for the AI blueprint. + + Replicates sanic-ext's build_spec startup hook (sanic_ext/extensions/openapi/blueprint.py) for + just our routes, so tool definitions work identically in production and under pytest.""" + spec = SpecificationBuilder() + store = OperationStore() + for uri, route_name, route_parameters, method_handlers, host in get_all_routes(api_app, '/docs'): + if not uri.startswith(AI_PATH_PREFIX): + continue + for method, handler in method_handlers: + if method in ('OPTIONS', 'HEAD', 'TRACE'): + continue + if handler not in store and (func := getattr(handler, '__func__', None)) and func in store: + handler = func + operation = store[handler] + if operation._exclude: + continue + operation._default['operationId'] = f'{method.lower()}~{route_name}' + for parameter in route_parameters: + if not any(p.fields['name'] == parameter.name for p in operation.parameters): + operation.parameter(parameter.name, parameter.cast, 'path') + operation._app = api_app + spec.operation(uri, method, operation) + return spec.build(api_app).serialize() + + +def _resolve_ref(spec: dict, schema: Optional[dict]) -> dict: + if not schema: + return {} + if ref := schema.get('$ref'): + name = ref.rsplit('/', 1)[-1] + return (spec.get('components', {}).get('schemas', {}) or {}).get(name, {}) + return schema + + +def _clean_property(schema: dict) -> dict: + """A lean JSON-schema property for the model: type (+items), nothing else.""" + schema = schema or {} + cleaned = {} + if type_ := schema.get('type'): + cleaned['type'] = type_ + if items := schema.get('items'): + cleaned['items'] = _clean_property(items) + return cleaned or {'type': 'string'} + + +_TOOL_DEFINITIONS: Optional[Dict[str, ToolDefinition]] = None + + +def get_ai_tool_definitions(refresh: bool = False) -> Dict[str, ToolDefinition]: + """All tools the agent loop can ever execute, keyed by name. Built lazily, cached. + + This map is also the hard allowlist: the executor can only reach /api/ai endpoints that + appear here, and manage endpoints never do.""" + global _TOOL_DEFINITIONS + if _TOOL_DEFINITIONS is not None and not refresh: + return _TOOL_DEFINITIONS + + spec = _build_ai_spec() + tools = {} + for path, operations in (spec.get('paths') or {}).items(): + if not path.startswith(AI_PATH_PREFIX) or path.startswith(EXCLUDED_PREFIX): + continue + for method, operation in operations.items(): + if method not in ('get', 'post'): + continue + name = operation.get('operationId') + description = operation.get('description') or operation.get('summary') or '' + if not name or '~' in name: + # Every tool endpoint must declare a deliberate @openapi.operation name. + logger.warning(f'AI endpoint without an operation name is not a tool: {method} {path}') + continue + + path_params, query_params = {}, {} + for parameter in operation.get('parameters') or []: + cleaned = _clean_property(parameter.get('schema')) + if parameter.get('in') == 'path': + path_params[parameter['name']] = cleaned + elif parameter.get('in') == 'query': + query_params[parameter['name']] = cleaned + + body_params = {} + content = (operation.get('requestBody') or {}).get('content') or {} + # sanic-ext serializes dataclass bodies under '*/*'. + body_schema = (content.get('application/json') or content.get('*/*') or {}).get('schema') + body_schema = _resolve_ref(spec, body_schema) + for prop_name, prop_schema in (body_schema.get('properties') or {}).items(): + body_params[prop_name] = _clean_property(_resolve_ref(spec, prop_schema)) + + tools[name] = ToolDefinition( + name=name, description=description, method=method, path=path, + path_params=path_params, query_params=query_params, body_params=body_params, + ) + + _TOOL_DEFINITIONS = tools + return tools + + +def get_mode_tools(mode_tools: tuple) -> List[dict]: + """OpenAI-format tool definitions for one mode's allowlist.""" + definitions = get_ai_tool_definitions() + return [definitions[name].as_openai_tool() for name in mode_tools if name in definitions] + + +async def call_ai_endpoint(method: str, path: str, query: dict, body: Optional[dict]) -> Tuple[int, dict]: + """Execute one blueprint endpoint over localhost. Patched in tests.""" + url = f'http://127.0.0.1:{API_PORT}{path}' + if query: + url = f'{url}?{urlencode(query)}' + if method == 'post': + async with aiohttp_post(url, json_=body or {}, timeout=60) as response: + return response.status, await response.json() + async with aiohttp_get(url, timeout=60) as response: + return response.status, await response.json() + + +def _validate_args(tool: ToolDefinition, args: dict) -> Optional[str]: + """Return a crisp, instructive error string for the model, or None when the args are usable.""" + known = {**tool.path_params, **tool.query_params, **tool.body_params} + if unknown := sorted(set(args) - set(known)): + return f'Unknown arguments {unknown}. {tool.name} accepts: {sorted(known)}.' + if missing := sorted(set(tool.path_params) - set(args)): + return f'Missing required arguments {missing} for {tool.name}.' + for name, value in args.items(): + expected = _TYPE_MAP.get(known[name].get('type')) + if expected in (int, float) and value is not None and not isinstance(value, (int, float)) \ + and not (isinstance(value, str) and value.lstrip('-').isdigit()): + return f'Argument {name} must be a {known[name]["type"]}, got {value!r}.' + # Path parameters land in the request URL; refuse anything that could escape the + # endpoint's own path (the allowlist must hold even against a crafted value). + if name in tool.path_params and isinstance(value, str): + segments = value.split('/') + if value.startswith('/') or '' in segments or '.' in segments or '..' in segments: + return f'Argument {name} must not contain path traversal.' + return None + + +async def execute_tool(mode_tools: tuple, name: str, args: dict) -> Tuple[bool, str]: + """Validate and execute one tool call. Returns (success, result text for the model). + + Every failure is a plain-language error string — small models retry well on crisp errors.""" + definitions = get_ai_tool_definitions() + if name not in definitions or name not in mode_tools: + available = sorted(set(mode_tools) & set(definitions)) + return False, f'Unknown tool {name!r}. Available tools: {available}.' + tool = definitions[name] + + args = args or {} + if error := _validate_args(tool, args): + return False, error + + path = tool.path + for param in tool.path_params: + # Encode everything but '/' (help-doc slugs contain it); traversal was rejected above. + path = path.replace(f'{{{param}}}', quote(str(args[param]), safe='/')) + # The substituted path must still be inside the endpoint the tool declares. + static_prefix = tool.path.split('{', 1)[0] + if not path.startswith(static_prefix) or '..' in path.split('/'): + return False, f'Arguments produced an invalid path for {tool.name}.' + query = {k: v for k, v in args.items() if k in tool.query_params and v is not None} + body = {k: v for k, v in args.items() if k in tool.body_params} if tool.method == 'post' else None + + try: + status, response = await call_ai_endpoint(tool.method, path, query, body) + except Exception as e: + logger.error(f'Tool {name} failed', exc_info=e) + return False, f'Tool {name} failed to execute; try different arguments or another tool.' + + text = json.dumps(response, default=str) + if len(text) > TOOL_RESULT_MAX_CHARS: + text = text[:TOOL_RESULT_MAX_CHARS] + '… (result truncated; use limit/offset to get less at a time)' + if status >= 400: + message = response.get('message') or response.get('error') or f'HTTP {status}' + return False, f'{tool.name} returned an error: {message}' + return True, text diff --git a/modules/ai/api.py b/modules/ai/api.py index 4a72ef94..7413d65e 100644 --- a/modules/ai/api.py +++ b/modules/ai/api.py @@ -117,6 +117,7 @@ async def get_video(_: Request, file_group_id: int): ' continue reading.', ) @openapi.operation('get_video_captions') +@openapi.parameter('offset', int, 'query') @openapi.response(HTTPStatus.OK, schema.AIPagedTextResponse) @openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse) async def get_video_captions(request: Request, file_group_id: int): @@ -156,6 +157,7 @@ async def search_archives(_: Request, body: schema.AIArchiveSearchRequest): ' again with offset=next_offset to continue reading.', ) @openapi.operation('get_archive_text') +@openapi.parameter('offset', int, 'query') @openapi.response(HTTPStatus.OK, schema.AIPagedTextResponse) @openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse) async def get_archive_text(request: Request, file_group_id: int): @@ -268,6 +270,8 @@ async def search_zims(request: Request, body: schema.AIZimSearchRequest): ' reading.', ) @openapi.operation('get_zim_entry') +@openapi.parameter('path', str, 'query') +@openapi.parameter('offset', int, 'query') @openapi.response(HTTPStatus.OK, schema.AIPagedTextResponse) @openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse) async def get_zim_entry(request: Request, zim_id: int): @@ -289,6 +293,7 @@ async def get_zim_entry(request: Request, zim_id: int): ' name with search_archives.', ) @openapi.operation('list_collections') +@openapi.parameter('kind', str, 'query') async def list_collections(request: Request): kind = request.args.get('kind') collections = search_collections(request.ctx.session, kind=kind) @@ -366,6 +371,7 @@ async def search_help(_: Request, body: schema.AIHelpSearchRequest): ' pages are paged: request again with offset=next_offset to continue reading.', ) @openapi.operation('get_help_doc') +@openapi.parameter('offset', int, 'query') @openapi.response(HTTPStatus.OK, schema.AIPagedTextResponse) @openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse) async def get_help_doc(request: Request, slug: str): @@ -466,6 +472,8 @@ async def ai_list_services(_: Request): ' Pass lines to read more history (at most 1000).', ) @openapi.operation('get_service_logs') +@openapi.parameter('lines', int, 'query') +@openapi.parameter('since', str, 'query') @openapi.response(HTTPStatus.BAD_GATEWAY, JSONErrorResponse) async def ai_get_service_logs(request: Request, name: str): try: @@ -584,6 +592,8 @@ async def manage_settings(_: Request, body: schema.AIManageSettingsRequest): ' request again with offset=next_offset to continue reading.', ) @openapi.operation('read_file') +@openapi.parameter('path', str, 'query') +@openapi.parameter('offset', int, 'query') @openapi.response(HTTPStatus.OK, schema.AIPagedTextResponse) @openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse) async def read_file(request: Request): diff --git a/modules/ai/chat.py b/modules/ai/chat.py new file mode 100644 index 00000000..43166e08 --- /dev/null +++ b/modules/ai/chat.py @@ -0,0 +1,200 @@ +"""POST /api/chat — the agent loop, streamed over SSE. + +Events (each `event: ` with a JSON `data:` line): + status {message} e.g. "Loading the model…" + token {content} one streamed answer fragment + tool_call {tool, args} the model is using a tool + tool_result {tool, success} the tool finished + done {content} the complete final answer + error {message} + +The loop: build the mode's system prompt and tool definitions -> stream from llama-server -> +execute tool calls against /api/ai only (the mode allowlist is enforced server-side) -> repeat, +bounded by ai.yaml's max_tool_calls -> stream the final answer. + +A streaming response outlives the request middleware, so nothing here may use +request.ctx.session; tools execute through the blueprint, which manages its own sessions.""" +import dataclasses +import json +import time +from http import HTTPStatus +from typing import AsyncGenerator, List, Optional + +from sanic import Blueprint, Request +from sanic_ext import validate +from sanic_ext.extensions.openapi import openapi + +from wrolpi.common import aiohttp_post, logger +from wrolpi.errors import APIError, ValidationError + +from . import agent, service +from .config import get_ai_config +from .modes import MODE_SUGGESTIONS, MODES + +chat_bp = Blueprint('AIChat', url_prefix='/api/chat') + +logger = logger.getChild(__name__) + +# The whole conversation turn (all tool calls and answers) must finish within this. +CHAT_TIMEOUT_SECONDS = 600 +FALLBACK_MESSAGE = "I couldn't complete that request. Try rephrasing, or a more specific question." + + +class AIDisabled(APIError): + code = 'AI_DISABLED' + summary = 'AI is not enabled. Enable it and select a model on the AI page.' + status_code = HTTPStatus.CONFLICT + + +@dataclasses.dataclass +class ChatRequest: + mode: str = '' + messages: List[dict] = dataclasses.field(default_factory=list) + model: Optional[str] = None # reserved; llama-server runs the single active model + + +@chat_bp.get('/modes') +@openapi.description('The available chat modes with their suggestion chips.') +async def get_modes(_: Request): + from wrolpi.api_utils import json_response + modes = [dict(name=name, suggestions=MODE_SUGGESTIONS.get(name, [])) for name in MODES] + return json_response(dict(modes=modes)) + + +async def llama_chat_stream(messages: List[dict], tools: List[dict]) -> AsyncGenerator[dict, None]: + """Stream one /v1/chat/completions call from llama-server. + + Yields {'type': 'token', 'content': str} for answer fragments, then exactly one + {'type': 'stop', 'tool_calls': [...]} where tool_calls is empty for a final answer. + Patched in tests.""" + body = dict(messages=messages, stream=True) + if tools: + body['tools'] = tools + + tool_calls = {} # index -> {id, name, arguments} + async with aiohttp_post(f'{service.LLAMA_URL}/v1/chat/completions', json_=body, + timeout=CHAT_TIMEOUT_SECONDS) as response: + if response.status != 200: + raise RuntimeError(f'llama-server returned HTTP {response.status}') + async for raw_line in response.content: + line = raw_line.decode('UTF-8', errors='replace').strip() + if not line.startswith('data:'): + continue + data = line[5:].strip() + if data == '[DONE]': + break + try: + chunk = json.loads(data) + except ValueError: + continue + choice = (chunk.get('choices') or [{}])[0] + delta = choice.get('delta') or {} + if content := delta.get('content'): + yield dict(type='token', content=content) + for tool_call in delta.get('tool_calls') or []: + index = tool_call.get('index', 0) + entry = tool_calls.setdefault(index, dict(id=None, name='', arguments='')) + if tool_call.get('id'): + entry['id'] = tool_call['id'] + function = tool_call.get('function') or {} + if function.get('name'): + entry['name'] = function['name'] + if function.get('arguments'): + entry['arguments'] += function['arguments'] + + yield dict(type='stop', tool_calls=[tool_calls[i] for i in sorted(tool_calls)]) + + +async def _send_event(response, event: str, data: dict): + await response.send(f'event: {event}\ndata: {json.dumps(data, default=str)}\n\n') + + +@chat_bp.post('/') +@openapi.definition( + summary='Chat with the local AI assistant (SSE stream)', + body=ChatRequest, +) +@validate(ChatRequest) +async def post_chat(request: Request, body: ChatRequest): + if body.mode not in MODES: + raise ValidationError(f'Unknown mode: {body.mode!r}. Modes: {sorted(MODES)}') + if not body.messages: + raise ValidationError('messages must not be empty') + config = get_ai_config() + if not config.enabled or not config.active_model: + raise AIDisabled() + + mode = MODES[body.mode] + tools = agent.get_mode_tools(mode['tools']) + max_tool_calls = config.max_tool_calls + deadline = time.time() + CHAT_TIMEOUT_SECONDS + + response = await request.respond(content_type='text/event-stream', headers={ + 'Cache-Control': 'no-cache', + # Ask proxies not to buffer the stream. + 'X-Accel-Buffering': 'no', + }) + + try: + if not await service.llama_healthy(): + await _send_event(response, 'status', dict(message='Loading the model…')) + if not await service.ensure_llama_running(): + await _send_event(response, 'error', dict(message='The AI service could not be started.')) + return + + messages = [dict(role='system', content=mode['system_prompt'])] + messages += [dict(role=m.get('role', 'user'), content=m.get('content', '')) for m in body.messages] + + final_answer = [] + for _ in range(max_tool_calls + 1): + if time.time() > deadline: + await _send_event(response, 'error', dict(message=FALLBACK_MESSAGE)) + return + + service.record_llm_activity() + tool_calls = [] + final_answer = [] + async for event in llama_chat_stream(messages, tools): + if event['type'] == 'token': + final_answer.append(event['content']) + await _send_event(response, 'token', dict(content=event['content'])) + elif event['type'] == 'stop': + tool_calls = event['tool_calls'] + + if not tool_calls: + # A final answer. + content = ''.join(final_answer).strip() + await _send_event(response, 'done', dict(content=content or FALLBACK_MESSAGE)) + return + + # Execute the tool calls, then loop for the model's next step. + messages.append(dict(role='assistant', content=''.join(final_answer) or None, tool_calls=[ + dict(id=i['id'] or f'call_{n}', type='function', + function=dict(name=i['name'], arguments=i['arguments'])) + for n, i in enumerate(tool_calls) + ])) + for n, tool_call in enumerate(tool_calls): + name = tool_call['name'] + try: + args = json.loads(tool_call['arguments']) if tool_call['arguments'] else {} + except ValueError: + args = None + await _send_event(response, 'tool_call', dict(tool=name, args=args)) + if args is None: + success, result = False, 'The tool arguments were not valid JSON. Try again.' + else: + success, result = await agent.execute_tool(mode['tools'], name, args) + await _send_event(response, 'tool_result', dict(tool=name, success=success)) + messages.append(dict(role='tool', tool_call_id=tool_call['id'] or f'call_{n}', + content=result)) + + # Tool-call budget exhausted without a final answer. + await _send_event(response, 'error', dict(message=FALLBACK_MESSAGE)) + except Exception as e: + logger.error('Chat failed', exc_info=e) + try: + await _send_event(response, 'error', dict(message='The AI failed to answer. Check the logs.')) + except Exception: + pass + finally: + await response.eof() diff --git a/modules/ai/modes.py b/modules/ai/modes.py new file mode 100644 index 00000000..c4aa05f6 --- /dev/null +++ b/modules/ai/modes.py @@ -0,0 +1,81 @@ +"""Chat modes: a mode is a system prompt plus a tool allowlist, enforced server-side in the agent +loop — the model cannot opt out. Tool names are the /api/ai blueprint's operationIds.""" + +COMMON_PROMPT = ( + 'You are the WROLPi assistant, running locally on the user\'s own offline WROLPi server.' + ' WROLPi is an off-grid digital library: videos, archived web pages, ebooks, Zim encyclopedias,' + ' maps, and inventories. You have no internet access.' + ' Answer briefly and plainly. Use the provided tools to find real information before answering;' + ' never invent library content or system state.' +) + +MODES = { + 'help': dict( + system_prompt=( + f'{COMMON_PROMPT}' + ' Your job is to teach the user how to use WROLPi.' + ' Search the help documentation with search_help, read pages with get_help_doc, and answer' + ' from what the docs actually say. Always include the link of every help page you used.' + ' If the docs do not cover the question, say so.' + ), + tools=( + 'search_help', + 'get_help_doc', + 'get_system_status', + ), + ), + 'research': dict( + system_prompt=( + f'{COMMON_PROMPT}' + ' Your job is to be the librarian: find and summarize content from the user\'s library.' + ' Search first, then read captions, page text, or entries when the user wants detail.' + ' ALWAYS include the WROLPi link for every item you mention, exactly as the tools return it.' + ' Never give external URLs. When a search returns a large total, narrow the query instead' + ' of paging. If nothing matches, say so.' + ), + tools=( + 'search_all', + 'search_videos', + 'get_video', + 'get_video_captions', + 'search_archives', + 'get_archive_text', + 'search_docs', + 'get_doc', + 'list_zims', + 'search_zims', + 'get_zim_entry', + 'list_collections', + 'list_inventories', + 'get_inventory', + 'read_file', + ), + ), + 'system': dict( + system_prompt=( + f'{COMMON_PROMPT}' + ' Your job is to troubleshoot the user\'s WROLPi.' + ' Start every diagnosis with get_system_status; read a failing service\'s logs with' + ' get_service_logs; check drives with list_disks.' + ' You CANNOT execute commands or change anything — you can only read.' + ' Only suggest commands that appear in the help documentation: find them with search_help,' + ' quote them exactly, and include the help page link.' + ), + tools=( + 'get_system_status', + 'list_services', + 'get_service_logs', + 'list_disks', + 'search_help', + 'get_help_doc', + ), + ), +} + +# Suggestion chips shown by the Chat tab; small models do much better when the opening message +# matches what the mode's prompt was tuned for. +MODE_SUGGESTIONS = { + 'help': ['How do I download videos?', 'How do I add my own files?', 'What are Tags?'], + 'research': ['Find videos about canning', 'What do I have about gardening?', 'Summarize my newest archives'], + 'system': ['Why is my drive full?', 'Is everything running?', 'Why do downloads fail?'], +} diff --git a/modules/ai/schema.py b/modules/ai/schema.py index 7d7390db..6f6485d1 100644 --- a/modules/ai/schema.py +++ b/modules/ai/schema.py @@ -1,5 +1,5 @@ import dataclasses -from typing import Optional +from typing import List, Optional @dataclasses.dataclass @@ -7,8 +7,8 @@ class AISearchRequest: search_str: Optional[str] = None limit: int = 5 offset: int = 0 - mimetypes: Optional[list] = None - tag_names: Optional[list] = None + mimetypes: Optional[List[str]] = None + tag_names: Optional[List[str]] = None @dataclasses.dataclass @@ -17,7 +17,7 @@ class AIVideoSearchRequest: limit: int = 5 offset: int = 0 channel_id: Optional[int] = None - tag_names: Optional[list] = None + tag_names: Optional[List[str]] = None @dataclasses.dataclass @@ -26,7 +26,7 @@ class AIArchiveSearchRequest: limit: int = 5 offset: int = 0 domain: Optional[str] = None - tag_names: Optional[list] = None + tag_names: Optional[List[str]] = None @dataclasses.dataclass @@ -38,7 +38,7 @@ class AIDocSearchRequest: mimetype: Optional[str] = None limit: int = 5 offset: int = 0 - tag_names: Optional[list] = None + tag_names: Optional[List[str]] = None @dataclasses.dataclass diff --git a/modules/ai/test/test_agent.py b/modules/ai/test/test_agent.py new file mode 100644 index 00000000..f63e1526 --- /dev/null +++ b/modules/ai/test/test_agent.py @@ -0,0 +1,135 @@ +"""Tests for tool-definition generation and tool execution.""" +import json +from unittest import mock +from urllib.parse import urlencode + +import pytest + +from modules.ai import agent +from modules.ai.modes import MODES + + +@pytest.fixture +async def tool_definitions(async_client): + # Any request finalizes the router, which the spec builder requires. + await async_client.get('/api/ai/zims') + return agent.get_ai_tool_definitions(refresh=True) + + +@pytest.mark.asyncio +async def test_tool_definitions_from_spec(tool_definitions): + """Tool definitions are generated from the blueprint's own OpenAPI operations.""" + tools = tool_definitions + # Every tool any mode allowlists must exist — this guards allowlists against drift. + for mode_name, mode in MODES.items(): + for tool_name in mode['tools']: + assert tool_name in tools, f'Mode {mode_name} allowlists missing tool {tool_name}' + + # Every tool has a model-facing description. + for name, tool in tools.items(): + assert tool.description, f'Tool {name} has no description' + + # Manage endpoints are never tools. + assert not any(t.path.startswith('/api/ai/manage') for t in tools.values()) + assert 'manage_catalog' not in tools and 'manage_settings' not in tools + # The chat endpoint itself is not a tool. + assert not any(t.path.startswith('/api/chat') for t in tools.values()) + + # Body, query, and path parameters all appear in the flat schema. + schema = tools['search_videos'].parameters_schema() + assert schema['properties']['search_str'] == {'type': 'string'} + assert schema['properties']['tag_names'] == {'type': 'array', 'items': {'type': 'string'}} + schema = tools['get_video_captions'].parameters_schema() + assert schema['required'] == ['file_group_id'] + assert schema['properties']['offset'] == {'type': 'integer'} + + # OpenAI tool format. + openai_tool = tools['search_all'].as_openai_tool() + assert openai_tool['type'] == 'function' + assert openai_tool['function']['name'] == 'search_all' + + +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) + + +@pytest.mark.asyncio +async def test_execute_tool(async_client, tool_definitions, video_factory, test_session): + """Tools execute against the blueprint; every failure is an instructive error string.""" + video_factory(title='pressure canning') + test_session.commit() + research_tools = MODES['research']['tools'] + + with _endpoint_shim(async_client): + # A real search through the blueprint. + success, result = await agent.execute_tool(research_tools, 'search_videos', + dict(search_str='canning')) + assert success is True + result = json.loads(result) + assert result['total'] == 1 + + # Unknown tool. + success, error = await agent.execute_tool(research_tools, 'delete_everything', {}) + assert success is False and 'Unknown tool' in error + + # Mode allowlist: research mode cannot read service logs even though the tool exists. + success, error = await agent.execute_tool(research_tools, 'get_service_logs', + dict(name='wrolpi-api')) + assert success is False and 'Unknown tool' in error + + # Unknown argument names the accepted ones. + success, error = await agent.execute_tool(research_tools, 'search_videos', dict(query='x')) + assert success is False and 'accepts' in error and 'search_str' in error + + # Missing required path parameter. + success, error = await agent.execute_tool(research_tools, 'get_video_captions', {}) + assert success is False and 'Missing required' in error + + # Wrong type. + success, error = await agent.execute_tool(research_tools, 'get_video_captions', + dict(file_group_id='banana')) + assert success is False and 'must be a integer' in error + + # An endpoint error comes back as a message, not an exception. + success, error = await agent.execute_tool(research_tools, 'get_video_captions', + dict(file_group_id=123456)) + assert success is False and 'error' in error.lower() + + +@pytest.mark.asyncio +async def test_execute_tool_path_traversal(async_client, tool_definitions): + """A crafted path parameter cannot escape the tool's endpoint (allowlist bypass).""" + help_tools = MODES['help']['tools'] + with _endpoint_shim(async_client): + for slug in ('../../manage/catalog', '../manage/settings', '/etc/passwd', + 'a/../../files/refresh', 'a//b', 'a/./b'): + success, error = await agent.execute_tool(help_tools, 'get_help_doc', dict(slug=slug)) + assert success is False, f'traversal not refused: {slug}' + assert 'traversal' in error or 'invalid path' in error + + # Legitimate nested slugs keep working (404 from the endpoint is fine, no traversal error). + success, error = await agent.execute_tool(help_tools, 'get_help_doc', dict(slug='system/logs')) + assert 'traversal' not in error and 'invalid path' not in (error or '') + + +@pytest.mark.asyncio +async def test_execute_tool_result_truncated(async_client, tool_definitions, archive_factory, test_session): + """A huge tool result is capped with a paging hint.""" + archive_factory(domain='example.com', title='big page', contents='word ' * 5_000) + test_session.commit() + from modules.archive.models import Archive + archive = test_session.query(Archive).one() + + with _endpoint_shim(async_client): + success, result = await agent.execute_tool(MODES['research']['tools'], 'get_archive_text', + dict(file_group_id=archive.file_group_id)) + assert success is True + assert len(result) <= agent.TOOL_RESULT_MAX_CHARS + 100 diff --git a/modules/ai/test/test_chat_api.py b/modules/ai/test/test_chat_api.py new file mode 100644 index 00000000..cdb8b13f --- /dev/null +++ b/modules/ai/test/test_chat_api.py @@ -0,0 +1,188 @@ +"""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'] diff --git a/wrolpi/root_api.py b/wrolpi/root_api.py index 14daf1f2..a2b3d3bd 100644 --- a/wrolpi/root_api.py +++ b/wrolpi/root_api.py @@ -14,6 +14,7 @@ from vininfo import Vin from vininfo.details._base import VinDetails from modules.ai.api import ai_bp +from modules.ai.chat import chat_bp from modules.archive.api import archive_bp from modules.docs.api import docs_bp from modules.flasher.api import flasher_bp @@ -55,6 +56,7 @@ api_bp = Blueprint('RootAPI', url_prefix='/api') # Blueprints order here defines what order they are displayed in OpenAPI Docs. api_app.blueprint(api_bp) api_app.blueprint(ai_bp) +api_app.blueprint(chat_bp) api_app.blueprint(archive_bp) api_app.blueprint(collection_bp) # Unified collection endpoints api_app.blueprint(config_bp)