Read attached chat images with a vision model #82
No files matched your search
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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.")
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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))
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
Reference in new issue
Block a user