diff options
Diffstat (limited to 'src/ollama_web/routes')
| -rw-r--r-- | src/ollama_web/routes/auth.py | 32 | ||||
| -rw-r--r-- | src/ollama_web/routes/chat.py | 57 | ||||
| -rw-r--r-- | src/ollama_web/routes/pages.py | 10 | ||||
| -rw-r--r-- | src/ollama_web/routes/sessions.py | 51 |
4 files changed, 128 insertions, 22 deletions
diff --git a/src/ollama_web/routes/auth.py b/src/ollama_web/routes/auth.py new file mode 100644 index 0000000..808c2ec --- /dev/null +++ b/src/ollama_web/routes/auth.py @@ -0,0 +1,32 @@ +"""Authentication routes.""" + +from __future__ import annotations + +from starlette.requests import Request +from starlette.responses import JSONResponse + +from ..auth import clear_auth_cookies, is_authenticated, set_auth_cookies + + +async def login(request: Request) -> JSONResponse: + try: + body = await request.json() + except Exception: # noqa: BLE001 + body = {} + pin = str(body.get("pin", "")) + if pin != request.app.state.settings.pin: + return JSONResponse({"authenticated": False, "error": "Invalid PIN"}, status_code=401) + response = JSONResponse({"authenticated": True}) + csrf = set_auth_cookies(response, request.app.state.settings.secret_key) + response.headers["X-CSRF-Token"] = csrf + return response + + +async def logout(request: Request) -> JSONResponse: + response = JSONResponse({"authenticated": False}) + clear_auth_cookies(response) + return response + + +async def status(request: Request) -> JSONResponse: + return JSONResponse({"authenticated": is_authenticated(request)}) diff --git a/src/ollama_web/routes/chat.py b/src/ollama_web/routes/chat.py index eb7a555..9e0eae6 100644 --- a/src/ollama_web/routes/chat.py +++ b/src/ollama_web/routes/chat.py @@ -3,6 +3,7 @@ from __future__ import annotations import functools +import html import inspect import json import logging @@ -14,9 +15,10 @@ from ollama._types import Image as OllamaImage from starlette.requests import Request from starlette.responses import StreamingResponse -from ..sessions import SessionStore +from ..config import settings +from ..sessions import SessionStore, is_valid_id from ..tools.fetch import pop_fetched_files -from ..tools.image import resize_image +from ..tools.helper.image import resize_image from ..tools.registry import ToolRegistry, default_registry logger = logging.getLogger("ollama_web.chat") @@ -31,12 +33,18 @@ _TOOL_SYSTEM_PROMPT = ( ) +def _limit_text(text: str, limit: int) -> str: + if len(text) <= limit: + return text + return text[:limit] + "\n…[truncated]" + + def _parse_messages(raw: list[dict[str, Any]]) -> list[Message]: """Convert raw JSON message dicts into ``ollama.Message`` objects.""" out: list[Message] = [] for m in raw: role = m.get("role", "user") - content = m.get("content", "") + content = _limit_text(str(m.get("content", "") or ""), 20000) name = m.get("name") images = m.get("images") kwargs: dict[str, Any] = {"role": role, "content": content} @@ -74,9 +82,9 @@ def _attachment_context(store: SessionStore, session_id: str, file_ids: list[str chat_file = store.get_file(session_id, file_id) if chat_file is None or chat_file.is_image: continue - text = chat_file.text or "" + text = _limit_text(chat_file.text or "", settings.max_attachment_text_chars) parts.append( - f"<attached-file name=\"{chat_file.name}\" id=\"{chat_file.id}\" " + f"<attached-file name=\"{html.escape(chat_file.name)}\" id=\"{chat_file.id}\" " f"source=\"{chat_file.source}\">\n{text}\n</attached-file>" ) if not parts: @@ -127,7 +135,9 @@ async def _chat_event_stream( think = payload.get("think") store: SessionStore = request.app.state.session_store - file_ids: list[str] = list(payload.get("file_ids", [])) + file_ids: list[str] = [ + str(file_id) for file_id in payload.get("file_ids", []) if is_valid_id(str(file_id)) + ] user_content = "" assistant_text = "" @@ -140,7 +150,11 @@ async def _chat_event_stream( raw_messages = payload.get("messages", []) if raw_messages: - user_content = str(raw_messages[-1].get("content", "") or "") + user_content = _limit_text( + str(raw_messages[-1].get("content", "") or ""), + request.app.state.settings.max_message_chars, + ) + raw_messages[-1]["content"] = user_content attachment_ctx = "" if session_id and file_ids: @@ -192,7 +206,10 @@ async def _chat_event_stream( messages, model, think=think, host=host, registry=registry ): if event.get("type") == "delta": - assistant_text += event.get("content", "") + assistant_text = _limit_text( + assistant_text + event.get("content", ""), + request.app.state.settings.max_message_chars, + ) yield event except Exception as exc: # noqa: BLE001 @@ -201,7 +218,10 @@ async def _chat_event_stream( 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 対応モデルを選択してください。" + assistant_text = ( + f"[ERROR] {error_msg}\n\n" + "画像を含むメッセージを送信する場合、vision 対応モデルを選択してください。" + ) finally: # Persist user and assistant messages to the session even when the @@ -241,6 +261,16 @@ async def chat(request: Request) -> StreamingResponse: """POST /api/chat: stream chat events as Server-Sent Events.""" try: body = await request.body() + if len(body) > request.app.state.settings.max_message_chars * 4: + data = json.dumps( + {"type": "error", "message": "Request body too large"}, + ensure_ascii=False, + ) + return StreamingResponse( + _single_event(data), + media_type="text/event-stream", + status_code=413, + ) payload = json.loads(body.decode("utf-8")) except Exception as exc: logger.error("failed to parse request json: %s", exc) @@ -252,7 +282,12 @@ async def chat(request: Request) -> StreamingResponse: session_id = payload.get("session_id") msg_count = len(payload.get("messages", [])) - logger.info("chat request model=%s session=%s messages=%d", payload.get("model"), session_id, 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, session_id), media_type="text/event-stream", @@ -261,4 +296,4 @@ async def chat(request: Request) -> StreamingResponse: async def _single_event(data: str) -> AsyncIterator[bytes]: """Yield a single SSE event and finish.""" - yield f"data: {data}\n\n".encode()
\ No newline at end of file + yield f"data: {data}\n\n".encode() diff --git a/src/ollama_web/routes/pages.py b/src/ollama_web/routes/pages.py index 8e1d084..52af9e8 100644 --- a/src/ollama_web/routes/pages.py +++ b/src/ollama_web/routes/pages.py @@ -25,3 +25,13 @@ async def index(request: Request) -> Response: {"models": models, "default_model": default_model}, media_type="text/html; charset=utf-8", ) + + +async def login(request: Request) -> Response: + """Render the PIN login page.""" + return templates.TemplateResponse( + request, + "login.html", + {}, + media_type="text/html; charset=utf-8", + ) diff --git a/src/ollama_web/routes/sessions.py b/src/ollama_web/routes/sessions.py index 8640ec3..7b01705 100644 --- a/src/ollama_web/routes/sessions.py +++ b/src/ollama_web/routes/sessions.py @@ -2,20 +2,24 @@ from __future__ import annotations +import json import logging -from typing import Any +from typing import Any, cast +from starlette.datastructures import UploadFile from starlette.requests import Request from starlette.responses import JSONResponse -from ..sessions import SessionStore, _IMAGE_MIMES, _IMAGE_SUFFIXES -from ..tools.pdf import extract_pdf_text +from ..sessions import _IMAGE_MIMES, _IMAGE_SUFFIXES, SessionStore, is_valid_id +from ..tools.helper.pdf import extract_pdf_text logger = logging.getLogger("ollama_web.sessions") +_UPLOAD_CHUNK_SIZE = 1024 * 1024 + def _get_store(request: Request) -> SessionStore: - return request.app.state.session_store + return cast(SessionStore, request.app.state.session_store) def _is_image_file(name: str, mime: str) -> bool: @@ -36,7 +40,8 @@ async def list_sessions(request: Request) -> JSONResponse: async def create_session(request: Request) -> JSONResponse: store = _get_store(request) - body = await request.json() if await request.body() else {} + raw = await request.body() + body = json.loads(raw.decode("utf-8")) if raw else {} title = body.get("title") session = store.create(title=title) return JSONResponse(session) @@ -45,6 +50,8 @@ async def create_session(request: Request) -> JSONResponse: async def get_session(request: Request) -> JSONResponse: store = _get_store(request) session_id = request.path_params["session_id"] + if not is_valid_id(session_id): + return JSONResponse({"error": "Invalid session id"}, status_code=400) session = store.get(session_id) if session is None: return JSONResponse({"error": "Session not found"}, status_code=404) @@ -54,6 +61,8 @@ async def get_session(request: Request) -> JSONResponse: async def delete_session(request: Request) -> JSONResponse: store = _get_store(request) session_id = request.path_params["session_id"] + if not is_valid_id(session_id): + return JSONResponse({"error": "Invalid session id"}, status_code=400) ok = store.delete(session_id) if not ok: return JSONResponse({"error": "Session not found"}, status_code=404) @@ -63,6 +72,8 @@ async def delete_session(request: Request) -> JSONResponse: async def upload_file(request: Request) -> JSONResponse: store = _get_store(request) session_id = request.path_params["session_id"] + if not is_valid_id(session_id): + return JSONResponse({"error": "Invalid session id"}, status_code=400) session = store.get(session_id) if session is None: return JSONResponse({"error": "Session not found"}, status_code=404) @@ -79,22 +90,38 @@ async def upload_file(request: Request) -> JSONResponse: for upload in files: if upload is None: continue + if not hasattr(upload, "read"): + results.append({"name": "unknown", "error": "Invalid upload"}) + continue + upload_file = cast(UploadFile, upload) try: - data = await upload.read() + chunks: list[bytes] = [] + total = 0 + while True: + chunk = await upload_file.read(_UPLOAD_CHUNK_SIZE) + if not chunk: + break + total += len(chunk) + if total > max_size: + raise ValueError( + f"File exceeds {request.app.state.settings.max_upload_mb} MB limit" + ) + chunks.append(chunk) + data = b"".join(chunks) except Exception as exc: # noqa: BLE001 - results.append({"name": getattr(upload, "filename", "unknown"), "error": str(exc)}) + results.append({"name": getattr(upload_file, "filename", "unknown"), "error": str(exc)}) continue if len(data) > max_size: results.append( { - "name": getattr(upload, "filename", "unknown"), + "name": getattr(upload_file, "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" + name = getattr(upload_file, "filename", "uploaded-file") or "uploaded-file" + content_type = getattr(upload_file, "content_type", None) or "application/octet-stream" text = "" if _is_image_file(name, content_type): # Images are passed to ollama via Message.images (base64), not as text. @@ -128,7 +155,9 @@ 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"] + if not is_valid_id(session_id) or not is_valid_id(file_id): + return JSONResponse({"error": "Invalid id"}, status_code=400) 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 + return JSONResponse({"deleted": True}) |