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

1 files changed, 36 insertions, 1 deletions

diff --git a/src/ollama_web/sessions.py b/src/ollama_web/sessions.py
index bfe659c..7cc06ad 100644
--- a/src/ollama_web/sessions.py
+++ b/src/ollama_web/sessions.py
@@ -17,6 +17,20 @@ def _now() -> str:
     return datetime.now(timezone.utc).isoformat()
 
 
+_IMAGE_MIMES: set[str] = {
+    "image/png",
+    "image/jpeg",
+    "image/gif",
+    "image/bmp",
+    "image/webp",
+    "image/tiff",
+    "image/x-png",
+    "image/jpg",
+}
+
+_IMAGE_SUFFIXES: set[str] = {".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", ".tif", ".tiff"}
+
+
 @dataclass
 class ChatFile:
     id: str
@@ -50,6 +64,16 @@ class ChatFile:
             source=data.get("source", "upload"),
         )
 
+    @property
+    def is_image(self) -> bool:
+        name_lower = self.name.lower()
+        if self.mime.lower() in _IMAGE_MIMES:
+            return True
+        for suffix in _IMAGE_SUFFIXES:
+            if name_lower.endswith(suffix):
+                return True
+        return False
+
 
 @dataclass
 class ChatMessage:
@@ -245,4 +269,15 @@ class SessionStore:
         for f in session.get("files", []):
             if f["id"] == file_id:
                 return ChatFile.from_dict(f)
-        return None
\ No newline at end of file
+        return None
+
+    def get_file_data(self, session_id: str, file_id: str) -> bytes | None:
+        chat_file = self.get_file(session_id, file_id)
+        if chat_file is None:
+            return None
+        abs_path = self.data_dir / chat_file.path
+        try:
+            return abs_path.read_bytes()
+        except Exception:  # noqa: BLE001
+            return None
+