llm-gemini 0.26__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.
- {llm_gemini-0.26 → llm_gemini-0.26.1}/PKG-INFO +1 -1
- {llm_gemini-0.26 → llm_gemini-0.26.1}/llm_gemini.egg-info/PKG-INFO +1 -1
- {llm_gemini-0.26 → llm_gemini-0.26.1}/llm_gemini.py +25 -0
- {llm_gemini-0.26 → llm_gemini-0.26.1}/pyproject.toml +1 -1
- llm_gemini-0.26.1/tests/test_gemini.py +583 -0
- llm_gemini-0.26/tests/test_gemini.py +0 -274
- {llm_gemini-0.26 → llm_gemini-0.26.1}/LICENSE +0 -0
- {llm_gemini-0.26 → llm_gemini-0.26.1}/README.md +0 -0
- {llm_gemini-0.26 → llm_gemini-0.26.1}/llm_gemini.egg-info/SOURCES.txt +0 -0
- {llm_gemini-0.26 → llm_gemini-0.26.1}/llm_gemini.egg-info/dependency_links.txt +0 -0
- {llm_gemini-0.26 → llm_gemini-0.26.1}/llm_gemini.egg-info/entry_points.txt +0 -0
- {llm_gemini-0.26 → llm_gemini-0.26.1}/llm_gemini.egg-info/requires.txt +0 -0
- {llm_gemini-0.26 → llm_gemini-0.26.1}/llm_gemini.egg-info/top_level.txt +0 -0
- {llm_gemini-0.26 → llm_gemini-0.26.1}/setup.cfg +0 -0
@@ -207,6 +207,11 @@ def cleanup_schema(schema, in_properties=False):
|
|
207
207
|
"Gemini supports only a subset of JSON schema"
|
208
208
|
keys_to_remove = ("$schema", "additionalProperties", "title")
|
209
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
|
+
|
210
215
|
if isinstance(schema, dict):
|
211
216
|
# Only remove keys if we're not inside a 'properties' block.
|
212
217
|
if not in_properties:
|
@@ -224,6 +229,26 @@ def cleanup_schema(schema, in_properties=False):
|
|
224
229
|
return schema
|
225
230
|
|
226
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
|
+
|
227
252
|
class _SharedGemini:
|
228
253
|
needs_key = "gemini"
|
229
254
|
key_env_var = "LLM_GEMINI_KEY"
|
@@ -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,274 +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_resolved_model():
|
247
|
-
model = llm.get_model("gemini-flash-latest")
|
248
|
-
response = model.prompt("hi", key=GEMINI_API_KEY)
|
249
|
-
response.text()
|
250
|
-
assert response.resolved_model == "gemini-2.5-flash-preview-09-2025"
|
251
|
-
|
252
|
-
|
253
|
-
@pytest.mark.vcr
|
254
|
-
def test_tools():
|
255
|
-
model = llm.get_model("gemini-2.0-flash")
|
256
|
-
names = ["Charles", "Sammy"]
|
257
|
-
chain_response = model.chain(
|
258
|
-
"Two names for a pet pelican",
|
259
|
-
tools=[
|
260
|
-
llm.Tool.function(lambda: names.pop(0), name="pelican_name_generator"),
|
261
|
-
],
|
262
|
-
key=GEMINI_API_KEY,
|
263
|
-
)
|
264
|
-
text = chain_response.text()
|
265
|
-
assert text == "Okay, here are two names for a pet pelican: Charles and Sammy.\n"
|
266
|
-
# This one did three
|
267
|
-
assert len(chain_response._responses) == 3
|
268
|
-
first, second, third = chain_response._responses
|
269
|
-
assert len(first.tool_calls()) == 1
|
270
|
-
assert first.tool_calls()[0].name == "pelican_name_generator"
|
271
|
-
assert len(second.tool_calls()) == 1
|
272
|
-
assert second.tool_calls()[0].name == "pelican_name_generator"
|
273
|
-
assert second.prompt.tool_results[0].output == "Charles"
|
274
|
-
assert third.prompt.tool_results[0].output == "Sammy"
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|