Skip to content

Commit 434f5e9

Browse files
houriantouser
hourianto
authored and
user
committed
fix: make the SSE parser properly ignore non-data chunks
1 parent 2324732 commit 434f5e9

10 files changed

+339
-404
lines changed

‎google/genai/_api_client.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,7 @@ def segments(self):
191191
# we must strip before parsing.
192192
if chunk.startswith(b'data: '):
193193
chunk = chunk[len(b'data: ') :]
194-
yield json.loads(str(chunk, 'utf-8'))
194+
yield json.loads(str(chunk, 'utf-8'))
195195

196196
async def async_segments(self) -> AsyncIterator[Any]:
197197
if isinstance(self.response_stream, list):
@@ -211,7 +211,7 @@ async def async_segments(self) -> AsyncIterator[Any]:
211211
# which we must strip before parsing.
212212
if chunk.startswith('data: '):
213213
chunk = chunk[len('data: ') :]
214-
yield json.loads(chunk)
214+
yield json.loads(chunk)
215215
else:
216216
raise ValueError(
217217
'Error parsing streaming response.'

‎google/genai/_automatic_function_calling_util.py

+1-11
Original file line numberDiff line numberDiff line change
@@ -46,14 +46,6 @@ def _is_builtin_primitive_or_compound(
4646
return annotation in _py_builtin_type_to_schema_type.keys()
4747

4848

49-
def _raise_for_any_of_if_mldev(schema: types.Schema):
50-
if schema.any_of:
51-
raise ValueError(
52-
'AnyOf is not supported in function declaration schema for'
53-
' the Gemini API.'
54-
)
55-
56-
5749
def _raise_for_default_if_mldev(schema: types.Schema):
5850
if schema.default is not None:
5951
raise ValueError(
@@ -64,7 +56,6 @@ def _raise_for_default_if_mldev(schema: types.Schema):
6456

6557
def _raise_if_schema_unsupported(api_option: Literal['VERTEX_AI', 'GEMINI_API'], schema: types.Schema):
6658
if api_option == 'GEMINI_API':
67-
_raise_for_any_of_if_mldev(schema)
6859
_raise_for_default_if_mldev(schema)
6960

7061

@@ -292,8 +283,7 @@ def _parse_schema_from_parameter(
292283
),
293284
func_name,
294285
)
295-
if api_option == 'VERTEX_AI':
296-
schema.required = _get_required_fields(schema)
286+
schema.required = _get_required_fields(schema)
297287
_raise_if_schema_unsupported(api_option, schema)
298288
return schema
299289
raise ValueError(

‎google/genai/_transformers.py

-4
Original file line numberDiff line numberDiff line change
@@ -617,10 +617,6 @@ def process_schema(
617617

618618
any_of = schema.get('anyOf', None)
619619
if any_of is not None:
620-
if client and not client.vertexai:
621-
raise ValueError(
622-
'AnyOf is not supported in the response schema for the Gemini API.'
623-
)
624620
for sub_schema in any_of:
625621
# $ref is present in any_of if the schema is a union of Pydantic classes
626622
ref_key = sub_schema.get('$ref', None)

‎google/genai/caches.py

+18-18
Original file line numberDiff line numberDiff line change
@@ -174,15 +174,9 @@ def _Schema_to_mldev(
174174
if getv(from_object, ['pattern']) is not None:
175175
raise ValueError('pattern parameter is not supported in Gemini API.')
176176

177-
if getv(from_object, ['minimum']) is not None:
178-
raise ValueError('minimum parameter is not supported in Gemini API.')
179-
180177
if getv(from_object, ['default']) is not None:
181178
raise ValueError('default parameter is not supported in Gemini API.')
182179

183-
if getv(from_object, ['any_of']) is not None:
184-
raise ValueError('any_of parameter is not supported in Gemini API.')
185-
186180
if getv(from_object, ['max_length']) is not None:
187181
raise ValueError('max_length parameter is not supported in Gemini API.')
188182

@@ -195,12 +189,12 @@ def _Schema_to_mldev(
195189
if getv(from_object, ['min_properties']) is not None:
196190
raise ValueError('min_properties parameter is not supported in Gemini API.')
197191

198-
if getv(from_object, ['maximum']) is not None:
199-
raise ValueError('maximum parameter is not supported in Gemini API.')
200-
201192
if getv(from_object, ['max_properties']) is not None:
202193
raise ValueError('max_properties parameter is not supported in Gemini API.')
203194

195+
if getv(from_object, ['any_of']) is not None:
196+
setv(to_object, ['anyOf'], getv(from_object, ['any_of']))
197+
204198
if getv(from_object, ['description']) is not None:
205199
setv(to_object, ['description'], getv(from_object, ['description']))
206200

@@ -216,9 +210,15 @@ def _Schema_to_mldev(
216210
if getv(from_object, ['max_items']) is not None:
217211
setv(to_object, ['maxItems'], getv(from_object, ['max_items']))
218212

213+
if getv(from_object, ['maximum']) is not None:
214+
setv(to_object, ['maximum'], getv(from_object, ['maximum']))
215+
219216
if getv(from_object, ['min_items']) is not None:
220217
setv(to_object, ['minItems'], getv(from_object, ['min_items']))
221218

219+
if getv(from_object, ['minimum']) is not None:
220+
setv(to_object, ['minimum'], getv(from_object, ['minimum']))
221+
222222
if getv(from_object, ['nullable']) is not None:
223223
setv(to_object, ['nullable'], getv(from_object, ['nullable']))
224224

@@ -253,15 +253,9 @@ def _Schema_to_vertex(
253253
if getv(from_object, ['pattern']) is not None:
254254
setv(to_object, ['pattern'], getv(from_object, ['pattern']))
255255

256-
if getv(from_object, ['minimum']) is not None:
257-
setv(to_object, ['minimum'], getv(from_object, ['minimum']))
258-
259256
if getv(from_object, ['default']) is not None:
260257
setv(to_object, ['default'], getv(from_object, ['default']))
261258

262-
if getv(from_object, ['any_of']) is not None:
263-
setv(to_object, ['anyOf'], getv(from_object, ['any_of']))
264-
265259
if getv(from_object, ['max_length']) is not None:
266260
setv(to_object, ['maxLength'], getv(from_object, ['max_length']))
267261

@@ -274,12 +268,12 @@ def _Schema_to_vertex(
274268
if getv(from_object, ['min_properties']) is not None:
275269
setv(to_object, ['minProperties'], getv(from_object, ['min_properties']))
276270

277-
if getv(from_object, ['maximum']) is not None:
278-
setv(to_object, ['maximum'], getv(from_object, ['maximum']))
279-
280271
if getv(from_object, ['max_properties']) is not None:
281272
setv(to_object, ['maxProperties'], getv(from_object, ['max_properties']))
282273

274+
if getv(from_object, ['any_of']) is not None:
275+
setv(to_object, ['anyOf'], getv(from_object, ['any_of']))
276+
283277
if getv(from_object, ['description']) is not None:
284278
setv(to_object, ['description'], getv(from_object, ['description']))
285279

@@ -295,9 +289,15 @@ def _Schema_to_vertex(
295289
if getv(from_object, ['max_items']) is not None:
296290
setv(to_object, ['maxItems'], getv(from_object, ['max_items']))
297291

292+
if getv(from_object, ['maximum']) is not None:
293+
setv(to_object, ['maximum'], getv(from_object, ['maximum']))
294+
298295
if getv(from_object, ['min_items']) is not None:
299296
setv(to_object, ['minItems'], getv(from_object, ['min_items']))
300297

298+
if getv(from_object, ['minimum']) is not None:
299+
setv(to_object, ['minimum'], getv(from_object, ['minimum']))
300+
301301
if getv(from_object, ['nullable']) is not None:
302302
setv(to_object, ['nullable'], getv(from_object, ['nullable']))
303303

‎google/genai/models.py

+18-18
Original file line numberDiff line numberDiff line change
@@ -175,15 +175,9 @@ def _Schema_to_mldev(
175175
if getv(from_object, ['pattern']) is not None:
176176
raise ValueError('pattern parameter is not supported in Gemini API.')
177177

178-
if getv(from_object, ['minimum']) is not None:
179-
raise ValueError('minimum parameter is not supported in Gemini API.')
180-
181178
if getv(from_object, ['default']) is not None:
182179
raise ValueError('default parameter is not supported in Gemini API.')
183180

184-
if getv(from_object, ['any_of']) is not None:
185-
raise ValueError('any_of parameter is not supported in Gemini API.')
186-
187181
if getv(from_object, ['max_length']) is not None:
188182
raise ValueError('max_length parameter is not supported in Gemini API.')
189183

@@ -196,12 +190,12 @@ def _Schema_to_mldev(
196190
if getv(from_object, ['min_properties']) is not None:
197191
raise ValueError('min_properties parameter is not supported in Gemini API.')
198192

199-
if getv(from_object, ['maximum']) is not None:
200-
raise ValueError('maximum parameter is not supported in Gemini API.')
201-
202193
if getv(from_object, ['max_properties']) is not None:
203194
raise ValueError('max_properties parameter is not supported in Gemini API.')
204195

196+
if getv(from_object, ['any_of']) is not None:
197+
setv(to_object, ['anyOf'], getv(from_object, ['any_of']))
198+
205199
if getv(from_object, ['description']) is not None:
206200
setv(to_object, ['description'], getv(from_object, ['description']))
207201

@@ -217,9 +211,15 @@ def _Schema_to_mldev(
217211
if getv(from_object, ['max_items']) is not None:
218212
setv(to_object, ['maxItems'], getv(from_object, ['max_items']))
219213

214+
if getv(from_object, ['maximum']) is not None:
215+
setv(to_object, ['maximum'], getv(from_object, ['maximum']))
216+
220217
if getv(from_object, ['min_items']) is not None:
221218
setv(to_object, ['minItems'], getv(from_object, ['min_items']))
222219

220+
if getv(from_object, ['minimum']) is not None:
221+
setv(to_object, ['minimum'], getv(from_object, ['minimum']))
222+
223223
if getv(from_object, ['nullable']) is not None:
224224
setv(to_object, ['nullable'], getv(from_object, ['nullable']))
225225

@@ -254,15 +254,9 @@ def _Schema_to_vertex(
254254
if getv(from_object, ['pattern']) is not None:
255255
setv(to_object, ['pattern'], getv(from_object, ['pattern']))
256256

257-
if getv(from_object, ['minimum']) is not None:
258-
setv(to_object, ['minimum'], getv(from_object, ['minimum']))
259-
260257
if getv(from_object, ['default']) is not None:
261258
setv(to_object, ['default'], getv(from_object, ['default']))
262259

263-
if getv(from_object, ['any_of']) is not None:
264-
setv(to_object, ['anyOf'], getv(from_object, ['any_of']))
265-
266260
if getv(from_object, ['max_length']) is not None:
267261
setv(to_object, ['maxLength'], getv(from_object, ['max_length']))
268262

@@ -275,12 +269,12 @@ def _Schema_to_vertex(
275269
if getv(from_object, ['min_properties']) is not None:
276270
setv(to_object, ['minProperties'], getv(from_object, ['min_properties']))
277271

278-
if getv(from_object, ['maximum']) is not None:
279-
setv(to_object, ['maximum'], getv(from_object, ['maximum']))
280-
281272
if getv(from_object, ['max_properties']) is not None:
282273
setv(to_object, ['maxProperties'], getv(from_object, ['max_properties']))
283274

275+
if getv(from_object, ['any_of']) is not None:
276+
setv(to_object, ['anyOf'], getv(from_object, ['any_of']))
277+
284278
if getv(from_object, ['description']) is not None:
285279
setv(to_object, ['description'], getv(from_object, ['description']))
286280

@@ -296,9 +290,15 @@ def _Schema_to_vertex(
296290
if getv(from_object, ['max_items']) is not None:
297291
setv(to_object, ['maxItems'], getv(from_object, ['max_items']))
298292

293+
if getv(from_object, ['maximum']) is not None:
294+
setv(to_object, ['maximum'], getv(from_object, ['maximum']))
295+
299296
if getv(from_object, ['min_items']) is not None:
300297
setv(to_object, ['minItems'], getv(from_object, ['min_items']))
301298

299+
if getv(from_object, ['minimum']) is not None:
300+
setv(to_object, ['minimum'], getv(from_object, ['minimum']))
301+
302302
if getv(from_object, ['nullable']) is not None:
303303
setv(to_object, ['nullable'], getv(from_object, ['nullable']))
304304

0 commit comments

Comments
 (0)