WIP

Author: collindutter

griptape/drivers/prompt/google_prompt_driver.py

@@ -5,7 +5,6 @@
 from typing import TYPE_CHECKING, Optional
 
 from attrs import Factory, define, field
-from google.genai.types import GenerateContentResponse
 from pydantic import BaseModel
 
 from griptape.artifacts import ActionArtifact, TextArtifact
@@ -35,7 +34,7 @@
     from collections.abc import Iterator
 
     from google.genai import Client
-    from google.genai.types import Content, ContentDict, Part
+    from google.genai.types import Content, ContentDict, GenerateContentResponse, Part, Tool
 
     from griptape.drivers.prompt.base_prompt_driver import StructuredOutputStrategy
     from griptape.tools import BaseTool
@@ -97,16 +96,14 @@ def try_run(self, prompt_stack: PromptStack) -> Message:
     def try_stream(self, prompt_stack: PromptStack) -> Iterator[DeltaMessage]:
         params = self._base_params(prompt_stack)
         logging.debug(params)
-        response = self.client.models.generate_content_stream(
-            **params,
-        )
+        response = self.client.models.generate_content_stream(**params)
 
         prompt_token_count = None
         for chunk in response:
             logger.debug(chunk.model_dump())
             usage_metadata = chunk.usage_metadata
 
-            content = self.__to_prompt_stack_delta_message_content(chunk.parts[0]) if chunk.parts else None
+            content = self.__to_prompt_stack_delta_message_content(chunk)
             # Only want to output the prompt token count once since it is static each chunk
             if prompt_token_count is None:
                 yield DeltaMessage(
@@ -135,33 +132,35 @@ def _base_params(self, prompt_stack: PromptStack) -> dict:
         params = {
             "model": self.model,
             "contents": self.__to_google_messages(prompt_stack),
-            "config": types.GenerateContentConfig(
-                **{
-                    "stop_sequences": [] if self.use_native_tools else self.tokenizer.stop_sequences,
-                    "max_output_tokens": self.max_tokens,
-                    "temperature": self.temperature,
-                    "top_p": self.top_p,
-                    "top_k": self.top_k,
-                    "system_instruction": system_instruction,
-                    **self.extra_params,
-                },
-            ),
+        }
+
+        config = {
+            "stop_sequences": [] if self.use_native_tools else self.tokenizer.stop_sequences,
+            "max_output_tokens": self.max_tokens,
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "top_k": self.top_k,
+            "system_instruction": system_instruction,
+            "automatic_function_calling": types.AutomaticFunctionCallingConfig(disable=True),
+            **self.extra_params,
         }
 
         if (
             self.structured_output_strategy == "native"
             and isinstance(prompt_stack.output_schema, type)
             and issubclass(prompt_stack.output_schema, BaseModel)
         ):
-            params["config"].response_schema = prompt_stack.output_schema
+            config["response_schema"] = prompt_stack.output_schema
 
         if prompt_stack.tools and self.use_native_tools:
-            params["tool_config"] = {"function_calling_config": {"mode": self.tool_choice}}
+            config["tool_config"] = {"function_calling_config": {"mode": self.tool_choice}}
 
             if prompt_stack.output_schema is not None and self.structured_output_strategy == "tool":
-                params["tool_config"]["function_calling_config"]["mode"] = "auto"
+                config["tool_config"]["function_calling_config"]["mode"] = "auto"
+
+            config["tools"] = self.__to_google_tools(prompt_stack.tools)
 
-            params["tools"] = self.__to_google_tools(prompt_stack.tools)
+        params["config"] = types.GenerateContentConfig(**config)
 
         return params
 
@@ -182,7 +181,7 @@ def __to_google_role(self, message: Message) -> str:
             return "model"
         return "user"
 
-    def __to_google_tools(self, tools: list[BaseTool]) -> list[dict]:
+    def __to_google_tools(self, tools: list[BaseTool]) -> list[Tool]:
         types = import_optional_dependency("google.genai.types")
 
         tool_declarations = []
@@ -194,7 +193,7 @@ def __to_google_tools(self, tools: list[BaseTool]) -> list[dict]:
                     schema = schema["properties"]["values"]
 
                 schema = remove_key_in_dict_recursively(schema, "additionalProperties")
-                tool_declaration = types.FunctionDeclaration(
+                function_declaration = types.FunctionDeclaration(
                     name=tool.to_native_tool_name(activity),
                     description=tool.activity_description(activity),
                     **(
@@ -209,8 +208,9 @@ def __to_google_tools(self, tools: list[BaseTool]) -> list[dict]:
                         else {}
                     ),
                 )
+                google_tool = types.Tool(function_declarations=[function_declaration])
 
-                tool_declarations.append(tool_declaration)
+                tool_declarations.append(google_tool)
 
         return tool_declarations
 
@@ -259,11 +259,11 @@ def __to_prompt_stack_message_content(self, content: GenerateContentResponse) ->
 
         return []
 
-    def __to_prompt_stack_delta_message_content(self, content: Part) -> BaseDeltaMessageContent:
+    def __to_prompt_stack_delta_message_content(self, content: GenerateContentResponse) -> BaseDeltaMessageContent:
         if content.text:
             return TextDeltaMessageContent(content.text)
-        if content.function_call:
-            function_call = content.function_call
+        if content.function_calls:
+            function_call = content.function_calls[0]
 
             args = function_call.args
             return ActionCallDeltaMessageContent(

griptape/utils/dict_utils.py

@@ -12,6 +12,8 @@ def remove_null_values_in_dict_recursively(d: dict) -> dict:
 def remove_key_in_dict_recursively(d: dict, key: str) -> dict:
     if isinstance(d, dict):
         return {k: remove_key_in_dict_recursively(v, key) for k, v in d.items() if k != key}
+    if isinstance(d, list):
+        return [remove_key_in_dict_recursively(v, key) for v in d]
     return d
 
 

tests/unit/drivers/prompt/test_google_prompt_driver.py

@@ -1,8 +1,7 @@
 from unittest.mock import MagicMock, Mock
 
 import pytest
-from google.generativeai.protos import FunctionCall, FunctionResponse, Part
-from google.generativeai.types import ContentDict, GenerationConfig
+from google.genai.types import ContentDict, FunctionCall, FunctionResponse, GenerationConfig, Part
 from google.protobuf.json_format import MessageToDict
 from schema import Schema
 
@@ -51,7 +50,7 @@ class TestGooglePromptDriver:
 
     @pytest.fixture()
     def mock_generative_model(self, mocker):
-        mock_generative_model = mocker.patch("google.generativeai.GenerativeModel")
+        mock_generative_model = mocker.patch("google.genai.GenerativeModel")
         mocker.patch("google.protobuf.json_format.MessageToDict").return_value = {
             "args": {"foo": "bar"},
         }
@@ -71,7 +70,7 @@ def mock_generative_model(self, mocker):
 
     @pytest.fixture()
     def mock_stream_generative_model(self, mocker):
-        mock_generative_model = mocker.patch("google.generativeai.GenerativeModel")
+        mock_generative_model = mocker.patch("google.genai.GenerativeModel")
         mocker.patch("google.protobuf.json_format.MessageToDict").return_value = {
             "args": {"foo": "bar"},
         }