diff options
Diffstat (limited to 'src/ollama_web/routes')
| -rw-r--r-- | src/ollama_web/routes/chat.py | 124 | ||||
| -rw-r--r-- | src/ollama_web/routes/sessions.py | 44 |
2 files changed, 132 insertions, 36 deletions
diff --git a/src/ollama_web/routes/chat.py b/src/ollama_web/routes/chat.py index d94b567..748b090 100644 --- a/src/ollama_web/routes/chat.py +++ b/src/ollama_web/routes/chat.py @@ -24,7 +24,7 @@ from ..settings_store import load_app_settings from ..tools.fetch import pop_fetched_files from ..tools.helper.image import resize_image from ..tools.pdf import make_session_pdf_tool -from ..tools.registry import ToolRegistry, default_registry +from ..tools.registry import ToolArtifact, ToolRegistry, default_registry logger = logging.getLogger("ollama_web.chat") @@ -107,7 +107,25 @@ async def _build_registry( logger.info("collected %d MCP tool(s): %s", len(mcp_tools), tool_names) if mcp_tools: - executor = make_mcp_executor() + artifact_handler = None + if session_id is not None and store is not None: + + def persist_artifact(full_name: str, artifact: ToolArtifact) -> None: + if not artifact.for_user: + return + chat_file = store.add_file( + session_id=session_id, + name=artifact.name, + data=artifact.data, + mime=artifact.mime, + text="", + source=full_name, + ) + artifact.file_id = chat_file.id + + artifact_handler = persist_artifact + + executor = make_mcp_executor(artifact_handler=artifact_handler) for tool_def in mcp_tools: fn = tool_def.get("function", {}) name = fn.get("name") @@ -127,21 +145,23 @@ def _attachment_context( file_ids: list[str], *, total_limit: int, + language: str | None = None, ) -> str: if total_limit <= 0: return "" chat_files = [] for file_id in file_ids: chat_file = store.get_file(session_id, file_id) - if chat_file is not None and not chat_file.is_image: + if chat_file is not None and not chat_file.is_image and not chat_file.is_audio: chat_files.append(chat_file) if not chat_files: return "" parts: list[str] = [] - prefix = get_prompt("attachment_prefix") + prefix = get_prompt("attachment_prefix", lang=language) paging_hint = get_prompt( "attachment_paging", + lang=language, fallback=( "PDF previews may have unread content. Use pdf_to_text with the attachment URI " "and next_token until has_more is absent." @@ -150,9 +170,11 @@ def _attachment_context( used = len(prefix) + len(paging_hint) + 4 per_file = max(256, min(4000, (total_limit - used) // max(1, len(chat_files)))) pdf_tool = make_session_pdf_tool( - lambda file_id: store.get_file_data(session_id, file_id) - if store.get_file(session_id, file_id) is not None - else None + lambda file_id: ( + store.get_file_data(session_id, file_id) + if store.get_file(session_id, file_id) is not None + else None + ) ) for chat_file in chat_files: @@ -174,6 +196,20 @@ def _attachment_context( return "\n\n".join([prefix, paging_hint, *parts]) +def _has_image_attachments(store: SessionStore, session_id: str, file_ids: list[str]) -> bool: + return any( + (chat_file := store.get_file(session_id, file_id)) is not None and chat_file.is_image + for file_id in file_ids + ) + + +def _has_audio_attachments(store: SessionStore, session_id: str, file_ids: list[str]) -> bool: + return any( + (chat_file := store.get_file(session_id, file_id)) is not None and chat_file.is_audio + for file_id in file_ids + ) + + def _collect_images(store: SessionStore, session_id: str, file_ids: list[str]) -> list[bytes]: images: list[bytes] = [] for file_id in file_ids: @@ -251,6 +287,7 @@ async def _chat_event_stream( ] user_content = "" assistant_text = "" + assistant_file_ids: list[str] = [] last_done_event: dict[str, Any] | None = None try: @@ -260,7 +297,7 @@ async def _chat_event_stream( if session is not None: session_messages = _parse_messages(session.get("messages", [])) - raw_messages = payload.get("messages", []) + raw_messages = list(payload.get("messages", [])) if raw_messages: user_content = _limit_text( str(raw_messages[-1].get("content", "") or ""), @@ -268,6 +305,29 @@ async def _chat_event_stream( ) raw_messages[-1]["content"] = user_content + # An attachment-only submission still needs a user message to carry + # either image bytes or the localized unsupported-vision notice. + if file_ids and not raw_messages: + raw_messages.append({"role": "user", "content": ""}) + + host = request.app.state.settings.ollama_host + + from ollama_web import llm # local import to avoid circular at module load + + capabilities = llm.get_model_capabilities(model, host=host) + logger.info( + "chat stream model=%s capabilities=%s", + model, + sorted(capabilities), + ) + + # Fallback: some cloud/proxy models support tools even though ollama + # does not advertise the capability. Enable only tools; vision remains + # opt-in because sending image bytes to an unknown model is not safe. + if "tools" not in capabilities and ":cloud" in model: + capabilities.add("tools") + logger.info("enabled tools capability fallback for cloud model %s", model) + attachment_ctx = "" if session_id and file_ids: attachment_budget = max( @@ -279,6 +339,7 @@ async def _chat_event_stream( session_id, file_ids, total_limit=attachment_budget, + language=language, ) if attachment_ctx and user_content: @@ -289,11 +350,22 @@ async def _chat_event_stream( elif attachment_ctx: raw_messages.append({"role": "user", "content": attachment_ctx}) - # Attach image bytes to the latest user message per ollama spec. + # Attach image bytes only when Ollama explicitly advertises vision. + # Other models receive a localized fixed notice, never base64 text. if session_id and file_ids and raw_messages: - images = _collect_images(store, session_id, file_ids) - if images: - raw_messages[-1]["images"] = images + notices: list[str] = [] + if _has_image_attachments(store, session_id, file_ids): + if "vision" in capabilities: + images = _collect_images(store, session_id, file_ids) + if images: + raw_messages[-1]["images"] = images + else: + notices.append(get_prompt("vision_unsupported_attachment", lang=language)) + if _has_audio_attachments(store, session_id, file_ids): + notices.append(get_prompt("audio_unsupported_attachment", lang=language)) + if notices: + current = str(raw_messages[-1].get("content", "") or "") + raw_messages[-1]["content"] = f"{current}\n\n" + "\n\n".join(notices) messages = _parse_messages(raw_messages) if session_messages: @@ -305,10 +377,6 @@ async def _chat_event_stream( if not any(getattr(m, "role", None) == "system" for m in messages): messages.insert(0, Message(role="system", content=system_prompt)) - host = request.app.state.settings.ollama_host - - from ollama_web import llm # local import to avoid circular at module load - registry = await _build_registry(session_id, store) total_image_bytes = sum( @@ -323,20 +391,6 @@ async def _chat_event_stream( total_image_bytes, ) - capabilities = llm.get_model_capabilities(model, host=host) - logger.info( - "chat stream model=%s capabilities=%s", - model, - sorted(capabilities), - ) - - # Fallback: some cloud/proxy models support tools even though ollama - # does not advertise the capability. Enable tools when the model name - # hints at a cloud endpoint and tools are not already reported. - if "tools" not in capabilities and ":cloud" in model: - capabilities.add("tools") - logger.info("enabled tools capability fallback for cloud model %s", model) - async for event in llm.astream_chat_with_tools( messages, model, @@ -355,6 +409,11 @@ async def _chat_event_stream( if event.get("type") == "done": event["num_ctx"] = num_ctx last_done_event = event + if event.get("type") == "tool_end": + for media in event.get("media", []): + file_id = str(media.get("id", "")) + if is_valid_id(file_id) and file_id not in assistant_file_ids: + assistant_file_ids.append(file_id) yield event except Exception as exc: # noqa: BLE001 @@ -370,19 +429,20 @@ async def _chat_event_stream( 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: + if session_id and (user_content or file_ids): store.add_message( session_id=session_id, role="user", content=user_content, attachments=file_ids, ) - if assistant_text: + if assistant_text or assistant_file_ids: token_count = _token_count_from_done(last_done_event, num_ctx) store.add_message( session_id=session_id, role="assistant", content=assistant_text, + attachments=assistant_file_ids, token_count=token_count, ) diff --git a/src/ollama_web/routes/sessions.py b/src/ollama_web/routes/sessions.py index 73f79d6..5d1c63a 100644 --- a/src/ollama_web/routes/sessions.py +++ b/src/ollama_web/routes/sessions.py @@ -8,9 +8,15 @@ from typing import Any, cast from starlette.datastructures import UploadFile from starlette.requests import Request -from starlette.responses import JSONResponse - -from ..sessions import _IMAGE_MIMES, _IMAGE_SUFFIXES, SessionStore, is_valid_id +from starlette.responses import FileResponse, JSONResponse, Response + +from ..sessions import ( + _AUDIO_MIMES, + _IMAGE_MIMES, + _IMAGE_SUFFIXES, + SessionStore, + is_valid_id, +) from ..tools.helper.pdf import extract_pdf_page logger = logging.getLogger("ollama_web.sessions") @@ -124,7 +130,10 @@ async def upload_file(request: Request) -> JSONResponse: 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. + # Images are passed to capable models via Message.images, not as text. + text = "" + elif content_type.lower().startswith("audio/"): + # Audio is UI-only until Ollama audio input support is explicit. text = "" elif content_type == "application/pdf" or name.lower().endswith(".pdf"): # Store only a bounded preview. The chat-scoped PDF tool reads the @@ -173,6 +182,33 @@ async def delete_file(request: Request) -> JSONResponse: return JSONResponse({"deleted": True}) +async def get_file_content(request: Request) -> Response: + """Serve a stored image or audio file from the authenticated origin.""" + 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) + chat_file = store.get_file(session_id, file_id) + path = store.get_file_path(session_id, file_id) + if chat_file is None or path is None: + return JSONResponse({"error": "File not found"}, status_code=404) + + mime = chat_file.mime.lower() + # Trust the stored, allowlisted MIME type rather than the filename. This + # prevents an HTML upload named ``*.png`` from being served inline. + if mime not in _IMAGE_MIMES and mime not in _AUDIO_MIMES: + return JSONResponse({"error": "File is not previewable"}, status_code=415) + + return FileResponse( + path, + media_type=mime, + filename=chat_file.name, + content_disposition_type="inline", + headers={"Cache-Control": "private, no-store"}, + ) + + async def clear_messages(request: Request) -> JSONResponse: store = _get_store(request) session_id = request.path_params["session_id"] |