Search grounding

google/generativeai/types/content_types.py

@@ -71,6 +71,27 @@
     "FunctionLibraryType",
 ]
 
+Mode = protos.DynamicRetrievalConfig.Mode
+
+ModeOptions = Union[int, str, Mode]
+
+_MODE: dict[ModeOptions, Mode] = {
+    Mode.MODE_UNSPECIFIED: Mode.MODE_UNSPECIFIED,
+    0: Mode.MODE_UNSPECIFIED,
+    "mode_unspecified": Mode.MODE_UNSPECIFIED,
+    "unspecified": Mode.MODE_UNSPECIFIED,
+    Mode.MODE_DYNAMIC: Mode.MODE_DYNAMIC,
+    1: Mode.MODE_DYNAMIC,
+    "mode_dynamic": Mode.MODE_DYNAMIC,
+    "dynamic": Mode.MODE_DYNAMIC,
+}
+
+
+def to_mode(x: ModeOptions) -> Mode:
+    if isinstance(x, str):
+        x = x.lower()
+    return _MODE[x]
+
 
 def pil_to_blob(img):
     # When you load an image with PIL you get a subclass of PIL.Image
@@ -650,12 +671,56 @@ 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
+    elif isinstance(gsr, Iterable) and not isinstance(gsr, Mapping):
+        # Handle list of protos.Tool(...) and list of protos.GoogleSearchRetrieval
+        return gsr
+    elif isinstance(gsr, Mapping):
+        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"],
+                    }
+                }
+            }
+            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"],
+                    }
+                }
+            }
+        return gsr
+    else:
+        raise TypeError(
+            "Invalid input type. Expected an instance of `genai.GoogleSearchRetrieval`.\n"
+            f"However, received an object of type: {type(gsr)}.\n"
+            f"Object Value: {gsr}"
+        )
+
+
 class Tool:
-    """A wrapper for `protos.Tool`, Contains a collection of related `FunctionDeclaration` objects."""
+    """A wrapper for `protos.Tool`, Contains a collection of related `FunctionDeclaration` objects,
+    protos.CodeExecution object, and protos.GoogleSearchRetrieval object."""
 
     def __init__(
         self,
+        *,
         function_declarations: Iterable[FunctionDeclarationType] | None = None,
+        google_search_retrieval: Union[protos.GoogleSearchRetrieval, str] | None = None,
         code_execution: protos.CodeExecution | None = None,
     ):
         # The main path doesn't use this but is seems useful.
@@ -674,15 +739,32 @@ def __init__(
             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)}
+                    }
+                }
+            else:
+                _make_google_search_retrieval(google_search_retrieval)
+
         self._proto = protos.Tool(
             function_declarations=[_encode_fd(fd) for fd in self._function_declarations],
+            google_search_retrieval=google_search_retrieval,
             code_execution=code_execution,
         )
 
+        print(self._proto.google_search_retrieval)
+
     @property
     def function_declarations(self) -> list[FunctionDeclaration | protos.FunctionDeclaration]:
         return self._function_declarations
 
+    @property
+    def google_search_retrieval(self) -> protos.GoogleSearchRetrieval:
+        return self._proto.google_search_retrieval
+
     @property
     def code_execution(self) -> protos.CodeExecution:
         return self._proto.code_execution
@@ -711,7 +793,7 @@ class ToolDict(TypedDict):
 
 
 ToolType = Union[
-    Tool, protos.Tool, ToolDict, Iterable[FunctionDeclarationType], FunctionDeclarationType
+    str, Tool, protos.Tool, ToolDict, Iterable[FunctionDeclarationType], FunctionDeclarationType
 ]
 
 
@@ -723,20 +805,41 @@ def _make_tool(tool: ToolType) -> Tool:
             code_execution = tool.code_execution
         else:
             code_execution = None
-        return Tool(function_declarations=tool.function_declarations, code_execution=code_execution)
+
+        if "google_search_retrieval" in 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,
+            code_execution=code_execution,
+        )
     elif isinstance(tool, dict):
-        if "function_declarations" in tool or "code_execution" in tool:
+        if (
+            "function_declarations" in tool
+            or "google_search_retrieval" in tool
+            or "code_execution" in tool
+        ):
             return Tool(**tool)
         else:
             fd = tool
             return Tool(function_declarations=[protos.FunctionDeclaration(**fd)])
     elif isinstance(tool, str):
         if tool.lower() == "code_execution":
             return Tool(code_execution=protos.CodeExecution())
+        # Check to see if one of the mode enums matches
+        elif tool.lower() == "google_search_retrieval":
+            return Tool(google_search_retrieval=protos.GoogleSearchRetrieval())
         else:
-            raise ValueError("The only string that can be passed as a tool is 'code_execution'.")
+            raise ValueError(
+                "The only string that can be passed as a tool is 'code_execution', or one of the specified values for the `mode` parameter for google_search_retrieval."
+            )
     elif isinstance(tool, protos.CodeExecution):
         return Tool(code_execution=tool)
+    elif isinstance(tool, protos.GoogleSearchRetrieval):
+        return Tool(google_search_retrieval=tool)
     elif isinstance(tool, Iterable):
         return Tool(function_declarations=tool)
     else:
@@ -792,7 +895,7 @@ def to_proto(self):
 
 def _make_tools(tools: ToolsType) -> list[Tool]:
     if isinstance(tools, str):
-        if tools.lower() == "code_execution":
+        if tools.lower() == "code_execution" or tools.lower() == "google_search_retrieval":
             return [_make_tool(tools)]
         else:
             raise ValueError("The only string that can be passed as a tool is 'code_execution'.")

tests/test_content.py

@@ -433,6 +433,79 @@ def test_code_execution(self, tools):
             t = content_types._make_tool(tools)  # Pass code execution into tools
             self.assertIsInstance(t.code_execution, protos.CodeExecution)
 
+    @parameterized.named_parameters(
+        ["string", "google_search_retrieval"],
+        ["empty_dictionary", {"google_search_retrieval": {}}],
+        ["empty_dictionary_with_dynamic_retrieval_config", {"google_search_retrieval": {"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",
+            {
+                "google_search_retrieval": {
+                    "dynamic_retrieval_config": {"mode": "unspecified", "dynamic_threshold": 0.5}
+                }
+            },
+        ],
+        [
+            "dictionary_without_dynamic_retrieval_config",
+            {"google_search_retrieval": {"mode": "unspecified", "dynamic_threshold": 0.5}},
+        ],
+        [
+            "proto_object",
+            protos.GoogleSearchRetrieval(
+                dynamic_retrieval_config=protos.DynamicRetrievalConfig(
+                    mode="MODE_UNSPECIFIED", dynamic_threshold=0.5
+                )
+            ),
+        ],
+        [
+            "proto_passed_in",
+            protos.Tool(
+                google_search_retrieval=protos.GoogleSearchRetrieval(
+                    dynamic_retrieval_config=protos.DynamicRetrievalConfig(
+                        mode="MODE_UNSPECIFIED", dynamic_threshold=0.5
+                    )
+                )
+            ),
+        ],
+        [
+            "proto_object_list",
+            [
+                protos.GoogleSearchRetrieval(
+                    dynamic_retrieval_config=protos.DynamicRetrievalConfig(
+                        mode="MODE_UNSPECIFIED", dynamic_threshold=0.5
+                    )
+                )
+            ],
+        ],
+        [
+            "proto_passed_in_list",
+            [
+                protos.Tool(
+                    google_search_retrieval=protos.GoogleSearchRetrieval(
+                        dynamic_retrieval_config=protos.DynamicRetrievalConfig(
+                            mode="MODE_UNSPECIFIED", dynamic_threshold=0.5
+                        )
+                    )
+                )
+            ],
+        ],
+    )
+    def test_search_grounding(self, tools):
+        if isinstance(tools, Iterable):
+            t = content_types._make_tools(tools)
+            self.assertIsInstance(t[0].google_search_retrieval, protos.GoogleSearchRetrieval)
+        else:
+            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):
         def a():
             pass