x

langchain-ai · Feb 7, 2024 · 96aeb2b · 96aeb2b
1 parent f53c557
commit 96aeb2b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 20 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,22 +115,7 @@ def get_generation_info(
                 }
                 for rating in candidate.safety_ratings
             ],
-            "citation_metadata": (
-                {
-                    "citations": [
-                        {
-                            "start_index": citation.start_index,
-                            "end_index": citation.end_index,
-                            "publication_date": citation.publication_date,
-                            "title": citation.title,
-                            "uri": citation.uri,
-                        }
-                        for citation in candidate.citation_metadata.citations
-                    ]
-                }
-                if candidate.citation_metadata
-                else None
-            ),
+            "citation_metadata": proto.Message.to_dict(candidate.citation_metadata),
         }
     # https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text-chat#response_body
     else:

diff --git a/libs/partners/google-vertexai/tests/integration_tests/test_tools.py b/libs/partners/google-vertexai/tests/integration_tests/test_tools.py
@@ -68,7 +68,7 @@ def test_tools() -> None:
     llm_with_tools = llm.bind(functions=tools)
 
     agent: Any = (
-        {  # type: ignore[var-annotated]
+        {
             "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"]