google-vertexai[patch]: serializable citation metadata, release 0.0.4 (…

…#17145) was breaking in langserve before
langchain-ai · Feb 7, 2024 · 52be84a · 52be84a
1 parent 19ff81e
commit 52be84a
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 12 deletions.
diff --git a/libs/partners/google-vertexai/langchain_google_vertexai/_utils.py b/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:

diff --git a/libs/partners/google-vertexai/pyproject.toml b/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"

diff --git a/libs/partners/google-vertexai/tests/integration_tests/test_tools.py b/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"]

diff --git a/libs/partners/google-vertexai/tests/unit_tests/test_chat_models.py b/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()