diff options
Diffstat (limited to 'src/ollama_web/routes')
| -rw-r--r-- | src/ollama_web/routes/chat.py | 184 | ||||
| -rw-r--r-- | src/ollama_web/routes/sessions.py | 18 |
2 files changed, 141 insertions, 61 deletions
diff --git a/src/ollama_web/routes/chat.py b/src/ollama_web/routes/chat.py index 8b6c302..08b330a 100644 --- a/src/ollama_web/routes/chat.py +++ b/src/ollama_web/routes/chat.py @@ -10,11 +10,13 @@ from collections.abc import AsyncIterator 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 ..sessions import SessionStore from ..tools.fetch import pop_fetched_files +from ..tools.image import resize_image from ..tools.registry import ToolRegistry, default_registry logger = logging.getLogger("ollama_web.chat") @@ -32,7 +34,8 @@ def _parse_messages(raw: list[dict[str, Any]]) -> list[Message]: if name: kwargs["name"] = name if images: - kwargs["images"] = 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 @@ -60,7 +63,7 @@ def _attachment_context(store: SessionStore, session_id: str, file_ids: list[str parts: list[str] = [] for file_id in file_ids: chat_file = store.get_file(session_id, file_id) - if chat_file is None: + if chat_file is None or chat_file.is_image: continue text = chat_file.text or "" parts.append( @@ -72,88 +75,151 @@ def _attachment_context(store: SessionStore, session_id: str, file_ids: list[str return "\n\n".join(["以下はユーザーが添付したファイルの内容です：", *parts]) +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 + + 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") store: SessionStore = request.app.state.session_store file_ids: list[str] = list(payload.get("file_ids", [])) - - session_messages: list[Message] = [] - if session_id: - session = store.get(session_id) - if session is not None: - session_messages = _parse_messages(session.get("messages", [])) - user_content = "" - raw_messages = payload.get("messages", []) - if raw_messages: - user_content = str(raw_messages[-1].get("content", "") or "") + assistant_text = "" - attachment_ctx = "" - if session_id and file_ids: - attachment_ctx = _attachment_context(store, session_id, file_ids) + 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", [])) - if attachment_ctx and user_content: - # Inject attachment context into the latest user message. - augmented = f"{user_content}\n\n{attachment_ctx}" + raw_messages = payload.get("messages", []) if raw_messages: - raw_messages[-1]["content"] = augmented - elif attachment_ctx: - raw_messages.append({"role": "user", "content": attachment_ctx}) + user_content = str(raw_messages[-1].get("content", "") or "") - messages = _parse_messages(raw_messages) - if session_messages: - messages = session_messages + messages + attachment_ctx = "" + if session_id and file_ids: + attachment_ctx = _attachment_context(store, session_id, file_ids) - host = request.app.state.settings.ollama_host + 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}) - from ollama_web import llm # local import to avoid circular at module load + # Attach image bytes to the latest user message per ollama spec. + if session_id and file_ids and raw_messages: + images = _collect_images(store, session_id, file_ids) + if images: + raw_messages[-1]["images"] = images - registry = _session_aware_registry(session_id) + messages = _parse_messages(raw_messages) + if session_messages: + messages = session_messages + messages - assistant_text = "" - finished = False - async for event in llm.astream_chat_with_tools(messages, model, think=think, host=host, registry=registry): - if event.get("type") == "delta": - assistant_text += event.get("content", "") - elif event.get("type") == "done": - finished = True - data = json.dumps(event, ensure_ascii=False) - yield f"data: {data}\n\n".encode() + host = request.app.state.settings.ollama_host + + from ollama_web import llm # local import to avoid circular at module load - # Persist user and assistant messages to the session. - if session_id and user_content: - store.add_message( - session_id=session_id, - role="user", - content=user_content, - attachments=file_ids, + registry = _session_aware_registry(session_id) + + 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, ) - if assistant_text: - store.add_message( - session_id=session_id, - role="assistant", - content=assistant_text, - ) - # 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( + async for event in llm.astream_chat_with_tools( + messages, model, think=think, host=host, registry=registry + ): + if event.get("type") == "delta": + assistant_text += event.get("content", "") + 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画像を含むメッセージを送信する場合、vision 対応モデルを選択してください。" + + 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: + store.add_message( session_id=session_id, - name=f["name"], - data=f["data"], - mime=f["mime"], - text=f["text"], - source="fetch", + role="user", + content=user_content, + attachments=file_ids, ) + if assistant_text: + store.add_message( + session_id=session_id, + role="assistant", + content=assistant_text, + ) + + # 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: diff --git a/src/ollama_web/routes/sessions.py b/src/ollama_web/routes/sessions.py index 8980960..8640ec3 100644 --- a/src/ollama_web/routes/sessions.py +++ b/src/ollama_web/routes/sessions.py @@ -8,7 +8,7 @@ from typing import Any from starlette.requests import Request from starlette.responses import JSONResponse -from ..sessions import SessionStore +from ..sessions import SessionStore, _IMAGE_MIMES, _IMAGE_SUFFIXES from ..tools.pdf import extract_pdf_text logger = logging.getLogger("ollama_web.sessions") @@ -18,6 +18,17 @@ def _get_store(request: Request) -> SessionStore: return request.app.state.session_store +def _is_image_file(name: str, mime: str) -> bool: + mime_lower = mime.lower() + name_lower = name.lower() + if mime_lower in _IMAGE_MIMES: + return True + for suffix in _IMAGE_SUFFIXES: + if name_lower.endswith(suffix): + return True + return False + + async def list_sessions(request: Request) -> JSONResponse: store = _get_store(request) return JSONResponse({"sessions": store.list_sessions()}) @@ -85,7 +96,10 @@ async def upload_file(request: Request) -> JSONResponse: name = getattr(upload, "filename", "uploaded-file") or "uploaded-file" content_type = getattr(upload, "content_type", None) or "application/octet-stream" text = "" - if content_type == "application/pdf" or name.lower().endswith(".pdf"): + if _is_image_file(name, content_type): + # Images are passed to ollama via Message.images (base64), not as text. + text = "" + elif content_type == "application/pdf" or name.lower().endswith(".pdf"): text = extract_pdf_text(data) else: # Try decode as text for text-like uploads. |