Read attached chat images with a vision model (#82)
Deploy Beta / unit-tests (push) Successful in 10s
Unit Tests / test (push) Successful in 9s
Deploy Beta / docker (push) Successful in 23s
Deploy Beta / deploy-beta (push) Successful in 1m22s

## Summary
- Image attachments (PNG, JPEG, WebP, GIF) skip the document parser and go to a vision-capable Ollama model (`OLLAMA_MODEL_VISION`, default `gemma4:latest`).
- The image is stored on the prompt. `PromptSerializer` returns `file_type` and base64 `image` so history can redraw it.
- Unsupported image types and files over 8 MB return a websocket error and are not saved.
- A later text turn does not reuse a previous image as a spreadsheet.

Pairs with `chat_web_app` `feat/128-vision-image-prompt` ([#128](ai_ml_operations/chat_web_app#128)).

## Test plan
- [ ] `manage.py test chat_backend.tests.test_vision chat_backend.tests.test_consumers.GraphNodeTestCase`
- [ ] Send a prompt plus PNG over `ws/chat_again/` and confirm the vision model is called.
- [ ] `GET conversation_details` includes `image` and `file_type` for that turn.
- [ ] CSV upload still follows the document path.

Reviewed-on: #82
This commit was merged in pull request #82.
This commit is contained in:
westfarn committed 2026-10-20 06:13:31 -07:00
1 parent 32f57c2376
commit fc06202a1f
11 files changed
+384 -41

No files matched your search

+2
View File
@@ -16,6 +16,8 @@ OLLAMA_BASE_URL=http://127.0.0.1:11434
# OLLAMA_MODEL=gpt-oss:20b # OLLAMA_MODEL=gpt-oss:20b
# OLLAMA_MODEL_THINKING=gpt-oss:20b # OLLAMA_MODEL_THINKING=gpt-oss:20b
# OLLAMA_MODEL_FAST=gemma4:latest # 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_MODEL_UTILITY=llama3.2
# OLLAMA_EMBED_MODEL=nomic-embed-text # OLLAMA_EMBED_MODEL=nomic-embed-text
# OLLAMA_NUM_CTX_THINKING=16384 # OLLAMA_NUM_CTX_THINKING=16384
+2
View File
@@ -37,6 +37,8 @@ OLLAMA_BASE_URL=http://10.0.0.128:11434
OLLAMA_MODEL=gpt-oss:20b OLLAMA_MODEL=gpt-oss:20b
OLLAMA_MODEL_THINKING=gpt-oss:20b OLLAMA_MODEL_THINKING=gpt-oss:20b
OLLAMA_MODEL_FAST=gemma4:latest 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_MODEL_UTILITY=llama3.2
OLLAMA_EMBED_MODEL=nomic-embed-text OLLAMA_EMBED_MODEL=nomic-embed-text
OLLAMA_NUM_CTX_THINKING=16384 OLLAMA_NUM_CTX_THINKING=16384
+30 -26
View File
@@ -1,5 +1,4 @@
import json import json
import base64
import logging import logging
import pandas as pd import pandas as pd
from datetime import datetime from datetime import datetime
@@ -14,6 +13,7 @@ from langchain_core.messages import HumanMessage, AIMessage
from langchain_community.vectorstores import Chroma from langchain_community.vectorstores import Chroma
from langchain_ollama import OllamaEmbeddings from langchain_ollama import OllamaEmbeddings
from chat_backend.ollama_config import ( from chat_backend.ollama_config import (
ROLE_VISION,
ollama_embeddings_kwargs, ollama_embeddings_kwargs,
ollama_model_for_role, ollama_model_for_role,
resolve_chat_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.moderation_classifier import moderation_classifier, ModerationLabel
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
from .services.data_analysis_service import AsyncDataAnalysisService 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.grounded_chat import prepare_grounded_chat
from .services.status_context import ( from .services.status_context import (
emit_status, 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 # Get the very first prompt in the conversation that has a file
prompt_with_file = await Prompt.objects.filter( prompt_with_file = await Prompt.objects.filter(
conversation_id=conversation_id 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: if prompt_with_file and prompt_with_file.file:
# Opening a DatabaseStorage file hits the DB, so read inside the thread. # Opening a DatabaseStorage file hits the DB, so read inside the thread.
@@ -414,29 +417,19 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
decoded_file = None decoded_file = None
if file: if file:
decoded_file = base64.b64decode(file) decoded_file, file_type, attachment_error = prepare_attachment(
logger.debug(decoded_file) file, file_type
# The `altered_message` should only be created if a file exists )
# and you want to pass its content directly to the classifier. if attachment_error:
# Here, we'll let the classifier decide based on the user's prompt await self.send_json_message(
# and then handle the file content separately. json.dumps(
altered_message = message {
if "csv" in file_type: "type": "error",
file_type = "csv" "content": attachment_error,
#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) return
#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"
logger.info(f'received: "{message}" for conversation {conversation_id}') logger.info(f'received: "{message}" for conversation {conversation_id}')
@@ -476,6 +469,15 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
decoded_file = input_dict.get("decoded_file") decoded_file = input_dict.get("decoded_file")
file_type = input_dict.get("file_type") 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 # Feature Flag + plan gate: Image Generation
if prompt_type == PromptType.IMAGE_GENERATION: if prompt_type == PromptType.IMAGE_GENERATION:
if not getattr(settings, "ALLOW_IMAGE_GENERATION", False): if not getattr(settings, "ALLOW_IMAGE_GENERATION", False):
@@ -603,7 +605,9 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
# messages = messages[:-1] + [HumanMessage(content=altered_message)] # messages = messages[:-1] + [HumanMessage(content=altered_message)]
# I'll add it to the input_dict if needed. # 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_metric = await create_prompt_metric(
prompt_instance.id, prompt_instance.id,
prompt_instance.message, prompt_instance.message,
+41 -12
View File
@@ -1,5 +1,4 @@
import json import json
import base64
import logging import logging
import pandas as pd import pandas as pd
from datetime import datetime 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.moderation_classifier import moderation_classifier, ModerationLabel
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
from .services.data_analysis_service import AsyncDataAnalysisService 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.grounded_chat import prepare_grounded_chat
from .services.status_context import ( from .services.status_context import (
emit_status, emit_status,
@@ -36,7 +36,11 @@ from .services.status_context import (
set_status_emitter, set_status_emitter,
) )
from .services.ws_frames import citations_frame, status_frame 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 ( from .utils import (
TokenUsageCollector, TokenUsageCollector,
aiter_text_chunks, aiter_text_chunks,
@@ -202,7 +206,9 @@ async def get_conversation_file_async(conversation_id):
try: try:
prompt_with_file = await Prompt.objects.filter( prompt_with_file = await Prompt.objects.filter(
conversation_id=conversation_id 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: if prompt_with_file and prompt_with_file.file:
# Opening a DatabaseStorage file hits the DB, so read inside the thread. # Opening a DatabaseStorage file hits the DB, so read inside the thread.
@@ -250,6 +256,10 @@ async def classification_node(state: ChatState) -> ChatState:
prompt_type = await PROMPT_CLASSIFIER.classify_async(msg) 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 # Override logic
if decoded_file and (prompt_type == PromptType.DATA_ANALYSIS or 'analyze' in msg.lower() or 'data' in msg.lower()): if decoded_file and (prompt_type == PromptType.DATA_ANALYSIS or 'analyze' in msg.lower() or 'data' in msg.lower()):
prompt_type = PromptType.DATA_ANALYSIS prompt_type = PromptType.DATA_ANALYSIS
@@ -270,6 +280,17 @@ async def generation_node(state: ChatState) -> ChatState:
decoded_file = state.get("decoded_file") decoded_file = state.get("decoded_file")
file_type = state.get("file_type") 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 # Feature Flag + plan gate: Image Generation
if prompt_type == PromptType.IMAGE_GENERATION: if prompt_type == PromptType.IMAGE_GENERATION:
if not getattr(settings, "ALLOW_IMAGE_GENERATION", False): if not getattr(settings, "ALLOW_IMAGE_GENERATION", False):
@@ -485,13 +506,19 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
print("Conversation ID: ", conversation_id) print("Conversation ID: ", conversation_id)
decoded_file = None decoded_file = None
if file: if file:
decoded_file = base64.b64decode(file) decoded_file, file_type, attachment_error = prepare_attachment(
if "csv" in file_type: file_type = "csv" file, file_type
elif "xmlformats-officedocument" in file_type: file_type = "xlsx" )
elif "word" in file_type: file_type = "docx" if attachment_error:
elif "pdf" in file_type: file_type = "pdf" await self.send_json_message(
elif "text" in file_type: file_type = "txt" json.dumps(
else: file_type = "Not Sure" {
"type": "error",
"content": attachment_error,
}
)
)
return
# Pre-fetch messages and file # Pre-fetch messages and file
messages, prompt_instance = await get_messages( messages, prompt_instance = await get_messages(
@@ -501,7 +528,9 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
if not decoded_file: if not decoded_file:
decoded_file, file_type = await get_conversation_file_async(conversation_id) 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_metric = await create_prompt_metric(
prompt_instance.id, prompt_instance.id,
prompt_instance.message, prompt_instance.message,
@@ -556,7 +585,7 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
citations = [] citations = []
final_model = resolved_model final_model = resolved_model
if should_use_agent(message): if should_use_agent(message) and not is_image_type(file_type):
try: try:
await enforce_feature_gate(chat_user, "agentic_tasks") await enforce_feature_gate(chat_user, "agentic_tasks")
except FeatureNotAllowed as exc: except FeatureNotAllowed as exc:
+9 -2
View File
@@ -17,6 +17,8 @@ ROLE_EMBED = "embed"
# independent plan steps concurrently on a smaller/cheaper model. # independent plan steps concurrently on a smaller/cheaper model.
ROLE_ORCHESTRATOR = "orchestrator" ROLE_ORCHESTRATOR = "orchestrator"
ROLE_SUBAGENT = "subagent" ROLE_SUBAGENT = "subagent"
# Image + prompt turns (#128). Never falls back to the text chat model.
ROLE_VISION = "vision"
_VALID_ROLES = { _VALID_ROLES = {
ROLE_THINKING, ROLE_THINKING,
@@ -25,6 +27,7 @@ _VALID_ROLES = {
ROLE_EMBED, ROLE_EMBED,
ROLE_ORCHESTRATOR, ROLE_ORCHESTRATOR,
ROLE_SUBAGENT, ROLE_SUBAGENT,
ROLE_VISION,
} }
@@ -56,6 +59,7 @@ def ollama_model_for_role(role: str) -> str:
ROLE_EMBED: "OLLAMA_EMBED_MODEL", ROLE_EMBED: "OLLAMA_EMBED_MODEL",
ROLE_ORCHESTRATOR: "OLLAMA_MODEL_ORCHESTRATOR", ROLE_ORCHESTRATOR: "OLLAMA_MODEL_ORCHESTRATOR",
ROLE_SUBAGENT: "OLLAMA_MODEL_SUBAGENT", ROLE_SUBAGENT: "OLLAMA_MODEL_SUBAGENT",
ROLE_VISION: "OLLAMA_MODEL_VISION",
}[role] }[role]
role_default = { role_default = {
ROLE_THINKING: "gpt-oss:20b", ROLE_THINKING: "gpt-oss:20b",
@@ -64,14 +68,15 @@ def ollama_model_for_role(role: str) -> str:
ROLE_EMBED: "nomic-embed-text", ROLE_EMBED: "nomic-embed-text",
ROLE_ORCHESTRATOR: "gpt-oss:20b", ROLE_ORCHESTRATOR: "gpt-oss:20b",
ROLE_SUBAGENT: "llama3.2", ROLE_SUBAGENT: "llama3.2",
ROLE_VISION: "gemma4:latest",
}[role] }[role]
configured = getattr(settings, role_setting, None) configured = getattr(settings, role_setting, None)
if configured: if configured:
return configured return configured
# Embeddings must never silently fall back to a chat model (#62). # Embeddings and vision must never silently fall back to a text chat model.
if role == ROLE_EMBED: if role in (ROLE_EMBED, ROLE_VISION):
return role_default return role_default
legacy = getattr(settings, "OLLAMA_MODEL", None) 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) return int(getattr(settings, "OLLAMA_NUM_CTX_SUBAGENT", 8192) or 8192)
if role == ROLE_ORCHESTRATOR: if role == ROLE_ORCHESTRATOR:
return int(getattr(settings, "OLLAMA_NUM_CTX_ORCHESTRATOR", 16384) or 16384) 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) return int(getattr(settings, "OLLAMA_NUM_CTX_THINKING", 16384) or 16384)
+9 -1
View File
@@ -257,6 +257,7 @@ class PromptSerializer(serializers.ModelSerializer):
tokens_in = serializers.SerializerMethodField() tokens_in = serializers.SerializerMethodField()
tokens_out = serializers.SerializerMethodField() tokens_out = serializers.SerializerMethodField()
feedback = serializers.SerializerMethodField() feedback = serializers.SerializerMethodField()
image = serializers.SerializerMethodField()
class Meta: class Meta:
model = Prompt model = Prompt
@@ -269,8 +270,10 @@ class PromptSerializer(serializers.ModelSerializer):
"tokens_out", "tokens_out",
"citations", "citations",
"feedback", "feedback",
"file_type",
"image",
) )
read_only_fields = ("citations", "feedback") read_only_fields = ("citations", "feedback", "file_type", "image")
def _token_pair(self, obj): def _token_pair(self, obj):
cache = self.context.setdefault("_prompt_token_cache", {}) cache = self.context.setdefault("_prompt_token_cache", {})
@@ -314,6 +317,11 @@ class PromptSerializer(serializers.ModelSerializer):
self.context["_prompt_feedback_by_id"] = by_prompt self.context["_prompt_feedback_by_id"] = by_prompt
return by_prompt.get(obj.id) 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: def validate_message(self, value: str) -> str:
if value is None or not str(value).strip(): if value is None or not str(value).strip():
raise serializers.ValidationError("Message text cannot be empty.") raise serializers.ValidationError("Message text cannot be empty.")
+186
View File
@@ -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.")
@@ -23,6 +23,7 @@ STATUS_LABELS: dict[str, str] = {
"reading_sources": "Reading sources", "reading_sources": "Reading sources",
"retrieving_docs": "Searching your documents", "retrieving_docs": "Searching your documents",
"analysing": "Analysing your file", "analysing": "Analysing your file",
"reading_image": "Reading your image",
"refining": "Refining the answer", "refining": "Refining the answer",
"writing": "", # tokens streaming — no label "writing": "", # tokens streaming — no label
} }
@@ -355,6 +355,42 @@ class GraphNodeTestCase(TransactionTestCase):
self.prompt.message, b"a,b\n1,2\n", "csv" 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): async def test_generation_node_routes_rag_prompts_with_workspace(self):
with mock.patch.object(consumers_graph, "AsyncRAGService") as service: with mock.patch.object(consumers_graph, "AsyncRAGService") as service:
service.return_value.generate_response.return_value = "generator" service.return_value.generate_response.return_value = "generator"
+64
View File
@@ -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))
+4
View File
@@ -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. # sub-agents run independent plan steps on a smaller/cheaper model.
OLLAMA_MODEL_ORCHESTRATOR = env("OLLAMA_MODEL_ORCHESTRATOR", "") or OLLAMA_MODEL OLLAMA_MODEL_ORCHESTRATOR = env("OLLAMA_MODEL_ORCHESTRATOR", "") or OLLAMA_MODEL
OLLAMA_MODEL_SUBAGENT = env("OLLAMA_MODEL_SUBAGENT", "") or OLLAMA_MODEL_UTILITY 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( OLLAMA_NUM_CTX_ORCHESTRATOR = int(
env("OLLAMA_NUM_CTX_ORCHESTRATOR", "16384") or "16384" env("OLLAMA_NUM_CTX_ORCHESTRATOR", "16384") or "16384"
) )