機能追加(セッション、左ペインでのセッション管理、ファイル取得ツール、pdfテキスト化ツール)

2 files changed, 240 insertions, 9 deletions

diff --git a/src/ollama_web/routes/chat.py b/src/ollama_web/routes/chat.py
index 2612279..8b6c302 100644
--- a/src/ollama_web/routes/chat.py
+++ b/src/ollama_web/routes/chat.py
@@ -2,7 +2,10 @@
 
 from __future__ import annotations
 
+import functools
+import inspect
 import json
+import logging
 from collections.abc import AsyncIterator
 from typing import Any
 
@@ -10,6 +13,12 @@ from ollama import Message
 from starlette.requests import Request
 from starlette.responses import StreamingResponse
 
+from ..sessions import SessionStore
+from ..tools.fetch import pop_fetched_files
+from ..tools.registry import ToolRegistry, default_registry
+
+logger = logging.getLogger("ollama_web.chat")
+
 
 def _parse_messages(raw: list[dict[str, Any]]) -> list[Message]:
     """Convert raw JSON message dicts into ``ollama.Message`` objects."""
@@ -28,27 +37,127 @@ def _parse_messages(raw: list[dict[str, Any]]) -> list[Message]:
     return out
 
 
-async def _event_stream(request: Request, payload: dict[str, Any]) -> AsyncIterator[bytes]:
+def _session_aware_registry(session_id: str | None) -> ToolRegistry:
+    """Return a registry whose search_and_fetch receives the session id."""
+    reg = default_registry()
+
+    original = reg.get("search_and_fetch")
+    if original is not None and session_id is not None:
+
+        @functools.wraps(original)
+        def wrapped(*args: Any, **kwargs: Any) -> Any:
+            sig = inspect.signature(original)
+            bound = sig.bind(*args, **kwargs)
+            bound.arguments["session_id"] = session_id
+            return original(*bound.args, **bound.kwargs)
+
+        reg._tools["search_and_fetch"] = wrapped  # noqa: SLF001
+
+    return reg
+
+
+def _attachment_context(store: SessionStore, session_id: str, file_ids: list[str]) -> str:
+    parts: list[str] = []
+    for file_id in file_ids:
+        chat_file = store.get_file(session_id, file_id)
+        if chat_file is None:
+            continue
+        text = chat_file.text or ""
+        parts.append(
+            f"<attached-file name=\"{chat_file.name}\" id=\"{chat_file.id}\" "
+            f"source=\"{chat_file.source}\">\n{text}\n</attached-file>"
+        )
+    if not parts:
+        return ""
+    return "\n\n".join(["以下はユーザーが添付したファイルの内容です：", *parts])
+
+
+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."""
     model = payload.get("model") or request.app.state.settings.default_model
-    raw_messages = payload.get("messages", [])
     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 "")
+
+    attachment_ctx = ""
+    if session_id and file_ids:
+        attachment_ctx = _attachment_context(store, session_id, file_ids)
+
+    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})
+
     messages = _parse_messages(raw_messages)
+    if session_messages:
+        messages = session_messages + 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):
+    registry = _session_aware_registry(session_id)
+
+    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()
 
+    # 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,
+        )
+        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",
+            )
+
 
 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"))
@@ -59,14 +168,16 @@ async def chat(request: Request) -> StreamingResponse:
             _single_event(data),
             media_type="text/event-stream",
         )
+
+    session_id = payload.get("session_id")
     msg_count = len(payload.get("messages", []))
-    logger.info("chat request model=%s messages=%d", payload.get("model"), msg_count)
+    logger.info("chat request model=%s session=%s messages=%d", payload.get("model"), session_id, msg_count)
     return StreamingResponse(
-        _event_stream(request, payload),
+        _event_stream(request, payload, session_id),
         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()
+    yield f"data: {data}\n\n".encode()
\ No newline at end of file
diff --git a/src/ollama_web/routes/sessions.py b/src/ollama_web/routes/sessions.py
new file mode 100644
index 0000000..8980960
--- /dev/null
+++ b/src/ollama_web/routes/sessions.py
@@ -0,0 +1,120 @@
+"""Session management API routes."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+
+from ..sessions import SessionStore
+from ..tools.pdf import extract_pdf_text
+
+logger = logging.getLogger("ollama_web.sessions")
+
+
+def _get_store(request: Request) -> SessionStore:
+    return request.app.state.session_store
+
+
+async def list_sessions(request: Request) -> JSONResponse:
+    store = _get_store(request)
+    return JSONResponse({"sessions": store.list_sessions()})
+
+
+async def create_session(request: Request) -> JSONResponse:
+    store = _get_store(request)
+    body = await request.json() if await request.body() else {}
+    title = body.get("title")
+    session = store.create(title=title)
+    return JSONResponse(session)
+
+
+async def get_session(request: Request) -> JSONResponse:
+    store = _get_store(request)
+    session_id = request.path_params["session_id"]
+    session = store.get(session_id)
+    if session is None:
+        return JSONResponse({"error": "Session not found"}, status_code=404)
+    return JSONResponse(session)
+
+
+async def delete_session(request: Request) -> JSONResponse:
+    store = _get_store(request)
+    session_id = request.path_params["session_id"]
+    ok = store.delete(session_id)
+    if not ok:
+        return JSONResponse({"error": "Session not found"}, status_code=404)
+    return JSONResponse({"deleted": True})
+
+
+async def upload_file(request: Request) -> JSONResponse:
+    store = _get_store(request)
+    session_id = request.path_params["session_id"]
+    session = store.get(session_id)
+    if session is None:
+        return JSONResponse({"error": "Session not found"}, status_code=404)
+
+    max_size = request.app.state.settings.max_upload_mb * 1024 * 1024
+    try:
+        form = await request.form()
+    except Exception as exc:  # noqa: BLE001
+        logger.error("failed to parse multipart form: %s", exc)
+        return JSONResponse({"error": f"Invalid form data: {exc}"}, status_code=400)
+
+    files = form.getlist("files") if hasattr(form, "getlist") else [form.get("files")]
+    results: list[dict[str, Any]] = []
+    for upload in files:
+        if upload is None:
+            continue
+        try:
+            data = await upload.read()
+        except Exception as exc:  # noqa: BLE001
+            results.append({"name": getattr(upload, "filename", "unknown"), "error": str(exc)})
+            continue
+        if len(data) > max_size:
+            results.append(
+                {
+                    "name": getattr(upload, "filename", "unknown"),
+                    "error": f"File exceeds {request.app.state.settings.max_upload_mb} MB limit",
+                }
+            )
+            continue
+
+        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"):
+            text = extract_pdf_text(data)
+        else:
+            # Try decode as text for text-like uploads.
+            try:
+                text = data.decode("utf-8")
+            except UnicodeDecodeError:
+                try:
+                    text = data.decode("cp932")
+                except UnicodeDecodeError:
+                    text = data.decode("utf-8", errors="replace")
+
+        chat_file = store.add_file(
+            session_id=session_id,
+            name=name,
+            data=data,
+            mime=content_type,
+            text=text,
+            source="upload",
+        )
+        results.append(chat_file.to_dict())
+
+    return JSONResponse({"files": results})
+
+
+async def delete_file(request: Request) -> JSONResponse:
+    store = _get_store(request)
+    session_id = request.path_params["session_id"]
+    file_id = request.path_params["file_id"]
+    ok = store.remove_file(session_id, file_id)
+    if not ok:
+        return JSONResponse({"error": "File not found"}, status_code=404)
+    return JSONResponse({"deleted": True})
\ No newline at end of file