Skip to content

Commit

Permalink
fix: Pass along model version in GenerateContentResponse (#621)
Browse files Browse the repository at this point in the history
* 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>
  • Loading branch information
Annhiluc and vishal-dharm authored Nov 26, 2024
1 parent 4c00257 commit b8772ed
Show file tree
Hide file tree
Showing 3 changed files with 14 additions and 1 deletion.
10 changes: 10 additions & 0 deletions google/generativeai/types/generation_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -359,10 +359,16 @@ def _join_chunks(chunks: Iterable[protos.GenerateContentResponse]):
else:
usage_metadata = None

if "model_version" in chunks[-1]:
model_version = chunks[-1].model_version
else:
model_version = None

return protos.GenerateContentResponse(
candidates=_join_candidate_lists(c.candidates for c in chunks),
prompt_feedback=_join_prompt_feedbacks(c.prompt_feedback for c in chunks),
usage_metadata=usage_metadata,
model_version=model_version,
)


Expand Down Expand Up @@ -539,6 +545,10 @@ def prompt_feedback(self):
def usage_metadata(self):
return self._result.usage_metadata

@property
def model_version(self):
return self._result.model_version

def __str__(self) -> str:
if self._done:
_iterator = "None"
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ def get_version():
release_status = "Development Status :: 5 - Production/Stable"

dependencies = [
"google-ai-generativelanguage==0.6.10",
"google-ai-generativelanguage==0.6.13",
"google-api-core",
"google-api-python-client",
"google-auth>=2.15.0", # 2.15 adds API key auth support
Expand Down
3 changes: 3 additions & 0 deletions tests/test_generation.py
Original file line number Diff line number Diff line change
Expand Up @@ -493,6 +493,8 @@ def test_join_chunks(self):
prompt_token_count=5
)

chunks[-1].model_version = "gemini-1.5-flash-002"

result = generation_types._join_chunks(chunks)

expected = protos.GenerateContentResponse(
Expand All @@ -509,6 +511,7 @@ def test_join_chunks(self):
],
},
"usage_metadata": {"prompt_token_count": 5},
"model_version": "gemini-1.5-flash-002",
},
)

Expand Down

0 comments on commit b8772ed

Please sign in to comment.