llm-gemini 0.25__tar.gz → 0.26.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: llm-gemini
3
- Version: 0.25
3
+ Version: 0.26.1
4
4
  Summary: LLM plugin to access Google's Gemini family of models
5
5
  Author: Simon Willison
6
6
  License-Expression: Apache-2.0
@@ -10,7 +10,7 @@ Project-URL: Issues, https://github.com/simonw/llm-gemini/issues
10
10
  Project-URL: CI, https://github.com/simonw/llm-gemini/actions
11
11
  Description-Content-Type: text/markdown
12
12
  License-File: LICENSE
13
- Requires-Dist: llm>=0.26
13
+ Requires-Dist: llm>=0.27
14
14
  Requires-Dist: httpx
15
15
  Requires-Dist: ijson
16
16
  Provides-Extra: test
@@ -75,6 +75,9 @@ result = runner.invoke(cli.cli, ["models", "-q", "gemini/"])
75
75
  lines = reversed(result.output.strip().split("\n"))
76
76
  to_output = []
77
77
  NOTES = {
78
+ "gemini/gemini-flash-latest": "Latest Gemini Flash",
79
+ "gemini/gemini-flash-lite-latest": "Latest Gemini Flash Lite",
80
+ "gemini/gemini-2.5-flash": "Gemini 2.5 Flash",
78
81
  "gemini/gemini-2.5-pro": "Gemini 2.5 Pro",
79
82
  "gemini/gemini-2.5-flash": "Gemini 2.5 Flash",
80
83
  "gemini/gemini-2.5-flash-lite": "Gemini 2.5 Flash Lite",
@@ -93,6 +96,10 @@ for line in lines:
93
96
  )
94
97
  cog.out("\n".join(to_output))
95
98
  ]]] -->
99
+ - `gemini/gemini-2.5-flash-lite-preview-09-2025`
100
+ - `gemini/gemini-2.5-flash-preview-09-2025`
101
+ - `gemini/gemini-flash-lite-latest`: Latest Gemini Flash Lite
102
+ - `gemini/gemini-flash-latest`: Latest Gemini Flash
96
103
  - `gemini/gemini-2.5-flash-lite`: Gemini 2.5 Flash Lite
97
104
  - `gemini/gemini-2.5-pro`: Gemini 2.5 Pro
98
105
  - `gemini/gemini-2.5-flash`: Gemini 2.5 Flash
@@ -52,6 +52,9 @@ result = runner.invoke(cli.cli, ["models", "-q", "gemini/"])
52
52
  lines = reversed(result.output.strip().split("\n"))
53
53
  to_output = []
54
54
  NOTES = {
55
+ "gemini/gemini-flash-latest": "Latest Gemini Flash",
56
+ "gemini/gemini-flash-lite-latest": "Latest Gemini Flash Lite",
57
+ "gemini/gemini-2.5-flash": "Gemini 2.5 Flash",
55
58
  "gemini/gemini-2.5-pro": "Gemini 2.5 Pro",
56
59
  "gemini/gemini-2.5-flash": "Gemini 2.5 Flash",
57
60
  "gemini/gemini-2.5-flash-lite": "Gemini 2.5 Flash Lite",
@@ -70,6 +73,10 @@ for line in lines:
70
73
  )
71
74
  cog.out("\n".join(to_output))
72
75
  ]]] -->
76
+ - `gemini/gemini-2.5-flash-lite-preview-09-2025`
77
+ - `gemini/gemini-2.5-flash-preview-09-2025`
78
+ - `gemini/gemini-flash-lite-latest`: Latest Gemini Flash Lite
79
+ - `gemini/gemini-flash-latest`: Latest Gemini Flash
73
80
  - `gemini/gemini-2.5-flash-lite`: Gemini 2.5 Flash Lite
74
81
  - `gemini/gemini-2.5-pro`: Gemini 2.5 Pro
75
82
  - `gemini/gemini-2.5-flash`: Gemini 2.5 Flash
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: llm-gemini
3
- Version: 0.25
3
+ Version: 0.26.1
4
4
  Summary: LLM plugin to access Google's Gemini family of models
5
5
  Author: Simon Willison
6
6
  License-Expression: Apache-2.0
@@ -10,7 +10,7 @@ Project-URL: Issues, https://github.com/simonw/llm-gemini/issues
10
10
  Project-URL: CI, https://github.com/simonw/llm-gemini/actions
11
11
  Description-Content-Type: text/markdown
12
12
  License-File: LICENSE
13
- Requires-Dist: llm>=0.26
13
+ Requires-Dist: llm>=0.27
14
14
  Requires-Dist: httpx
15
15
  Requires-Dist: ijson
16
16
  Provides-Extra: test
@@ -75,6 +75,9 @@ result = runner.invoke(cli.cli, ["models", "-q", "gemini/"])
75
75
  lines = reversed(result.output.strip().split("\n"))
76
76
  to_output = []
77
77
  NOTES = {
78
+ "gemini/gemini-flash-latest": "Latest Gemini Flash",
79
+ "gemini/gemini-flash-lite-latest": "Latest Gemini Flash Lite",
80
+ "gemini/gemini-2.5-flash": "Gemini 2.5 Flash",
78
81
  "gemini/gemini-2.5-pro": "Gemini 2.5 Pro",
79
82
  "gemini/gemini-2.5-flash": "Gemini 2.5 Flash",
80
83
  "gemini/gemini-2.5-flash-lite": "Gemini 2.5 Flash Lite",
@@ -93,6 +96,10 @@ for line in lines:
93
96
  )
94
97
  cog.out("\n".join(to_output))
95
98
  ]]] -->
99
+ - `gemini/gemini-2.5-flash-lite-preview-09-2025`
100
+ - `gemini/gemini-2.5-flash-preview-09-2025`
101
+ - `gemini/gemini-flash-lite-latest`: Latest Gemini Flash Lite
102
+ - `gemini/gemini-flash-latest`: Latest Gemini Flash
96
103
  - `gemini/gemini-2.5-flash-lite`: Gemini 2.5 Flash Lite
97
104
  - `gemini/gemini-2.5-pro`: Gemini 2.5 Pro
98
105
  - `gemini/gemini-2.5-flash`: Gemini 2.5 Flash
@@ -1,4 +1,4 @@
1
- llm>=0.26
1
+ llm>=0.27
2
2
  httpx
3
3
  ijson
4
4
 
@@ -45,6 +45,10 @@ GOOGLE_SEARCH_MODELS = {
45
45
  "gemini-2.5-pro",
46
46
  "gemini-2.5-flash",
47
47
  "gemini-2.5-flash-lite",
48
+ "gemini-flash-latest",
49
+ "gemini-flash-lite-latest",
50
+ "gemini-2.5-flash-preview-09-2025",
51
+ "gemini-2.5-flash-lite-preview-09-2025",
48
52
  }
49
53
 
50
54
  # Older Google models used google_search_retrieval instead of google_search
@@ -70,6 +74,10 @@ THINKING_BUDGET_MODELS = {
70
74
  "gemini-2.5-pro",
71
75
  "gemini-2.5-flash",
72
76
  "gemini-2.5-flash-lite",
77
+ "gemini-flash-latest",
78
+ "gemini-flash-lite-latest",
79
+ "gemini-2.5-flash-preview-09-2025",
80
+ "gemini-2.5-flash-lite-preview-09-2025",
73
81
  }
74
82
 
75
83
  NO_VISION_MODELS = {"gemma-3-1b-it", "gemma-3n-e4b-it"}
@@ -156,6 +164,11 @@ def register_models(register):
156
164
  "gemini-2.5-pro",
157
165
  # 22nd July 2025:
158
166
  "gemini-2.5-flash-lite",
167
+ # 25th Spetember 2025:
168
+ "gemini-flash-latest",
169
+ "gemini-flash-lite-latest",
170
+ "gemini-2.5-flash-preview-09-2025",
171
+ "gemini-2.5-flash-lite-preview-09-2025",
159
172
  ):
160
173
  can_google_search = model_id in GOOGLE_SEARCH_MODELS
161
174
  can_thinking_budget = model_id in THINKING_BUDGET_MODELS
@@ -194,6 +207,11 @@ def cleanup_schema(schema, in_properties=False):
194
207
  "Gemini supports only a subset of JSON schema"
195
208
  keys_to_remove = ("$schema", "additionalProperties", "title")
196
209
 
210
+ # First pass: resolve $ref references using $defs
211
+ if isinstance(schema, dict) and "$defs" in schema:
212
+ defs = schema.pop("$defs")
213
+ _resolve_refs(schema, defs)
214
+
197
215
  if isinstance(schema, dict):
198
216
  # Only remove keys if we're not inside a 'properties' block.
199
217
  if not in_properties:
@@ -211,6 +229,26 @@ def cleanup_schema(schema, in_properties=False):
211
229
  return schema
212
230
 
213
231
 
232
+ def _resolve_refs(schema, defs):
233
+ """Recursively resolve $ref references in schema using definitions."""
234
+ if isinstance(schema, dict):
235
+ if "$ref" in schema:
236
+ # Extract the reference path (e.g., "#/$defs/Dog" -> "Dog")
237
+ ref_path = schema.pop("$ref")
238
+ if ref_path.startswith("#/$defs/"):
239
+ def_name = ref_path.split("/")[-1]
240
+ if def_name in defs:
241
+ # Replace the $ref with the actual definition
242
+ schema.update(copy.deepcopy(defs[def_name]))
243
+
244
+ # Recursively resolve refs in nested structures
245
+ for value in schema.values():
246
+ _resolve_refs(value, defs)
247
+ elif isinstance(schema, list):
248
+ for item in schema:
249
+ _resolve_refs(item, defs)
250
+
251
+
214
252
  class _SharedGemini:
215
253
  needs_key = "gemini"
216
254
  key_env_var = "LLM_GEMINI_KEY"
@@ -543,6 +581,8 @@ class GeminiPro(_SharedGemini, llm.KeyModel):
543
581
  gathered.append(event)
544
582
  events.clear()
545
583
  response.response_json = gathered[-1]
584
+ resolved_model = gathered[-1]["modelVersion"]
585
+ response.set_resolved_model(resolved_model)
546
586
  self.set_usage(response)
547
587
 
548
588
 
@@ -1,13 +1,13 @@
1
1
  [project]
2
2
  name = "llm-gemini"
3
- version = "0.25"
3
+ version = "0.26.1"
4
4
  description = "LLM plugin to access Google's Gemini family of models"
5
5
  readme = "README.md"
6
6
  authors = [{name = "Simon Willison"}]
7
7
  license = "Apache-2.0"
8
8
  classifiers = []
9
9
  dependencies = [
10
- "llm>=0.26",
10
+ "llm>=0.27",
11
11
  "httpx",
12
12
  "ijson"
13
13
  ]
@@ -0,0 +1,583 @@
1
+ from click.testing import CliRunner
2
+ import llm
3
+ from llm.cli import cli
4
+ import nest_asyncio
5
+ import json
6
+ import os
7
+ import pytest
8
+ import pydantic
9
+ from pydantic import BaseModel
10
+ from typing import List, Optional
11
+ from llm_gemini import cleanup_schema
12
+
13
+ nest_asyncio.apply()
14
+
15
+ GEMINI_API_KEY = os.environ.get("PYTEST_GEMINI_API_KEY", None) or "gm-..."
16
+
17
+
18
+ @pytest.mark.vcr
19
+ @pytest.mark.asyncio
20
+ async def test_prompt():
21
+ model = llm.get_model("gemini-1.5-flash-latest")
22
+ response = model.prompt("Name for a pet pelican, just the name", key=GEMINI_API_KEY)
23
+ assert str(response) == "Percy\n"
24
+ assert response.response_json == {
25
+ "candidates": [
26
+ {
27
+ "finishReason": "STOP",
28
+ "safetyRatings": [
29
+ {
30
+ "category": "HARM_CATEGORY_HATE_SPEECH",
31
+ "probability": "NEGLIGIBLE",
32
+ },
33
+ {
34
+ "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
35
+ "probability": "NEGLIGIBLE",
36
+ },
37
+ {
38
+ "category": "HARM_CATEGORY_HARASSMENT",
39
+ "probability": "NEGLIGIBLE",
40
+ },
41
+ {
42
+ "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
43
+ "probability": "NEGLIGIBLE",
44
+ },
45
+ ],
46
+ }
47
+ ],
48
+ "modelVersion": "gemini-1.5-flash-latest",
49
+ }
50
+ assert response.token_details == {
51
+ "candidatesTokenCount": 2,
52
+ "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 9}],
53
+ "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 2}],
54
+ }
55
+ assert response.input_tokens == 9
56
+ assert response.output_tokens == 2
57
+
58
+ # And try it async too
59
+ async_model = llm.get_async_model("gemini-1.5-flash-latest")
60
+ response = await async_model.prompt(
61
+ "Name for a pet pelican, just the name", key=GEMINI_API_KEY
62
+ )
63
+ text = await response.text()
64
+ assert text == "Percy\n"
65
+
66
+
67
+ @pytest.mark.vcr
68
+ @pytest.mark.asyncio
69
+ async def test_prompt_with_pydantic_schema():
70
+ class Dog(pydantic.BaseModel):
71
+ name: str
72
+ age: int
73
+ bio: str
74
+
75
+ class Dogs(BaseModel):
76
+ dogs: List[Dog]
77
+
78
+ model = llm.get_model("gemini-1.5-flash-latest")
79
+ response = model.prompt(
80
+ "Invent a cool dog", key=GEMINI_API_KEY, schema=Dog, stream=False
81
+ )
82
+ assert json.loads(response.text()) == {
83
+ "age": 3,
84
+ "bio": "A fluffy Samoyed with exceptional intelligence and a love for belly rubs. He's mastered several tricks, including fetching the newspaper and opening doors.",
85
+ "name": "Cloud",
86
+ }
87
+ assert response.response_json == {
88
+ "candidates": [
89
+ {
90
+ "finishReason": "STOP",
91
+ "safetyRatings": [
92
+ {
93
+ "category": "HARM_CATEGORY_HATE_SPEECH",
94
+ "probability": "NEGLIGIBLE",
95
+ },
96
+ {
97
+ "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
98
+ "probability": "NEGLIGIBLE",
99
+ },
100
+ {
101
+ "category": "HARM_CATEGORY_HARASSMENT",
102
+ "probability": "NEGLIGIBLE",
103
+ },
104
+ {
105
+ "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
106
+ "probability": "NEGLIGIBLE",
107
+ },
108
+ ],
109
+ }
110
+ ],
111
+ "modelVersion": "gemini-1.5-flash-latest",
112
+ }
113
+ assert response.input_tokens == 10
114
+
115
+
116
+ @pytest.mark.vcr
117
+ @pytest.mark.asyncio
118
+ async def test_prompt_with_multiple_dogs():
119
+ class Dog(pydantic.BaseModel):
120
+ name: str
121
+ age: int
122
+ bio: str
123
+
124
+ class Dogs(BaseModel):
125
+ dogs: List[Dog]
126
+
127
+ model = llm.get_model("gemini-2.0-flash")
128
+ response = model.prompt(
129
+ "Invent 3 cool dogs", key=GEMINI_API_KEY, schema=Dogs, stream=False
130
+ )
131
+ result = json.loads(response.text())
132
+
133
+ # Verify we got 3 dogs
134
+ assert "dogs" in result
135
+ assert len(result["dogs"]) == 3
136
+
137
+ # Verify each dog has the required fields
138
+ for dog in result["dogs"]:
139
+ assert "name" in dog
140
+ assert "age" in dog
141
+ assert "bio" in dog
142
+ assert isinstance(dog["name"], str)
143
+ assert isinstance(dog["age"], int)
144
+ assert isinstance(dog["bio"], str)
145
+
146
+
147
+ @pytest.mark.vcr
148
+ @pytest.mark.parametrize(
149
+ "model_id",
150
+ (
151
+ "gemini-embedding-exp-03-07",
152
+ "gemini-embedding-exp-03-07-128",
153
+ "gemini-embedding-exp-03-07-512",
154
+ ),
155
+ )
156
+ def test_embedding(model_id, monkeypatch):
157
+ monkeypatch.setenv("LLM_GEMINI_KEY", GEMINI_API_KEY)
158
+ model = llm.get_embedding_model(model_id)
159
+ response = model.embed("Some text goes here")
160
+ expected_length = 3072
161
+ if model_id.endswith("-128"):
162
+ expected_length = 128
163
+ elif model_id.endswith("-512"):
164
+ expected_length = 512
165
+ assert len(response) == expected_length
166
+
167
+
168
+ @pytest.mark.parametrize(
169
+ "schema,expected",
170
+ [
171
+ # Test 1: Top-level keys removal
172
+ (
173
+ {
174
+ "$schema": "http://json-schema.org/draft-07/schema#",
175
+ "title": "Example Schema",
176
+ "additionalProperties": False,
177
+ "type": "object",
178
+ },
179
+ {"type": "object"},
180
+ ),
181
+ # Test 2: Preserve keys within a "properties" block
182
+ (
183
+ {
184
+ "type": "object",
185
+ "properties": {
186
+ "authors": {"type": "string"},
187
+ "title": {"type": "string"},
188
+ "reference": {"type": "string"},
189
+ "year": {"type": "string"},
190
+ },
191
+ "title": "This should be removed from the top-level",
192
+ },
193
+ {
194
+ "type": "object",
195
+ "properties": {
196
+ "authors": {"type": "string"},
197
+ "title": {"type": "string"},
198
+ "reference": {"type": "string"},
199
+ "year": {"type": "string"},
200
+ },
201
+ },
202
+ ),
203
+ # Test 3: Nested keys outside and inside properties block
204
+ (
205
+ {
206
+ "definitions": {
207
+ "info": {
208
+ "title": "Info title", # should be removed because it's not inside a "properties" block
209
+ "description": "A description",
210
+ "properties": {
211
+ "name": {
212
+ "title": "Name Title",
213
+ "type": "string",
214
+ }, # title here should be preserved
215
+ "$schema": {
216
+ "type": "string"
217
+ }, # should be preserved as it's within properties
218
+ },
219
+ }
220
+ },
221
+ "$schema": "http://example.com/schema",
222
+ },
223
+ {
224
+ "definitions": {
225
+ "info": {
226
+ "description": "A description",
227
+ "properties": {
228
+ "name": {"title": "Name Title", "type": "string"},
229
+ "$schema": {"type": "string"},
230
+ },
231
+ }
232
+ }
233
+ },
234
+ ),
235
+ # Test 4: List of schemas
236
+ (
237
+ [
238
+ {
239
+ "$schema": "http://json-schema.org/draft-07/schema#",
240
+ "type": "object",
241
+ },
242
+ {"title": "Should be removed", "type": "array"},
243
+ ],
244
+ [{"type": "object"}, {"type": "array"}],
245
+ ),
246
+ ],
247
+ )
248
+ def test_cleanup_schema(schema, expected):
249
+ # Use a deep copy so the original test data remains unchanged.
250
+ result = cleanup_schema(schema)
251
+ assert result == expected
252
+
253
+
254
+ # Tests for $ref resolution - patterns that now work with nested models
255
+ @pytest.mark.parametrize(
256
+ "schema,expected",
257
+ [
258
+ # Test 1: Direct model reference (Person with Address)
259
+ (
260
+ {
261
+ "properties": {
262
+ "name": {"type": "string"},
263
+ "address": {"$ref": "#/$defs/Address"},
264
+ },
265
+ "required": ["name", "address"],
266
+ "type": "object",
267
+ "$defs": {
268
+ "Address": {
269
+ "properties": {
270
+ "street": {"type": "string"},
271
+ "city": {"type": "string"},
272
+ },
273
+ "required": ["street", "city"],
274
+ "type": "object",
275
+ }
276
+ },
277
+ },
278
+ {
279
+ "properties": {
280
+ "name": {"type": "string"},
281
+ "address": {
282
+ "properties": {
283
+ "street": {"type": "string"},
284
+ "city": {"type": "string"},
285
+ },
286
+ "required": ["street", "city"],
287
+ "type": "object",
288
+ },
289
+ },
290
+ "required": ["name", "address"],
291
+ "type": "object",
292
+ },
293
+ ),
294
+ # Test 2: List of models (Dogs with List[Dog])
295
+ (
296
+ {
297
+ "properties": {
298
+ "dogs": {"items": {"$ref": "#/$defs/Dog"}, "type": "array"}
299
+ },
300
+ "required": ["dogs"],
301
+ "type": "object",
302
+ "$defs": {
303
+ "Dog": {
304
+ "properties": {
305
+ "name": {"type": "string"},
306
+ "age": {"type": "integer"},
307
+ },
308
+ "required": ["name", "age"],
309
+ "type": "object",
310
+ }
311
+ },
312
+ },
313
+ {
314
+ "properties": {
315
+ "dogs": {
316
+ "items": {
317
+ "properties": {
318
+ "name": {"type": "string"},
319
+ "age": {"type": "integer"},
320
+ },
321
+ "required": ["name", "age"],
322
+ "type": "object",
323
+ },
324
+ "type": "array",
325
+ }
326
+ },
327
+ "required": ["dogs"],
328
+ "type": "object",
329
+ },
330
+ ),
331
+ # Test 3: Optional model field
332
+ (
333
+ {
334
+ "properties": {
335
+ "name": {"type": "string"},
336
+ "employer": {
337
+ "anyOf": [{"$ref": "#/$defs/Company"}, {"type": "null"}]
338
+ },
339
+ },
340
+ "required": ["name"],
341
+ "type": "object",
342
+ "$defs": {
343
+ "Company": {
344
+ "properties": {"company_name": {"type": "string"}},
345
+ "required": ["company_name"],
346
+ "type": "object",
347
+ }
348
+ },
349
+ },
350
+ {
351
+ "properties": {
352
+ "name": {"type": "string"},
353
+ "employer": {
354
+ "anyOf": [
355
+ {
356
+ "properties": {"company_name": {"type": "string"}},
357
+ "required": ["company_name"],
358
+ "type": "object",
359
+ },
360
+ {"type": "null"},
361
+ ]
362
+ },
363
+ },
364
+ "required": ["name"],
365
+ "type": "object",
366
+ },
367
+ ),
368
+ # Test 4: Nested composition (Customer -> List[Order] -> List[Item])
369
+ (
370
+ {
371
+ "properties": {
372
+ "name": {"type": "string"},
373
+ "orders": {"items": {"$ref": "#/$defs/Order"}, "type": "array"},
374
+ },
375
+ "required": ["name", "orders"],
376
+ "type": "object",
377
+ "$defs": {
378
+ "Order": {
379
+ "properties": {
380
+ "items": {
381
+ "items": {"$ref": "#/$defs/Item"},
382
+ "type": "array",
383
+ }
384
+ },
385
+ "required": ["items"],
386
+ "type": "object",
387
+ },
388
+ "Item": {
389
+ "properties": {
390
+ "product_name": {"type": "string"},
391
+ "quantity": {"type": "integer"},
392
+ },
393
+ "required": ["product_name", "quantity"],
394
+ "type": "object",
395
+ },
396
+ },
397
+ },
398
+ {
399
+ "properties": {
400
+ "name": {"type": "string"},
401
+ "orders": {
402
+ "items": {
403
+ "properties": {
404
+ "items": {
405
+ "items": {
406
+ "properties": {
407
+ "product_name": {"type": "string"},
408
+ "quantity": {"type": "integer"},
409
+ },
410
+ "required": ["product_name", "quantity"],
411
+ "type": "object",
412
+ },
413
+ "type": "array",
414
+ }
415
+ },
416
+ "required": ["items"],
417
+ "type": "object",
418
+ },
419
+ "type": "array",
420
+ },
421
+ },
422
+ "required": ["name", "orders"],
423
+ "type": "object",
424
+ },
425
+ ),
426
+ ],
427
+ )
428
+ def test_cleanup_schema_with_refs(schema, expected):
429
+ """Test that $ref resolution works for various nested model patterns."""
430
+ import copy
431
+
432
+ result = cleanup_schema(copy.deepcopy(schema))
433
+ assert result == expected
434
+
435
+
436
+ @pytest.mark.vcr
437
+ def test_nested_model_direct_reference():
438
+ """Test Pattern 1: Direct model reference (Person with Address)"""
439
+
440
+ class Address(BaseModel):
441
+ street: str
442
+ city: str
443
+
444
+ class Person(BaseModel):
445
+ name: str
446
+ address: Address
447
+
448
+ model = llm.get_model("gemini-2.0-flash")
449
+ response = model.prompt(
450
+ "Create a person named Alice living in San Francisco",
451
+ key=GEMINI_API_KEY,
452
+ schema=Person,
453
+ stream=False,
454
+ )
455
+ result = json.loads(response.text())
456
+ assert "name" in result
457
+ assert "address" in result
458
+ assert "street" in result["address"]
459
+ assert "city" in result["address"]
460
+
461
+
462
+ @pytest.mark.vcr
463
+ def test_nested_model_list():
464
+ """Test Pattern 2: List of models (already covered by test_prompt_with_multiple_dogs)"""
465
+ pass # Covered by test_prompt_with_multiple_dogs
466
+
467
+
468
+ @pytest.mark.vcr
469
+ def test_nested_model_optional():
470
+ """Test Pattern 3: Optional model field"""
471
+
472
+ class Company(BaseModel):
473
+ company_name: str
474
+
475
+ class Person(BaseModel):
476
+ name: str
477
+ employer: Optional[Company]
478
+
479
+ model = llm.get_model("gemini-2.0-flash")
480
+ response = model.prompt(
481
+ "Create a person named Bob who works at TechCorp",
482
+ key=GEMINI_API_KEY,
483
+ schema=Person,
484
+ stream=False,
485
+ )
486
+ result = json.loads(response.text())
487
+ assert "name" in result
488
+ assert "employer" in result
489
+ if result["employer"] is not None:
490
+ assert "company_name" in result["employer"]
491
+
492
+
493
+ @pytest.mark.vcr
494
+ def test_nested_model_deep_composition():
495
+ """Test Pattern 4: Nested composition (Customer -> Orders -> Items)"""
496
+
497
+ class Item(BaseModel):
498
+ product_name: str
499
+ quantity: int
500
+
501
+ class Order(BaseModel):
502
+ items: List[Item]
503
+
504
+ class Customer(BaseModel):
505
+ name: str
506
+ orders: List[Order]
507
+
508
+ model = llm.get_model("gemini-2.0-flash")
509
+ response = model.prompt(
510
+ "Create a customer named Carol with 2 orders, each containing 2 items",
511
+ key=GEMINI_API_KEY,
512
+ schema=Customer,
513
+ stream=False,
514
+ )
515
+ result = json.loads(response.text())
516
+ assert "name" in result
517
+ assert "orders" in result
518
+ assert len(result["orders"]) > 0
519
+ for order in result["orders"]:
520
+ assert "items" in order
521
+ assert len(order["items"]) > 0
522
+ for item in order["items"]:
523
+ assert "product_name" in item
524
+ assert "quantity" in item
525
+
526
+
527
+ @pytest.mark.vcr
528
+ def test_cli_gemini_models(tmpdir, monkeypatch):
529
+ user_dir = tmpdir / "llm.datasette.io"
530
+ user_dir.mkdir()
531
+ monkeypatch.setenv("LLM_USER_PATH", str(user_dir))
532
+ # With no key set should error nicely
533
+ runner = CliRunner()
534
+ result = runner.invoke(cli, ["gemini", "models"])
535
+ assert result.exit_code == 1
536
+ assert (
537
+ "Error: You must set the LLM_GEMINI_KEY environment variable or use --key\n"
538
+ == result.output
539
+ )
540
+ # Try again with --key
541
+ result2 = runner.invoke(cli, ["gemini", "models", "--key", GEMINI_API_KEY])
542
+ assert result2.exit_code == 0
543
+ assert "gemini-1.5-flash-latest" in result2.output
544
+ # And with --method
545
+ result3 = runner.invoke(
546
+ cli, ["gemini", "models", "--key", GEMINI_API_KEY, "--method", "embedContent"]
547
+ )
548
+ assert result3.exit_code == 0
549
+ models = json.loads(result3.output)
550
+ for model in models:
551
+ assert "embedContent" in model["supportedGenerationMethods"]
552
+
553
+
554
+ @pytest.mark.vcr
555
+ def test_resolved_model():
556
+ model = llm.get_model("gemini-flash-latest")
557
+ response = model.prompt("hi", key=GEMINI_API_KEY)
558
+ response.text()
559
+ assert response.resolved_model == "gemini-2.5-flash-preview-09-2025"
560
+
561
+
562
+ @pytest.mark.vcr
563
+ def test_tools():
564
+ model = llm.get_model("gemini-2.0-flash")
565
+ names = ["Charles", "Sammy"]
566
+ chain_response = model.chain(
567
+ "Two names for a pet pelican",
568
+ tools=[
569
+ llm.Tool.function(lambda: names.pop(0), name="pelican_name_generator"),
570
+ ],
571
+ key=GEMINI_API_KEY,
572
+ )
573
+ text = chain_response.text()
574
+ assert text == "Okay, here are two names for a pet pelican: Charles and Sammy.\n"
575
+ # This one did three
576
+ assert len(chain_response._responses) == 3
577
+ first, second, third = chain_response._responses
578
+ assert len(first.tool_calls()) == 1
579
+ assert first.tool_calls()[0].name == "pelican_name_generator"
580
+ assert len(second.tool_calls()) == 1
581
+ assert second.tool_calls()[0].name == "pelican_name_generator"
582
+ assert second.prompt.tool_results[0].output == "Charles"
583
+ assert third.prompt.tool_results[0].output == "Sammy"
@@ -1,266 +0,0 @@
1
- from click.testing import CliRunner
2
- import llm
3
- from llm.cli import cli
4
- import nest_asyncio
5
- import json
6
- import os
7
- import pytest
8
- import pydantic
9
- from llm_gemini import cleanup_schema
10
-
11
- nest_asyncio.apply()
12
-
13
- GEMINI_API_KEY = os.environ.get("PYTEST_GEMINI_API_KEY", None) or "gm-..."
14
-
15
-
16
- @pytest.mark.vcr
17
- @pytest.mark.asyncio
18
- async def test_prompt():
19
- model = llm.get_model("gemini-1.5-flash-latest")
20
- response = model.prompt("Name for a pet pelican, just the name", key=GEMINI_API_KEY)
21
- assert str(response) == "Percy\n"
22
- assert response.response_json == {
23
- "candidates": [
24
- {
25
- "finishReason": "STOP",
26
- "safetyRatings": [
27
- {
28
- "category": "HARM_CATEGORY_HATE_SPEECH",
29
- "probability": "NEGLIGIBLE",
30
- },
31
- {
32
- "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
33
- "probability": "NEGLIGIBLE",
34
- },
35
- {
36
- "category": "HARM_CATEGORY_HARASSMENT",
37
- "probability": "NEGLIGIBLE",
38
- },
39
- {
40
- "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
41
- "probability": "NEGLIGIBLE",
42
- },
43
- ],
44
- }
45
- ],
46
- "modelVersion": "gemini-1.5-flash-latest",
47
- }
48
- assert response.token_details == {
49
- "candidatesTokenCount": 2,
50
- "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 9}],
51
- "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 2}],
52
- }
53
- assert response.input_tokens == 9
54
- assert response.output_tokens == 2
55
-
56
- # And try it async too
57
- async_model = llm.get_async_model("gemini-1.5-flash-latest")
58
- response = await async_model.prompt(
59
- "Name for a pet pelican, just the name", key=GEMINI_API_KEY
60
- )
61
- text = await response.text()
62
- assert text == "Percy\n"
63
-
64
-
65
- @pytest.mark.vcr
66
- @pytest.mark.asyncio
67
- async def test_prompt_with_pydantic_schema():
68
- class Dog(pydantic.BaseModel):
69
- name: str
70
- age: int
71
- bio: str
72
-
73
- model = llm.get_model("gemini-1.5-flash-latest")
74
- response = model.prompt(
75
- "Invent a cool dog", key=GEMINI_API_KEY, schema=Dog, stream=False
76
- )
77
- assert json.loads(response.text()) == {
78
- "age": 3,
79
- "bio": "A fluffy Samoyed with exceptional intelligence and a love for belly rubs. He's mastered several tricks, including fetching the newspaper and opening doors.",
80
- "name": "Cloud",
81
- }
82
- assert response.response_json == {
83
- "candidates": [
84
- {
85
- "finishReason": "STOP",
86
- "safetyRatings": [
87
- {
88
- "category": "HARM_CATEGORY_HATE_SPEECH",
89
- "probability": "NEGLIGIBLE",
90
- },
91
- {
92
- "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
93
- "probability": "NEGLIGIBLE",
94
- },
95
- {
96
- "category": "HARM_CATEGORY_HARASSMENT",
97
- "probability": "NEGLIGIBLE",
98
- },
99
- {
100
- "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
101
- "probability": "NEGLIGIBLE",
102
- },
103
- ],
104
- }
105
- ],
106
- "modelVersion": "gemini-1.5-flash-latest",
107
- }
108
- assert response.input_tokens == 10
109
-
110
-
111
- @pytest.mark.vcr
112
- @pytest.mark.parametrize(
113
- "model_id",
114
- (
115
- "gemini-embedding-exp-03-07",
116
- "gemini-embedding-exp-03-07-128",
117
- "gemini-embedding-exp-03-07-512",
118
- ),
119
- )
120
- def test_embedding(model_id, monkeypatch):
121
- monkeypatch.setenv("LLM_GEMINI_KEY", GEMINI_API_KEY)
122
- model = llm.get_embedding_model(model_id)
123
- response = model.embed("Some text goes here")
124
- expected_length = 3072
125
- if model_id.endswith("-128"):
126
- expected_length = 128
127
- elif model_id.endswith("-512"):
128
- expected_length = 512
129
- assert len(response) == expected_length
130
-
131
-
132
- @pytest.mark.parametrize(
133
- "schema,expected",
134
- [
135
- # Test 1: Top-level keys removal
136
- (
137
- {
138
- "$schema": "http://json-schema.org/draft-07/schema#",
139
- "title": "Example Schema",
140
- "additionalProperties": False,
141
- "type": "object",
142
- },
143
- {"type": "object"},
144
- ),
145
- # Test 2: Preserve keys within a "properties" block
146
- (
147
- {
148
- "type": "object",
149
- "properties": {
150
- "authors": {"type": "string"},
151
- "title": {"type": "string"},
152
- "reference": {"type": "string"},
153
- "year": {"type": "string"},
154
- },
155
- "title": "This should be removed from the top-level",
156
- },
157
- {
158
- "type": "object",
159
- "properties": {
160
- "authors": {"type": "string"},
161
- "title": {"type": "string"},
162
- "reference": {"type": "string"},
163
- "year": {"type": "string"},
164
- },
165
- },
166
- ),
167
- # Test 3: Nested keys outside and inside properties block
168
- (
169
- {
170
- "definitions": {
171
- "info": {
172
- "title": "Info title", # should be removed because it's not inside a "properties" block
173
- "description": "A description",
174
- "properties": {
175
- "name": {
176
- "title": "Name Title",
177
- "type": "string",
178
- }, # title here should be preserved
179
- "$schema": {
180
- "type": "string"
181
- }, # should be preserved as it's within properties
182
- },
183
- }
184
- },
185
- "$schema": "http://example.com/schema",
186
- },
187
- {
188
- "definitions": {
189
- "info": {
190
- "description": "A description",
191
- "properties": {
192
- "name": {"title": "Name Title", "type": "string"},
193
- "$schema": {"type": "string"},
194
- },
195
- }
196
- }
197
- },
198
- ),
199
- # Test 4: List of schemas
200
- (
201
- [
202
- {
203
- "$schema": "http://json-schema.org/draft-07/schema#",
204
- "type": "object",
205
- },
206
- {"title": "Should be removed", "type": "array"},
207
- ],
208
- [{"type": "object"}, {"type": "array"}],
209
- ),
210
- ],
211
- )
212
- def test_cleanup_schema(schema, expected):
213
- # Use a deep copy so the original test data remains unchanged.
214
- result = cleanup_schema(schema)
215
- assert result == expected
216
-
217
-
218
- @pytest.mark.vcr
219
- def test_cli_gemini_models(tmpdir, monkeypatch):
220
- user_dir = tmpdir / "llm.datasette.io"
221
- user_dir.mkdir()
222
- monkeypatch.setenv("LLM_USER_PATH", str(user_dir))
223
- # With no key set should error nicely
224
- runner = CliRunner()
225
- result = runner.invoke(cli, ["gemini", "models"])
226
- assert result.exit_code == 1
227
- assert (
228
- "Error: You must set the LLM_GEMINI_KEY environment variable or use --key\n"
229
- == result.output
230
- )
231
- # Try again with --key
232
- result2 = runner.invoke(cli, ["gemini", "models", "--key", GEMINI_API_KEY])
233
- assert result2.exit_code == 0
234
- assert "gemini-1.5-flash-latest" in result2.output
235
- # And with --method
236
- result3 = runner.invoke(
237
- cli, ["gemini", "models", "--key", GEMINI_API_KEY, "--method", "embedContent"]
238
- )
239
- assert result3.exit_code == 0
240
- models = json.loads(result3.output)
241
- for model in models:
242
- assert "embedContent" in model["supportedGenerationMethods"]
243
-
244
-
245
- @pytest.mark.vcr
246
- def test_tools():
247
- model = llm.get_model("gemini-2.0-flash")
248
- names = ["Charles", "Sammy"]
249
- chain_response = model.chain(
250
- "Two names for a pet pelican",
251
- tools=[
252
- llm.Tool.function(lambda: names.pop(0), name="pelican_name_generator"),
253
- ],
254
- key=GEMINI_API_KEY,
255
- )
256
- text = chain_response.text()
257
- assert text == "Okay, here are two names for a pet pelican: Charles and Sammy.\n"
258
- # This one did three
259
- assert len(chain_response._responses) == 3
260
- first, second, third = chain_response._responses
261
- assert len(first.tool_calls()) == 1
262
- assert first.tool_calls()[0].name == "pelican_name_generator"
263
- assert len(second.tool_calls()) == 1
264
- assert second.tool_calls()[0].name == "pelican_name_generator"
265
- assert second.prompt.tool_results[0].output == "Charles"
266
- assert third.prompt.tool_results[0].output == "Sammy"
File without changes
File without changes