Add basic SI and tool_config support (#257)

* Add basic SI and tool_config support

* fix tool_config=None

* Fix format.

* Fix types

Author: MarkDaoust

google/generativeai/generative_models.py

@@ -73,6 +73,8 @@ def __init__(
         safety_settings: safety_types.SafetySettingOptions | None = None,
         generation_config: generation_types.GenerationConfigType | None = None,
         tools: content_types.FunctionLibraryType | None = None,
+        tool_config: content_types.ToolConfigType | None = None,
+        system_instructions: content_types.ContentType | None = None,
     ):
         if "/" not in model_name:
             model_name = "models/" + model_name
@@ -83,6 +85,16 @@ def __init__(
         self._generation_config = generation_types.to_generation_config_dict(generation_config)
         self._tools = content_types.to_function_library(tools)
 
+        if tool_config is None:
+            self._tool_config = None
+        else:
+            self._tool_config = content_types.to_tool_config(tool_config)
+
+        if system_instructions is None:
+            self._system_instructions = None
+        else:
+            self._system_instructions = content_types.to_content(system_instructions)
+
         self._client = None
         self._async_client = None
 
@@ -110,6 +122,7 @@ def _prepare_request(
         generation_config: generation_types.GenerationConfigType | None = None,
         safety_settings: safety_types.SafetySettingOptions | None = None,
         tools: content_types.FunctionLibraryType | None,
+        tool_config: content_types.ToolConfigType | None,
     ) -> glm.GenerateContentRequest:
         """Creates a `glm.GenerateContentRequest` from raw inputs."""
         if not contents:
@@ -119,6 +132,11 @@ def _prepare_request(
         if tools_lib is not None:
             tools_lib = tools_lib.to_proto()
 
+        if tool_config is None:
+            tool_config = self._tool_config
+        else:
+            tool_config = content_types.to_tool_config(tool_config)
+
         contents = content_types.to_contents(contents)
 
         generation_config = generation_types.to_generation_config_dict(generation_config)
@@ -136,6 +154,8 @@ def _prepare_request(
             generation_config=merged_gc,
             safety_settings=merged_ss,
             tools=tools_lib,
+            tool_config=tool_config,
+            system_instructions=self._system_instructions,
         )
 
     def _get_tools_lib(
@@ -154,6 +174,7 @@ def generate_content(
         safety_settings: safety_types.SafetySettingOptions | None = None,
         stream: bool = False,
         tools: content_types.FunctionLibraryType | None = None,
+        tool_config: content_types.ToolConfigType | None = None,
         request_options: dict[str, Any] | None = None,
     ) -> generation_types.GenerateContentResponse:
         """A multipurpose function to generate responses from the model.
@@ -215,6 +236,7 @@ def generate_content(
             generation_config=generation_config,
             safety_settings=safety_settings,
             tools=tools,
+            tool_config=tool_config,
         )
         if self._client is None:
             self._client = client.get_default_generative_client()
@@ -252,6 +274,7 @@ async def generate_content_async(
         safety_settings: safety_types.SafetySettingOptions | None = None,
         stream: bool = False,
         tools: content_types.FunctionLibraryType | None = None,
+        tool_config: content_types.ToolConfigType | None = None,
         request_options: dict[str, Any] | None = None,
     ) -> generation_types.AsyncGenerateContentResponse:
         """The async version of `GenerativeModel.generate_content`."""
@@ -260,6 +283,7 @@ async def generate_content_async(
             generation_config=generation_config,
             safety_settings=safety_settings,
             tools=tools,
+            tool_config=tool_config,
         )
         if self._async_client is None:
             self._async_client = client.get_default_generative_async_client()
@@ -388,6 +412,7 @@ def send_message(
         safety_settings: safety_types.SafetySettingOptions = None,
         stream: bool = False,
         tools: content_types.FunctionLibraryType | None = None,
+        tool_config: content_types.ToolConfigType | None = None,
     ) -> generation_types.GenerateContentResponse:
         """Sends the conversation history with the added message and returns the model's response.
 
@@ -446,6 +471,7 @@ def send_message(
             safety_settings=safety_settings,
             stream=stream,
             tools=tools_lib,
+            tool_config=tool_config,
         )
 
         self._check_response(response=response, stream=stream)
@@ -529,6 +555,7 @@ async def send_message_async(
         safety_settings: safety_types.SafetySettingOptions = None,
         stream: bool = False,
         tools: content_types.FunctionLibraryType | None = None,
+        tool_config: content_types.ToolConfigType | None = None,
     ) -> generation_types.AsyncGenerateContentResponse:
         """The async version of `ChatSession.send_message`."""
         if self.enable_automatic_function_calling and stream:
@@ -557,6 +584,7 @@ async def send_message_async(
             safety_settings=safety_settings,
             stream=stream,
             tools=tools_lib,
+            tool_config=tool_config,
         )
 
         self._check_response(response=response, stream=stream)

google/generativeai/types/content_types.py

@@ -661,3 +661,69 @@ def to_function_library(lib: FunctionLibraryType | None) -> FunctionLibrary | No
         return lib
     else:
         return FunctionLibrary(tools=lib)
+
+
+FunctionCallingMode = glm.FunctionCallingConfig.Mode
+
+# fmt: off
+_FUNCTION_CALLING_MODE = {
+    1: FunctionCallingMode.AUTO,
+    FunctionCallingMode.AUTO: FunctionCallingMode.AUTO,
+    "mode_auto": FunctionCallingMode.AUTO,
+    "auto": FunctionCallingMode.AUTO,
+
+    2: FunctionCallingMode.ANY,
+    FunctionCallingMode.ANY: FunctionCallingMode.ANY,
+    "mode_any": FunctionCallingMode.ANY,
+    "any": FunctionCallingMode.ANY,
+
+    3: FunctionCallingMode.NONE,
+    FunctionCallingMode.NONE: FunctionCallingMode.NONE,
+    "mode_none": FunctionCallingMode.NONE,
+    "none": FunctionCallingMode.NONE,
+}
+# fmt: on
+
+FunctionCallingModeType = Union[FunctionCallingMode, str, int]
+
+
+def to_function_calling_mode(x: FunctionCallingModeType) -> FunctionCallingMode:
+    if isinstance(x, str):
+        x = x.lower()
+    return _FUNCTION_CALLING_MODE[x]
+
+
+class FunctionCallingConfigDict(TypedDict):
+    mode: FunctionCallingModeType
+    allowed_function_names: list[str]
+
+
+FunctionCallingConfigType = Union[FunctionCallingConfigDict, glm.FunctionCallingConfig]
+
+
+def to_function_calling_config(obj: FunctionCallingConfigType) -> glm.FunctionCallingConfig:
+    if isinstance(obj, (FunctionCallingMode, str, int)):
+        obj = {"mode": to_function_calling_mode(obj)}
+
+    return glm.FunctionCallingConfig(obj)
+
+
+class ToolConfigDict:
+    function_calling_config: FunctionCallingConfigType
+
+
+ToolConfigType = Union[ToolConfigDict, glm.ToolConfig]
+
+
+def to_tool_config(obj: ToolConfigType) -> glm.ToolConfig:
+    if isinstance(obj, glm.ToolConfig):
+        return obj
+    elif isinstance(obj, dict):
+        fcc = obj.pop("function_calling_config")
+        fcc = to_function_calling_config(fcc)
+        obj["function_calling_config"] = fcc
+        return glm.ToolConfig(**obj)
+    else:
+        raise TypeError(
+            f"Could not convert input to `glm.ToolConfig`: \n'" f"  type: {type(obj)}\n", obj
+        )