path: root/src/ollama_web/routes/chat.py

"""SSE streaming chat endpoint."""

from __future__ import annotations

import functools
import html
import inspect
import json
import logging
from collections.abc import AsyncIterator
from datetime import datetime
from typing import Any

from ollama import Message
from ollama._types import Image as OllamaImage
from starlette.requests import Request
from starlette.responses import StreamingResponse

from ..i18n import t
from ..mcp import collect_mcp_tools, make_mcp_executor
from ..prompts import get_prompt
from ..sessions import SessionStore, is_valid_id
from ..settings_store import load_app_settings
from ..tools.fetch import pop_fetched_files
from ..tools.helper.image import resize_image
from ..tools.pdf import make_session_pdf_tool
from ..tools.registry import ToolArtifact, ToolRegistry, default_registry

logger = logging.getLogger("ollama_web.chat")


def _limit_text(text: str, limit: int) -> str:
    if len(text) <= limit:
        return text
    return text[:limit] + "\n…[truncated]"


def _compose_system_prompt(language: str, custom_prompt: str) -> str:
    """Combine the localized built-in tool prompt, current time, and user prompt."""
    tool_prompt = get_prompt("tool_system", lang=language)
    time_line = get_prompt("current_time", lang=language).format(
        time=datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    )
    parts = [tool_prompt, time_line]
    if custom_prompt:
        parts.append(custom_prompt)
    return "\n\n".join(parts)


def _parse_messages(raw: list[dict[str, Any]]) -> list[Message]:
    """Convert raw JSON message dicts into ``ollama.Message`` objects."""
    out: list[Message] = []
    for m in raw:
        role = m.get("role", "user")
        content = _limit_text(str(m.get("content", "") or ""), 20000)
        name = m.get("name")
        images = m.get("images")
        kwargs: dict[str, Any] = {"role": role, "content": content}
        if name:
            kwargs["name"] = name
        if images:
            # ollama expects a sequence of ollama._types.Image objects.
            kwargs["images"] = [OllamaImage(value=img) for img in images]
        out.append(Message(**kwargs))
    return out


async def _build_registry(
    session_id: str | None, store: SessionStore | None = None
) -> ToolRegistry:
    """Return a registry with built-in tools and any configured MCP tools."""
    reg = default_registry()

    original = reg.get("search_and_fetch")
    if original is not None and session_id is not None:

        @functools.wraps(original)
        def wrapped(*args: Any, **kwargs: Any) -> Any:
            sig = inspect.signature(original)
            bound = sig.bind(*args, **kwargs)
            bound.arguments["session_id"] = session_id
            return original(*bound.args, **bound.kwargs)

        reg._tools["search_and_fetch"] = wrapped  # noqa: SLF001

    if session_id is not None and store is not None:

        def load_pdf_attachment(file_id: str) -> bytes | None:
            chat_file = store.get_file(session_id, file_id)
            if chat_file is None or chat_file.is_image:
                return None
            if chat_file.mime.lower() != "application/pdf" and not chat_file.name.lower().endswith(
                ".pdf"
            ):
                return None
            return store.get_file_data(session_id, file_id)

        reg._tools["pdf_to_text"] = make_session_pdf_tool(load_pdf_attachment)  # noqa: SLF001

    try:
        mcp_tools = await collect_mcp_tools()
    except Exception as exc:  # noqa: BLE001
        logger.warning("failed to collect MCP tools: %s", exc)
        mcp_tools = []

    tool_names = [t.get("function", {}).get("name") for t in mcp_tools]
    logger.info("collected %d MCP tool(s): %s", len(mcp_tools), tool_names)

    if mcp_tools:
        artifact_handler = None
        if session_id is not None and store is not None:

            def persist_artifact(full_name: str, artifact: ToolArtifact) -> None:
                if not artifact.for_user:
                    return
                chat_file = store.add_file(
                    session_id=session_id,
                    name=artifact.name,
                    data=artifact.data,
                    mime=artifact.mime,
                    text="",
                    source=full_name,
                )
                artifact.file_id = chat_file.id

            artifact_handler = persist_artifact

        executor = make_mcp_executor(artifact_handler=artifact_handler)
        for tool_def in mcp_tools:
            fn = tool_def.get("function", {})
            name = fn.get("name")
            if not name:
                continue
            try:
                reg.register_mcp_tool(name, tool_def, executor)
            except ValueError as exc:
                logger.warning("skipping duplicate MCP tool %s: %s", name, exc)

    return reg


def _attachment_context(
    store: SessionStore,
    session_id: str,
    file_ids: list[str],
    *,
    total_limit: int,
    language: str | None = None,
) -> str:
    if total_limit <= 0:
        return ""
    chat_files = []
    for file_id in file_ids:
        chat_file = store.get_file(session_id, file_id)
        if chat_file is not None and not chat_file.is_image and not chat_file.is_audio:
            chat_files.append(chat_file)
    if not chat_files:
        return ""

    parts: list[str] = []
    prefix = get_prompt("attachment_prefix", lang=language)
    paging_hint = get_prompt(
        "attachment_paging",
        lang=language,
        fallback=(
            "PDF previews may have unread content. Use pdf_to_text with the attachment URI "
            "and next_token until has_more is absent."
        ),
    )
    used = len(prefix) + len(paging_hint) + 4
    per_file = max(256, min(4000, (total_limit - used) // max(1, len(chat_files))))
    pdf_tool = make_session_pdf_tool(
        lambda file_id: (
            store.get_file_data(session_id, file_id)
            if store.get_file(session_id, file_id) is not None
            else None
        )
    )

    for chat_file in chat_files:
        source_uri = f"attachment:{chat_file.id}"
        if chat_file.mime.lower() == "application/pdf" or chat_file.name.lower().endswith(".pdf"):
            text = pdf_tool(source_uri, max_chars=per_file)
        else:
            text = _limit_text(chat_file.text or "", per_file)
        block = (
            f'<attached-file name="{html.escape(chat_file.name)}" id="{chat_file.id}" '
            f'source="{chat_file.source}" uri="{source_uri}">\n{text}\n</attached-file>'
        )
        if used + len(block) + 2 > total_limit:
            break
        parts.append(block)
        used += len(block) + 2
    if not parts:
        return ""
    return "\n\n".join([prefix, paging_hint, *parts])


def _has_image_attachments(store: SessionStore, session_id: str, file_ids: list[str]) -> bool:
    return any(
        (chat_file := store.get_file(session_id, file_id)) is not None and chat_file.is_image
        for file_id in file_ids
    )


def _has_audio_attachments(store: SessionStore, session_id: str, file_ids: list[str]) -> bool:
    return any(
        (chat_file := store.get_file(session_id, file_id)) is not None and chat_file.is_audio
        for file_id in file_ids
    )


def _collect_images(store: SessionStore, session_id: str, file_ids: list[str]) -> list[bytes]:
    images: list[bytes] = []
    for file_id in file_ids:
        chat_file = store.get_file(session_id, file_id)
        if chat_file is None or not chat_file.is_image:
            continue
        data = store.get_file_data(session_id, file_id)
        if data is None:
            continue
        try:
            resized = resize_image(data, name=chat_file.name, max_dimension=1024, quality=85)
            images.append(resized)
        except Exception:  # noqa: BLE001
            logger.warning(
                "failed to resize image %s, falling back to original bytes", chat_file.name
            )
            images.append(data)
    return images


def _token_count_from_done(
    done_event: dict[str, Any] | None, num_ctx: int
) -> dict[str, int] | None:
    if not done_event:
        return None
    prompt_eval_count = done_event.get("prompt_eval_count")
    eval_count = done_event.get("eval_count")
    if prompt_eval_count is None or eval_count is None:
        return None
    return {
        "prompt_eval_count": int(prompt_eval_count),
        "eval_count": int(eval_count),
        "num_ctx": num_ctx,
        "total_count": int(prompt_eval_count) + int(eval_count),
    }


async def _event_stream(
    request: Request,
    payload: dict[str, Any],
    session_id: str | None,
) -> AsyncIterator[bytes]:
    """Yield SSE events from the tool-augmented chat loop."""
    # Send a keepalive comment immediately so the browser knows the connection is alive.
    yield b":\n\n"
    async for chunk in _chat_event_stream(request, payload, session_id):
        data = json.dumps(chunk, ensure_ascii=False)
        yield f"data: {data}\n\n".encode()


async def _chat_event_stream(
    request: Request,
    payload: dict[str, Any],
    session_id: str | None,
) -> AsyncIterator[dict[str, Any]]:
    """Core chat loop yielding events. Exceptions are converted to SSE error events."""
    model = payload.get("model") or request.app.state.settings.default_model
    think = payload.get("think")
    runtime_settings = request.app.state.settings
    app_settings = load_app_settings(runtime_settings.data_dir, runtime_settings.language)
    language = str(app_settings["ui"]["language"])
    ollama_settings = app_settings["ollama"]
    ollama_options = dict(ollama_settings["options"])
    custom_prompt = str(ollama_settings["system_prompt"])
    system_prompt = _compose_system_prompt(language, custom_prompt)
    num_ctx_value = ollama_options.get("num_ctx")
    try:
        num_ctx = int(num_ctx_value) if num_ctx_value is not None else 0
    except (TypeError, ValueError):
        num_ctx = 0

    store: SessionStore = request.app.state.session_store
    file_ids: list[str] = [
        str(file_id) for file_id in payload.get("file_ids", []) if is_valid_id(str(file_id))
    ]
    user_content = ""
    assistant_text = ""
    assistant_file_ids: list[str] = []
    last_done_event: dict[str, Any] | None = None

    try:
        session_messages: list[Message] = []
        if session_id:
            session = store.get(session_id)
            if session is not None:
                session_messages = _parse_messages(session.get("messages", []))

        raw_messages = list(payload.get("messages", []))
        if raw_messages:
            user_content = _limit_text(
                str(raw_messages[-1].get("content", "") or ""),
                request.app.state.settings.max_message_chars,
            )
            raw_messages[-1]["content"] = user_content

        # An attachment-only submission still needs a user message to carry
        # either image bytes or the localized unsupported-vision notice.
        if file_ids and not raw_messages:
            raw_messages.append({"role": "user", "content": ""})

        host = request.app.state.settings.ollama_host

        from ollama_web import llm  # local import to avoid circular at module load

        capabilities = llm.get_model_capabilities(model, host=host)
        logger.info(
            "chat stream model=%s capabilities=%s",
            model,
            sorted(capabilities),
        )

        # Fallback: some cloud/proxy models support tools even though ollama
        # does not advertise the capability. Enable only tools; vision remains
        # opt-in because sending image bytes to an unknown model is not safe.
        if "tools" not in capabilities and ":cloud" in model:
            capabilities.add("tools")
            logger.info("enabled tools capability fallback for cloud model %s", model)

        attachment_ctx = ""
        if session_id and file_ids:
            attachment_budget = max(
                0,
                request.app.state.settings.max_message_chars - len(user_content) - 256,
            )
            attachment_ctx = _attachment_context(
                store,
                session_id,
                file_ids,
                total_limit=attachment_budget,
                language=language,
            )

        if attachment_ctx and user_content:
            # Inject attachment context into the latest user message.
            augmented = f"{user_content}\n\n{attachment_ctx}"
            if raw_messages:
                raw_messages[-1]["content"] = augmented
        elif attachment_ctx:
            raw_messages.append({"role": "user", "content": attachment_ctx})

        # Attach image bytes only when Ollama explicitly advertises vision.
        # Other models receive a localized fixed notice, never base64 text.
        if session_id and file_ids and raw_messages:
            notices: list[str] = []
            if _has_image_attachments(store, session_id, file_ids):
                if "vision" in capabilities:
                    images = _collect_images(store, session_id, file_ids)
                    if images:
                        raw_messages[-1]["images"] = images
                else:
                    notices.append(get_prompt("vision_unsupported_attachment", lang=language))
            if _has_audio_attachments(store, session_id, file_ids):
                notices.append(get_prompt("audio_unsupported_attachment", lang=language))
            if notices:
                current = str(raw_messages[-1].get("content", "") or "")
                raw_messages[-1]["content"] = f"{current}\n\n" + "\n\n".join(notices)

        messages = _parse_messages(raw_messages)
        if session_messages:
            messages = session_messages + messages

        # Prepend a system prompt to encourage efficient tool usage when tools
        # are available. Avoid duplicate injection if the session already has
        # a system message.
        if not any(getattr(m, "role", None) == "system" for m in messages):
            messages.insert(0, Message(role="system", content=system_prompt))

        registry = await _build_registry(session_id, store)

        total_image_bytes = sum(
            len(img.value) for img in (getattr(messages[-1], "images", None) or [])
        )
        logger.info(
            "chat stream model=%s session=%s messages=%d images=%d image_bytes=%d",
            model,
            session_id,
            len(messages),
            len(getattr(messages[-1], "images", []) or []),
            total_image_bytes,
        )

        async for event in llm.astream_chat_with_tools(
            messages,
            model,
            think=think,
            host=host,
            registry=registry,
            capabilities=capabilities,
            options=ollama_options,
            language=language,
        ):
            if event.get("type") == "delta":
                assistant_text = _limit_text(
                    assistant_text + event.get("content", ""),
                    request.app.state.settings.max_message_chars,
                )
            if event.get("type") == "done":
                event["num_ctx"] = num_ctx
                last_done_event = event
            if event.get("type") == "tool_end":
                for media in event.get("media", []):
                    file_id = str(media.get("id", ""))
                    if is_valid_id(file_id) and file_id not in assistant_file_ids:
                        assistant_file_ids.append(file_id)
            yield event

    except Exception as exc:  # noqa: BLE001
        logger.exception("chat stream failed")
        error_msg = f"Stream failed: {exc}"
        yield {"type": "error", "message": error_msg}
        # Persist the error as the assistant message so it survives reloads
        # and re-renders in the chat pane.
        assistant_text = (
            f"[ERROR] {error_msg}\n\n{t('errors.vision_model_required', lang=language)}"
        )

    finally:
        # Persist user and assistant messages to the session even when the
        # upstream call failed, so the UI keeps the user message.
        if session_id and (user_content or file_ids):
            store.add_message(
                session_id=session_id,
                role="user",
                content=user_content,
                attachments=file_ids,
            )
            if assistant_text or assistant_file_ids:
                token_count = _token_count_from_done(last_done_event, num_ctx)
                store.add_message(
                    session_id=session_id,
                    role="assistant",
                    content=assistant_text,
                    attachments=assistant_file_ids,
                    token_count=token_count,
                )

        # Persist fetched files from tools like search_and_fetch.
        if session_id:
            fetched = pop_fetched_files(session_id)
            for f in fetched:
                store.add_file(
                    session_id=session_id,
                    name=f["name"],
                    data=f["data"],
                    mime=f["mime"],
                    text=f["text"],
                    source="fetch",
                )

        # Always emit a terminal done event so the client closes cleanly.
        yield {"type": "done"}


async def chat(request: Request) -> StreamingResponse:
    """POST /api/chat: stream chat events as Server-Sent Events."""
    try:
        body = await request.body()
        if len(body) > request.app.state.settings.max_message_chars * 4:
            data = json.dumps(
                {"type": "error", "message": "Request body too large"},
                ensure_ascii=False,
            )
            return StreamingResponse(
                _single_event(data),
                media_type="text/event-stream",
                status_code=413,
            )
        payload = json.loads(body.decode("utf-8"))
    except Exception as exc:
        logger.error("failed to parse request json: %s", exc)
        data = json.dumps({"type": "error", "message": f"Invalid JSON: {exc}"}, ensure_ascii=False)
        return StreamingResponse(
            _single_event(data),
            media_type="text/event-stream",
        )

    session_id = payload.get("session_id")
    msg_count = len(payload.get("messages", []))
    logger.info(
        "chat request model=%s session=%s messages=%d",
        payload.get("model"),
        session_id,
        msg_count,
    )
    return StreamingResponse(
        _event_stream(request, payload, session_id),
        media_type="text/event-stream",
    )


async def _single_event(data: str) -> AsyncIterator[bytes]:
    """Yield a single SSE event and finish."""
    yield f"data: {data}\n\n".encode()