vertexai[patch], genai[patch]: use tool name explicitly in with struc… (#341)

Author: baskaryan

libs/genai/langchain_google_genai/_function_utils.py

@@ -304,17 +304,17 @@ def _tool_choice_to_tool_config(
 ) -> _ToolConfigDict:
     allowed_function_names: Optional[List[str]] = None
     if tool_choice is True or tool_choice == "any":
-        mode = "any"
+        mode = "ANY"
         allowed_function_names = all_names
     elif tool_choice == "auto":
-        mode = "auto"
+        mode = "AUTO"
     elif tool_choice == "none":
-        mode = "none"
+        mode = "NONE"
     elif isinstance(tool_choice, str):
-        mode = "any"
+        mode = "ANY"
         allowed_function_names = [tool_choice]
     elif isinstance(tool_choice, list):
-        mode = "any"
+        mode = "ANY"
         allowed_function_names = tool_choice
     elif isinstance(tool_choice, dict):
         if "mode" in tool_choice:
@@ -334,7 +334,7 @@ def _tool_choice_to_tool_config(
         raise ValueError(f"Unrecognized tool choice format:\n\n{tool_choice=}")
     return _ToolConfigDict(
         function_calling_config={
-            "mode": mode,
+            "mode": mode.upper(),
             "allowed_function_names": allowed_function_names,
         }
     )

libs/genai/langchain_google_genai/chat_models.py

@@ -1189,7 +1189,7 @@ def bind_tools(
 
     @property
     def _supports_tool_choice(self) -> bool:
-        return "gemini-1.5" in self.model
+        return "gemini-1.5-pro" in self.model
 
 
 def _get_tool_name(

libs/genai/tests/integration_tests/test_chat_models.py

@@ -430,15 +430,20 @@ def my_tool(name: str, age: int) -> None:
     assert tool_call_chunk["args"] == '{"age": 27.0, "name": "Erick"}'
 
 
-def test_chat_google_genai_function_calling_with_structured_output() -> None:
+# Test with model that supports tool choice (gemini 1.5) and one that doesn't
+# (gemini 1).
+@pytest.mark.parametrize("model_name", [_MODEL, "models/gemini-1.5-pro-001"])
+def test_chat_google_genai_function_calling_with_structured_output(
+    model_name: str,
+) -> None:
     class MyModel(BaseModel):
         name: str
         age: int
 
     safety = {
         HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_ONLY_HIGH
     }
-    llm = ChatGoogleGenerativeAI(model=_MODEL, safety_settings=safety)
+    llm = ChatGoogleGenerativeAI(model=model_name, safety_settings=safety)
     model = llm.with_structured_output(MyModel)
     message = HumanMessage(content="My name is Erick and I am 27 years old")
 

libs/genai/tests/unit_tests/test_function_utils.py

@@ -101,7 +101,7 @@ def test_format_dict_to_genai_function() -> None:
 def test__tool_choice_to_tool_config(choice: Any) -> None:
     expected = _ToolConfigDict(
         function_calling_config={
-            "mode": "any",
+            "mode": "ANY",
             "allowed_function_names": ["foo"],
         },
     )

libs/vertexai/langchain_google_vertexai/chat_models.py

@@ -115,6 +115,7 @@
     _ToolChoiceType,
     _ToolsType,
     _format_to_gapic_tool,
+    _ToolType,
 )
 
 logger = logging.getLogger(__name__)
@@ -1601,7 +1602,8 @@ class AnswerWithJustification(BaseModel):
             parser = JsonOutputToolsParser(first_tool_only=True) | RunnableGenerator(
                 _yield_args
             )
-        llm = self.bind_tools([schema], tool_choice=self._is_gemini_advanced)
+        tool_choice = _get_tool_name(schema) if self._is_gemini_advanced else None
+        llm = self.bind_tools([schema], tool_choice=tool_choice)
         if include_raw:
             parser_with_fallback = RunnablePassthrough.assign(
                 parsed=itemgetter("raw") | parser, parsing_error=lambda _: None
@@ -1751,3 +1753,8 @@ def _get_usage_metadata_non_gemini(raw_metadata: dict) -> Optional[UsageMetadata
             output_tokens=output_tokens,
             total_tokens=input_tokens + output_tokens,
         )
+
+
+def _get_tool_name(tool: _ToolType) -> str:
+    vertexai_tool = _format_to_gapic_tool([tool])
+    return [f.name for f in vertexai_tool.function_declarations][0]

libs/vertexai/langchain_google_vertexai/functions_utils.py

@@ -53,9 +53,8 @@ class _ToolDictLike(TypedDict):
     retrieval: Optional[_RetrievalLike]
 
 
-_ToolsType = Sequence[
-    Union[gapic.Tool, vertexai.Tool, _ToolDictLike, _FunctionDeclarationLike]
-]
+_ToolType = Union[gapic.Tool, vertexai.Tool, _ToolDictLike, _FunctionDeclarationLike]
+_ToolsType = Sequence[_ToolType]
 
 _ALLOWED_SCHEMA_FIELDS = []
 _ALLOWED_SCHEMA_FIELDS.extend([f.name for f in gapic.Schema()._pb.DESCRIPTOR.fields])

libs/vertexai/langchain_google_vertexai/model_garden.py

@@ -350,7 +350,8 @@ def with_structured_output(
     ) -> Runnable[LanguageModelInput, Union[Dict, BaseModel]]:
         """Model wrapper that returns outputs formatted to match the given schema."""
 
-        llm = self.bind_tools([schema], tool_choice="any")
+        tool_name = convert_to_anthropic_tool(schema)["name"]
+        llm = self.bind_tools([schema], tool_choice=tool_name)
         if isinstance(schema, type) and issubclass(schema, BaseModel):
             output_parser = ToolsOutputParser(
                 first_tool_only=True, pydantic_schemas=[schema]