diff --git a/google/generativeai/types/generation_types.py b/google/generativeai/types/generation_types.py index 45ca6b21b..42571bbe7 100644 --- a/google/generativeai/types/generation_types.py +++ b/google/generativeai/types/generation_types.py @@ -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, ) @@ -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" diff --git a/setup.py b/setup.py index c61393765..5323906aa 100644 --- a/setup.py +++ b/setup.py @@ -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 diff --git a/tests/test_generation.py b/tests/test_generation.py index a1461e8b5..67b1544b6 100644 --- a/tests/test_generation.py +++ b/tests/test_generation.py @@ -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( @@ -509,6 +511,7 @@ def test_join_chunks(self): ], }, "usage_metadata": {"prompt_token_count": 5}, + "model_version": "gemini-1.5-flash-002", }, )