Read attached chat images with a vision model #82

Merged
westfarn merged 1 commits from feat/128-vision-image-prompt into master 2026-10-20 06:13:32 -07:00
11 changed files with 387 additions and 44 deletions
Showing only changes of commit c21f052d5b - Show all commits

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_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
+2
View File
@@ -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
+31 -27
View File
@@ -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,
+43 -14
View File
@@ -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:
+9 -2
View File
@@ -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)
+9 -1
View File
@@ -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.")
+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",
"retrieving_docs": "Searching your documents",
"analysing": "Analysing your file",
"reading_image": "Reading your image",
"refining": "Refining the answer",
"writing": "", # tokens streaming — no label
}
@@ -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"
+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.
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"
)