mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-08-14 00:47:52 +00:00
feat(vision): add image + screen capture input for vision models (#486)
OpenJarvis can run vision-capable local models (gemma3, qwen2.5-vl), but the CLI had no way to send them a picture -- the Ollama engine only serialized text. This adds end-to-end image input. What's new - `jarvis ask -i/--image <file>` attaches one or more images to the query. - `jarvis ask -S/--screen` captures the primary monitor (dependency-free on Windows via .NET; mss/Pillow fallback elsewhere). - Vision auto-routes to direct-to-engine mode; with an explicit --agent it warns rather than silently dropping the image. - Privacy guard: warns before sending an image to a non-local engine, keeping OpenJarvis local-first by default. - Context-window default raised 8k -> 16k (JARVIS_NUM_CTX) so an image plus a conversation fit. Implementation - Message.images carries base64 data; messages_to_dicts() forwards it to Ollama's /api/chat "images" field. Text-only messages are unchanged. - GuardrailsEngine preserves images when it rewrites a flagged message. Tests (tests/test_vision.py, 6/6 pass, ruff-clean) - payload forwarding, text path untouched, num_ctx override, guardrail image preservation. Verified on AMD RX 9070 XT (Ollama/Vulkan, 100% GPU) with gemma3:4b: solid-color image, file image, and live screen capture all described. Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Jon Saad-Falcon <jonsaadfalcon@gmail.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
Jon Saad-Falcon
parent
7ba334b5f0
commit
79e23719d4
@@ -8,6 +8,19 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
|
|||||||
|
|
||||||
## [Unreleased]
|
## [Unreleased]
|
||||||
|
|
||||||
|
### Added
|
||||||
|
|
||||||
|
**Vision input for `jarvis ask`** — attach images to a query with
|
||||||
|
`-i`/`--image` (repeatable) or capture the current screen with
|
||||||
|
`-S`/`--screen`, for vision-capable models such as `gemma3:4b`. Images flow
|
||||||
|
through `Message.images` into Ollama's `/api/chat` `images` field; text-only
|
||||||
|
requests are unaffected. A privacy guard warns before any image is sent to a
|
||||||
|
non-local engine, and the security guardrail now preserves images when it
|
||||||
|
sanitizes a flagged prompt. Screen capture uses the built-in Windows .NET
|
||||||
|
stack with `mss`/`Pillow` fallbacks on other platforms. Adds the
|
||||||
|
`JARVIS_NUM_CTX` environment variable to tune the Ollama context window
|
||||||
|
(default `16384`).
|
||||||
|
|
||||||
## [1.0.2] - 2026-05-24
|
## [1.0.2] - 2026-05-24
|
||||||
|
|
||||||
A patch release that fixes a packaging bug which broke the v1.0.1
|
A patch release that fixes a packaging bug which broke the v1.0.1
|
||||||
|
|||||||
@@ -66,6 +66,8 @@ jarvis ask "What is the capital of France?"
|
|||||||
| `--no-context` | flag | off | Disable memory context injection |
|
| `--no-context` | flag | off | Disable memory context injection |
|
||||||
| `-a`, `--agent AGENT` | string | none | Agent to use (`simple`, `orchestrator`) |
|
| `-a`, `--agent AGENT` | string | none | Agent to use (`simple`, `orchestrator`) |
|
||||||
| `--tools TOOLS` | string | none | Comma-separated tool names to enable |
|
| `--tools TOOLS` | string | none | Comma-separated tool names to enable |
|
||||||
|
| `-i`, `--image PATH` | path | none | Image file for a vision model (e.g. `gemma3:4b`); repeatable |
|
||||||
|
| `-S`, `--screen` | flag | off | Capture the current screen and send it to the vision model |
|
||||||
|
|
||||||
### Direct Mode vs Agent Mode
|
### Direct Mode vs Agent Mode
|
||||||
|
|
||||||
@@ -105,6 +107,39 @@ jarvis ask --no-context "Tell me about Python"
|
|||||||
jarvis ask --max-tokens 2048 "Write a detailed essay about AI"
|
jarvis ask --max-tokens 2048 "Write a detailed essay about AI"
|
||||||
```
|
```
|
||||||
|
|
||||||
|
### Vision Input
|
||||||
|
|
||||||
|
Vision-capable models (such as `gemma3:4b`) can read images alongside your
|
||||||
|
text prompt. Attach one or more image files with `-i`/`--image`, or capture
|
||||||
|
the current screen with `-S`/`--screen`:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Ask about a local image
|
||||||
|
jarvis ask -i screenshot.png "What is shown in this image?"
|
||||||
|
|
||||||
|
# Send multiple images (the flag is repeatable)
|
||||||
|
jarvis ask -i chart-a.png -i chart-b.png "Compare these two charts"
|
||||||
|
|
||||||
|
# Capture the current screen and ask about it
|
||||||
|
jarvis ask --screen "Summarize what's on my screen"
|
||||||
|
```
|
||||||
|
|
||||||
|
Vision runs in **direct mode** only. If you also pass `--agent`, the image is
|
||||||
|
ignored and a note is printed — re-run with `--agent ""` to force direct mode.
|
||||||
|
|
||||||
|
The Ollama context window can be tuned for large images or long prompts with
|
||||||
|
the `JARVIS_NUM_CTX` environment variable (default `16384`):
|
||||||
|
|
||||||
|
```bash
|
||||||
|
JARVIS_NUM_CTX=8192 jarvis ask --screen "What's on my screen?"
|
||||||
|
```
|
||||||
|
|
||||||
|
!!! note "Keep vision on-device"
|
||||||
|
Images are sensitive. OpenJarvis prints a privacy warning before sending
|
||||||
|
an image to a non-local engine, so a screenshot never leaves your machine
|
||||||
|
unnoticed. Use a local engine (e.g. `ollama` with `gemma3:4b`) to keep
|
||||||
|
vision fully local.
|
||||||
|
|
||||||
### JSON Output Format
|
### JSON Output Format
|
||||||
|
|
||||||
When using `--json` in **direct mode**, the output includes:
|
When using `--json` in **direct mode**, the output includes:
|
||||||
|
|||||||
@@ -0,0 +1,79 @@
|
|||||||
|
"""Screen capture for vision input (``jarvis ask --screen``).
|
||||||
|
|
||||||
|
Captures the primary monitor to a temporary PNG so it can be handed to a
|
||||||
|
vision-capable model. On Windows this uses the built-in .NET
|
||||||
|
``System.Drawing`` stack (no third-party dependency). Other platforms fall
|
||||||
|
back to ``mss`` or ``Pillow`` if installed.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
|
||||||
|
# PowerShell: capture the PRIMARY monitor (more legible for a vision model
|
||||||
|
# than a downscaled multi-monitor grab). {path} is filled in with forward
|
||||||
|
# slashes, which .NET accepts on Windows and which avoids backslash escaping.
|
||||||
|
_PS_CAPTURE = """
|
||||||
|
Add-Type -AssemblyName System.Windows.Forms, System.Drawing
|
||||||
|
$b = [System.Windows.Forms.Screen]::PrimaryScreen.Bounds
|
||||||
|
$bmp = New-Object System.Drawing.Bitmap($b.Width, $b.Height)
|
||||||
|
$g = [System.Drawing.Graphics]::FromImage($bmp)
|
||||||
|
$g.CopyFromScreen($b.X, $b.Y, 0, 0, $bmp.Size)
|
||||||
|
$bmp.Save("{path}", [System.Drawing.Imaging.ImageFormat]::Png)
|
||||||
|
$g.Dispose(); $bmp.Dispose()
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
def capture_screen_to_temp() -> str:
|
||||||
|
"""Capture the screen to a temp PNG and return its absolute path.
|
||||||
|
|
||||||
|
Raises ``RuntimeError`` with actionable guidance if capture fails or the
|
||||||
|
platform has no available backend.
|
||||||
|
"""
|
||||||
|
fd, path = tempfile.mkstemp(prefix="jarvis_screen_", suffix=".png")
|
||||||
|
os.close(fd)
|
||||||
|
|
||||||
|
if sys.platform.startswith("win"):
|
||||||
|
script = _PS_CAPTURE.replace("{path}", path.replace("\\", "/"))
|
||||||
|
proc = subprocess.run(
|
||||||
|
["powershell", "-NoProfile", "-NonInteractive", "-Command", script],
|
||||||
|
capture_output=True,
|
||||||
|
text=True,
|
||||||
|
timeout=30,
|
||||||
|
)
|
||||||
|
if (
|
||||||
|
proc.returncode != 0
|
||||||
|
or not os.path.exists(path)
|
||||||
|
or not os.path.getsize(path)
|
||||||
|
):
|
||||||
|
raise RuntimeError(
|
||||||
|
"screen capture failed: "
|
||||||
|
+ (proc.stderr.strip() or "empty image written")
|
||||||
|
)
|
||||||
|
return path
|
||||||
|
|
||||||
|
# Non-Windows: optional backends.
|
||||||
|
try:
|
||||||
|
import mss # type: ignore
|
||||||
|
|
||||||
|
with mss.mss() as sct:
|
||||||
|
sct.shot(mon=-1, output=path)
|
||||||
|
return path
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
from PIL import ImageGrab # type: ignore
|
||||||
|
|
||||||
|
ImageGrab.grab().save(path)
|
||||||
|
return path
|
||||||
|
except Exception as exc: # noqa: BLE001
|
||||||
|
raise RuntimeError(
|
||||||
|
"screen capture on this platform needs 'mss' or 'Pillow' "
|
||||||
|
"(try: pip install mss)"
|
||||||
|
) from exc
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ["capture_screen_to_temp"]
|
||||||
@@ -2,6 +2,7 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import base64
|
||||||
import json as json_mod
|
import json as json_mod
|
||||||
import logging
|
import logging
|
||||||
import sys
|
import sys
|
||||||
@@ -619,6 +620,21 @@ def _print_profile(
|
|||||||
"(default: ~/.openjarvis/knowledge.db)."
|
"(default: ~/.openjarvis/knowledge.db)."
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
@click.option(
|
||||||
|
"-i",
|
||||||
|
"--image",
|
||||||
|
"image_paths",
|
||||||
|
multiple=True,
|
||||||
|
type=click.Path(exists=True, dir_okay=False),
|
||||||
|
help="Image file for a vision model (e.g. gemma3). Repeatable.",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"-S",
|
||||||
|
"--screen",
|
||||||
|
"capture_screen",
|
||||||
|
is_flag=True,
|
||||||
|
help="Capture the current screen and send it to the vision model.",
|
||||||
|
)
|
||||||
@click.option(
|
@click.option(
|
||||||
"--persona",
|
"--persona",
|
||||||
"persona_name",
|
"persona_name",
|
||||||
@@ -645,6 +661,8 @@ def ask(
|
|||||||
research_mode: bool,
|
research_mode: bool,
|
||||||
knowledge_db: str | None,
|
knowledge_db: str | None,
|
||||||
persona_name: str | None,
|
persona_name: str | None,
|
||||||
|
image_paths: tuple[str, ...] = (),
|
||||||
|
capture_screen: bool = False,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Ask Jarvis a question."""
|
"""Ask Jarvis a question."""
|
||||||
quiet = (ctx.obj or {}).get("quiet", False) or output_json
|
quiet = (ctx.obj or {}).get("quiet", False) or output_json
|
||||||
@@ -652,6 +670,27 @@ def ask(
|
|||||||
console = Console(stderr=True)
|
console = Console(stderr=True)
|
||||||
query_text = " ".join(query)
|
query_text = " ".join(query)
|
||||||
|
|
||||||
|
# Vision: collect base64 images from --image files and/or --screen.
|
||||||
|
image_b64: list[str] = []
|
||||||
|
for _img_path in image_paths:
|
||||||
|
try:
|
||||||
|
with open(_img_path, "rb") as _fh:
|
||||||
|
image_b64.append(base64.b64encode(_fh.read()).decode("ascii"))
|
||||||
|
except OSError as exc:
|
||||||
|
console.print(f"[red]Could not read image {_img_path}: {exc}[/red]")
|
||||||
|
sys.exit(1)
|
||||||
|
if capture_screen:
|
||||||
|
try:
|
||||||
|
from openjarvis.cli._screen import capture_screen_to_temp
|
||||||
|
|
||||||
|
_shot = capture_screen_to_temp()
|
||||||
|
with open(_shot, "rb") as _fh:
|
||||||
|
image_b64.append(base64.b64encode(_fh.read()).decode("ascii"))
|
||||||
|
logger.debug("Captured screen to %s", _shot)
|
||||||
|
except Exception as exc: # noqa: BLE001
|
||||||
|
console.print(f"[red]Screen capture failed:[/red] {exc}")
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
wall_start = time.monotonic() if enable_profile else None
|
wall_start = time.monotonic() if enable_profile else None
|
||||||
|
|
||||||
# Load config
|
# Load config
|
||||||
@@ -671,11 +710,26 @@ def ask(
|
|||||||
# Without this fallback, `[agent].default_system_prompt` and the
|
# Without this fallback, `[agent].default_system_prompt` and the
|
||||||
# SOUL.md / MEMORY.md / USER.md persona system are silently bypassed for
|
# SOUL.md / MEMORY.md / USER.md persona system are silently bypassed for
|
||||||
# the most common command (`jarvis ask "..."`).
|
# the most common command (`jarvis ask "..."`).
|
||||||
|
agent_explicitly_set = agent_name is not None
|
||||||
if agent_name is None:
|
if agent_name is None:
|
||||||
configured_default = (config.agent.default_agent or "").strip()
|
configured_default = (config.agent.default_agent or "").strip()
|
||||||
if configured_default:
|
if configured_default:
|
||||||
agent_name = configured_default
|
agent_name = configured_default
|
||||||
|
|
||||||
|
# Vision flows only through direct-to-engine mode. If an image/screenshot
|
||||||
|
# was supplied without an explicit --agent, route to direct mode so the
|
||||||
|
# picture reaches the model; if an agent was explicitly requested, say
|
||||||
|
# plainly that the image is being skipped rather than dropping it silently.
|
||||||
|
if image_b64:
|
||||||
|
if not agent_explicitly_set:
|
||||||
|
agent_name = ""
|
||||||
|
else:
|
||||||
|
console.print(
|
||||||
|
"[yellow]Note:[/yellow] --image/--screen only works in direct "
|
||||||
|
"mode; the image is ignored with --agent set. Re-run with "
|
||||||
|
'`--agent ""` to use vision.'
|
||||||
|
)
|
||||||
|
|
||||||
# Track whether the user explicitly set --max-tokens
|
# Track whether the user explicitly set --max-tokens
|
||||||
user_set_max_tokens = max_tokens is not None
|
user_set_max_tokens = max_tokens is not None
|
||||||
|
|
||||||
@@ -871,6 +925,27 @@ def ask(
|
|||||||
return
|
return
|
||||||
|
|
||||||
# Direct-to-engine mode (no agent)
|
# Direct-to-engine mode (no agent)
|
||||||
|
# Privacy guard: a screenshot/image is sensitive, and OpenJarvis is
|
||||||
|
# local-first. If the active engine isn't local, warn before the image
|
||||||
|
# leaves the machine rather than silently uploading it to a third party.
|
||||||
|
_LOCAL_ENGINES = {
|
||||||
|
"ollama",
|
||||||
|
"llamacpp",
|
||||||
|
"vllm",
|
||||||
|
"sglang",
|
||||||
|
"exo",
|
||||||
|
"nexa",
|
||||||
|
"uzu",
|
||||||
|
"apple_fm",
|
||||||
|
"gemma_cpp",
|
||||||
|
}
|
||||||
|
if image_b64 and engine_name not in _LOCAL_ENGINES:
|
||||||
|
console.print(
|
||||||
|
f"[yellow]Privacy warning:[/yellow] sending {len(image_b64)} "
|
||||||
|
f"image(s) to a non-local engine ('{engine_name}'). The image will "
|
||||||
|
"leave this machine. Use a local engine (e.g. ollama) to keep "
|
||||||
|
"vision on-device."
|
||||||
|
)
|
||||||
messages = [Message(role=Role.USER, content=query_text)]
|
messages = [Message(role=Role.USER, content=query_text)]
|
||||||
|
|
||||||
# Memory-augmented context injection
|
# Memory-augmented context injection
|
||||||
@@ -897,6 +972,15 @@ def ask(
|
|||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
logger.debug("Failed to inject memory context: %s", exc)
|
logger.debug("Failed to inject memory context: %s", exc)
|
||||||
|
|
||||||
|
# Vision: attach images to the final user message *after* any context
|
||||||
|
# injection (which may rebuild the list). messages_to_dicts() forwards
|
||||||
|
# the "images" field to Ollama's /api/chat.
|
||||||
|
if image_b64:
|
||||||
|
for _m in reversed(messages):
|
||||||
|
if _m.role == Role.USER:
|
||||||
|
_m.images = image_b64
|
||||||
|
break
|
||||||
|
|
||||||
# Generate (InstrumentedEngine handles telemetry + energy recording)
|
# Generate (InstrumentedEngine handles telemetry + energy recording)
|
||||||
try:
|
try:
|
||||||
with console.status("[bold green]Generating...[/bold green]"):
|
with console.status("[bold green]Generating...[/bold green]"):
|
||||||
|
|||||||
@@ -68,6 +68,10 @@ class Message:
|
|||||||
tool_calls: Optional[List[ToolCall]] = None
|
tool_calls: Optional[List[ToolCall]] = None
|
||||||
tool_call_id: Optional[str] = None
|
tool_call_id: Optional[str] = None
|
||||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||||
|
# Base64-encoded image data for vision-capable models (e.g. gemma3,
|
||||||
|
# qwen2.5-vl). Forwarded to Ollama's /api/chat "images" field; None or
|
||||||
|
# empty for text-only messages (the common case).
|
||||||
|
images: Optional[List[str]] = None
|
||||||
|
|
||||||
|
|
||||||
@dataclass(slots=True)
|
@dataclass(slots=True)
|
||||||
|
|||||||
@@ -34,6 +34,10 @@ def messages_to_dicts(messages: Sequence[Message]) -> List[Dict[str, Any]]:
|
|||||||
]
|
]
|
||||||
if m.tool_call_id:
|
if m.tool_call_id:
|
||||||
d["tool_call_id"] = m.tool_call_id
|
d["tool_call_id"] = m.tool_call_id
|
||||||
|
# Vision: forward base64 images to the engine. Ollama's /api/chat
|
||||||
|
# accepts an "images" array on a message; text messages skip this.
|
||||||
|
if getattr(m, "images", None):
|
||||||
|
d["images"] = list(m.images)
|
||||||
out.append(d)
|
out.append(d)
|
||||||
return out
|
return out
|
||||||
|
|
||||||
|
|||||||
@@ -23,6 +23,19 @@ from openjarvis.engine._stubs import StreamChunk
|
|||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _default_num_ctx() -> int:
|
||||||
|
"""Default context window (tokens). Override with ``JARVIS_NUM_CTX``.
|
||||||
|
|
||||||
|
Raised above Ollama's 4k default so an image (which costs many tokens)
|
||||||
|
plus a real conversation fit. 16k is comfortable for small models on a
|
||||||
|
typical consumer GPU.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
return int(os.environ.get("JARVIS_NUM_CTX", "16384"))
|
||||||
|
except ValueError:
|
||||||
|
return 16384
|
||||||
|
|
||||||
|
|
||||||
@EngineRegistry.register("ollama")
|
@EngineRegistry.register("ollama")
|
||||||
class OllamaEngine(InferenceEngine):
|
class OllamaEngine(InferenceEngine):
|
||||||
"""Ollama backend via its native HTTP API."""
|
"""Ollama backend via its native HTTP API."""
|
||||||
@@ -73,7 +86,7 @@ class OllamaEngine(InferenceEngine):
|
|||||||
"options": {
|
"options": {
|
||||||
"temperature": temperature,
|
"temperature": temperature,
|
||||||
"num_predict": max_tokens,
|
"num_predict": max_tokens,
|
||||||
"num_ctx": kwargs.get("num_ctx", 8192),
|
"num_ctx": kwargs.get("num_ctx", _default_num_ctx()),
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
# Disable extended thinking by default (Qwen3.5 etc.).
|
# Disable extended thinking by default (Qwen3.5 etc.).
|
||||||
@@ -189,7 +202,7 @@ class OllamaEngine(InferenceEngine):
|
|||||||
"options": {
|
"options": {
|
||||||
"temperature": temperature,
|
"temperature": temperature,
|
||||||
"num_predict": max_tokens,
|
"num_predict": max_tokens,
|
||||||
"num_ctx": kwargs.get("num_ctx", 8192),
|
"num_ctx": kwargs.get("num_ctx", _default_num_ctx()),
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
# Mirror generate()'s default: disable extended thinking unless the
|
# Mirror generate()'s default: disable extended thinking unless the
|
||||||
@@ -268,7 +281,7 @@ class OllamaEngine(InferenceEngine):
|
|||||||
"options": {
|
"options": {
|
||||||
"temperature": temperature,
|
"temperature": temperature,
|
||||||
"num_predict": max_tokens,
|
"num_predict": max_tokens,
|
||||||
"num_ctx": kwargs.get("num_ctx", 8192),
|
"num_ctx": kwargs.get("num_ctx", _default_num_ctx()),
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
if "think" not in kwargs:
|
if "think" not in kwargs:
|
||||||
|
|||||||
@@ -192,6 +192,7 @@ class GuardrailsEngine(InferenceEngine):
|
|||||||
tool_calls=msg.tool_calls,
|
tool_calls=msg.tool_calls,
|
||||||
tool_call_id=msg.tool_call_id,
|
tool_call_id=msg.tool_call_id,
|
||||||
metadata=msg.metadata,
|
metadata=msg.metadata,
|
||||||
|
images=msg.images,
|
||||||
)
|
)
|
||||||
messages = processed
|
messages = processed
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,128 @@
|
|||||||
|
"""CLI-level regression tests for ``jarvis ask`` vision input.
|
||||||
|
|
||||||
|
The unit tests in ``tests/test_vision.py`` cover the ``Message.images`` ->
|
||||||
|
``messages_to_dicts`` serialization contract in isolation. These tests lock
|
||||||
|
the *end-to-end CLI wiring*: that ``--image`` reads a file, base64-encodes it,
|
||||||
|
attaches it to the final user ``Message``, and that the bytes actually reach
|
||||||
|
``engine.generate()`` -- and that the local-first privacy guard fires only for
|
||||||
|
non-local engines.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import importlib
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from click.testing import CliRunner
|
||||||
|
|
||||||
|
from openjarvis.cli import cli
|
||||||
|
from openjarvis.core.config import JarvisConfig
|
||||||
|
from openjarvis.core.types import Role
|
||||||
|
|
||||||
|
# Import the module (not the Click command attribute) so we can monkeypatch
|
||||||
|
# the names it looks up at call time.
|
||||||
|
_ask_mod = importlib.import_module("openjarvis.cli.ask")
|
||||||
|
|
||||||
|
# A minimal but valid 1x1 PNG so ``click.Path(exists=True)`` is satisfied and
|
||||||
|
# the bytes are deterministic.
|
||||||
|
_PNG_BYTES = base64.b64decode(
|
||||||
|
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk"
|
||||||
|
"+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class _RecordingEngine:
|
||||||
|
"""A fake engine that records the messages handed to ``generate()``."""
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.engine_id = "mock"
|
||||||
|
self.received: list[Any] = []
|
||||||
|
|
||||||
|
def health(self) -> bool:
|
||||||
|
return True
|
||||||
|
|
||||||
|
def list_models(self) -> list[str]:
|
||||||
|
return ["test-model"]
|
||||||
|
|
||||||
|
def generate(self, messages, *, model=None, **kwargs):
|
||||||
|
# Capture the exact Message objects the CLI built so the test can
|
||||||
|
# assert the image bytes reached the engine boundary.
|
||||||
|
self.received = list(messages)
|
||||||
|
return {
|
||||||
|
"content": "a 1x1 pixel",
|
||||||
|
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||||
|
"model": "test-model",
|
||||||
|
"finish_reason": "stop",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _patch_ask(monkeypatch, tmp_path: Path, *, engine_name: str) -> _RecordingEngine:
|
||||||
|
"""Wire ``jarvis ask`` to a recording engine reported under ``engine_name``."""
|
||||||
|
cfg = JarvisConfig()
|
||||||
|
cfg.telemetry.db_path = str(tmp_path / "telemetry.db")
|
||||||
|
# Keep memory context out of the picture so the user message we inspect is
|
||||||
|
# the one the CLI built directly from the query + image.
|
||||||
|
cfg.agent.context_from_memory = False
|
||||||
|
monkeypatch.setattr(_ask_mod, "load_config", lambda: cfg)
|
||||||
|
|
||||||
|
engine = _RecordingEngine()
|
||||||
|
monkeypatch.setattr(_ask_mod, "get_engine", lambda *a, **kw: (engine_name, engine))
|
||||||
|
monkeypatch.setattr(_ask_mod, "discover_engines", lambda c: [(engine_name, engine)])
|
||||||
|
monkeypatch.setattr(
|
||||||
|
_ask_mod, "discover_models", lambda e: {engine_name: ["test-model"]}
|
||||||
|
)
|
||||||
|
return engine
|
||||||
|
|
||||||
|
|
||||||
|
def _write_png(tmp_path: Path) -> tuple[Path, str]:
|
||||||
|
img = tmp_path / "pixel.png"
|
||||||
|
img.write_bytes(_PNG_BYTES)
|
||||||
|
return img, base64.b64encode(_PNG_BYTES).decode("ascii")
|
||||||
|
|
||||||
|
|
||||||
|
def test_image_reaches_engine_payload(monkeypatch, tmp_path: Path) -> None:
|
||||||
|
engine = _patch_ask(monkeypatch, tmp_path, engine_name="ollama")
|
||||||
|
img, expected_b64 = _write_png(tmp_path)
|
||||||
|
|
||||||
|
result = CliRunner().invoke(
|
||||||
|
cli,
|
||||||
|
["ask", "-i", str(img), "--no-context", "--agent", "", "describe this"],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.exit_code == 0, result.output
|
||||||
|
# The CLI must have routed to direct mode and called the engine.
|
||||||
|
assert engine.received, "engine.generate() was never called"
|
||||||
|
user_msgs = [m for m in engine.received if m.role == Role.USER]
|
||||||
|
assert user_msgs, "no USER message reached the engine"
|
||||||
|
assert user_msgs[-1].images == [expected_b64]
|
||||||
|
|
||||||
|
|
||||||
|
def test_privacy_warning_for_non_local_engine(monkeypatch, tmp_path: Path) -> None:
|
||||||
|
engine = _patch_ask(monkeypatch, tmp_path, engine_name="openai")
|
||||||
|
img, expected_b64 = _write_png(tmp_path)
|
||||||
|
|
||||||
|
result = CliRunner().invoke(
|
||||||
|
cli,
|
||||||
|
["ask", "-i", str(img), "--no-context", "--agent", "", "describe this"],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.exit_code == 0, result.output
|
||||||
|
assert "Privacy warning" in result.output
|
||||||
|
# The warning is informational; the image must still be delivered.
|
||||||
|
user_msgs = [m for m in engine.received if m.role == Role.USER]
|
||||||
|
assert user_msgs and user_msgs[-1].images == [expected_b64]
|
||||||
|
|
||||||
|
|
||||||
|
def test_no_privacy_warning_for_local_engine(monkeypatch, tmp_path: Path) -> None:
|
||||||
|
_patch_ask(monkeypatch, tmp_path, engine_name="ollama")
|
||||||
|
img, _ = _write_png(tmp_path)
|
||||||
|
|
||||||
|
result = CliRunner().invoke(
|
||||||
|
cli,
|
||||||
|
["ask", "-i", str(img), "--no-context", "--agent", "", "describe this"],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.exit_code == 0, result.output
|
||||||
|
assert "Privacy warning" not in result.output
|
||||||
@@ -0,0 +1,94 @@
|
|||||||
|
"""Tests for vision input support: ``Message.images`` -> Ollama payload.
|
||||||
|
|
||||||
|
These cover the data-flow contract that makes vision work end to end:
|
||||||
|
a ``Message`` can carry base64 images, the engine serializer forwards them
|
||||||
|
to Ollama's ``/api/chat`` ``images`` field, and text-only messages are
|
||||||
|
completely unaffected. The security guardrail must preserve images when it
|
||||||
|
rewrites a flagged message.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
import openjarvis.engine.ollama as ollama_mod
|
||||||
|
from openjarvis.core.types import Message, Role
|
||||||
|
from openjarvis.engine._base import messages_to_dicts
|
||||||
|
|
||||||
|
|
||||||
|
def test_message_defaults_to_no_images() -> None:
|
||||||
|
assert Message(role=Role.USER, content="hi").images is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_messages_to_dicts_omits_images_for_text() -> None:
|
||||||
|
dicts = messages_to_dicts([Message(role=Role.USER, content="hi")])
|
||||||
|
assert "images" not in dicts[0]
|
||||||
|
|
||||||
|
|
||||||
|
def test_messages_to_dicts_forwards_images() -> None:
|
||||||
|
b64 = "aGVsbG8=" # "hello"
|
||||||
|
dicts = messages_to_dicts(
|
||||||
|
[Message(role=Role.USER, content="what is this?", images=[b64])]
|
||||||
|
)
|
||||||
|
assert dicts[0]["role"] == "user"
|
||||||
|
assert dicts[0]["content"] == "what is this?"
|
||||||
|
assert dicts[0]["images"] == [b64]
|
||||||
|
|
||||||
|
|
||||||
|
def test_messages_to_dicts_empty_images_treated_as_text() -> None:
|
||||||
|
dicts = messages_to_dicts([Message(role=Role.USER, content="hi", images=[])])
|
||||||
|
assert "images" not in dicts[0]
|
||||||
|
|
||||||
|
|
||||||
|
def test_default_num_ctx_default_and_override(monkeypatch) -> None:
|
||||||
|
monkeypatch.delenv("JARVIS_NUM_CTX", raising=False)
|
||||||
|
assert ollama_mod._default_num_ctx() == 16384
|
||||||
|
|
||||||
|
monkeypatch.setenv("JARVIS_NUM_CTX", "8000")
|
||||||
|
assert ollama_mod._default_num_ctx() == 8000
|
||||||
|
|
||||||
|
# A non-integer override must fall back to the safe default, not crash.
|
||||||
|
monkeypatch.setenv("JARVIS_NUM_CTX", "not-an-int")
|
||||||
|
assert ollama_mod._default_num_ctx() == 16384
|
||||||
|
|
||||||
|
|
||||||
|
def test_guardrails_preserves_images_when_sanitizing() -> None:
|
||||||
|
"""A flagged message gets rewritten; its image must survive the rewrite."""
|
||||||
|
from openjarvis.security.guardrails import GuardrailsEngine
|
||||||
|
|
||||||
|
class _RecordingEngine:
|
||||||
|
"""Captures the messages the guardrail forwards to the real engine."""
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.received: list[Message] = []
|
||||||
|
|
||||||
|
def generate(self, messages, *, model, **kwargs):
|
||||||
|
self.received = list(messages)
|
||||||
|
return {"content": "ok"}
|
||||||
|
|
||||||
|
class _AlwaysFlag:
|
||||||
|
"""A scanner that flags everything, forcing the sanitize rewrite path."""
|
||||||
|
|
||||||
|
def scan(self, text: str):
|
||||||
|
finding = SimpleNamespace(
|
||||||
|
pattern_name="test",
|
||||||
|
threat_level=SimpleNamespace(value="low"),
|
||||||
|
description="always flags",
|
||||||
|
)
|
||||||
|
return SimpleNamespace(findings=[finding])
|
||||||
|
|
||||||
|
def redact(self, text: str) -> str:
|
||||||
|
return text
|
||||||
|
|
||||||
|
engine = _RecordingEngine()
|
||||||
|
guarded = GuardrailsEngine(
|
||||||
|
engine,
|
||||||
|
scanners=[_AlwaysFlag()],
|
||||||
|
scan_input=True,
|
||||||
|
scan_output=False,
|
||||||
|
)
|
||||||
|
msg = Message(role=Role.USER, content="suspicious", images=["aGVsbG8="])
|
||||||
|
|
||||||
|
guarded.generate([msg], model="x")
|
||||||
|
|
||||||
|
assert engine.received[0].images == ["aGVsbG8="]
|
||||||
Reference in New Issue
Block a user