diff --git a/.claude/commands/vlmrun-gateway-cli.md b/.claude/commands/vlmrun-gateway-cli.md new file mode 100644 index 0000000..ab28601 --- /dev/null +++ b/.claude/commands/vlmrun-gateway-cli.md @@ -0,0 +1,115 @@ +--- +description: Natural-language interface to the `vlmrun gw` gateway CLI — say what you want done to a file/URL/text and it runs the right command +argument-hint: ", e.g. 'embed this image ~/data/image.jpg' or 'transcribe ~/clip.mp3'" +--- + +# vlmrun gateway — do what I asked + +Request: **$ARGUMENTS** + +Translate the request above into the correct `vlmrun gw` command, run it, and +show the result. You are a thin natural-language front-end to an existing CLI — +choose the subcommand, model, and method, execute, and report. Do not modify the +CLI or write code. + +## 1. Parse the request + +From the request text pull out: + +- **Inputs** — local paths (`~/…`, `./…`, absolute), `http(s)://…` URLs, and any + literal text to act on ("embed the phrase 'blue parrot'"). Expand `~`. +- **Intent** — what to do with them (see the mapping below). + +If there is no input and the intent needs one, or the request is genuinely +ambiguous, ask one short clarifying question instead of guessing. A missing +model is not ambiguity — pick a sensible default. + +## 2. Map intent → subcommand, model, method + +Model ids drift, so confirm your choice against `vlmrun gw models` (and +`vlmrun gw models ` for its methods) before running. The defaults below +are current; if one is absent from the catalog, pick another model whose `task` +and `methods` fit. Always pass the full `/` id. + +**Transcription** — audio, or "transcribe" a video (its audio track). +Triggers: transcribe, subtitles/captions, "what is said". Extensions: mp3, wav, +m4a, flac, ogg; or a video when the intent is transcription. +→ `vlmrun gw transcribe -m nvidia/parakeet-tdt-0.6b-v3` + - subtitles → `-f srt` (or `-f vtt`); timestamps/detail → `-f verbose_json` + - a hosted URL → `--url ` (no positional file); language hint → `-l en` + +**Embedding** — turn an image, text, or video into a vector. +Triggers: embed, vectorize, "embedding for", similarity/index. +→ `vlmrun gw embed -m qwen/qwen3-vl-embedding-2b` + - literal text → `-t "…"` (repeatable); each file/`-t` is its own vector + - "embed the image together with this caption" → add `--join` (one file max) + - "give me N dims" → `--dimensions N`; scripting → `--json` for full vectors + +**Document parsing** — a PDF/DOCX, or an image of a document/form/receipt. +Triggers: parse, read, "to markdown", extract text, OCR, layout/structure. +→ `vlmrun gw chat -m --method ` + - to markdown / clean reading order → `-m zai-org/glm-ocr --method markdown` + - text + confidence + polygons (JSON) → `-m paddleocr/pp-ocrv6 --method ocr` + - just text-region boxes → `-m paddleocr/pp-ocrv6 --method detect` + - layout blocks / sections → `-m rednote-hilab/dots.mocr --method parse_layout` + - OCR knobs: `--method-params '{"lang":"en","score_threshold":0.5}'` (pp-ocrv6) + +**Image understanding / VQA** — a free-form question or description of an image. +Triggers: describe, caption, "what/how many/what color/is there…", summarize. +→ `vlmrun gw chat -p "" -m qwen/qwen3.5-0.8b` + (Use the OCR models above only when the intent is text extraction, not a + question about the scene.) + +Notes that matter: + +- Images are sent as `image_url`, documents as `document_url` — the CLI decides + from the file's real bytes, so a mislabelled extension is fine; you don't set + this. +- Only the OCR/document `chat` models are strict about needing a file. VQA takes + `-p`; embedding takes `-t`. +- If the user names a model or method explicitly, honor it over the defaults. + +## 3. Run it + +Before executing, make sure a stale dev environment isn't shadowing the +configured key: `unset VLMRUN_API_KEY VLMRUN_BASE_URL VLMRUN_GATEWAY_URL` (the +CLI reads the key from `~/.vlmrun/config.toml`). Then: + +1. Show the exact `vlmrun gw …` command you're about to run (one line). +2. Run it. Default to the human-readable panel; add `--json` only if the user + asked for raw output or is clearly scripting. +3. If it fails: a `model not found` error prints the valid ids — pick the right + one and retry once. An `Error: Unknown method …` means the method is wrong + for that model — drop `--method` (use its default) or pick a listed one. A + `model not found`/wrong-task pairing means you chose the wrong model for the + task; re-map and retry. Don't loop more than a couple of times — if still + stuck, report what you tried. + +## 4. Report + +State what you ran and give the result: the transcript, the extracted +text/markdown, the VQA answer, or for embeddings a one-line summary (how many +vectors, their dimension) rather than dumping raw floats unless asked. If you +picked a model/method the user didn't specify, say which and why in one line so +they can steer next time. + +Chat/OCR responses carry `usage.cost` (USD) — the panel footer shows it, and +`--json` puts it under `usage.cost`. Report the cost when it's relevant, and for +a batch over many files sum it into a total. + +## Examples + +- `"embed this image ~/data/image.jpg"` + → `vlmrun gw embed ~/data/image.jpg -m qwen/qwen3-vl-embedding-2b` +- `"what's in ~/photos/street.jpg?"` + → `vlmrun gw chat ~/photos/street.jpg -p "What's in this image?" -m qwen/qwen3.5-0.8b` +- `"extract the text from ~/scans/receipt.png"` + → `vlmrun gw chat ~/scans/receipt.png -m paddleocr/pp-ocrv6 --method ocr` +- `"parse ~/docs/contract.pdf to markdown"` + → `vlmrun gw chat ~/docs/contract.pdf -m zai-org/glm-ocr --method markdown` +- `"get the layout blocks of ~/forms/intake.jpg"` + → `vlmrun gw chat ~/forms/intake.jpg -m rednote-hilab/dots.mocr --method parse_layout` +- `"transcribe ~/calls/standup.mp4 as subtitles"` + → `vlmrun gw transcribe ~/calls/standup.mp4 -m nvidia/parakeet-tdt-0.6b-v3 -f srt` +- `"embed the caption 'a blue parrot' with ~/img/parrot.jpg as one vector"` + → `vlmrun gw embed ~/img/parrot.jpg -t "a blue parrot" --join -m qwen/qwen3-vl-embedding-2b` diff --git a/tests/test_gateway.py b/tests/test_gateway.py new file mode 100644 index 0000000..15c7d77 --- /dev/null +++ b/tests/test_gateway.py @@ -0,0 +1,1111 @@ +"""Tests for the OpenAI-compatible gateway resource and `vlmrun gw` CLI.""" + +from __future__ import annotations + +import json + +import pytest +from rich.markdown import Markdown +from rich.text import Text +from typer.testing import CliRunner + +from vlmrun.cli.cli import app +from vlmrun.cli._cli import gateway as gw +from vlmrun.client.gateway import Gateway +from vlmrun.constants import DEFAULT_GATEWAY_URL + +# Minimal real file headers, so mime sniffing sees what it would in the wild. +PNG_BYTES = b"\x89PNG\r\n\x1a\n" + b"\x00" * 8 +WEBP_BYTES = b"RIFF" + b"\x00\x00\x00\x00" + b"WEBP" + b"\x00" * 4 +MP4_BYTES = b"\x00\x00\x00\x20" + b"ftyp" + b"isom" + b"\x00" * 4 + +# --------------------------------------------------------------------------- +# Concrete fakes (per CLAUDE.md: no MagicMock) +# --------------------------------------------------------------------------- + + +class FakeUsage: + def __init__( + self, prompt_tokens: int = 10, completion_tokens: int = 20, cost=None + ) -> None: + self.prompt_tokens = prompt_tokens + self.completion_tokens = completion_tokens + self.total_tokens = prompt_tokens + completion_tokens + self.cost = cost + + def model_dump(self) -> dict: + return { + "prompt_tokens": self.prompt_tokens, + "completion_tokens": self.completion_tokens, + "total_tokens": self.total_tokens, + "cost": self.cost, + } + + +class FakeMessage: + def __init__(self, content: str) -> None: + self.content = content + + +class FakeChoice: + def __init__(self, content: str) -> None: + self.message = FakeMessage(content) + + +class FakeResponse: + def __init__(self, content: str, cost=None) -> None: + self.choices = [FakeChoice(content)] + self.usage = FakeUsage(cost=cost) + + +class FakeDelta: + def __init__(self, content: str) -> None: + self.content = content + + +class FakeStreamChoice: + def __init__(self, content: str) -> None: + self.delta = FakeDelta(content) + + +class FakeChunk: + def __init__(self, content: str, usage=None) -> None: + self.choices = [FakeStreamChoice(content)] + self.usage = usage + + +class FakeCompletions: + def __init__(self) -> None: + self.calls: list[dict] = [] + self.content = "Hello world" + self.cost = None + + def create(self, model, messages, stream=False, **kwargs): + self.calls.append( + {"model": model, "messages": messages, "stream": stream, **kwargs} + ) + if stream: + return iter( + [ + FakeChunk(self.content[:6]), + FakeChunk(self.content[6:], usage=FakeUsage(cost=self.cost)), + ] + ) + return FakeResponse(self.content, cost=self.cost) + + +class FakeModel: + """Mimics an OpenAI ``Model`` object as the gateway actually returns it. + + Mirrors a real ``GET /v1/openai/models`` payload: methods/aliases/task and + capabilities, and no pricing fields. + """ + + def __init__(self, id: str, **extra) -> None: + self._data = { + "id": id, + "object": "model", + "owned_by": "vlm-run", + "aliases": [], + "methods": [], + "default_method": "", + "extra_body_help": "", + "capabilities": {"supported_input_types": []}, + "task": "chat", + **extra, + } + + def model_dump(self) -> dict: + return dict(self._data) + + +class FakeEmbeddingItem: + def __init__(self, index: int) -> None: + self.object = "embedding" + self.index = index + self.embedding = [0.1, 0.2, 0.3, 0.4] + + +class FakeEmbeddingResponse: + def __init__(self, n: int) -> None: + self.data = [FakeEmbeddingItem(i) for i in range(n)] + self.usage = FakeUsage(prompt_tokens=5, completion_tokens=0) + + def model_dump(self) -> dict: + return { + "data": [ + {"object": e.object, "index": e.index, "embedding": e.embedding} + for e in self.data + ] + } + + +class FakeEmbeddings: + def __init__(self) -> None: + self.calls: list[dict] = [] + + def create(self, model, input, **kwargs): + self.calls.append({"model": model, "input": input, **kwargs}) + return FakeEmbeddingResponse(len(input)) + + +class FakeTranscription: + def __init__(self, text: str) -> None: + self.text = text + + def model_dump(self) -> dict: + return {"text": self.text} + + +class FakeTranscriptions: + def __init__(self) -> None: + self.calls: list[dict] = [] + + def create(self, model, file, **kwargs): + # `file` is a handle or tuple; record only what it resolves to. + name = getattr(file, "name", None) or ( + file[0] if isinstance(file, tuple) else str(file) + ) + self.calls.append({"model": model, "file": str(name), **kwargs}) + return FakeTranscription("hello from audio") + + +class FakeGateway: + def __init__(self, healthy: bool = True) -> None: + self.base_url = "https://gateway.vlm.run/v1" + self._healthy = healthy + self.completions = FakeCompletions() + self.embeddings = FakeEmbeddings() + self.transcriptions = FakeTranscriptions() + + def health(self) -> bool: + return self._healthy + + def models(self) -> list: + return [ + FakeModel( + "zai-org/glm-ocr", + aliases=["glm-ocr"], + methods=["ocr", "markdown"], + default_method="ocr", + extra_body_help='{"method":"ocr"} | {"method":"markdown"}' + " | document_url PDF (markdown per page)", + capabilities={ + "supported_input_types": ["text", "image_url", "document_url"] + }, + ), + FakeModel( + "paddleocr/pp-ocrv6", + aliases=["pp-ocrv6"], + methods=["ocr", "detect", "markdown"], + default_method="ocr", + extra_body_help='{"method":"ocr","method_params":' + '{"lang":"en","score_threshold":0.5}}', + capabilities={ + "supported_input_types": ["text", "image_url", "document_url"] + }, + ), + ] + + +class FakeClient: + """Concrete stand-in for VLMRun used as the CLI context object.""" + + def __init__(self, api_key=None, base_url=None, healthy: bool = True) -> None: + self.api_key = api_key or "test-key" + self.base_url = base_url or "https://api.vlm.run/v1" + self.timeout = 120.0 + self.max_retries = 1 + self.gateway = FakeGateway(healthy=healthy) + + +@pytest.fixture +def runner(): + return CliRunner() + + +@pytest.fixture +def patched_cli(monkeypatch): + """Patch the CLI's VLMRun factory + credentials so ctx.obj is a FakeClient.""" + monkeypatch.setenv("VLMRUN_API_KEY", "test-key") + holder = {} + + def _factory(api_key=None, base_url=None): + healthy = holder.get("healthy", True) + client = FakeClient(api_key=api_key, base_url=base_url, healthy=healthy) + if "content" in holder: + client.gateway.completions.content = holder["content"] + if "cost" in holder: + client.gateway.completions.cost = holder["cost"] + holder["client"] = client + return client + + monkeypatch.setattr("vlmrun.cli.cli.VLMRun", _factory) + return holder + + +# --------------------------------------------------------------------------- +# Client resource: Gateway +# --------------------------------------------------------------------------- + + +class _MiniClient: + def __init__(self, timeout=120.0) -> None: + self.api_key = "sk-test" + self.timeout = timeout + self.max_retries = 3 + + +class TestGatewayResource: + def test_timeout_raises_floor_at_default(self): + # The 120s default is bumped to 600s for slow gateway calls. + assert Gateway(_MiniClient(timeout=120.0))._timeout() == 600.0 + + def test_timeout_respects_explicit_short(self): + # A user's fail-fast timeout must not be silently widened. + assert Gateway(_MiniClient(timeout=5.0))._timeout() == 5.0 + + def test_timeout_respects_explicit_long(self): + assert Gateway(_MiniClient(timeout=900.0))._timeout() == 900.0 + + def test_timeout_none_stays_none(self): + assert Gateway(_MiniClient(timeout=None))._timeout() is None + + def test_default_base_url(self, monkeypatch): + monkeypatch.delenv("VLMRUN_GATEWAY_URL", raising=False) + g = Gateway(_MiniClient()) + assert g.base_url == DEFAULT_GATEWAY_URL + assert g.openai_base_url == f"{DEFAULT_GATEWAY_URL}/openai" + + def test_env_override(self, monkeypatch): + monkeypatch.setenv("VLMRUN_GATEWAY_URL", "https://gw.example.com/v1/") + g = Gateway(_MiniClient()) + assert g.base_url == "https://gw.example.com/v1" + + def test_param_override(self, monkeypatch): + monkeypatch.setenv("VLMRUN_GATEWAY_URL", "https://env.example.com/v1") + g = Gateway(_MiniClient(), base_url="https://param.example.com/v1") + assert g.base_url == "https://param.example.com/v1" + + def test_models_delegates_to_openai(self): + g = Gateway(_MiniClient()) + + class _Models: + def list(self): + return iter(["a", "b", "c"]) + + class _OpenAI: + models = _Models() + + # cached_property stored in instance __dict__ takes precedence. + g.__dict__["_openai"] = _OpenAI() + assert g.models() == ["a", "b", "c"] + + def test_health_dedicated_endpoint(self, monkeypatch): + g = Gateway(_MiniClient()) + + class _Resp: + status_code = 200 + is_success = True + + monkeypatch.setattr("httpx.get", lambda *a, **k: _Resp()) + assert g.health() is True + + def test_health_falls_back_to_models_on_404(self, monkeypatch): + g = Gateway(_MiniClient()) + + class _Resp: + status_code = 404 + is_success = False + + monkeypatch.setattr("httpx.get", lambda *a, **k: _Resp()) + + class _Models: + def list(self): + return iter([1]) + + class _OpenAI: + models = _Models() + + g.__dict__["_openai"] = _OpenAI() + assert g.health() is True + + def test_health_false_on_connection_error(self, monkeypatch): + g = Gateway(_MiniClient()) + + def _boom(*a, **k): + raise RuntimeError("no network") + + monkeypatch.setattr("httpx.get", _boom) + + class _Models: + def list(self): + raise RuntimeError("still down") + + class _OpenAI: + models = _Models() + + g.__dict__["_openai"] = _OpenAI() + assert g.health() is False + + +# --------------------------------------------------------------------------- +# CLI helper functions +# --------------------------------------------------------------------------- + + +class TestHelpers: + def test_guess_mime(self, tmp_path): + assert gw._guess_mime(tmp_path / "a.pdf") == "application/pdf" + assert gw._guess_mime(tmp_path / "a.png") == "image/png" + + def test_content_part_type(self, tmp_path): + assert gw._content_part_type(tmp_path / "a.pdf") == "document_url" + assert gw._content_part_type(tmp_path / "a.docx") == "document_url" + assert gw._content_part_type(tmp_path / "a.png") == "image_url" + assert gw._content_part_type(tmp_path / "a.jpg") == "image_url" + # Unidentifiable content still falls back to file_url. + assert gw._content_part_type(tmp_path / "a.bin") == "file_url" + + def test_encode_document_part(self, tmp_path): + f = tmp_path / "doc.pdf" + f.write_bytes(b"%PDF-1.7 fake") + part = gw._encode_file_part(f) + assert part["type"] == "document_url" + url = part["document_url"]["url"] + assert url.startswith("data:application/pdf;base64,") + + def test_encode_image_part(self, tmp_path): + f = tmp_path / "img.png" + f.write_bytes(PNG_BYTES) + part = gw._encode_file_part(f) + # Images go as image_url: file_url is routed through the gateway's + # document/PDF path and 400s on a plain image. + assert part["type"] == "image_url" + url = part["image_url"]["url"] + assert url.startswith("data:image/png;base64,") + + def test_encode_part_sniffs_mislabelled_extension(self, tmp_path): + # A WebP that claims to be a .jpg — the extension must not win, or the + # gateway misroutes it and fails. + f = tmp_path / "actually-webp.jpg" + f.write_bytes(WEBP_BYTES) + part = gw._encode_file_part(f) + assert part["type"] == "image_url" + assert part["image_url"]["url"].startswith("data:image/webp;base64,") + + def test_encode_part_sniffs_pdf_without_extension(self, tmp_path): + f = tmp_path / "nameless" + f.write_bytes(b"%PDF-1.4 fake") + part = gw._encode_file_part(f) + assert part["type"] == "document_url" + assert part["document_url"]["url"].startswith("data:application/pdf;base64,") + + def test_guess_mime_falls_back_to_extension(self, tmp_path): + f = tmp_path / "img.png" + f.write_bytes(b"not-a-real-png") + assert gw._guess_mime(f) == "image/png" + + def test_guess_mime_unknown(self, tmp_path): + f = tmp_path / "mystery.bin" + f.write_bytes(b"\x00\x01\x02\x03") + assert gw._guess_mime(f) == "application/octet-stream" + + def test_sniff_mime_signatures(self): + assert gw._sniff_mime(PNG_BYTES) == "image/png" + assert gw._sniff_mime(WEBP_BYTES) == "image/webp" + assert gw._sniff_mime(b"\xff\xd8\xff\xe0") == "image/jpeg" + assert gw._sniff_mime(b"%PDF-1.4") == "application/pdf" + assert gw._sniff_mime(b"nonsense") is None + + def test_build_messages_with_prompt(self, tmp_path): + f = tmp_path / "img.png" + f.write_bytes(PNG_BYTES) + messages = gw._build_messages([f], "describe") + assert len(messages) == 1 + content = messages[0]["content"] + assert content[0]["type"] == "image_url" + assert content[-1] == {"type": "text", "text": "describe"} + + def test_build_messages_mixed_files(self, tmp_path): + img = tmp_path / "img.png" + doc = tmp_path / "doc.pdf" + img.write_bytes(PNG_BYTES) + doc.write_bytes(b"%PDF fake") + content = gw._build_messages([img, doc], None)[0]["content"] + assert [p["type"] for p in content] == ["image_url", "document_url"] + + def test_parse_extra_json_and_string(self): + parsed = gw._parse_extra(["temperature=0", "max_tokens=4096", "label=hello"]) + assert parsed == {"temperature": 0, "max_tokens": 4096, "label": "hello"} + + def test_parse_extra_invalid(self): + with pytest.raises(Exception): + gw._parse_extra(["nonsense"]) + + def test_parse_response_format_shorthands(self): + assert gw._parse_response_format("text") == {"type": "text"} + assert gw._parse_response_format("json") == {"type": "json_object"} + assert gw._parse_response_format("json_object") == {"type": "json_object"} + + def test_parse_response_format_json_object(self): + schema = '{"type": "json_schema", "json_schema": {"name": "x"}}' + assert gw._parse_response_format(schema) == { + "type": "json_schema", + "json_schema": {"name": "x"}, + } + + def test_parse_response_format_invalid(self): + with pytest.raises(Exception): + gw._parse_response_format("yaml") + with pytest.raises(Exception): + gw._parse_response_format('{"no": "type key"}') + + def test_openai_create_params_introspection(self): + params = gw._openai_create_params() + # Standard OpenAI fields are accepted by create() ... + assert {"temperature", "max_tokens", "stream", "extra_body"} <= params + # ... gateway-specific ones are not, and must ride in extra_body. + assert not ({"method", "method_params", "document_dpi"} & params) + + def test_openai_create_params_missing_dep_raises_dependency_error( + self, monkeypatch + ): + # A missing openai package must surface DependencyError (with install + # hints), not a raw ImportError. + from vlmrun.client.exceptions import DependencyError + + def _raise(): + raise DependencyError( + message="OpenAI SDK is not installed", + suggestion="pip install openai", + error_type="missing_dependency", + ) + + monkeypatch.setattr(gw, "_require_openai", _raise) + gw._openai_create_params.cache_clear() + with pytest.raises(DependencyError): + gw._openai_create_params() + gw._openai_create_params.cache_clear() + + def test_split_create_kwargs_routes_gateway_fields(self): + kwargs, body = gw._split_create_kwargs( + {"temperature": 0, "method": "ocr", "document_dpi": 200} + ) + assert kwargs == {"temperature": 0} + assert body == {"method": "ocr", "document_dpi": 200} + + def test_split_create_kwargs_merges_explicit_extra_body(self): + kwargs, body = gw._split_create_kwargs( + {"extra_body": {"method": "ocr", "document_dpi": 72}, "method": "markdown"} + ) + assert kwargs == {} + # Routed keys win over the explicit extra_body payload. + assert body == {"method": "markdown", "document_dpi": 72} + + def test_renderable_plain_text_for_markup_and_json(self): + # OCR output is -wrapped; Markdown would render it as HTML + # and drop it entirely. + assert isinstance(gw._renderable('hi'), Text) + assert isinstance(gw._renderable('{"text": "hi"}'), Text) + assert isinstance(gw._renderable("\n hi"), Text) + + def test_renderable_markdown_for_prose(self): + assert isinstance(gw._renderable("# Heading\n\nsome text"), Markdown) + + def test_format_cost(self): + assert gw._format_cost(0.001508) == "$0.001508" + assert gw._format_cost(0.0482) == "$0.0482" + assert gw._format_cost(3.26e-06) == "$0.000003" + assert gw._format_cost(1e-9) == "<$0.000001" # real but tiny, never "$0" + assert gw._format_cost(0) == "$0" + assert gw._format_cost(None) is None + assert gw._format_cost("nan-ish") is None + + def test_content_error_detects_error_payload(self): + assert ( + gw._content_error('{"error": "Unknown method \'zzz\'"}') + == "Unknown method 'zzz'" + ) + + def test_content_error_ignores_normal_output(self): + # OCR JSON lines, document-wrapped text, and prose are not errors. + assert gw._content_error('{"text": "Arizona", "score": 1.0}') is None + assert gw._content_error('hi') is None + assert gw._content_error("plain transcript text") is None + # An object that merely contains an error key alongside data is not it. + assert gw._content_error('{"error": "x", "text": "y"}') is None + + def test_format_methods_marks_default(self): + out = gw._format_methods( + {"methods": ["ocr", "detect"], "default_method": "ocr"} + ) + assert "[bold]ocr[/bold]*" in out + assert "detect" in out and "detect*" not in out + + def test_format_methods_empty(self): + assert gw._format_methods({}) == "-" + + def test_format_inputs_strips_url_suffix(self): + out = gw._format_inputs( + {"capabilities": {"supported_input_types": ["text", "image_url"]}} + ) + assert out == "text, image" + + def test_parse_extra_body_help_splits_json_and_prose(self): + examples, notes = gw._parse_extra_body_help( + '{"method":"ocr"} | {"method":"ocr","method_params":{"lang":"en"}}' + " | document_url PDF (markdown per page)" + ) + assert examples == [ + {"method": "ocr"}, + {"method": "ocr", "method_params": {"lang": "en"}}, + ] + assert notes == ["document_url PDF (markdown per page)"] + + def test_parse_extra_body_help_empty(self): + assert gw._parse_extra_body_help("") == ([], []) + + def test_example_command_renders_method_and_params(self): + cmd = gw._example_command( + "pp-ocrv6", {"method": "ocr", "method_params": {"lang": "en"}}, "doc.pdf" + ) + assert cmd == ( + "vlmrun gw chat doc.pdf -m pp-ocrv6 --method ocr " + '--method-params \'{"lang": "en"}\'' + ) + + def test_example_command_routes_other_fields_to_extra(self): + cmd = gw._example_command("glm-ocr", {"document_dpi": 200}, "doc.pdf") + assert cmd == "vlmrun gw chat doc.pdf -m glm-ocr -e document_dpi=200" + + def test_sample_input_prefers_document(self): + assert ( + gw._sample_input( + { + "capabilities": { + "supported_input_types": ["image_url", "document_url"] + } + } + ) + == "doc.pdf" + ) + assert ( + gw._sample_input({"capabilities": {"supported_input_types": ["image_url"]}}) + == "img.jpg" + ) + + +# --------------------------------------------------------------------------- +# CLI commands +# --------------------------------------------------------------------------- + + +class TestGatewayCLI: + @pytest.mark.parametrize("alias", ["gw", "gateway"]) + def test_both_aliases_registered(self, runner, patched_cli, alias): + result = runner.invoke(app, [alias, "health"]) + assert result.exit_code == 0 + assert "healthy" in result.stdout.lower() + + def test_health_ok(self, runner, patched_cli): + result = runner.invoke(app, ["gw", "health"]) + assert result.exit_code == 0 + assert "healthy" in result.stdout.lower() + + def test_health_unreachable(self, runner, patched_cli): + patched_cli["healthy"] = False + result = runner.invoke(app, ["gw", "health"]) + assert result.exit_code == 1 + assert "unreachable" in result.stdout.lower() + + def test_models_table(self, runner, patched_cli): + result = runner.invoke(app, ["gw", "models"]) + assert result.exit_code == 0 + assert "zai-org/glm-ocr" in result.stdout + assert "paddleocr/pp-ocrv6" in result.stdout + # Methods are listed, with the default marked. + assert "detect" in result.stdout + assert "ocr*" in result.stdout + + def test_models_json(self, runner, patched_cli): + result = runner.invoke(app, ["gw", "models", "--json"]) + assert result.exit_code == 0 + data = json.loads(result.stdout) + ids = {m["id"] for m in data} + assert ids == {"zai-org/glm-ocr", "paddleocr/pp-ocrv6"} + + def test_model_detail_by_alias(self, runner, patched_cli): + result = runner.invoke(app, ["gw", "models", "pp-ocrv6"]) + assert result.exit_code == 0 + assert "paddleocr/pp-ocrv6" in result.stdout + # Detail view is scoped to the one model. + assert "zai-org/glm-ocr" not in result.stdout + assert "detect" in result.stdout + + def test_model_detail_by_id_shows_notes(self, runner, patched_cli): + result = runner.invoke(app, ["gw", "models", "zai-org/glm-ocr"]) + assert result.exit_code == 0 + # Prose fragments of extra_body_help surface as notes. + assert "markdown per page" in result.stdout + + def test_model_detail_unknown_model(self, runner, patched_cli): + result = runner.invoke(app, ["gw", "models", "nope"]) + assert result.exit_code == 1 + assert "not found" in result.stdout.lower() + + def test_model_detail_json_emits_runnable_commands(self, runner, patched_cli): + result = runner.invoke(app, ["gw", "models", "pp-ocrv6", "--json"]) + assert result.exit_code == 0 + entry = json.loads(result.stdout) + # Detail --json is a single object, not the catalog list. + assert entry["id"] == "paddleocr/pp-ocrv6" + assert entry["default_method"] == "ocr" + assert entry["methods"] == ["ocr", "detect", "markdown"] + assert entry["commands"] == [ + "vlmrun gw chat doc.pdf -m paddleocr/pp-ocrv6 --method ocr " + '--method-params \'{"lang": "en", "score_threshold": 0.5}\'' + ] + + def test_models_list_json_still_returns_catalog(self, runner, patched_cli): + result = runner.invoke(app, ["gw", "models", "--json"]) + assert result.exit_code == 0 + assert isinstance(json.loads(result.stdout), list) + + def test_methods_command_is_gone(self, runner, patched_cli): + result = runner.invoke(app, ["gw", "methods"]) + assert result.exit_code != 0 + + +class TestGatewayEmbed: + def test_embed_text(self, runner, patched_cli): + result = runner.invoke(app, ["gw", "embed", "-t", "hello", "-m", "emb"]) + assert result.exit_code == 0, result.stdout + call = patched_cli["client"].gateway.embeddings.calls[-1] + # Plain text rides as a bare string. + assert call["input"] == ["hello"] + + def test_embed_image_nests_content_parts(self, runner, patched_cli, tmp_path): + """Each item must be a *list* of parts; a flat parts list is rejected.""" + img = tmp_path / "a.png" + img.write_bytes(PNG_BYTES) + result = runner.invoke(app, ["gw", "embed", str(img), "-m", "emb"]) + assert result.exit_code == 0, result.stdout + call = patched_cli["client"].gateway.embeddings.calls[-1] + assert len(call["input"]) == 1 + item = call["input"][0] + assert isinstance(item, list) + assert item[0]["type"] == "image_url" + assert item[0]["image_url"]["url"].startswith("data:image/png;base64,") + + def test_embed_video_uses_video_url_part(self, runner, patched_cli, tmp_path): + vid = tmp_path / "clip.mp4" + vid.write_bytes(MP4_BYTES) + result = runner.invoke(app, ["gw", "embed", str(vid), "-m", "emb"]) + assert result.exit_code == 0, result.stdout + call = patched_cli["client"].gateway.embeddings.calls[-1] + assert call["input"][0][0]["type"] == "video_url" + + def test_embed_files_and_text_are_separate_vectors( + self, runner, patched_cli, tmp_path + ): + a, b = tmp_path / "a.png", tmp_path / "b.png" + a.write_bytes(PNG_BYTES) + b.write_bytes(PNG_BYTES) + result = runner.invoke( + app, ["gw", "embed", str(a), str(b), "-t", "cap", "-m", "emb"] + ) + assert result.exit_code == 0, result.stdout + call = patched_cli["client"].gateway.embeddings.calls[-1] + # Three independent items — batching two images into one item 500s. + assert len(call["input"]) == 3 + assert call["input"][2] == "cap" + + def test_embed_join_combines_into_one_vector(self, runner, patched_cli, tmp_path): + img = tmp_path / "a.png" + img.write_bytes(PNG_BYTES) + result = runner.invoke( + app, ["gw", "embed", str(img), "-t", "cap", "--join", "-m", "emb"] + ) + assert result.exit_code == 0, result.stdout + call = patched_cli["client"].gateway.embeddings.calls[-1] + assert len(call["input"]) == 1 + assert [p["type"] for p in call["input"][0]] == ["image_url", "text"] + + def test_embed_join_rejects_multiple_files(self, runner, patched_cli, tmp_path): + a, b = tmp_path / "a.png", tmp_path / "b.png" + a.write_bytes(PNG_BYTES) + b.write_bytes(PNG_BYTES) + result = runner.invoke( + app, ["gw", "embed", str(a), str(b), "--join", "-m", "emb"] + ) + assert result.exit_code == 1 + assert "at most one file" in result.stdout + + def test_embed_requires_input(self, runner, patched_cli): + result = runner.invoke(app, ["gw", "embed", "-m", "emb"]) + assert result.exit_code == 1 + assert "at least one file or --text" in result.stdout + + def test_embed_rejects_non_image_file(self, runner, patched_cli, tmp_path): + # A PDF (or any non-image/video) must be rejected client-side rather + # than sent as a mislabelled image_url the gateway would fail on. + doc = tmp_path / "doc.pdf" + doc.write_bytes(b"%PDF-1.4 fake") + result = runner.invoke(app, ["gw", "embed", str(doc), "-m", "emb"]) + assert result.exit_code == 1 + assert "images and video only" in result.stdout + assert not patched_cli["client"].gateway.embeddings.calls + + def test_embed_dimensions_passed_through(self, runner, patched_cli): + result = runner.invoke( + app, ["gw", "embed", "-t", "hi", "-m", "emb", "--dimensions", "64"] + ) + assert result.exit_code == 0, result.stdout + assert patched_cli["client"].gateway.embeddings.calls[-1]["dimensions"] == 64 + + def test_embed_json(self, runner, patched_cli): + result = runner.invoke(app, ["gw", "embed", "-t", "hi", "-m", "emb", "--json"]) + assert result.exit_code == 0, result.stdout + assert len(json.loads(result.stdout)["data"][0]["embedding"]) == 4 + + +class TestGatewayTranscribe: + def test_transcribe_file(self, runner, patched_cli, tmp_path): + audio = tmp_path / "clip.mp3" + audio.write_bytes(b"fake-audio") + result = runner.invoke(app, ["gw", "transcribe", str(audio), "-m", "asr"]) + assert result.exit_code == 0, result.stdout + assert "hello from audio" in result.stdout + call = patched_cli["client"].gateway.transcriptions.calls[-1] + assert call["response_format"] == "json" + + def test_transcribe_format_and_hints(self, runner, patched_cli, tmp_path): + audio = tmp_path / "clip.mp3" + audio.write_bytes(b"fake-audio") + result = runner.invoke( + app, + [ + "gw", + "transcribe", + str(audio), + "-m", + "asr", + "-f", + "srt", + "-l", + "en", + "-p", + "nouns", + ], + ) + assert result.exit_code == 0, result.stdout + call = patched_cli["client"].gateway.transcriptions.calls[-1] + assert call["response_format"] == "srt" + assert call["language"] == "en" + assert call["prompt"] == "nouns" + + def test_transcribe_url_rides_in_extra_body(self, runner, patched_cli): + result = runner.invoke( + app, ["gw", "transcribe", "--url", "https://x/a.mp3", "-m", "asr"] + ) + assert result.exit_code == 0, result.stdout + call = patched_cli["client"].gateway.transcriptions.calls[-1] + assert call["extra_body"] == {"url": "https://x/a.mp3"} + + def test_transcribe_requires_input(self, runner, patched_cli): + result = runner.invoke(app, ["gw", "transcribe", "-m", "asr"]) + assert result.exit_code == 1 + assert "audio file or --url" in result.stdout + + def test_transcribe_rejects_file_and_url(self, runner, patched_cli, tmp_path): + audio = tmp_path / "clip.mp3" + audio.write_bytes(b"fake-audio") + result = runner.invoke( + app, + ["gw", "transcribe", str(audio), "--url", "https://x/a.mp3", "-m", "asr"], + ) + assert result.exit_code == 1 + assert "not both" in result.stdout + + def test_transcribe_bad_format(self, runner, patched_cli, tmp_path): + audio = tmp_path / "clip.mp3" + audio.write_bytes(b"fake-audio") + result = runner.invoke( + app, ["gw", "transcribe", str(audio), "-m", "asr", "-f", "bogus"] + ) + assert result.exit_code == 1 + assert "unknown --format" in result.stdout.lower() + + def test_chat_requires_file(self, runner, patched_cli): + result = runner.invoke(app, ["gw", "chat", "-m", "glm-ocr"]) + assert result.exit_code == 1 + assert "at least one input file" in result.stdout.lower() + + def test_chat_with_file_json(self, runner, patched_cli, tmp_path): + f = tmp_path / "doc.pdf" + f.write_bytes(b"%PDF fake") + result = runner.invoke( + app, + ["gw", "chat", str(f), "-m", "glm-ocr", "--no-stream", "--json"], + ) + assert result.exit_code == 0, result.stdout + out = json.loads(result.stdout) + assert out["model"] == "glm-ocr" + assert out["content"] == "Hello world" + + def test_chat_streaming_default(self, runner, patched_cli, tmp_path): + f = tmp_path / "img.png" + f.write_bytes(b"fakepng") + result = runner.invoke(app, ["gw", "chat", str(f), "-m", "paddle-ocrv6"]) + assert result.exit_code == 0, result.stdout + assert "Hello world" in result.stdout + + def test_chat_sends_document_and_image_urls(self, runner, patched_cli, tmp_path): + pdf = tmp_path / "doc.pdf" + img = tmp_path / "scan.png" + pdf.write_bytes(b"%PDF fake") + img.write_bytes(PNG_BYTES) + result = runner.invoke( + app, + ["gw", "chat", str(pdf), str(img), "-m", "glm-ocr", "--no-stream"], + ) + assert result.exit_code == 0, result.stdout + call = patched_cli["client"].gateway.completions.calls[-1] + content = call["messages"][0]["content"] + assert content[0]["type"] == "document_url" + assert content[0]["document_url"]["url"].startswith( + "data:application/pdf;base64," + ) + assert content[1]["type"] == "image_url" + assert content[1]["image_url"]["url"].startswith("data:image/png;base64,") + + def test_chat_image_only_does_not_stream(self, runner, patched_cli, tmp_path): + """Images must not stream: the gateway returns a non-SSE body that an + SSE reader drains to empty.""" + img = tmp_path / "img.png" + img.write_bytes(PNG_BYTES) + result = runner.invoke(app, ["gw", "chat", str(img), "-m", "pp-ocrv6"]) + assert result.exit_code == 0, result.stdout + assert patched_cli["client"].gateway.completions.calls[-1]["stream"] is False + assert "Hello world" in result.stdout + + def test_chat_text_only_does_not_stream(self, runner, patched_cli): + result = runner.invoke( + app, ["gw", "chat", "-m", "qwen/qwen3.5-0.8b", "-p", "hi"] + ) + assert result.exit_code == 0, result.stdout + assert patched_cli["client"].gateway.completions.calls[-1]["stream"] is False + + def test_chat_document_streams_by_default(self, runner, patched_cli, tmp_path): + pdf = tmp_path / "doc.pdf" + pdf.write_bytes(b"%PDF fake") + result = runner.invoke(app, ["gw", "chat", str(pdf), "-m", "glm-ocr"]) + assert result.exit_code == 0, result.stdout + assert patched_cli["client"].gateway.completions.calls[-1]["stream"] is True + + def test_chat_document_no_stream_flag_respected( + self, runner, patched_cli, tmp_path + ): + pdf = tmp_path / "doc.pdf" + pdf.write_bytes(b"%PDF fake") + result = runner.invoke(app, ["gw", "chat", str(pdf), "-m", "glm-ocr", "-ns"]) + assert result.exit_code == 0, result.stdout + assert patched_cli["client"].gateway.completions.calls[-1]["stream"] is False + + def test_chat_multiple_files_and_extra(self, runner, patched_cli, tmp_path): + f1 = tmp_path / "a.pdf" + f2 = tmp_path / "b.pdf" + f1.write_bytes(b"%PDF a") + f2.write_bytes(b"%PDF b") + result = runner.invoke( + app, + [ + "gw", + "chat", + str(f1), + str(f2), + "-m", + "paddle-ocrv6", + "-e", + "temperature=0", + "--no-stream", + "--json", + ], + ) + assert result.exit_code == 0, result.stdout + out = json.loads(result.stdout) + assert out["content"] == "Hello world" + + def test_chat_method_and_params_sent_via_extra_body( + self, runner, patched_cli, tmp_path + ): + f = tmp_path / "img.png" + f.write_bytes(b"fakepng") + result = runner.invoke( + app, + [ + "gateway", + "chat", + str(f), + "-m", + "pp-ocrv6", + "--method", + "ocr", + "--method-params", + '{"lang": "en", "score_threshold": 0.9}', + "--no-stream", + ], + ) + assert result.exit_code == 0, result.stdout + call = patched_cli["client"].gateway.completions.calls[-1] + assert call["extra_body"] == { + "method": "ocr", + "method_params": {"lang": "en", "score_threshold": 0.9}, + } + # method must not leak into create()'s own kwargs. + assert "method" not in call + + def test_chat_response_format_sent_as_top_level_kwarg( + self, runner, patched_cli, tmp_path + ): + f = tmp_path / "img.png" + f.write_bytes(b"fakepng") + result = runner.invoke( + app, + [ + "gw", + "chat", + str(f), + "-m", + "pp-ocrv6", + "--response-format", + "json_object", + "--no-stream", + ], + ) + assert result.exit_code == 0, result.stdout + call = patched_cli["client"].gateway.completions.calls[-1] + # A standard OpenAI field, so it rides top-level, not in extra_body. + assert call["response_format"] == {"type": "json_object"} + assert "response_format" not in call.get("extra_body", {}) + + def test_chat_response_format_invalid(self, runner, patched_cli, tmp_path): + f = tmp_path / "img.png" + f.write_bytes(b"fakepng") + result = runner.invoke( + app, + ["gw", "chat", str(f), "-m", "pp-ocrv6", "--response-format", "yaml"], + ) + assert result.exit_code == 1 + assert "response-format" in result.stdout.lower() + + def test_chat_extra_routes_gateway_field_to_extra_body( + self, runner, patched_cli, tmp_path + ): + f = tmp_path / "img.png" + f.write_bytes(b"fakepng") + result = runner.invoke( + app, + [ + "gw", + "chat", + str(f), + "-m", + "pp-ocrv6", + "-e", + "document_dpi=200", + "-e", + "temperature=0", + "--no-stream", + ], + ) + assert result.exit_code == 0, result.stdout + call = patched_cli["client"].gateway.completions.calls[-1] + assert call["extra_body"] == {"document_dpi": 200} + assert call["temperature"] == 0 + + def test_chat_no_extra_body_when_unused(self, runner, patched_cli, tmp_path): + f = tmp_path / "img.png" + f.write_bytes(b"fakepng") + result = runner.invoke( + app, ["gw", "chat", str(f), "-m", "pp-ocrv6", "--no-stream"] + ) + assert result.exit_code == 0, result.stdout + assert "extra_body" not in patched_cli["client"].gateway.completions.calls[-1] + + def test_chat_invalid_method_params(self, runner, patched_cli, tmp_path): + f = tmp_path / "img.png" + f.write_bytes(b"fakepng") + result = runner.invoke( + app, ["gw", "chat", str(f), "-m", "pp-ocrv6", "--method-params", "nope"] + ) + assert result.exit_code == 1 + assert "valid json" in result.stdout.lower() + + def test_chat_method_params_must_be_object(self, runner, patched_cli, tmp_path): + f = tmp_path / "img.png" + f.write_bytes(b"fakepng") + result = runner.invoke( + app, ["gw", "chat", str(f), "-m", "pp-ocrv6", "--method-params", "[1, 2]"] + ) + assert result.exit_code == 1 + assert "json object" in result.stdout.lower() + + def test_chat_renders_document_wrapped_ocr_output( + self, runner, patched_cli, tmp_path + ): + """OCR output must survive rendering (regression: Markdown ate the tags).""" + patched_cli["content"] = ( + '\n\nDLN B58471293\n\n' + ) + f = tmp_path / "img.png" + f.write_bytes(b"fakepng") + result = runner.invoke( + app, ["gw", "chat", str(f), "-m", "pp-ocrv6", "--no-stream"] + ) + assert result.exit_code == 0, result.stdout + assert "DLN B58471293" in result.stdout + assert " |----------|-------------| | `VLMRUN_API_KEY` | Your VLM Run API key (required) | | `VLMRUN_CACHE_DIR` | Custom cache directory (default: `~/.vlmrun/cache/artifacts`) | +| `VLMRUN_GATEWAY_URL` | Override the model gateway base URL (default: `https://gateway.vlm.run/v1`) | ## How It Works @@ -218,6 +219,68 @@ vlmrun models list vlmrun fine-tuning create --model base_model --training-file training_file_id ``` +### Gateway (`vlmrun gateway` / `vlmrun gw`) - OpenAI-compatible models + +The gateway (`https://gateway.vlm.run/v1`) exposes third-party OCR, +vision-language, embedding and transcription models (e.g. `zai-org/glm-ocr`, +`paddleocr/pp-ocrv6`, `qwen/qwen3.5-0.8b`) through an OpenAI-compatible API, +authenticated with the same `VLMRUN_API_KEY`. `vlmrun gateway` and `vlmrun gw` +are the same command. + +Unlike `vlmrun chat` (which calls the Orion agent), the gateway is a raw +passthrough to the underlying models: input files are sent inline as base64 +`data:` URLs — documents as `document_url` content parts, images as +`image_url`, and `file_url` as a fallback for anything else. **Most chat/OCR +models do not accept text-only input**, so at least one file is required. + +Model ids are the full `/` shown by `vlmrun gw models`; the short +aliases listed there (e.g. `glm-ocr`, `pp-ocrv6`) also work. + +```bash +# Health check +vlmrun gw health + +# List models (task + methods); detail one model's methods, params and examples +vlmrun gw models +vlmrun gw models paddleocr/pp-ocrv6 +vlmrun gw models --json + +# Parse a document (PDF -> text/markdown) +vlmrun gw chat document.pdf -m zai-org/glm-ocr + +# Multiple documents / OCR an image +vlmrun gw chat doc1.pdf doc2.pdf -m paddleocr/pp-ocrv6 +vlmrun gw chat scan.jpg -m paddleocr/pp-ocrv6 + +# Select a model method (see `vlmrun gw models `) +vlmrun gw chat scan.jpg -m paddleocr/pp-ocrv6 --method detect +vlmrun gw chat scan.jpg -m paddleocr/pp-ocrv6 --method ocr \ + --method-params '{"lang": "en", "score_threshold": 0.5}' + +# Prompt a model that supports text input +vlmrun gw chat image.jpg -p "describe this image" -m qwen/qwen3.5-0.8b + +# Forward extra completion kwargs as key=value (JSON-parsed) +vlmrun gw chat document.pdf -m zai-org/glm-ocr -e temperature=0 -e max_tokens=4096 + +# Embed text, images or video (each input -> one vector; --join for a joint one) +vlmrun gw embed -t "a blue parrot" -m qwen/qwen3-vl-embedding-2b +vlmrun gw embed photo.jpg -m qwen/qwen3-vl-embedding-2b +vlmrun gw embed photo.jpg -t "caption" --join -m qwen/qwen3-vl-embedding-2b + +# Transcribe audio, or a video's audio track +vlmrun gw transcribe clip.mp3 -m nvidia/parakeet-tdt-0.6b-v3 +vlmrun gw transcribe clip.mp4 -m nvidia/parakeet-tdt-0.6b-v3 -f srt +``` + +| Command | Description | +|---------|-------------| +| `vlmrun gw health` | Check gateway reachability | +| `vlmrun gw models [MODEL]` | List models (task + methods), or detail one model's methods, params and example commands (`--json` for raw output) | +| `vlmrun gw chat FILES... -m MODEL` | Run a model over one or more files (`--method`/`--method-params`, `-p` prompt, `-e key=value` extras, `--no-stream`, `--json`) | +| `vlmrun gw embed [FILES...] -m MODEL` | Embed text (`-t`), images or video (`--join`, `--dimensions`, `--json`) | +| `vlmrun gw transcribe AUDIO -m MODEL` | Transcribe audio or a video's audio track (`-f` format, `--language`, `--prompt`, `--url`, `--json`) | + ### Predictions ```bash diff --git a/vlmrun/cli/_cli/gateway.py b/vlmrun/cli/_cli/gateway.py new file mode 100644 index 0000000..4c6fb3c --- /dev/null +++ b/vlmrun/cli/_cli/gateway.py @@ -0,0 +1,1127 @@ +"""Gateway commands for the VLM Run CLI. + +Talk to OpenAI-compatible OCR / VLM models hosted behind the VLM Run gateway +(``https://gateway.vlm.run/v1``), authenticating with the same +``VLMRUN_API_KEY`` used everywhere else. + +Unlike ``vlmrun chat`` (which uploads to the Files API and calls the Orion +agent), the gateway is a raw passthrough to third-party models. Inputs are +inlined as base64 ``data:`` URLs in the message content: documents use +``document_url`` content parts, images use ``image_url``, and ``file_url`` is +the fallback for anything unidentifiable. Most models (especially OCR models +such as ``zai-org/glm-ocr`` and ``paddleocr/pp-ocrv6``) do not accept +text-only input. + +Commands: ``health``, ``models`` (list or detail one model), ``chat``, +``embed`` (embeddings) and ``transcribe`` (audio transcriptions). +""" + +from __future__ import annotations + +import base64 +import inspect +import json +import mimetypes +import time +from functools import lru_cache +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +import typer +from rich.console import Console +from rich.markdown import Markdown +from rich.panel import Panel +from rich.table import Table +from rich.text import Text +from rich.tree import Tree +from rich import box + +from vlmrun.client import VLMRun +from vlmrun.client.gateway import _require_openai +from vlmrun.cli._cli.chat import ( + TimedStatus, + format_file_size, + handle_api_errors, +) +from vlmrun.constants import SUPPORTED_DOCUMENT_FILETYPES + +console = Console() + +CHAT_HELP = """Run OCR / VLM models on the VLM Run gateway. + +\b +EXAMPLES: + vlmrun gw chat doc.pdf -m zai-org/glm-ocr + vlmrun gw chat a.pdf b.pdf -m paddleocr/pp-ocrv6 + vlmrun gw chat img.jpg -m paddleocr/pp-ocrv6 + vlmrun gw chat img.jpg -p "describe this image" -m qwen/qwen3.5-0.8b + vlmrun gw chat doc.pdf -m zai-org/glm-ocr -e temperature=0 -e max_tokens=4096 + +\b +METHODS: + Each model exposes methods with a default. Run `vlmrun gw models ` for + its methods, params, and copy-pasteable example commands. + vlmrun gw chat img.jpg -m paddleocr/pp-ocrv6 --method detect + vlmrun gw chat img.jpg -m paddleocr/pp-ocrv6 --method ocr \\ + --method-params '{"lang": "en", "score_threshold": 0.5}' + +\b +NOTES: + Model ids are the full `/` shown by `vlmrun gw models`; short + aliases (e.g. `glm-ocr`) also work. + Most gateway models (e.g. OCR models) require at least one input file and do + not accept text-only prompts. Use -p only for models that support it. +""" + +GATEWAY_HELP = """OCR, VLM, embedding and transcription models on the VLM Run gateway. + +An OpenAI-compatible passthrough to third-party models, authenticated with the +same VLMRUN_API_KEY as the rest of the CLI. `vlmrun gateway` and `vlmrun gw` +are the same command. + +\b +Start here: + vlmrun gw models See what is available (task + methods per model) + vlmrun gw models Methods, params and copy-pasteable examples +""" + +app = typer.Typer( + help=GATEWAY_HELP, + add_completion=False, + no_args_is_help=True, +) + + +# Magic-byte signatures, checked before the filename extension. Extensions lie +# (a .jpg that is really WebP is common), and the gateway trusts the media type +# we declare in the data URL, so a wrong one makes it misroute the file. +_MAGIC_SIGNATURES: Tuple[Tuple[bytes, str], ...] = ( + (b"\xff\xd8\xff", "image/jpeg"), + (b"\x89PNG\r\n\x1a\n", "image/png"), + (b"GIF87a", "image/gif"), + (b"GIF89a", "image/gif"), + (b"BM", "image/bmp"), + (b"II*\x00", "image/tiff"), + (b"MM\x00*", "image/tiff"), + (b"%PDF", "application/pdf"), +) + + +def _sniff_mime(data: bytes) -> Optional[str]: + """MIME type from a file's magic bytes, or None if unrecognized.""" + if data[:4] == b"RIFF" and data[8:12] == b"WEBP": + return "image/webp" + if data[:4] == b"RIFF" and data[8:12] == b"AVI ": + return "video/x-msvideo" + if data[4:8] == b"ftyp": + return "video/mp4" + for signature, mime in _MAGIC_SIGNATURES: + if data.startswith(signature): + return mime + return None + + +def _guess_mime(path: Path, data: Optional[bytes] = None) -> str: + """Best-effort MIME type for a local file, preferring its actual content.""" + if data is None: + try: + with path.open("rb") as fh: + data = fh.read(16) + except OSError: + data = b"" + sniffed = _sniff_mime(data) + if sniffed: + return sniffed + mime, _ = mimetypes.guess_type(str(path)) + return mime or "application/octet-stream" + + +def _content_part_type(path: Path, mime: Optional[str] = None) -> str: + """Content-part type for a file. + + Documents (``.pdf``, ``.doc``, ``.docx``) are sent as ``document_url`` and + images as ``image_url``. ``file_url`` is the fallback for anything we cannot + identify: the gateway routes it through its document/PDF path, which fails + outright on a plain image. + """ + if path.suffix.lower() in SUPPORTED_DOCUMENT_FILETYPES: + return "document_url" + mime = mime or _guess_mime(path) + if mime == "application/pdf": + return "document_url" + if mime.startswith("image/"): + return "image_url" + return "file_url" + + +def _encode_file_part(path: Path) -> Dict[str, Any]: + """Encode a local file as a gateway data-URL content part. + + The gateway accepts files inline as base64 ``data:`` URLs under a + ``document_url`` (documents), ``image_url`` (images) or ``file_url`` + (anything else) content part. + """ + data = path.read_bytes() + b64 = base64.b64encode(data).decode("ascii") + mime = _guess_mime(path, data) + key = _content_part_type(path, mime) + return { + "type": key, + key: {"url": f"data:{mime};base64,{b64}"}, + } + + +def _parse_response_format(value: str) -> Dict[str, Any]: + """Parse ``--response-format`` into an OpenAI ``response_format`` object. + + Accepts the shorthands ``text`` and ``json_object`` (with ``json`` as an + alias), or a full JSON object for advanced cases (e.g. ``json_schema``). + Exits with a clear message on anything else. + """ + stripped = value.strip() + aliases = { + "text": {"type": "text"}, + "json": {"type": "json_object"}, + "json_object": {"type": "json_object"}, + } + if stripped in aliases: + return aliases[stripped] + if stripped.startswith("{"): + try: + parsed = json.loads(stripped) + except (json.JSONDecodeError, ValueError) as e: + console.print(f"[red]Error:[/] --response-format must be valid JSON: {e}") + raise typer.Exit(1) + if not isinstance(parsed, dict) or "type" not in parsed: + console.print( + "[red]Error:[/] --response-format JSON must be an object with a " + "'type' key, e.g. '{\"type\":\"json_object\"}'." + ) + raise typer.Exit(1) + return parsed + console.print( + f"[red]Error:[/] Unknown --response-format '{value}'. Use 'text', " + "'json_object', or a JSON object with a 'type' key." + ) + raise typer.Exit(1) + + +def _build_messages(files: List[Path], prompt: Optional[str]) -> List[Dict[str, Any]]: + """Build a single OpenAI-style user message from files + optional prompt.""" + content: List[Dict[str, Any]] = [_encode_file_part(f) for f in files] + if prompt: + content.append({"type": "text", "text": prompt}) + return [{"role": "user", "content": content}] + + +def _parse_extra(pairs: Optional[List[str]]) -> Dict[str, Any]: + """Parse repeatable ``key=value`` options into create() kwargs. + + Values are parsed as JSON when possible (so ``temperature=0.2`` becomes a + float and ``stop=["\\n"]`` becomes a list), else kept as strings. + """ + extra: Dict[str, Any] = {} + for pair in pairs or []: + if "=" not in pair: + console.print( + f"[red]Error:[/] Invalid --extra value '{pair}'. Use key=value." + ) + raise typer.Exit(1) + key, _, raw = pair.partition("=") + key = key.strip() + try: + value: Any = json.loads(raw) + except (json.JSONDecodeError, ValueError): + value = raw + extra[key] = value + return extra + + +@lru_cache(maxsize=1) +def _openai_create_params() -> frozenset: + """Parameter names accepted by the OpenAI SDK's ``chat.completions.create()``. + + Introspected rather than hardcoded so the split below tracks whatever + version of the ``openai`` package is installed. + """ + # Route a missing dependency through the SDK's DependencyError (with install + # hints) instead of surfacing a raw ImportError. Reuses _require_openai so + # the install message lives in one place. + _require_openai() + from openai.resources.chat.completions import Completions + + sig = inspect.signature(Completions.create) + names = { + p.name + for p in sig.parameters.values() + if p.kind in (p.KEYWORD_ONLY, p.POSITIONAL_OR_KEYWORD) + } + return frozenset(names - {"self"}) + + +def _split_create_kwargs( + extra: Dict[str, Any], +) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """Split user-supplied kwargs into OpenAI create() kwargs and extra_body. + + ``create()`` has an explicit signature and rejects unknown keywords, so + gateway-specific fields (``method``, ``document_dpi``, ...) must travel in + ``extra_body`` to reach the server as top-level request-body fields. + """ + known = _openai_create_params() + kwargs: Dict[str, Any] = {} + body: Dict[str, Any] = {} + for key, value in extra.items(): + if key in known: + kwargs[key] = value + else: + body[key] = value + + # An explicit -e extra_body={...} merges with the routed fields. + explicit = kwargs.pop("extra_body", None) + if isinstance(explicit, dict): + body = {**explicit, **body} + return kwargs, body + + +def _format_methods(model: Dict[str, Any]) -> str: + """Render a model's methods, marking the default with ``*``.""" + methods = model.get("methods") or [] + default = model.get("default_method") or "" + if not methods: + return "-" + return ", ".join(f"[bold]{m}[/bold]*" if m == default else m for m in methods) + + +def _format_inputs(model: Dict[str, Any]) -> str: + """Render the input types a model accepts, minus the ``_url`` noise.""" + caps = model.get("capabilities") or {} + types = caps.get("supported_input_types") or [] + if not types: + return "-" + return ", ".join(t.removesuffix("_url") for t in types) + + +def _model_dicts(client: VLMRun) -> List[Dict[str, Any]]: + """Fetch gateway models and normalize them to plain dicts.""" + with handle_api_errors(): + model_objs = client.gateway.models() + + rows: List[Dict[str, Any]] = [] + for m in model_objs: + if hasattr(m, "model_dump"): + rows.append(m.model_dump()) + elif isinstance(m, dict): + rows.append(m) + else: + rows.append({"id": str(m)}) + return sorted(rows, key=lambda r: str(r.get("id", ""))) + + +def _parse_extra_body_help(help_text: str) -> Tuple[List[Dict[str, Any]], List[str]]: + """Split a model's ``extra_body_help`` into JSON examples and prose notes. + + The gateway packs this field with ``|``-separated fragments, some of which + are JSON extra_body payloads (e.g. ``{"method":"ocr"}``) and some of which + are free-text hints. Parsing it keeps the examples we print in sync with + whatever the gateway currently advertises. + """ + examples: List[Dict[str, Any]] = [] + notes: List[str] = [] + for fragment in (help_text or "").split("|"): + fragment = fragment.strip() + if not fragment: + continue + try: + parsed = json.loads(fragment) + except (json.JSONDecodeError, ValueError): + notes.append(fragment) + continue + if isinstance(parsed, dict): + examples.append(parsed) + else: + notes.append(fragment) + return examples, notes + + +def _example_command(model_id: str, payload: Dict[str, Any], sample: str) -> str: + """Render an extra_body example as a runnable `vlmrun gw chat` command.""" + parts = [f"vlmrun gw chat {sample} -m {model_id}"] + method = payload.get("method") + if method: + parts.append(f"--method {method}") + params = payload.get("method_params") + if isinstance(params, dict): + parts.append(f"--method-params '{json.dumps(params)}'") + for key, value in payload.items(): + if key in ("method", "method_params"): + continue + parts.append(f"-e {key}={json.dumps(value)}") + return " ".join(parts) + + +def _sample_input(model: Dict[str, Any]) -> str: + """Pick a plausible sample filename for a model's example commands.""" + caps = model.get("capabilities") or {} + types = caps.get("supported_input_types") or [] + if "document_url" in types: + return "doc.pdf" + if "image_url" in types: + return "img.jpg" + if "video_url" in types: + return "clip.mp4" + return "input.bin" + + +@app.command() +def health(ctx: typer.Context) -> None: + """Check gateway health.""" + client: VLMRun = ctx.obj + with TimedStatus("Checking gateway...", console=console): + ok = client.gateway.health() + + if ok: + console.print( + Panel( + f"[green]Gateway is healthy[/green]\n[dim]{client.gateway.base_url}[/dim]", + title="[green]OK[/green]", + title_align="left", + border_style="green", + ) + ) + else: + console.print( + Panel( + f"[red]Gateway is unreachable[/red]\n[dim]{client.gateway.base_url}[/dim]", + title="[red]Unhealthy[/red]", + title_align="left", + border_style="red", + ) + ) + raise typer.Exit(1) + + +MODELS_HELP = """List gateway models, or detail one model. + +\b +EXAMPLES: + vlmrun gw models List every model with its methods. + vlmrun gw models paddleocr/pp-ocrv6 Methods, params and examples for one model. + vlmrun gw models --json Raw model catalog. +""" + + +def _model_detail(row: Dict[str, Any]) -> Panel: + """Render one model's methods, params, notes and example commands.""" + model_id = str(row.get("id", "-")) + examples, notes = _parse_extra_body_help(row.get("extra_body_help", "")) + sample = _sample_input(row) + + tree = Tree("", guide_style="dim", hide_root=True) + tree.add(f"[dim]task[/dim] {row.get('task', '-')}") + tree.add(f"[dim]methods[/dim] {_format_methods(row)}") + tree.add(f"[dim]inputs[/dim] {_format_inputs(row)}") + aliases = row.get("aliases") or [] + if aliases: + tree.add(f"[dim]aliases[/dim] {', '.join(aliases)}") + for note in notes: + tree.add(f"[dim]note[/dim] {note}") + if examples: + branch = tree.add("[dim]examples[/dim]") + for ex in examples: + branch.add(Text(_example_command(model_id, ex, sample), style="cyan")) + + return Panel( + tree, + title=f"[bold cyan]{model_id}[/bold cyan]", + title_align="left", + border_style="blue", + padding=(0, 1), + ) + + +def _model_detail_json(row: Dict[str, Any]) -> Dict[str, Any]: + """JSON form of a model's detail view, including runnable commands.""" + examples, notes = _parse_extra_body_help(row.get("extra_body_help", "")) + sample = _sample_input(row) + return { + "id": row.get("id"), + "aliases": row.get("aliases") or [], + "task": row.get("task"), + "methods": row.get("methods") or [], + "default_method": row.get("default_method") or None, + "supported_input_types": (row.get("capabilities") or {}).get( + "supported_input_types" + ) + or [], + "extra_body_examples": examples, + "notes": notes, + "commands": [ + _example_command(str(row.get("id")), ex, sample) for ex in examples + ], + } + + +@app.command(help=MODELS_HELP, context_settings={"max_content_width": 120}) +def models( + ctx: typer.Context, + model: Optional[str] = typer.Argument( + None, + help="Model id or alias. Shows that model's methods, params and examples.", + ), + output_json: bool = typer.Option(False, "--json", "-j", help="Output raw JSON."), +) -> None: + """List gateway models, or detail one model.""" + client: VLMRun = ctx.obj + rows = _model_dicts(client) + + if model: + wanted = model.strip() + match = [ + r + for r in rows + if wanted == str(r.get("id")) or wanted in (r.get("aliases") or []) + ] + if not match: + console.print( + f"[red]Error:[/] Model '{model}' not found on the gateway. " + "Run `vlmrun gw models` to list available models." + ) + raise typer.Exit(1) + if output_json: + print(json.dumps(_model_detail_json(match[0]), indent=2, default=str)) + return + console.print(_model_detail(match[0])) + return + + if output_json: + print(json.dumps(rows, indent=2, default=str)) + return + + table = Table( + show_header=True, + header_style="bold white", + box=box.SIMPLE_HEAVY, + padding=(0, 1), + ) + table.add_column("MODEL", style="bold cyan", no_wrap=True) + table.add_column("TASK", style="dim", no_wrap=True) + table.add_column("METHODS") + + for row in rows: + # Aliases are omitted here to keep method names from truncating at 80 + # columns; the per-model detail view lists them. + table.add_row( + str(row.get("id", "-")), + str(row.get("task", "-")), + _format_methods(row), + ) + + console.print( + Panel( + table, + title="[bold]Gateway Models[/bold]", + title_align="left", + subtitle=f"[dim]{len(rows)} model(s) · [bold]*[/bold] = default method · `vlmrun gw models ` for examples[/dim]", + subtitle_align="right", + border_style="blue", + padding=(0, 1), + ) + ) + + +@app.command(help=CHAT_HELP, context_settings={"max_content_width": 120}) +def chat( + ctx: typer.Context, + files: List[Path] = typer.Argument( + None, + help="Input document/image file(s) to process. Repeatable.", + exists=True, + readable=True, + ), + model: str = typer.Option( + ..., + "--model", + "-m", + help="Gateway model id, full / or alias (see `vlmrun gw models`).", + ), + prompt: Optional[str] = typer.Option( + None, + "--prompt", + "-p", + help="Optional text prompt (only for models that support text input).", + ), + method: Optional[str] = typer.Option( + None, + "--method", + "-M", + help="Model method, e.g. ocr, detect, markdown. Defaults to the model's default_method.", + ), + method_params: Optional[str] = typer.Option( + None, + "--method-params", + help='JSON object of method arguments, e.g. \'{"lang": "en"}\'.', + ), + response_format: Optional[str] = typer.Option( + None, + "--response-format", + help=( + "Ask the MODEL to constrain its output: 'text', 'json_object' (JSON " + 'mode), or a JSON object like \'{"type":"json_schema",...}\'. ' + "Sent to the gateway as `response_format`; not yet honored server-side. " + "(Distinct from --json, which formats the CLI's own output.)" + ), + ), + extra: Optional[List[str]] = typer.Option( + None, + "--extra", + "-e", + help="Extra create() kwarg as key=value (repeatable), e.g. -e temperature=0.", + ), + no_stream: bool = typer.Option( + False, "--no-stream", "-ns", help="Disable streaming." + ), + output_json: bool = typer.Option(False, "--json", "-j", help="Output raw JSON."), + timeout: Optional[float] = typer.Option( + None, "--timeout", help="Request timeout in seconds." + ), +) -> None: + """Run a gateway model over one or more documents/images.""" + client: VLMRun = ctx.obj + + if not files and not prompt: + console.print( + "[red]Error:[/] Provide at least one input file. " + "Most gateway models do not accept text-only input." + ) + raise typer.Exit(1) + + files = files or [] + create_kwargs, extra_body = _split_create_kwargs(_parse_extra(extra)) + if timeout is not None: + create_kwargs["timeout"] = timeout + + if method: + extra_body["method"] = method + if method_params: + try: + parsed_params = json.loads(method_params) + except (json.JSONDecodeError, ValueError) as e: + console.print(f"[red]Error:[/] --method-params must be valid JSON: {e}") + raise typer.Exit(1) + if not isinstance(parsed_params, dict): + console.print("[red]Error:[/] --method-params must be a JSON object.") + raise typer.Exit(1) + extra_body["method_params"] = parsed_params + + if response_format: + # A standard OpenAI create() field, so it rides as a top-level kwarg. + create_kwargs["response_format"] = _parse_response_format(response_format) + + if extra_body: + create_kwargs["extra_body"] = extra_body + + # Show the files being processed. + if files and not output_json: + tree = Tree("", guide_style="dim", hide_root=True) + for f in files: + size_str = format_file_size(f.stat().st_size) + tree.add(f"{f.name} [dim]({size_str})[/dim]") + console.print( + Panel( + tree, + title=f"Processing {len(files)} file(s) [dim]({model})[/dim]", + title_align="left", + border_style="dim", + ) + ) + + messages = _build_messages(files, prompt) + start_time = time.time() + status_msg = f"Processing ([bold]{model}[/bold])..." + + # The gateway only streams document (PDF) requests: it emits one SSE chunk + # per page. Image- and text-only requests ignore `stream` and return a + # single plain chat.completion body still labelled text/event-stream, so an + # SSE reader finds no events and yields empty content. Stream only when a + # document is actually present. + has_document = any( + part.get("type") == "document_url" + for message in messages + for part in message["content"] + ) + if not has_document or no_stream: + if output_json: + with handle_api_errors(): + response = client.gateway.completions.create( + model=model, messages=messages, stream=False, **create_kwargs + ) + else: + with ( + TimedStatus(status_msg, console=console), + handle_api_errors(), + ): + response = client.gateway.completions.create( + model=model, messages=messages, stream=False, **create_kwargs + ) + latency_s = time.time() - start_time + content = response.choices[0].message.content or "" + usage = response.usage + else: + chunks: List[str] = [] + usage = None + + def _consume(stream) -> None: + nonlocal usage + for chunk in stream: + if ( + chunk.choices + and chunk.choices[0].delta + and chunk.choices[0].delta.content + ): + chunks.append(chunk.choices[0].delta.content) + if getattr(chunk, "usage", None): + usage = chunk.usage + + if output_json: + with handle_api_errors(): + _consume( + client.gateway.completions.create( + model=model, messages=messages, stream=True, **create_kwargs + ) + ) + else: + with ( + TimedStatus(status_msg, console=console), + handle_api_errors(), + ): + _consume( + client.gateway.completions.create( + model=model, messages=messages, stream=True, **create_kwargs + ) + ) + content = "".join(chunks) + latency_s = time.time() - start_time + + error = _content_error(content) + + if output_json: + out = { + "model": model, + "content": content, + "latency_s": latency_s, + "usage": usage.model_dump() if hasattr(usage, "model_dump") else usage, + } + print(json.dumps(out, indent=2, default=str)) + raise typer.Exit(1 if error else 0) + + if error: + console.print(f"[red]Error:[/] {error}") + raise typer.Exit(1) + + _print_output(content, model, latency_s, usage) + + +EMBED_HELP = """Embed text, images or video with a gateway embedding model. + +\b +EXAMPLES: + vlmrun gw embed -t "a blue parrot" -m qwen/qwen3-vl-embedding-2b + vlmrun gw embed photo.jpg -m qwen/qwen3-vl-embedding-2b + vlmrun gw embed a.jpg b.jpg -t "caption" -m qwen/qwen3-vl-embedding-2b + vlmrun gw embed photo.jpg -t "caption" --join -m qwen/qwen3-vl-embedding-2b + vlmrun gw embed -t "hi" -m qwen/qwen3-vl-embedding-2b --dimensions 64 + vlmrun gw embed photo.jpg -m qwen/qwen3-vl-embedding-2b --json # full vectors + +\b +NOTES: + Every file and every -t/--text is embedded as its own vector. Use --join to + embed them together as a single vector instead (e.g. an image plus its + caption); models embed at most one image per vector, so --join accepts at + most one file. + Video is accepted by the API but is not currently backed by any embedding + model: it returns the same vector regardless of the clip. +""" + +TRANSCRIBE_HELP = """Transcribe audio with a gateway transcription model. + +\b +EXAMPLES: + vlmrun gw transcribe clip.mp3 -m nvidia/parakeet-tdt-0.6b-v3 + vlmrun gw transcribe clip.mp4 -m nvidia/parakeet-tdt-0.6b-v3 # video's audio track + vlmrun gw transcribe clip.mp3 -m nvidia/parakeet-tdt-0.6b-v3 -f srt + vlmrun gw transcribe clip.mp3 -m nvidia/parakeet-tdt-0.6b-v3 --language en + vlmrun gw transcribe --url https://example.com/a.mp3 -m nvidia/parakeet-tdt-0.6b-v3 + +\b +NOTES: + Accepts audio files, or a video file whose audio track is transcribed. + Formats: json, text, verbose_json, srt, vtt. +""" + +TRANSCRIBE_FORMATS = ("json", "text", "verbose_json", "srt", "vtt") + + +def _embed_part(path: Path) -> Dict[str, Any]: + """Encode a local file as an embedding content part. + + Embedding models take images and video only; anything else (a PDF, a text + file) is rejected here rather than sent as a mislabelled ``image_url`` that + the gateway would fail on. + """ + data = path.read_bytes() + mime = _guess_mime(path, data) + if mime.startswith("video/"): + key = "video_url" + elif mime.startswith("image/"): + key = "image_url" + else: + console.print( + f"[red]Error:[/] Cannot embed '{path.name}': unsupported type " + f"'{mime}'. Embedding models accept images and video only " + "(use --text for text)." + ) + raise typer.Exit(1) + b64 = base64.b64encode(data).decode("ascii") + return {"type": key, key: {"url": f"data:{mime};base64,{b64}"}} + + +def _content_error(content: str) -> Optional[str]: + """Return the error message if the response body is a gateway error payload. + + On some 200 responses (e.g. an unknown ``--method``) the gateway ships an + ``{"error": "..."}`` object as the message content instead of a real + result. Detect that exact shape — a lone ``error`` key — so the CLI can fail + loudly instead of rendering it as a successful response. Normal outputs are + either ````-wrapped or ``{"text": ...}`` lines, so this never + misfires. + """ + stripped = content.strip() + if not stripped.startswith("{"): + return None + try: + parsed = json.loads(stripped) + except (json.JSONDecodeError, ValueError): + return None + if isinstance(parsed, dict) and list(parsed.keys()) == ["error"]: + return str(parsed["error"]) + return None + + +def _renderable(content: str): + """Pick a Rich renderable for gateway output. + + OCR responses are wrapped in ````/```` tags, and detect/ocr + methods emit JSON lines. Rich's Markdown renderer treats angle-bracket tags + as HTML and drops them, blanking the output, so only render as Markdown when + the payload does not start with markup or JSON. + """ + stripped = content.lstrip() + if stripped.startswith(("<", "{", "[")): + return Text(content) + return Markdown(content) + + +def _format_cost(cost: Any) -> Optional[str]: + """Format the gateway's per-request ``usage.cost`` (USD), or None.""" + try: + value = float(cost) + except (TypeError, ValueError): + return None + if value <= 0: + return "$0" + if value >= 0.01: + return f"${value:.4f}" + # Sub-cent: show up to 6 decimals without scientific notation, but never + # collapse a real cost to "$0". + formatted = f"${value:.6f}".rstrip("0").rstrip(".") + return formatted if formatted != "$0" else "<$0.000001" + + +def _print_output(content: str, model: str, latency_s: float, usage: Any) -> None: + """Render the gateway response in a Rich panel.""" + stats = [model] + if usage is not None: + total = getattr(usage, "total_tokens", None) + if total: + prompt_toks = getattr(usage, "prompt_tokens", 0) + completion_toks = getattr(usage, "completion_tokens", 0) + stats.append(f"P:{prompt_toks} / C:{completion_toks} / T:{total} tokens") + cost = _format_cost(getattr(usage, "cost", None)) + if cost: + stats.append(cost) + stats.append(f"{latency_s:.2f}s") + + console.print( + Panel( + _renderable(content) if content else "[dim](empty response)[/dim]", + title="[bold]Response[/bold]", + title_align="left", + subtitle=f"[dim][white]{' · '.join(stats)}[/white][/dim]", + subtitle_align="right", + border_style="blue", + padding=(1, 2), + ) + ) + + +@app.command(help=EMBED_HELP, context_settings={"max_content_width": 120}) +def embed( + ctx: typer.Context, + files: List[Path] = typer.Argument( + None, + help="Image/video file(s) to embed. Repeatable.", + exists=True, + readable=True, + ), + model: str = typer.Option( + ..., "--model", "-m", help="Embedding model id (see `vlmrun gw models`)." + ), + text: Optional[List[str]] = typer.Option( + None, + "--text", + "-t", + help="Text to embed (repeatable). Its own vector unless --join is set.", + ), + join: bool = typer.Option( + False, + "--join", + help="Embed all inputs together as one vector (max one file).", + ), + dimensions: Optional[int] = typer.Option( + None, "--dimensions", "-d", help="Truncate vectors to this many dimensions." + ), + output_json: bool = typer.Option( + False, "--json", "-j", help="Output raw JSON, including full vectors." + ), + timeout: Optional[float] = typer.Option( + None, "--timeout", help="Request timeout in seconds." + ), +) -> None: + """Embed text, images or video with a gateway embedding model.""" + client: VLMRun = ctx.obj + files = files or [] + texts = list(text or []) + + if not files and not texts: + console.print("[red]Error:[/] Provide at least one file or --text to embed.") + raise typer.Exit(1) + if join and len(files) > 1: + console.print( + "[red]Error:[/] --join accepts at most one file: embedding models take " + "a single image per vector. Drop --join to embed each file separately." + ) + raise typer.Exit(1) + + # `input` is a list whose items are each either a plain string or a *list* + # of content parts; a flat list of parts is rejected by the API. + inputs: List[Any] = [] + labels: List[str] = [] + if join: + parts = [_embed_part(f) for f in files] + parts += [{"type": "text", "text": t} for t in texts] + inputs.append(parts) + labels.append(" + ".join([f.name for f in files] + [f'"{t}"' for t in texts])) + else: + for f in files: + inputs.append([_embed_part(f)]) + labels.append(f.name) + for t in texts: + inputs.append(t) + labels.append(f'"{t}"') + + create_kwargs: Dict[str, Any] = {} + if dimensions is not None: + create_kwargs["dimensions"] = dimensions + if timeout is not None: + create_kwargs["timeout"] = timeout + + start_time = time.time() + if output_json: + with handle_api_errors(): + response = client.gateway.embeddings.create( + model=model, input=inputs, **create_kwargs + ) + else: + with ( + TimedStatus(f"Embedding ([bold]{model}[/bold])...", console=console), + handle_api_errors(), + ): + response = client.gateway.embeddings.create( + model=model, input=inputs, **create_kwargs + ) + latency_s = time.time() - start_time + + if output_json: + print( + json.dumps( + response.model_dump() if hasattr(response, "model_dump") else response, + indent=2, + default=str, + ) + ) + return + + table = Table( + show_header=True, + header_style="bold white", + box=box.SIMPLE_HEAVY, + padding=(0, 1), + ) + table.add_column("INPUT", style="bold cyan") + table.add_column("DIMS", justify="right") + table.add_column("PREVIEW", style="dim") + + for i, item in enumerate(response.data): + vector = item.embedding + label = labels[i] if i < len(labels) else str(i) + if isinstance(vector, str): # encoding_format=base64 + preview, dims = f"{vector[:28]}...", "-" + else: + preview = "[" + ", ".join(f"{v:+.3f}" for v in vector[:4]) + ", ...]" + dims = str(len(vector)) + table.add_row(label if len(label) <= 34 else label[:31] + "...", dims, preview) + + usage = getattr(response, "usage", None) + stats = [model] + total = getattr(usage, "total_tokens", None) + if total: + stats.append(f"T:{total} tokens") + stats.append(f"{latency_s:.2f}s") + + console.print( + Panel( + table, + title="[bold]Embeddings[/bold]", + title_align="left", + subtitle=f"[dim]{' · '.join(stats)}[/dim]", + subtitle_align="right", + border_style="blue", + padding=(0, 1), + ) + ) + + +@app.command(help=TRANSCRIBE_HELP, context_settings={"max_content_width": 120}) +def transcribe( + ctx: typer.Context, + file: Optional[Path] = typer.Argument( + None, + help="Audio file (or a video whose audio track is transcribed).", + exists=True, + readable=True, + ), + model: str = typer.Option( + ..., "--model", "-m", help="Transcription model id (see `vlmrun gw models`)." + ), + url: Optional[str] = typer.Option( + None, "--url", help="Hosted audio URL instead of a local file." + ), + response_format: str = typer.Option( + "json", + "--format", + "-f", + help=f"Response format: {', '.join(TRANSCRIBE_FORMATS)}.", + ), + language: Optional[str] = typer.Option( + None, "--language", "-l", help="ISO-639-1 language hint, e.g. en." + ), + prompt: Optional[str] = typer.Option( + None, "--prompt", "-p", help="Context to bias transcription (proper nouns)." + ), + output_json: bool = typer.Option(False, "--json", "-j", help="Output raw JSON."), + timeout: Optional[float] = typer.Option( + None, "--timeout", help="Request timeout in seconds." + ), +) -> None: + """Transcribe audio with a gateway transcription model.""" + client: VLMRun = ctx.obj + + if not file and not url: + console.print("[red]Error:[/] Provide an audio file or --url.") + raise typer.Exit(1) + if file and url: + console.print("[red]Error:[/] Provide either a file or --url, not both.") + raise typer.Exit(1) + if response_format not in TRANSCRIBE_FORMATS: + console.print( + f"[red]Error:[/] Unknown --format '{response_format}'. " + f"Choose from: {', '.join(TRANSCRIBE_FORMATS)}." + ) + raise typer.Exit(1) + + create_kwargs: Dict[str, Any] = {"response_format": response_format} + if language: + create_kwargs["language"] = language + if prompt: + create_kwargs["prompt"] = prompt + if timeout is not None: + create_kwargs["timeout"] = timeout + if url: + # `url` is a gateway extension to the OpenAI transcription form. + create_kwargs["extra_body"] = {"url": url} + + if file and not output_json: + console.print( + Panel( + f"{file.name} [dim]({format_file_size(file.stat().st_size)})[/dim]", + title=f"Transcribing [dim]({model})[/dim]", + title_align="left", + border_style="dim", + ) + ) + + start_time = time.time() + + def _create(): + if file: + with file.open("rb") as fh: + return client.gateway.transcriptions.create( + model=model, file=fh, **create_kwargs + ) + # The OpenAI SDK requires a `file`; the gateway reads `url` instead. + return client.gateway.transcriptions.create( + model=model, file=("audio.mp3", b"", "audio/mpeg"), **create_kwargs + ) + + if output_json: + with handle_api_errors(): + response = _create() + else: + with ( + TimedStatus(f"Transcribing ([bold]{model}[/bold])...", console=console), + handle_api_errors(), + ): + response = _create() + latency_s = time.time() - start_time + + text_out = response if isinstance(response, str) else getattr(response, "text", "") + + if output_json: + if hasattr(response, "model_dump"): + print(json.dumps(response.model_dump(), indent=2, default=str)) + else: + print( + json.dumps( + {"model": model, "text": text_out, "latency_s": latency_s}, + indent=2, + default=str, + ) + ) + return + + console.print( + Panel( + Text(text_out) if text_out else "[dim](empty transcript)[/dim]", + title="[bold]Transcript[/bold]", + title_align="left", + subtitle=f"[dim]{model} · {response_format} · {latency_s:.2f}s[/dim]", + subtitle_align="right", + border_style="blue", + padding=(1, 2), + ) + ) + + +if __name__ == "__main__": + app() diff --git a/vlmrun/cli/cli.py b/vlmrun/cli/cli.py index 5594c0b..8b3c7ce 100644 --- a/vlmrun/cli/cli.py +++ b/vlmrun/cli/cli.py @@ -18,6 +18,7 @@ from vlmrun.cli._cli.execute import EXECUTE_HELP, execute from vlmrun.cli._cli.executions import app as executions_app from vlmrun.cli._cli.files import app as files_app +from vlmrun.cli._cli.gateway import app as gateway_app from vlmrun.cli._cli.generate import GENERATE_HELP, generate from vlmrun.cli._cli.hub import app as hub_app from vlmrun.cli._cli.models import app as models_app @@ -128,6 +129,8 @@ def main( app.add_typer(predictions_app, name="predictions") app.add_typer(files_app, name="files") app.add_typer(hub_app, name="hub") +app.add_typer(gateway_app, name="gateway") +app.add_typer(gateway_app, name="gw") app.add_typer(models_app, name="models") app.add_typer(skills_app, name="skills") app.add_typer(artifacts_app, name="artifacts") diff --git a/vlmrun/client/client.py b/vlmrun/client/client.py index 7c811cc..2e13bab 100644 --- a/vlmrun/client/client.py +++ b/vlmrun/client/client.py @@ -22,6 +22,7 @@ ) from vlmrun.client.feedback import Feedback from vlmrun.client.agent import Agent +from vlmrun.client.gateway import Gateway from vlmrun.client.skills import Skills from vlmrun.client.executions import Executions from vlmrun.client.artifacts import Artifacts @@ -119,6 +120,7 @@ def __post_init__(self): self.video._requestor._timeout = 120.0 self.feedback = Feedback(self) self.agent = Agent(self) + self.gateway = Gateway(self) self.skills = Skills(self) self.executions = Executions(self) self.artifacts = Artifacts(self) diff --git a/vlmrun/client/gateway.py b/vlmrun/client/gateway.py new file mode 100644 index 0000000..7b37120 --- /dev/null +++ b/vlmrun/client/gateway.py @@ -0,0 +1,241 @@ +"""VLM Run OpenAI-compatible model gateway resource. + +The gateway (``https://gateway.vlm.run/v1``) exposes an OpenAI-compatible +surface for third-party OCR / vision-language models (e.g. ``glm-ocr``, +``paddle-ocrv6``, ``qwen3.6-0.8b``). It authenticates with the same +``VLMRUN_API_KEY`` used everywhere else in the SDK. + +This mirrors the :class:`~vlmrun.client.agent.Agent` completions pattern: +we point the OpenAI SDK at ``{gateway_url}/openai`` and reuse the familiar +chat-completions / models interface. +""" + +from __future__ import annotations + +import os +from functools import cached_property +from typing import Any, List, Optional + +from vlmrun.constants import DEFAULT_GATEWAY_URL +from vlmrun.client.exceptions import DependencyError +from vlmrun.types.abstract import VLMRunProtocol + + +def _require_openai(): + """Import the OpenAI SDK or raise a helpful :class:`DependencyError`.""" + try: + import openai # noqa: F401 + except ImportError as e: + raise DependencyError( + message="OpenAI SDK is not installed", + suggestion="Install it with `pip install vlmrun[openai]` or `pip install openai`", + error_type="missing_dependency", + ) from e + return openai + + +class Gateway: + """OpenAI-compatible model gateway resource for VLM Run. + + Provides access to third-party OCR / VLM models hosted behind the VLM Run + gateway using the standard OpenAI chat-completions and models interfaces. + + Attributes: + base_url: Gateway base URL (defaults to ``VLMRUN_GATEWAY_URL`` env var or + ``https://gateway.vlm.run/v1``). + """ + + def __init__( + self, client: "VLMRunProtocol", base_url: Optional[str] = None + ) -> None: + """Initialize the Gateway resource. + + Args: + client: VLM Run API client instance (provides the API key). + base_url: Optional gateway base URL override. Falls back to the + ``VLMRUN_GATEWAY_URL`` environment variable, then the default. + """ + self._client = client + self._base_url = ( + base_url or os.getenv("VLMRUN_GATEWAY_URL") or DEFAULT_GATEWAY_URL + ) + + @property + def base_url(self) -> str: + """Gateway base URL (without trailing slash).""" + return self._base_url.rstrip("/") + + @property + def openai_base_url(self) -> str: + """OpenAI-compatible base URL used by the OpenAI SDK.""" + return f"{self.base_url}/openai" + + def _timeout(self) -> Optional[float]: + # Gateway calls (especially multi-page PDF OCR) routinely exceed the + # client's 120s default, so raise the floor to 600s — but only when the + # user is still at that default. An explicit timeout (whether a longer + # deadline or a shorter fail-fast) is theirs to keep. + timeout = self._client.timeout + if timeout is None: + return None + if timeout == 120.0: + return 600.0 + return timeout + + @cached_property + def _openai(self): + """Synchronous OpenAI client pointed at the gateway.""" + openai = _require_openai() + return openai.OpenAI( + api_key=self._client.api_key, + base_url=self.openai_base_url, + timeout=self._timeout(), + max_retries=self._client.max_retries, + ) + + @cached_property + def _async_openai(self): + """Asynchronous OpenAI client pointed at the gateway.""" + openai = _require_openai() + return openai.AsyncOpenAI( + api_key=self._client.api_key, + base_url=self.openai_base_url, + timeout=self._timeout(), + max_retries=self._client.max_retries, + ) + + @cached_property + def completions(self): + """OpenAI-compatible chat completions interface (synchronous). + + Example: + ```python + from vlmrun import VLMRun + + client = VLMRun() + response = client.gateway.completions.create( + model="glm-ocr", + messages=[{"role": "user", "content": [ + {"type": "document_url", "document_url": {"url": "data:application/pdf;base64,..."}}, + ]}], + ) + ``` + + Raises: + DependencyError: If the ``openai`` package is not installed. + + Returns: + OpenAI Completions object configured for the VLM Run gateway. + """ + return self._openai.chat.completions + + @cached_property + def async_completions(self): + """OpenAI-compatible chat completions interface (asynchronous). + + Raises: + DependencyError: If the ``openai`` package is not installed. + + Returns: + OpenAI AsyncCompletions object configured for the VLM Run gateway. + """ + return self._async_openai.chat.completions + + @cached_property + def embeddings(self): + """OpenAI-compatible embeddings interface (synchronous). + + Note: + Multimodal input nests content parts one level deeper than plain + text: ``input`` is a list whose items are either a string or a + *list* of content parts. + + Example: + ```python + from vlmrun.client import VLMRun + + client = VLMRun() + response = client.gateway.embeddings.create( + model="qwen/qwen3-vl-embedding-2b", + input=[[{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}}]], + ) + ``` + + Raises: + DependencyError: If the ``openai`` package is not installed. + + Returns: + OpenAI Embeddings object configured for the VLM Run gateway. + """ + return self._openai.embeddings + + @cached_property + def transcriptions(self): + """OpenAI-compatible audio transcriptions interface (synchronous). + + Example: + ```python + from vlmrun.client import VLMRun + + client = VLMRun() + with open("clip.mp3", "rb") as fh: + response = client.gateway.transcriptions.create( + model="nvidia/parakeet-tdt-0.6b-v3", file=fh + ) + ``` + + Raises: + DependencyError: If the ``openai`` package is not installed. + + Returns: + OpenAI Transcriptions object configured for the VLM Run gateway. + """ + return self._openai.audio.transcriptions + + def models(self) -> List[Any]: + """List models available on the gateway. + + Returns the raw OpenAI ``Model`` objects. Gateway models carry extra + metadata (input/output pricing, modality support, etc.) beyond the + standard OpenAI fields; those are preserved on each object's + ``model_extra``. + + Raises: + DependencyError: If the ``openai`` package is not installed. + + Returns: + List of OpenAI ``Model`` objects. + """ + return list(self._openai.models.list()) + + def health(self) -> bool: + """Check gateway liveness. + + Attempts a ``GET {gateway}/health`` request and falls back to listing + models as a liveness probe if no dedicated health endpoint responds. + + Returns: + True if the gateway is reachable and authenticated, else False. + """ + # httpx is a hard dependency of the openai SDK, so it is always + # available whenever the gateway is usable. + import httpx + + headers = {"Authorization": f"Bearer {self._client.api_key}"} + try: + resp = httpx.get(f"{self.base_url}/health", headers=headers, timeout=30.0) + except Exception: + # No dedicated health route reachable — fall back to a real call. + try: + self.models() + return True + except Exception: + return False + + if resp.status_code == 404: + try: + self.models() + return True + except Exception: + return False + return resp.is_success diff --git a/vlmrun/constants.py b/vlmrun/constants.py index a6670bf..5a5941e 100644 --- a/vlmrun/constants.py +++ b/vlmrun/constants.py @@ -3,6 +3,10 @@ DEFAULT_BASE_URL = "https://api.vlm.run/v1" +# OpenAI-compatible model gateway (third-party OCR / VLM models). +# Override with the VLMRUN_GATEWAY_URL environment variable. +DEFAULT_GATEWAY_URL = "https://gateway.vlm.run/v1" + # Cache directories - use VLMRUN_CACHE_DIR env var if set, otherwise default to ~/.vlmrun/cache VLMRUN_HOME = Path.home() / ".vlmrun" VLMRUN_HOME.mkdir(parents=True, exist_ok=True) diff --git a/vlmrun/types/abstract.py b/vlmrun/types/abstract.py index 701a6a8..17da428 100644 --- a/vlmrun/types/abstract.py +++ b/vlmrun/types/abstract.py @@ -25,6 +25,7 @@ class VLMRunProtocol(Protocol): fine_tuning: Any feedback: Any agent: Any + gateway: Any requestor: Any artifacts: Any diff --git a/vlmrun/version.py b/vlmrun/version.py index 7bbb2ef..49e0fc1 100644 --- a/vlmrun/version.py +++ b/vlmrun/version.py @@ -1 +1 @@ -__version__ = "0.6.5" +__version__ = "0.7.0"