diff options
Diffstat (limited to 'src/ollama_web/sessions.py')
| -rw-r--r-- | src/ollama_web/sessions.py | 248 |
1 files changed, 248 insertions, 0 deletions
diff --git a/src/ollama_web/sessions.py b/src/ollama_web/sessions.py new file mode 100644 index 0000000..bfe659c --- /dev/null +++ b/src/ollama_web/sessions.py @@ -0,0 +1,248 @@ +"""Persistent session storage for ollama-web.""" + +from __future__ import annotations + +import json +import shutil +import uuid +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from .config import settings + + +def _now() -> str: + return datetime.now(timezone.utc).isoformat() + + +@dataclass +class ChatFile: + id: str + name: str + path: str + mime: str + size: int + text: str + source: str # "upload" | "fetch" + + def to_dict(self) -> dict[str, Any]: + return { + "id": self.id, + "name": self.name, + "path": self.path, + "mime": self.mime, + "size": self.size, + "text": self.text, + "source": self.source, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "ChatFile": + return cls( + id=data["id"], + name=data["name"], + path=data["path"], + mime=data.get("mime", "application/octet-stream"), + size=data.get("size", 0), + text=data.get("text", ""), + source=data.get("source", "upload"), + ) + + +@dataclass +class ChatMessage: + role: str + content: str + attachments: list[str] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + return { + "role": self.role, + "content": self.content, + "attachments": list(self.attachments), + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "ChatMessage": + return cls( + role=data.get("role", "user"), + content=data.get("content", ""), + attachments=list(data.get("attachments", [])), + ) + + +class SessionStore: + """Manage chat sessions as JSON files under the configured data directory.""" + + def __init__(self, data_dir: str | Path | None = None) -> None: + self.data_dir = Path(data_dir or settings.data_dir).resolve() + self.sessions_dir = self.data_dir / "sessions" + self.sessions_dir.mkdir(parents=True, exist_ok=True) + + def _session_dir(self, session_id: str) -> Path: + return self.sessions_dir / session_id + + def _session_file(self, session_id: str) -> Path: + return self._session_dir(session_id) / "session.json" + + def _files_dir(self, session_id: str) -> Path: + files_dir = self._session_dir(session_id) / "files" + files_dir.mkdir(parents=True, exist_ok=True) + return files_dir + + def list_sessions(self) -> list[dict[str, Any]]: + out: list[dict[str, Any]] = [] + if not self.sessions_dir.exists(): + return out + for entry in sorted(self.sessions_dir.iterdir(), key=lambda p: p.stat().st_mtime, reverse=True): + if not entry.is_dir(): + continue + session_file = entry / "session.json" + if not session_file.exists(): + continue + try: + data = json.loads(session_file.read_text(encoding="utf-8")) + out.append( + { + "id": data["id"], + "title": data.get("title", "Untitled"), + "updated_at": data.get("updated_at", data.get("created_at", "")), + } + ) + except Exception: # noqa: BLE001 + continue + return out + + def create(self, title: str | None = None) -> dict[str, Any]: + session_id = uuid.uuid4().hex + now = _now() + session = { + "id": session_id, + "title": title or "新しいセッション", + "created_at": now, + "updated_at": now, + "messages": [], + "files": [], + } + self._session_dir(session_id).mkdir(parents=True, exist_ok=True) + self._save(session_id, session) + return session + + def get(self, session_id: str) -> dict[str, Any] | None: + session_file = self._session_file(session_id) + if not session_file.exists(): + return None + try: + return json.loads(session_file.read_text(encoding="utf-8")) + except Exception: # noqa: BLE001 + return None + + def save(self, session: dict[str, Any]) -> None: + session["updated_at"] = _now() + self._save(session["id"], session) + + def _save(self, session_id: str, session: dict[str, Any]) -> None: + session_file = self._session_file(session_id) + session_file.parent.mkdir(parents=True, exist_ok=True) + session_file.write_text(json.dumps(session, ensure_ascii=False, indent=2), encoding="utf-8") + + def delete(self, session_id: str) -> bool: + session_dir = self._session_dir(session_id) + if not session_dir.exists(): + return False + try: + shutil.rmtree(session_dir) + return True + except Exception: # noqa: BLE001 + return False + + def add_message( + self, + session_id: str, + role: str, + content: str, + attachments: list[str] | None = None, + ) -> dict[str, Any] | None: + session = self.get(session_id) + if session is None: + return None + session["messages"].append( + ChatMessage( + role=role, + content=content, + attachments=list(attachments or []), + ).to_dict() + ) + if role == "user" and len(session["messages"]) <= 2: + first_user = next( + (m for m in session["messages"] if m["role"] == "user"), + None, + ) + if first_user: + title = first_user["content"].strip().replace("\n", " ")[:40] + session["title"] = title + ("…" if len(first_user["content"]) > 40 else "") + self.save(session) + return session + + def add_file( + self, + session_id: str, + name: str, + data: bytes, + mime: str, + text: str, + source: str = "upload", + ) -> ChatFile: + self._session_dir(session_id).mkdir(parents=True, exist_ok=True) + file_id = uuid.uuid4().hex + ext = Path(name).suffix + stored_name = f"{file_id}{ext}" + file_path = self._files_dir(session_id) / stored_name + file_path.write_bytes(data) + rel_path = str((self._session_dir(session_id) / "files" / stored_name).relative_to(self.data_dir)) + chat_file = ChatFile( + id=file_id, + name=name, + path=rel_path, + mime=mime, + size=len(data), + text=text, + source=source, + ) + session = self.get(session_id) + if session is None: + session = self.create() + session_id = session["id"] + session["files"].append(chat_file.to_dict()) + self.save(session) + return chat_file + + def remove_file(self, session_id: str, file_id: str) -> bool: + session = self.get(session_id) + if session is None: + return False + files = session.get("files", []) + target = next((f for f in files if f["id"] == file_id), None) + if target is None: + return False + files.remove(target) + session["files"] = files + try: + abs_path = self.data_dir / target["path"] + if abs_path.exists(): + abs_path.unlink() + except Exception: # noqa: BLE001 + pass + self.save(session) + return True + + def get_file(self, session_id: str, file_id: str) -> ChatFile | None: + session = self.get(session_id) + if session is None: + return None + for f in session.get("files", []): + if f["id"] == file_id: + return ChatFile.from_dict(f) + return None
\ No newline at end of file |