Update and add aditional test cases

Author: shilpakancharla

google/generativeai/types/content_types.py

@@ -670,8 +670,10 @@ def _encode_fd(fd: FunctionDeclaration | protos.FunctionDeclaration) -> protos.F
 
     return fd.to_proto()
 
+
 GoogleSearchRetrievalType = Union[protos.GoogleSearchRetrieval, dict[str, float]]
 
+
 def _make_google_search_retrieval(gsr: GoogleSearchRetrievalType):
     if isinstance(gsr, protos.GoogleSearchRetrieval):
         return gsr
@@ -682,13 +684,25 @@ def _make_google_search_retrieval(gsr: GoogleSearchRetrievalType):
         if "mode" in gsr["dynamic_retrieval_config"]:
             print(to_mode(gsr["dynamic_retrieval_config"]["mode"]))
             # Create proto object from dictionary
-            gsr = {"google_search_retrieval": {"dynamic_retrieval_config": {"mode": to_mode(gsr["dynamic_retrieval_config"]["mode"]),
-                                                                        "dynamic_threshold": gsr["dynamic_retrieval_config"]["dynamic_threshold"]}}}
+            gsr = {
+                "google_search_retrieval": {
+                    "dynamic_retrieval_config": {
+                        "mode": to_mode(gsr["dynamic_retrieval_config"]["mode"]),
+                        "dynamic_threshold": gsr["dynamic_retrieval_config"]["dynamic_threshold"],
+                    }
+                }
+            }
             print(gsr)
         elif "mode" in gsr.keys():
             # Create proto object from dictionary
-            gsr = {"google_search_retrieval": {"dynamic_retrieval_config": {"mode": to_mode(gsr["mode"]),
-                                                                        "dynamic_threshold": gsr["dynamic_threshold"]}}}
+            gsr = {
+                "google_search_retrieval": {
+                    "dynamic_retrieval_config": {
+                        "mode": to_mode(gsr["mode"]),
+                        "dynamic_threshold": gsr["dynamic_threshold"],
+                    }
+                }
+            }
         return gsr
     else:
         raise TypeError(
@@ -724,10 +738,14 @@ def __init__(
             # Consistent fields
             self._function_declarations = []
             self._index = {}
-        
+
         if google_search_retrieval:
             if isinstance(google_search_retrieval, str):
-                google_search_retrieval = {"google_search_retrieval" : {"dynamic_retrieval_config": {"mode": to_mode(google_search_retrieval)}}} 
+                google_search_retrieval = {
+                    "google_search_retrieval": {
+                        "dynamic_retrieval_config": {"mode": to_mode(google_search_retrieval)}
+                    }
+                }
             else:
                 _make_google_search_retrieval(google_search_retrieval)
 
@@ -792,7 +810,7 @@ def _make_tool(tool: ToolType) -> Tool:
             google_search_retrieval = tool.google_search_retrieval
         else:
             google_search_retrieval = None
-        
+
         return Tool(
             function_declarations=tool.function_declarations,
             google_search_retrieval=google_search_retrieval,

tests/test_content.py

@@ -435,6 +435,16 @@ def test_code_execution(self, tools):
 
     @parameterized.named_parameters(
         ["string", "google_search_retrieval"],
+        ["empty_dictionary", {"google_search_retrieval": {}}],
+        ["empty_dictionary_with_dynamic_retrieval_config", {"dynamic_retrieval_config": {}}],
+        [
+            "dictionary_with_mode_integer",
+            {"google_search_retrieval": {"dynamic_retrieval_config": {"mode": 0}}},
+        ],
+        [
+            "dictionary_with_mode_string",
+            {"google_search_retrieval": {"dynamic_retrieval_config": {"mode": "DYNAMIC"}}},
+        ],
         [
             "dictionary_with_dynamic_retrieval_config",
             {
@@ -493,7 +503,7 @@ def test_search_grounding(self, tools):
             t = content_types._make_tools(tools)
             self.assertIsInstance(t[0].google_search_retrieval, protos.GoogleSearchRetrieval)
         else:
-            t = content_types._make_tool(tools)  # Pass code execution into tools
+            t = content_types._make_tool(tools)  # Pass google_search_retrieval into tools
             self.assertIsInstance(t.google_search_retrieval, protos.GoogleSearchRetrieval)
 
     def test_two_fun_is_one_tool(self):