google-vertexai[patch]: function calling integration test

libs/partners/google-vertexai/langchain_google_vertexai/_utils.py

@@ -5,6 +5,7 @@
 from typing import Any, Callable, Dict, Optional, Union
 
 import google.api_core
+import proto  # type: ignore[import-untyped]
 from google.api_core.gapic_v1.client_info import ClientInfo
 from google.cloud import storage
 from langchain_core.callbacks import (
@@ -114,7 +115,11 @@ def get_generation_info(
                 }
                 for rating in candidate.safety_ratings
             ],
-            "citation_metadata": candidate.citation_metadata,
+            "citation_metadata": (
+                proto.Message.to_dict(candidate.citation_metadata)
+                if candidate.citation_metadata
+                else None
+            ),
         }
     # https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text-chat#response_body
     else:

libs/partners/google-vertexai/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langchain-google-vertexai"
-version = "0.0.3"
+version = "0.0.4"
 description = "An integration package connecting GoogleVertexAI and LangChain"
 authors = []
 readme = "README.md"

libs/partners/google-vertexai/tests/integration_tests/test_chat_models.py

@@ -1,4 +1,6 @@
 """Test ChatGoogleVertexAI chat model."""
+
+import json
 from typing import Optional, cast
 
 import pytest
@@ -9,6 +11,7 @@
     SystemMessage,
 )
 from langchain_core.outputs import ChatGeneration, LLMResult
+from langchain_core.pydantic_v1 import BaseModel
 
 from langchain_google_vertexai.chat_models import ChatVertexAI
 
@@ -220,3 +223,26 @@ def test_chat_vertexai_system_message(model_name: Optional[str]) -> None:
     response = model([system_message, message1, message2, message3])
     assert isinstance(response, AIMessage)
     assert isinstance(response.content, str)
+
+
+def test_chat_vertexai_gemini_function_calling() -> None:
+    class MyModel(BaseModel):
+        name: str
+        age: int
+
+    model = ChatVertexAI(model_name="gemini-pro").bind(functions=[MyModel])
+    message = HumanMessage(content="My name is Erick and I am 27 years old")
+    response = model.invoke([message])
+    assert isinstance(response, AIMessage)
+    assert isinstance(response.content, str)
+    assert response.content == ""
+    function_call = response.additional_kwargs.get("function_call")
+    assert function_call
+    assert function_call["name"] == "MyModel"
+    arguments_str = function_call.get("arguments")
+    assert arguments_str
+    arguments = json.loads(arguments_str)
+    assert arguments == {
+        "name": "Erick",
+        "age": 27.0,
+    }

libs/partners/google-vertexai/tests/integration_tests/test_tools.py

@@ -1,6 +1,6 @@
 import os
 import re
-from typing import List, Union
+from typing import Any, List, Union
 
 from langchain_core.agents import AgentAction, AgentActionMessageLog, AgentFinish
 from langchain_core.messages import AIMessageChunk
@@ -44,11 +44,11 @@ def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
 
 
 def test_tools() -> None:
-    from langchain.agents import AgentExecutor  # type: ignore[import-not-found]
-    from langchain.agents.format_scratchpad import (  # type: ignore[import-not-found]
+    from langchain.agents import AgentExecutor
+    from langchain.agents.format_scratchpad import (
         format_to_openai_function_messages,
     )
-    from langchain.chains import LLMMathChain  # type: ignore[import-not-found]
+    from langchain.chains import LLMMathChain
 
     llm = ChatVertexAI(model_name="gemini-pro")
     math_chain = LLMMathChain.from_llm(llm=llm)
@@ -67,8 +67,8 @@ def test_tools() -> None:
     )
     llm_with_tools = llm.bind(functions=tools)
 
-    agent = (
-        {  # type: ignore[var-annotated]
+    agent: Any = (
+        {
             "input": lambda x: x["input"],
             "agent_scratchpad": lambda x: format_to_openai_function_messages(
                 x["intermediate_steps"]
@@ -115,7 +115,7 @@ def test_multiple_tools() -> None:
     from langchain.agents import AgentExecutor
     from langchain.agents.format_scratchpad import format_to_openai_function_messages
     from langchain.chains import LLMMathChain
-    from langchain.utilities import (  # type: ignore[import-not-found]
+    from langchain.utilities import (
         GoogleSearchAPIWrapper,
     )
 
@@ -149,8 +149,8 @@ def test_multiple_tools() -> None:
     )
     llm_with_tools = llm.bind(functions=tools)
 
-    agent = (
-        {  # type: ignore[var-annotated]
+    agent: Any = (
+        {
             "input": lambda x: x["input"],
             "agent_scratchpad": lambda x: format_to_openai_function_messages(
                 x["intermediate_steps"]

libs/partners/google-vertexai/tests/unit_tests/test_chat_models.py

@@ -187,7 +187,7 @@ def test_default_params_gemini() -> None:
             StubGeminiResponse(
                 text="Goodbye",
                 content=Mock(parts=[Mock(function_call=None)]),
-                citation_metadata=Mock(),
+                citation_metadata=None,
             )
         ]
         mock_chat = MagicMock()