diff options
Diffstat (limited to 'src/ollama_web/routes/chat.py')
| -rw-r--r-- | src/ollama_web/routes/chat.py | 72 |
1 files changed, 72 insertions, 0 deletions
diff --git a/src/ollama_web/routes/chat.py b/src/ollama_web/routes/chat.py new file mode 100644 index 0000000..2612279 --- /dev/null +++ b/src/ollama_web/routes/chat.py @@ -0,0 +1,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() |
