From c21f052d5b0a33f890a86baf7f826cf9c33b9f40 Mon Sep 17 00:00:00 2001 From: Ryan Westfall Date: Tue, 20 Oct 2026 08:11:42 -0500 Subject: [PATCH] Read attached chat images with a vision model. Image uploads were treated as unknown documents and never shown to the model. PNG, JPEG, WebP, and GIF now go to Ollama vision, and history returns the image bytes. Co-authored-by: Cursor --- .env.example | 2 + .env.prod.example | 2 + llm_be/chat_backend/consumers.py | 58 +++--- llm_be/chat_backend/consumers_graph.py | 57 ++++-- llm_be/chat_backend/ollama_config.py | 11 +- llm_be/chat_backend/serializers.py | 10 +- llm_be/chat_backend/services/vision.py | 186 ++++++++++++++++++++ llm_be/chat_backend/services/ws_frames.py | 1 + llm_be/chat_backend/tests/test_consumers.py | 36 ++++ llm_be/chat_backend/tests/test_vision.py | 64 +++++++ llm_be/llm_be/settings.py | 4 + 11 files changed, 387 insertions(+), 44 deletions(-) create mode 100644 llm_be/chat_backend/services/vision.py create mode 100644 llm_be/chat_backend/tests/test_vision.py diff --git a/.env.example b/.env.example index af0f1a8..18f4744 100644 --- a/.env.example +++ b/.env.example @@ -16,6 +16,8 @@ OLLAMA_BASE_URL=http://127.0.0.1:11434 # OLLAMA_MODEL=gpt-oss:20b # OLLAMA_MODEL_THINKING=gpt-oss:20b # OLLAMA_MODEL_FAST=gemma4:latest +# Vision (image + prompt). Must be a vision-capable tag. Default gemma4:latest. +# OLLAMA_MODEL_VISION=gemma4:latest # OLLAMA_MODEL_UTILITY=llama3.2 # OLLAMA_EMBED_MODEL=nomic-embed-text # OLLAMA_NUM_CTX_THINKING=16384 diff --git a/.env.prod.example b/.env.prod.example index 85100c9..0c3d2a4 100644 --- a/.env.prod.example +++ b/.env.prod.example @@ -37,6 +37,8 @@ OLLAMA_BASE_URL=http://10.0.0.128:11434 OLLAMA_MODEL=gpt-oss:20b OLLAMA_MODEL_THINKING=gpt-oss:20b OLLAMA_MODEL_FAST=gemma4:latest +# Vision (image + prompt). Must be a vision-capable tag. +OLLAMA_MODEL_VISION=gemma4:latest OLLAMA_MODEL_UTILITY=llama3.2 OLLAMA_EMBED_MODEL=nomic-embed-text OLLAMA_NUM_CTX_THINKING=16384 diff --git a/llm_be/chat_backend/consumers.py b/llm_be/chat_backend/consumers.py index be63e77..20f3328 100644 --- a/llm_be/chat_backend/consumers.py +++ b/llm_be/chat_backend/consumers.py @@ -1,5 +1,4 @@ import json -import base64 import logging import pandas as pd from datetime import datetime @@ -14,6 +13,7 @@ from langchain_core.messages import HumanMessage, AIMessage from langchain_community.vectorstores import Chroma from langchain_ollama import OllamaEmbeddings from chat_backend.ollama_config import ( + ROLE_VISION, ollama_embeddings_kwargs, ollama_model_for_role, resolve_chat_role, @@ -38,6 +38,7 @@ from .services.title_generator import title_generator from .services.moderation_classifier import moderation_classifier, ModerationLabel from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType from .services.data_analysis_service import AsyncDataAnalysisService +from .services.vision import AsyncVisionService, is_image_type, prepare_attachment from .services.grounded_chat import prepare_grounded_chat from .services.status_context import ( emit_status, @@ -278,7 +279,9 @@ async def get_conversation_file_async(conversation_id): # Get the very first prompt in the conversation that has a file prompt_with_file = await Prompt.objects.filter( conversation_id=conversation_id - ).exclude(file='').order_by('created').afirst() + ).exclude(file='').exclude( + file_type__in=["png", "jpeg", "jpg", "webp", "gif"] + ).order_by('created').afirst() if prompt_with_file and prompt_with_file.file: # Opening a DatabaseStorage file hits the DB, so read inside the thread. @@ -414,29 +417,19 @@ class ChatConsumerAgain(AsyncWebsocketConsumer): decoded_file = None if file: - decoded_file = base64.b64decode(file) - logger.debug(decoded_file) - # The `altered_message` should only be created if a file exists - # and you want to pass its content directly to the classifier. - # Here, we'll let the classifier decide based on the user's prompt - # and then handle the file content separately. - altered_message = message - if "csv" in file_type: - file_type = "csv" - #altered_message = f"{message}\n The file type is csv and the file contents are: {decoded_file}" - elif "xmlformats-officedocument" in file_type: - file_type = "xlsx" - #df = pd.read_excel(decoded_file) - #altered_message = f"{message}\n The file type is xlsx and the file contents are: {df}" - elif "word" in file_type: - file_type = "docx" - elif "pdf" in file_type: - file_type = "pdf" - elif "text" in file_type: - file_type = "txt" - #altered_message = f"{message}\n The file type is txt and the file contents are: {decoded_file}" - else: - file_type = "Not Sure" + decoded_file, file_type, attachment_error = prepare_attachment( + file, file_type + ) + if attachment_error: + await self.send_json_message( + json.dumps( + { + "type": "error", + "content": attachment_error, + } + ) + ) + return logger.info(f'received: "{message}" for conversation {conversation_id}') @@ -475,7 +468,16 @@ class ChatConsumerAgain(AsyncWebsocketConsumer): conversation_id = input_dict["conversation_id"] decoded_file = input_dict.get("decoded_file") file_type = input_dict.get("file_type") - + + if is_image_type(file_type) and decoded_file: + service = AsyncVisionService() + input_dict["_resolved_model"] = service.model_name + input_dict["_citations"] = [] + await emit_status("reading_image") + return service.generate_response( + prompt_instance.message, decoded_file, file_type + ) + # Feature Flag + plan gate: Image Generation if prompt_type == PromptType.IMAGE_GENERATION: if not getattr(settings, "ALLOW_IMAGE_GENERATION", False): @@ -603,7 +605,9 @@ class ChatConsumerAgain(AsyncWebsocketConsumer): # messages = messages[:-1] + [HumanMessage(content=altered_message)] # I'll add it to the input_dict if needed. - resolved_model = ollama_model_for_role(resolve_chat_role(model)) + resolved_model = ollama_model_for_role( + ROLE_VISION if is_image_type(file_type) else resolve_chat_role(model) + ) prompt_metric = await create_prompt_metric( prompt_instance.id, prompt_instance.message, diff --git a/llm_be/chat_backend/consumers_graph.py b/llm_be/chat_backend/consumers_graph.py index 7b37333..90e4fec 100644 --- a/llm_be/chat_backend/consumers_graph.py +++ b/llm_be/chat_backend/consumers_graph.py @@ -1,5 +1,4 @@ import json -import base64 import logging import pandas as pd from datetime import datetime @@ -29,6 +28,7 @@ from .services.title_generator import title_generator from .services.moderation_classifier import moderation_classifier, ModerationLabel from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType from .services.data_analysis_service import AsyncDataAnalysisService +from .services.vision import AsyncVisionService, is_image_type, prepare_attachment from .services.grounded_chat import prepare_grounded_chat from .services.status_context import ( emit_status, @@ -36,7 +36,11 @@ from .services.status_context import ( set_status_emitter, ) from .services.ws_frames import citations_frame, status_frame -from chat_backend.ollama_config import ollama_model_for_role, resolve_chat_role +from chat_backend.ollama_config import ( + ROLE_VISION, + ollama_model_for_role, + resolve_chat_role, +) from .utils import ( TokenUsageCollector, aiter_text_chunks, @@ -202,7 +206,9 @@ async def get_conversation_file_async(conversation_id): try: prompt_with_file = await Prompt.objects.filter( conversation_id=conversation_id - ).exclude(file='').order_by('created').afirst() + ).exclude(file='').exclude( + file_type__in=["png", "jpeg", "jpg", "webp", "gif"] + ).order_by('created').afirst() if prompt_with_file and prompt_with_file.file: # Opening a DatabaseStorage file hits the DB, so read inside the thread. @@ -249,7 +255,11 @@ async def classification_node(state: ChatState) -> ChatState: decoded_file = state.get("decoded_file") prompt_type = await PROMPT_CLASSIFIER.classify_async(msg) - + + # Images are read by the vision model, not the spreadsheet parser. + if is_image_type(state.get("file_type")): + return {"prompt_type": PromptType.GENERAL_CHAT} + # Override logic if decoded_file and (prompt_type == PromptType.DATA_ANALYSIS or 'analyze' in msg.lower() or 'data' in msg.lower()): prompt_type = PromptType.DATA_ANALYSIS @@ -269,7 +279,18 @@ async def generation_node(state: ChatState) -> ChatState: conversation_id = state["conversation_id"] decoded_file = state.get("decoded_file") file_type = state.get("file_type") - + + if is_image_type(file_type) and decoded_file: + service = AsyncVisionService() + await emit_status("reading_image") + generator = service.generate_response( + prompt_instance.message, decoded_file, file_type + ) + return { + "response_generator": generator, + "resolved_model": service.model_name, + } + # Feature Flag + plan gate: Image Generation if prompt_type == PromptType.IMAGE_GENERATION: if not getattr(settings, "ALLOW_IMAGE_GENERATION", False): @@ -485,13 +506,19 @@ class ChatConsumerGraph(AsyncWebsocketConsumer): print("Conversation ID: ", conversation_id) decoded_file = None if file: - decoded_file = base64.b64decode(file) - if "csv" in file_type: file_type = "csv" - elif "xmlformats-officedocument" in file_type: file_type = "xlsx" - elif "word" in file_type: file_type = "docx" - elif "pdf" in file_type: file_type = "pdf" - elif "text" in file_type: file_type = "txt" - else: file_type = "Not Sure" + decoded_file, file_type, attachment_error = prepare_attachment( + file, file_type + ) + if attachment_error: + await self.send_json_message( + json.dumps( + { + "type": "error", + "content": attachment_error, + } + ) + ) + return # Pre-fetch messages and file messages, prompt_instance = await get_messages( @@ -501,7 +528,9 @@ class ChatConsumerGraph(AsyncWebsocketConsumer): if not decoded_file: decoded_file, file_type = await get_conversation_file_async(conversation_id) - resolved_model = ollama_model_for_role(resolve_chat_role(model)) + resolved_model = ollama_model_for_role( + ROLE_VISION if is_image_type(file_type) else resolve_chat_role(model) + ) prompt_metric = await create_prompt_metric( prompt_instance.id, prompt_instance.message, @@ -556,7 +585,7 @@ class ChatConsumerGraph(AsyncWebsocketConsumer): citations = [] final_model = resolved_model - if should_use_agent(message): + if should_use_agent(message) and not is_image_type(file_type): try: await enforce_feature_gate(chat_user, "agentic_tasks") except FeatureNotAllowed as exc: diff --git a/llm_be/chat_backend/ollama_config.py b/llm_be/chat_backend/ollama_config.py index c3486b4..43706e6 100644 --- a/llm_be/chat_backend/ollama_config.py +++ b/llm_be/chat_backend/ollama_config.py @@ -17,6 +17,8 @@ ROLE_EMBED = "embed" # independent plan steps concurrently on a smaller/cheaper model. ROLE_ORCHESTRATOR = "orchestrator" ROLE_SUBAGENT = "subagent" +# Image + prompt turns (#128). Never falls back to the text chat model. +ROLE_VISION = "vision" _VALID_ROLES = { ROLE_THINKING, @@ -25,6 +27,7 @@ _VALID_ROLES = { ROLE_EMBED, ROLE_ORCHESTRATOR, ROLE_SUBAGENT, + ROLE_VISION, } @@ -56,6 +59,7 @@ def ollama_model_for_role(role: str) -> str: ROLE_EMBED: "OLLAMA_EMBED_MODEL", ROLE_ORCHESTRATOR: "OLLAMA_MODEL_ORCHESTRATOR", ROLE_SUBAGENT: "OLLAMA_MODEL_SUBAGENT", + ROLE_VISION: "OLLAMA_MODEL_VISION", }[role] role_default = { ROLE_THINKING: "gpt-oss:20b", @@ -64,14 +68,15 @@ def ollama_model_for_role(role: str) -> str: ROLE_EMBED: "nomic-embed-text", ROLE_ORCHESTRATOR: "gpt-oss:20b", ROLE_SUBAGENT: "llama3.2", + ROLE_VISION: "gemma4:latest", }[role] configured = getattr(settings, role_setting, None) if configured: return configured - # Embeddings must never silently fall back to a chat model (#62). - if role == ROLE_EMBED: + # Embeddings and vision must never silently fall back to a text chat model. + if role in (ROLE_EMBED, ROLE_VISION): return role_default legacy = getattr(settings, "OLLAMA_MODEL", None) @@ -94,6 +99,8 @@ def ollama_num_ctx_for_role(role: str) -> int: return int(getattr(settings, "OLLAMA_NUM_CTX_SUBAGENT", 8192) or 8192) if role == ROLE_ORCHESTRATOR: return int(getattr(settings, "OLLAMA_NUM_CTX_ORCHESTRATOR", 16384) or 16384) + if role == ROLE_VISION: + return int(getattr(settings, "OLLAMA_NUM_CTX_VISION", 8192) or 8192) return int(getattr(settings, "OLLAMA_NUM_CTX_THINKING", 16384) or 16384) diff --git a/llm_be/chat_backend/serializers.py b/llm_be/chat_backend/serializers.py index 44b08fb..496b423 100644 --- a/llm_be/chat_backend/serializers.py +++ b/llm_be/chat_backend/serializers.py @@ -257,6 +257,7 @@ class PromptSerializer(serializers.ModelSerializer): tokens_in = serializers.SerializerMethodField() tokens_out = serializers.SerializerMethodField() feedback = serializers.SerializerMethodField() + image = serializers.SerializerMethodField() class Meta: model = Prompt @@ -269,8 +270,10 @@ class PromptSerializer(serializers.ModelSerializer): "tokens_out", "citations", "feedback", + "file_type", + "image", ) - read_only_fields = ("citations", "feedback") + read_only_fields = ("citations", "feedback", "file_type", "image") def _token_pair(self, obj): cache = self.context.setdefault("_prompt_token_cache", {}) @@ -314,6 +317,11 @@ class PromptSerializer(serializers.ModelSerializer): self.context["_prompt_feedback_by_id"] = by_prompt return by_prompt.get(obj.id) + def get_image(self, obj): + from chat_backend.services.vision import image_base64 + + return image_base64(obj) + def validate_message(self, value: str) -> str: if value is None or not str(value).strip(): raise serializers.ValidationError("Message text cannot be empty.") diff --git a/llm_be/chat_backend/services/vision.py b/llm_be/chat_backend/services/vision.py new file mode 100644 index 0000000..fc88b0a --- /dev/null +++ b/llm_be/chat_backend/services/vision.py @@ -0,0 +1,186 @@ +"""Image attachments on a chat turn (chat_web_app#128). + +Document uploads stay on the csv/xlsx/pdf path. An image is stored on the +prompt and sent to a vision-capable Ollama model with the user's text. +""" + +from __future__ import annotations + +import base64 +import json +import logging + +import httpx + +from chat_backend.ollama_config import ( + ROLE_VISION, + ollama_base_url, + ollama_model_for_role, + ollama_num_ctx_for_role, +) + +logger = logging.getLogger(__name__) + +MAX_IMAGE_BYTES = 8 * 1024 * 1024 + +IMAGE_MIME = { + "png": "image/png", + "jpeg": "image/jpeg", + "jpg": "image/jpeg", + "webp": "image/webp", + "gif": "image/gif", +} + +_IMAGE_EXTENSIONS = frozenset(IMAGE_MIME) + + +class VisionChunk: + """Stream piece that ``chunk_text`` and ``extract_token_usage`` both read.""" + + def __init__(self, text: str = "", info: dict | None = None): + self.content = text + self.generation_info = info or None + + +def canonical_file_type(file_type: str) -> str: + """Map a browser MIME (or short extension) to the stored file_type. + + Order matches the historical websocket normalizer so spreadsheets that + share the openxml MIME family still land on ``xlsx``. + """ + lowered = (file_type or "").lower() + if "csv" in lowered: + return "csv" + if "xmlformats-officedocument" in lowered: + return "xlsx" + if "word" in lowered: + return "docx" + if "pdf" in lowered: + return "pdf" + if "text" in lowered: + return "txt" + if "png" in lowered: + return "png" + if "jpeg" in lowered or "jpg" in lowered: + return "jpeg" + if "webp" in lowered: + return "webp" + if "gif" in lowered: + return "gif" + if lowered.startswith("image/"): + return "unsupported_image" + return "Not Sure" + + +def is_image_type(file_type: str | None) -> bool: + return (file_type or "").lower() in _IMAGE_EXTENSIONS + + +def image_mime(file_type: str | None) -> str: + return IMAGE_MIME.get((file_type or "").lower(), "application/octet-stream") + + +def prepare_attachment(file_b64: str, file_type: str) -> tuple[bytes, str, str | None]: + """Decode an upload. ``error`` is a user-facing string when the image is rejected.""" + try: + decoded = base64.b64decode(file_b64, validate=False) + except Exception: + return b"", "Not Sure", "Could not read the attached file." + + canonical = canonical_file_type(file_type) + if canonical == "unsupported_image": + return ( + decoded, + canonical, + "Unsupported image. Use PNG, JPEG, WebP, or GIF.", + ) + if is_image_type(canonical) and len(decoded) > MAX_IMAGE_BYTES: + return ( + decoded, + canonical, + "Image is too large. Maximum size is 8 MB.", + ) + return decoded, canonical, None + + +def image_base64(prompt) -> str | None: + """Base64 of a stored prompt image, or None for non-images.""" + if not is_image_type(getattr(prompt, "file_type", None)): + return None + if not prompt.file_exists(): + return None + try: + data = prompt.file.read() + except Exception: + logger.exception( + "failed to read prompt image prompt_id=%s", getattr(prompt, "id", None) + ) + return None + if not data: + return None + if isinstance(data, str): + data = data.encode() + return base64.b64encode(data).decode("ascii") + + +class AsyncVisionService: + """Stream an answer from Ollama's chat API with one image.""" + + def __init__(self): + self.model_name = ollama_model_for_role(ROLE_VISION) + self.num_ctx = ollama_num_ctx_for_role(ROLE_VISION) + + async def generate_response(self, query: str, image_bytes: bytes, file_type: str): + payload = { + "model": self.model_name, + "stream": True, + "options": {"num_ctx": self.num_ctx}, + "messages": [ + { + "role": "user", + "content": ( + "The user attached an image. Answer their question using what " + "is visible in the image, including any text in it.\n\n" + f"{query}" + ), + "images": [base64.b64encode(image_bytes).decode("ascii")], + } + ], + } + url = f"{ollama_base_url().rstrip('/')}/api/chat" + try: + async with httpx.AsyncClient(timeout=httpx.Timeout(180.0)) as client: + async with client.stream("POST", url, json=payload) as response: + if response.status_code >= 400: + body = await response.aread() + logger.error( + "vision model rejected request status=%s body=%s", + response.status_code, + body[:500], + ) + yield VisionChunk( + text="The vision model could not read that image." + ) + return + async for line in response.aiter_lines(): + if not line: + continue + try: + event = json.loads(line) + except json.JSONDecodeError: + continue + message = event.get("message") or {} + text = message.get("content") or "" + if text: + yield VisionChunk(text=text) + if event.get("done"): + info = { + key: event[key] + for key in ("prompt_eval_count", "eval_count") + if key in event + } + if info: + yield VisionChunk(info=info) + except httpx.HTTPError: + logger.exception("vision request failed model=%s", self.model_name) + yield VisionChunk(text="The vision model could not read that image.") diff --git a/llm_be/chat_backend/services/ws_frames.py b/llm_be/chat_backend/services/ws_frames.py index a47777b..9fdec50 100644 --- a/llm_be/chat_backend/services/ws_frames.py +++ b/llm_be/chat_backend/services/ws_frames.py @@ -23,6 +23,7 @@ STATUS_LABELS: dict[str, str] = { "reading_sources": "Reading sources", "retrieving_docs": "Searching your documents", "analysing": "Analysing your file", + "reading_image": "Reading your image", "refining": "Refining the answer", "writing": "", # tokens streaming — no label } diff --git a/llm_be/chat_backend/tests/test_consumers.py b/llm_be/chat_backend/tests/test_consumers.py index 4276baa..b5cf827 100644 --- a/llm_be/chat_backend/tests/test_consumers.py +++ b/llm_be/chat_backend/tests/test_consumers.py @@ -355,6 +355,42 @@ class GraphNodeTestCase(TransactionTestCase): self.prompt.message, b"a,b\n1,2\n", "csv" ) + async def test_generation_node_sends_images_to_vision(self): + with mock.patch.object(consumers_graph, "AsyncVisionService") as service: + service.return_value.model_name = "gemma4:latest" + service.return_value.generate_response.return_value = "generator" + + result = await consumers_graph.generation_node( + self._state( + prompt_type=PromptType.DATA_ANALYSIS, + message="what does this say?", + decoded_file=b"\x89PNG", + file_type="png", + ) + ) + + self.assertEqual(result["response_generator"], "generator") + self.assertEqual(result["resolved_model"], "gemma4:latest") + service.return_value.generate_response.assert_called_once_with( + self.prompt.message, b"\x89PNG", "png" + ) + + async def test_attached_image_is_not_data_analysis(self): + with mock.patch.object( + consumers_graph.PROMPT_CLASSIFIER, + "classify_async", + mock.AsyncMock(return_value=PromptType.DATA_ANALYSIS), + ): + result = await consumers_graph.classification_node( + self._state( + message="analyze this chart", + decoded_file=b"\x89PNG", + file_type="png", + ) + ) + + self.assertEqual(result["prompt_type"], PromptType.GENERAL_CHAT) + async def test_generation_node_routes_rag_prompts_with_workspace(self): with mock.patch.object(consumers_graph, "AsyncRAGService") as service: service.return_value.generate_response.return_value = "generator" diff --git a/llm_be/chat_backend/tests/test_vision.py b/llm_be/chat_backend/tests/test_vision.py new file mode 100644 index 0000000..ec64b5b --- /dev/null +++ b/llm_be/chat_backend/tests/test_vision.py @@ -0,0 +1,64 @@ +import base64 + +from django.core.files.uploadedfile import SimpleUploadedFile +from django.test import TestCase + +from chat_backend.services.vision import ( + canonical_file_type, + image_base64, + is_image_type, + prepare_attachment, +) +from chat_backend.tests.factories import make_conversation, make_prompt + + +class VisionAttachmentTestCase(TestCase): + def test_canonical_types(self): + self.assertEqual(canonical_file_type("image/png"), "png") + self.assertEqual(canonical_file_type("image/jpeg"), "jpeg") + self.assertEqual(canonical_file_type("image/jpg"), "jpeg") + self.assertEqual(canonical_file_type("image/webp"), "webp") + self.assertEqual(canonical_file_type("image/gif"), "gif") + self.assertEqual(canonical_file_type("text/plain"), "txt") + self.assertEqual(canonical_file_type("application/pdf"), "pdf") + self.assertEqual(canonical_file_type("image/bmp"), "unsupported_image") + self.assertEqual(canonical_file_type("image/svg+xml"), "unsupported_image") + + def test_prepare_attachment_rejects_oversize_and_unknown_images(self): + png = base64.b64encode(b"\x89PNG").decode() + decoded, kind, error = prepare_attachment(png, "image/png") + self.assertEqual(decoded, b"\x89PNG") + self.assertEqual(kind, "png") + self.assertIsNone(error) + self.assertTrue(is_image_type(kind)) + + huge = base64.b64encode(b"x" * (8 * 1024 * 1024 + 1)).decode() + _, _, error = prepare_attachment(huge, "image/png") + self.assertIn("too large", error) + + _, _, error = prepare_attachment(png, "image/bmp") + self.assertIn("Unsupported image", error) + + def test_document_types_unchanged(self): + raw = base64.b64encode(b"a,b\n1,2\n").decode() + decoded, kind, error = prepare_attachment(raw, "text/csv") + self.assertEqual(kind, "csv") + self.assertIsNone(error) + self.assertEqual(decoded, b"a,b\n1,2\n") + self.assertFalse(is_image_type(kind)) + + def test_image_base64_reads_stored_png(self): + prompt = make_prompt(make_conversation()) + prompt.file_type = "png" + prompt.file.save("shot.png", SimpleUploadedFile("shot.png", b"\x89PNG")) + prompt.save() + + self.assertEqual(image_base64(prompt), base64.b64encode(b"\x89PNG").decode()) + + def test_image_base64_skips_documents(self): + prompt = make_prompt(make_conversation()) + prompt.file_type = "csv" + prompt.file.save("data.csv", SimpleUploadedFile("data.csv", b"a,b")) + prompt.save() + + self.assertIsNone(image_base64(prompt)) diff --git a/llm_be/llm_be/settings.py b/llm_be/llm_be/settings.py index 33ccab7..4390015 100644 --- a/llm_be/llm_be/settings.py +++ b/llm_be/llm_be/settings.py @@ -169,6 +169,10 @@ OLLAMA_NUM_CTX_UTILITY = int(env("OLLAMA_NUM_CTX_UTILITY", "4096") or "4096") # sub-agents run independent plan steps on a smaller/cheaper model. OLLAMA_MODEL_ORCHESTRATOR = env("OLLAMA_MODEL_ORCHESTRATOR", "") or OLLAMA_MODEL OLLAMA_MODEL_SUBAGENT = env("OLLAMA_MODEL_SUBAGENT", "") or OLLAMA_MODEL_UTILITY +# Vision turns (image + prompt). Must be a vision-capable Ollama tag. +# Empty falls back to gemma4:latest inside ollama_model_for_role, not OLLAMA_MODEL. +OLLAMA_MODEL_VISION = env("OLLAMA_MODEL_VISION", "") or "" +OLLAMA_NUM_CTX_VISION = int(env("OLLAMA_NUM_CTX_VISION", "8192") or "8192") OLLAMA_NUM_CTX_ORCHESTRATOR = int( env("OLLAMA_NUM_CTX_ORCHESTRATOR", "16384") or "16384" )