summaryrefslogtreecommitdiff
path: root/src/ollama_web/routes/chat.py
diff options
context:
space:
mode:
Diffstat (limited to 'src/ollama_web/routes/chat.py')
-rw-r--r--src/ollama_web/routes/chat.py72
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()