Skip to content

Commit b8772ed

Browse files
fix: Pass along model version in GenerateContentResponse (#621)
* fix: Pass along model_version in GenerateContentResponse. * Revert autogenerated doc files from 94eb16e. * Fix 'argument list too long' error and add couple vision examples (#634) * Update google-ai-generativelanguage version in requirements. * Format updated generation_types and test using black. --------- Co-authored-by: Vishal Dharmadhikari <61256217+vishal-dharm@users.noreply.github.com>
1 parent 4c00257 commit b8772ed

File tree

3 files changed

+14
-1
lines changed

3 files changed

+14
-1
lines changed

google/generativeai/types/generation_types.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -359,10 +359,16 @@ def _join_chunks(chunks: Iterable[protos.GenerateContentResponse]):
359359
else:
360360
usage_metadata = None
361361

362+
if "model_version" in chunks[-1]:
363+
model_version = chunks[-1].model_version
364+
else:
365+
model_version = None
366+
362367
return protos.GenerateContentResponse(
363368
candidates=_join_candidate_lists(c.candidates for c in chunks),
364369
prompt_feedback=_join_prompt_feedbacks(c.prompt_feedback for c in chunks),
365370
usage_metadata=usage_metadata,
371+
model_version=model_version,
366372
)
367373

368374

@@ -539,6 +545,10 @@ def prompt_feedback(self):
539545
def usage_metadata(self):
540546
return self._result.usage_metadata
541547

548+
@property
549+
def model_version(self):
550+
return self._result.model_version
551+
542552
def __str__(self) -> str:
543553
if self._done:
544554
_iterator = "None"

setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ def get_version():
4242
release_status = "Development Status :: 5 - Production/Stable"
4343

4444
dependencies = [
45-
"google-ai-generativelanguage==0.6.10",
45+
"google-ai-generativelanguage==0.6.13",
4646
"google-api-core",
4747
"google-api-python-client",
4848
"google-auth>=2.15.0", # 2.15 adds API key auth support

tests/test_generation.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -493,6 +493,8 @@ def test_join_chunks(self):
493493
prompt_token_count=5
494494
)
495495

496+
chunks[-1].model_version = "gemini-1.5-flash-002"
497+
496498
result = generation_types._join_chunks(chunks)
497499

498500
expected = protos.GenerateContentResponse(
@@ -509,6 +511,7 @@ def test_join_chunks(self):
509511
],
510512
},
511513
"usage_metadata": {"prompt_token_count": 5},
514+
"model_version": "gemini-1.5-flash-002",
512515
},
513516
)
514517

0 commit comments

Comments
 (0)