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.
- {llm_gemini-0.25 → llm_gemini-0.26.1}/PKG-INFO +9 -2
- {llm_gemini-0.25 → llm_gemini-0.26.1}/README.md +7 -0
- {llm_gemini-0.25 → llm_gemini-0.26.1}/llm_gemini.egg-info/PKG-INFO +9 -2
- {llm_gemini-0.25 → llm_gemini-0.26.1}/llm_gemini.egg-info/requires.txt +1 -1
- {llm_gemini-0.25 → llm_gemini-0.26.1}/llm_gemini.py +40 -0
- {llm_gemini-0.25 → llm_gemini-0.26.1}/pyproject.toml +2 -2
- llm_gemini-0.26.1/tests/test_gemini.py +583 -0
- llm_gemini-0.25/tests/test_gemini.py +0 -266
- {llm_gemini-0.25 → llm_gemini-0.26.1}/LICENSE +0 -0
- {llm_gemini-0.25 → llm_gemini-0.26.1}/llm_gemini.egg-info/SOURCES.txt +0 -0
- {llm_gemini-0.25 → llm_gemini-0.26.1}/llm_gemini.egg-info/dependency_links.txt +0 -0
- {llm_gemini-0.25 → llm_gemini-0.26.1}/llm_gemini.egg-info/entry_points.txt +0 -0
- {llm_gemini-0.25 → llm_gemini-0.26.1}/llm_gemini.egg-info/top_level.txt +0 -0
- {llm_gemini-0.25 → llm_gemini-0.26.1}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
|
|
1
1
|
Metadata-Version: 2.4
|
2
2
|
Name: llm-gemini
|
3
|
-
Version: 0.
|
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.
|
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.
|
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.
|
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
|
@@ -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.
|
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.
|
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
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|