Add request options to chat. (#341)

* Add request options to chat

Change-Id: I6f7e4c980fd7e2a14fec4c3e2d837ad745c69c9a

* fix async

Change-Id: Ia224e9e8327443a9920ce5d9a877ebb8c272e583

* fix

Change-Id: I7eed70131346c7d7ffe435c8f6909f7eb3f7e9f7

* merge from main

Change-Id: I4b92a5bc25aa7bf11bfaf31aa6c029096f3e68bc

* add tests

Change-Id: I368315f220413ba9508012721e64093372555590

* format

Change-Id: I26c7fa1f040e7d1ea16068034d78fb9f6cc13db0

Author: MarkDaoust

google/generativeai/generative_models.py

@@ -444,6 +444,7 @@ def send_message(
         stream: bool = False,
         tools: content_types.FunctionLibraryType | None = None,
         tool_config: content_types.ToolConfigType | None = None,
+        request_options: helper_types.RequestOptionsType | None = None,
     ) -> generation_types.GenerateContentResponse:
         """Sends the conversation history with the added message and returns the model's response.
 
@@ -476,6 +477,9 @@ def send_message(
              safety_settings: Overrides for the model's safety settings.
              stream: If True, yield response chunks as they are generated.
         """
+        if request_options is None:
+            request_options = {}
+
         if self.enable_automatic_function_calling and stream:
             raise NotImplementedError(
                 "Unsupported configuration: The `google.generativeai` SDK currently does not support the combination of `stream=True` and `enable_automatic_function_calling=True`."
@@ -504,6 +508,7 @@ def send_message(
             stream=stream,
             tools=tools_lib,
             tool_config=tool_config,
+            request_options=request_options,
         )
 
         self._check_response(response=response, stream=stream)
@@ -516,6 +521,7 @@ def send_message(
                 safety_settings=safety_settings,
                 stream=stream,
                 tools_lib=tools_lib,
+                request_options=request_options,
             )
 
         self._last_sent = content
@@ -546,7 +552,15 @@ def _get_function_calls(self, response) -> list[glm.FunctionCall]:
         return function_calls
 
     def _handle_afc(
-        self, *, response, history, generation_config, safety_settings, stream, tools_lib
+        self,
+        *,
+        response,
+        history,
+        generation_config,
+        safety_settings,
+        stream,
+        tools_lib,
+        request_options,
     ) -> tuple[list[glm.Content], glm.Content, generation_types.BaseGenerateContentResponse]:
 
         while function_calls := self._get_function_calls(response):
@@ -572,6 +586,7 @@ def _handle_afc(
                 safety_settings=safety_settings,
                 stream=stream,
                 tools=tools_lib,
+                request_options=request_options,
             )
 
             self._check_response(response=response, stream=stream)
@@ -588,8 +603,12 @@ async def send_message_async(
         stream: bool = False,
         tools: content_types.FunctionLibraryType | None = None,
         tool_config: content_types.ToolConfigType | None = None,
+        request_options: helper_types.RequestOptionsType | None = None,
     ) -> generation_types.AsyncGenerateContentResponse:
         """The async version of `ChatSession.send_message`."""
+        if request_options is None:
+            request_options = {}
+
         if self.enable_automatic_function_calling and stream:
             raise NotImplementedError(
                 "Unsupported configuration: The `google.generativeai` SDK currently does not support the combination of `stream=True` and `enable_automatic_function_calling=True`."
@@ -618,6 +637,7 @@ async def send_message_async(
             stream=stream,
             tools=tools_lib,
             tool_config=tool_config,
+            request_options=request_options,
         )
 
         self._check_response(response=response, stream=stream)
@@ -630,6 +650,7 @@ async def send_message_async(
                 safety_settings=safety_settings,
                 stream=stream,
                 tools_lib=tools_lib,
+                request_options=request_options,
             )
 
         self._last_sent = content
@@ -638,7 +659,15 @@ async def send_message_async(
         return response
 
     async def _handle_afc_async(
-        self, *, response, history, generation_config, safety_settings, stream, tools_lib
+        self,
+        *,
+        response,
+        history,
+        generation_config,
+        safety_settings,
+        stream,
+        tools_lib,
+        request_options,
     ) -> tuple[list[glm.Content], glm.Content, generation_types.BaseGenerateContentResponse]:
 
         while function_calls := self._get_function_calls(response):
@@ -664,6 +693,7 @@ async def _handle_afc_async(
                 safety_settings=safety_settings,
                 stream=stream,
                 tools=tools_lib,
+                request_options=request_options,
             )
 
             self._check_response(response=response, stream=stream)

tests/test_generative_models.py

@@ -12,6 +12,8 @@
 from google.generativeai import generative_models
 from google.generativeai.types import content_types
 from google.generativeai.types import generation_types
+from google.generativeai.types import helper_types
+
 
 import PIL.Image
 
@@ -37,49 +39,63 @@ def simple_response(text: str) -> glm.GenerateContentResponse:
     return glm.GenerateContentResponse({"candidates": [{"content": simple_part(text)}]})
 
 
+class MockGenerativeServiceClient:
+    def __init__(self, test):
+        self.test = test
+        self.observed_requests = []
+        self.observed_kwargs = []
+        self.responses = collections.defaultdict(list)
+
+    def generate_content(
+        self,
+        request: glm.GenerateContentRequest,
+        **kwargs,
+    ) -> glm.GenerateContentResponse:
+        self.test.assertIsInstance(request, glm.GenerateContentRequest)
+        self.observed_requests.append(request)
+        self.observed_kwargs.append(kwargs)
+        response = self.responses["generate_content"].pop(0)
+        return response
+
+    def stream_generate_content(
+        self,
+        request: glm.GetModelRequest,
+        **kwargs,
+    ) -> Iterable[glm.GenerateContentResponse]:
+        self.observed_requests.append(request)
+        self.observed_kwargs.append(kwargs)
+        response = self.responses["stream_generate_content"].pop(0)
+        return response
+
+    def count_tokens(
+        self,
+        request: glm.CountTokensRequest,
+        **kwargs,
+    ) -> Iterable[glm.GenerateContentResponse]:
+        self.observed_requests.append(request)
+        self.observed_kwargs.append(kwargs)
+        response = self.responses["count_tokens"].pop(0)
+        return response
+
+
 class CUJTests(parameterized.TestCase):
     """Tests are in order with the design doc."""
 
-    def setUp(self):
-        self.client = unittest.mock.MagicMock()
+    @property
+    def observed_requests(self):
+        return self.client.observed_requests
 
-        client_lib._client_manager.clients["generative"] = self.client
-
-        def add_client_method(f):
-            name = f.__name__
-            setattr(self.client, name, f)
-            return f
+    @property
+    def observed_kwargs(self):
+        return self.client.observed_kwargs
 
-        self.observed_requests = []
-        self.responses = collections.defaultdict(list)
+    @property
+    def responses(self):
+        return self.client.responses
 
-        @add_client_method
-        def generate_content(
-            request: glm.GenerateContentRequest,
-            **kwargs,
-        ) -> glm.GenerateContentResponse:
-            self.assertIsInstance(request, glm.GenerateContentRequest)
-            self.observed_requests.append(request)
-            response = self.responses["generate_content"].pop(0)
-            return response
-
-        @add_client_method
-        def stream_generate_content(
-            request: glm.GetModelRequest,
-            **kwargs,
-        ) -> Iterable[glm.GenerateContentResponse]:
-            self.observed_requests.append(request)
-            response = self.responses["stream_generate_content"].pop(0)
-            return response
-
-        @add_client_method
-        def count_tokens(
-            request: glm.CountTokensRequest,
-            **kwargs,
-        ) -> Iterable[glm.GenerateContentResponse]:
-            self.observed_requests.append(request)
-            response = self.responses["count_tokens"].pop(0)
-            return response
+    def setUp(self):
+        self.client = MockGenerativeServiceClient(self)
+        client_lib._client_manager.clients["generative"] = self.client
 
     def test_hello(self):
         # Generate text from text prompt
@@ -451,7 +467,7 @@ def test_copy_history(self):
         chat1 = model.start_chat()
         chat1.send_message("hello1")
 
-        chat2 = copy.deepcopy(chat1)
+        chat2 = copy.copy(chat1)
         chat2.send_message("hello2")
 
         chat1.send_message("hello3")
@@ -810,7 +826,7 @@ def test_async_code_match(self, obj, aobj):
         )
 
         asource = re.sub(" *?# type: ignore", "", asource)
-        self.assertEqual(source, asource)
+        self.assertEqual(source, asource, f"error in {obj=}")
 
     def test_repr_for_unary_non_streamed_response(self):
         model = generative_models.GenerativeModel(model_name="gemini-pro")
@@ -1208,15 +1224,30 @@ def test_repr_for_system_instruction(self):
         self.assertIn("system_instruction='Be excellent.'", result)
 
     def test_count_tokens_called_with_request_options(self):
-        self.client.count_tokens = unittest.mock.MagicMock()
-        request = unittest.mock.ANY
+        self.responses["count_tokens"].append(glm.CountTokensResponse())
         request_options = {"timeout": 120}
 
-        self.responses["count_tokens"] = [glm.CountTokensResponse(total_tokens=7)]
         model = generative_models.GenerativeModel("gemini-pro-vision")
         model.count_tokens([{"role": "user", "parts": ["hello"]}], request_options=request_options)
 
-        self.client.count_tokens.assert_called_once_with(request, **request_options)
+        self.assertEqual(request_options, self.observed_kwargs[0])
+
+    def test_chat_with_request_options(self):
+        self.responses["generate_content"].append(
+            glm.GenerateContentResponse(
+                {
+                    "candidates": [{"finish_reason": "STOP"}],
+                }
+            )
+        )
+        request_options = {"timeout": 120}
+
+        model = generative_models.GenerativeModel("gemini-pro")
+        chat = model.start_chat()
+        chat.send_message("hello", request_options=helper_types.RequestOptions(**request_options))
+
+        request_options["retry"] = None
+        self.assertEqual(request_options, self.observed_kwargs[0])
 
 
 if __name__ == "__main__":