Skip to content

Commit 8fbcfec

Browse files
montasaurussimonw
andauthored
Add Image Support - thanks, @montasaurus
* add image support * bump python and llm versions * reset user directory on each test * add recordings to readme * Ensure compatible OpenAI library --------- Co-authored-by: Simon Willison <swillison@gmail.com>
1 parent b01ab71 commit 8fbcfec

10 files changed

Lines changed: 1780 additions & 46 deletions

File tree

.github/workflows/publish.yml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ jobs:
1212
runs-on: ubuntu-latest
1313
strategy:
1414
matrix:
15-
python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
15+
python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
1616
steps:
1717
- uses: actions/checkout@v4
1818
- name: Set up Python ${{ matrix.python-version }}
@@ -38,7 +38,7 @@ jobs:
3838
- name: Set up Python
3939
uses: actions/setup-python@v5
4040
with:
41-
python-version: "3.12"
41+
python-version: "3.13"
4242
cache: pip
4343
cache-dependency-path: pyproject.toml
4444
- name: Install dependencies

.github/workflows/test.yml

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ jobs:
1010
runs-on: ubuntu-latest
1111
strategy:
1212
matrix:
13-
python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
13+
python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
1414
steps:
1515
- uses: actions/checkout@v4
1616
- name: Set up Python ${{ matrix.python-version }}
@@ -25,4 +25,3 @@ jobs:
2525
- name: Run tests
2626
run: |
2727
pytest
28-

README.md

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,13 @@ Now you can prompt Claude using:
5454
```bash
5555
cat llm_openrouter.py | llm -m claude -s 'write some pytest tests for this'
5656
```
57+
58+
Images are supported too, for some models:
59+
```bash
60+
llm -m openrouter/anthropic/claude-3.5-sonnet 'describe this image' -a https://static.simonwillison.net/static/2024/pelicans.jpg
61+
llm -m openrouter/anthropic/claude-3-haiku 'extract text' -a page.png
62+
```
63+
5764
## Development
5865

5966
To set up this plugin locally, first checkout the code. Then create a new virtual environment:
@@ -64,9 +71,13 @@ source venv/bin/activate
6471
```
6572
Now install the dependencies and test dependencies:
6673
```bash
67-
pip install -e '.[test]'
74+
llm install -e '.[test]'
6875
```
6976
To run the tests:
7077
```bash
7178
pytest
7279
```
80+
To add new recordings and snapshots, run:
81+
```bash
82+
pytest --record-mode=once --inline-snapshot=create
83+
```

llm_openrouter.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,12 @@ def register_models(register):
2929
if not key:
3030
return
3131
for model_definition in get_openrouter_models():
32+
supports_images = get_supports_images(model_definition)
3233
register(
3334
OpenRouterChat(
3435
model_id="openrouter/{}".format(model_definition["id"]),
3536
model_name=model_definition["id"],
37+
vision=supports_images,
3638
api_base="https://openrouter.ai/api/v1",
3739
headers={"HTTP-Referer": "https://llm.datasette.io/", "X-Title": "LLM"},
3840
)
@@ -78,3 +80,14 @@ def fetch_cached_json(url, path, cache_timeout):
7880
raise DownloadError(
7981
f"Failed to download data and no cache is available at {path}"
8082
)
83+
84+
85+
def get_supports_images(model_definition):
86+
try:
87+
# e.g. `text->text` or `text+image->text`
88+
modality = model_definition["architecture"]["modality"]
89+
90+
input_modalities = modality.split("->")[0].split("+")
91+
return "image" in input_modalities
92+
except Exception:
93+
return False

pyproject.toml

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,9 @@ classifiers = [
99
"License :: OSI Approved :: Apache Software License"
1010
]
1111
dependencies = [
12-
"llm>=0.8",
13-
"httpx"
12+
"llm>=0.17",
13+
"httpx",
14+
"openai>=1.57.0",
1415
]
1516

1617
[project.urls]
@@ -23,4 +24,4 @@ CI = "https://github.com/simonw/llm-openrouter/actions"
2324
openrouter = "llm_openrouter"
2425

2526
[project.optional-dependencies]
26-
test = ["pytest", "pytest-httpx"]
27+
test = ["pytest", "pytest-recording", "inline-snapshot"]

tests/cassettes/test_llm_openrouter/test_image_prompt.yaml

Lines changed: 608 additions & 0 deletions
Large diffs are not rendered by default.

tests/cassettes/test_llm_openrouter/test_llm_models.yaml

Lines changed: 483 additions & 0 deletions
Large diffs are not rendered by default.

tests/cassettes/test_llm_openrouter/test_prompt.yaml

Lines changed: 591 additions & 0 deletions
Large diffs are not rendered by default.

tests/conftest.py

Lines changed: 7 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,5 @@
1-
import json
2-
import llm
31
import pytest
42

5-
DUMMY_MODELS = {
6-
"data": [
7-
{
8-
"id": "openai/gpt-3.5-turbo",
9-
"pricing": {"prompt": "0.0000015", "completion": "0.000002"},
10-
"context_length": 4095,
11-
"per_request_limits": {
12-
"prompt_tokens": "2871318",
13-
"completion_tokens": "2153488",
14-
},
15-
},
16-
{
17-
"id": "anthropic/claude-2",
18-
"pricing": {"prompt": "0.00001102", "completion": "0.00003268"},
19-
"context_length": 100000,
20-
"per_request_limits": {
21-
"prompt_tokens": "390832",
22-
"completion_tokens": "131792",
23-
},
24-
},
25-
]
26-
}
27-
283

294
@pytest.fixture
305
def user_path(tmpdir):
@@ -36,7 +11,11 @@ def user_path(tmpdir):
3611
@pytest.fixture(autouse=True)
3712
def env_setup(monkeypatch, user_path):
3813
monkeypatch.setenv("LLM_USER_PATH", str(user_path))
39-
# Write out the models.json file
40-
(llm.user_dir() / "openrouter_models.json").write_text(
41-
json.dumps(DUMMY_MODELS), "utf-8"
14+
monkeypatch.setenv(
15+
"LLM_OPENROUTER_KEY",
16+
"sk-...",
17+
)
18+
monkeypatch.setenv(
19+
"OPENROUTER_KEY",
20+
"sk-...",
4221
)

tests/test_llm_openrouter.py

Lines changed: 59 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,71 @@
1+
import llm
2+
import pytest
13
from click.testing import CliRunner
4+
from inline_snapshot import snapshot
25
from llm.cli import cli
3-
import json
4-
import pytest
6+
7+
TINY_PNG = (
8+
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\xa6\x00\x00\x01\x1a"
9+
b"\x02\x03\x00\x00\x00\xe6\x99\xc4^\x00\x00\x00\tPLTE\xff\xff\xff"
10+
b"\x00\xff\x00\xfe\x01\x00\x12t\x01J\x00\x00\x00GIDATx\xda\xed\xd81\x11"
11+
b"\x000\x08\xc0\xc0.]\xea\xaf&Q\x89\x04V\xe0>\xf3+\xc8\x91Z\xf4\xa2\x08EQ\x14E"
12+
b"Q\x14EQ\x14EQ\xd4B\x91$I3\xbb\xbf\x08EQ\x14EQ\x14EQ\x14E\xd1\xa5"
13+
b"\xd4\x17\x91\xc6\x95\x05\x15\x0f\x9f\xc5\t\x9f\xa4\x00\x00\x00\x00IEND\xaeB`"
14+
b"\x82"
15+
)
16+
17+
18+
@pytest.mark.vcr
19+
def test_prompt():
20+
model = llm.get_model("openrouter/openai/gpt-4o")
21+
response = model.prompt("Two names for a pet pelican, be brief")
22+
assert str(response) == snapshot("Gully or Skipper")
23+
response_dict = dict(response.response_json)
24+
response_dict.pop("id") # differs between requests
25+
assert response_dict == snapshot(
26+
{
27+
"content": "Gully or Skipper",
28+
"role": "assistant",
29+
"finish_reason": "stop",
30+
"usage": {"completion_tokens": 5, "prompt_tokens": 17, "total_tokens": 22},
31+
"object": "chat.completion.chunk",
32+
"model": "openai/gpt-4o",
33+
"created": 1731200404,
34+
}
35+
)
536

637

7-
@pytest.mark.parametrize("set_key", (False, True))
8-
def test_llm_models(set_key, user_path):
38+
@pytest.mark.vcr
39+
def test_llm_models():
940
runner = CliRunner()
10-
if set_key:
11-
(user_path / "keys.json").write_text(json.dumps({"openrouter": "x"}), "utf-8")
1241
result = runner.invoke(cli, ["models", "list"])
1342
assert result.exit_code == 0, result.output
1443
fragments = (
1544
"OpenRouter: openrouter/openai/gpt-3.5-turbo",
1645
"OpenRouter: openrouter/anthropic/claude-2",
1746
)
1847
for fragment in fragments:
19-
if set_key:
20-
assert fragment in result.output
21-
else:
22-
assert fragment not in result.output
48+
assert fragment in result.output
49+
50+
51+
@pytest.mark.vcr
52+
def test_image_prompt():
53+
model = llm.get_model("openrouter/anthropic/claude-3.5-sonnet")
54+
response = model.prompt(
55+
"Describe image in three words",
56+
attachments=[llm.Attachment(content=TINY_PNG)],
57+
)
58+
assert str(response) == snapshot("Red and green")
59+
response_dict = response.response_json
60+
response_dict.pop("id") # differs between requests
61+
assert response_dict == snapshot(
62+
{
63+
"content": "Red and green",
64+
"role": "assistant",
65+
"finish_reason": "end_turn",
66+
"usage": {"completion_tokens": 7, "prompt_tokens": 82, "total_tokens": 89},
67+
"object": "chat.completion.chunk",
68+
"model": "anthropic/claude-3.5-sonnet",
69+
"created": 1731200406,
70+
}
71+
)

0 commit comments

Comments
 (0)