vision対応モデル向けに添付画像の処理を変更

1 files changed, 125 insertions, 59 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: