summaryrefslogtreecommitdiff
path: root/src/ollama_web/routes/chat.py
blob: 2612279c63ab98ff70d474353f198c7cf4acd7bf (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
"""SSE streaming chat endpoint."""

from __future__ import annotations

import json
from collections.abc import AsyncIterator
from typing import Any

from ollama import Message
from starlette.requests import Request
from starlette.responses import StreamingResponse


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 = m.get("content", "")
        name = m.get("name")
        images = m.get("images")
        kwargs: dict[str, Any] = {"role": role, "content": content}
        if name:
            kwargs["name"] = name
        if images:
            kwargs["images"] = images
        out.append(Message(**kwargs))
    return out


async def _event_stream(request: Request, payload: dict[str, Any]) -> AsyncIterator[bytes]:
    """Yield SSE events from the tool-augmented chat loop."""
    model = payload.get("model") or request.app.state.settings.default_model
    raw_messages = payload.get("messages", [])
    think = payload.get("think")

    messages = _parse_messages(raw_messages)
    host = request.app.state.settings.ollama_host

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

    async for event in llm.astream_chat_with_tools(messages, model, think=think, host=host):
        data = json.dumps(event, ensure_ascii=False)
        yield f"data: {data}\n\n".encode()


async def chat(request: Request) -> StreamingResponse:
    """POST /api/chat: stream chat events as Server-Sent Events."""
    import logging

    logger = logging.getLogger("ollama_web.chat")
    try:
        body = await request.body()
        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",
        )
    msg_count = len(payload.get("messages", []))
    logger.info("chat request model=%s messages=%d", payload.get("model"), msg_count)
    return StreamingResponse(
        _event_stream(request, payload),
        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()