added grounding tools (#224)

langchain-ai · May 16, 2024 · 8793ee8 · 8793ee8
1 parent 481ed09
commit 8793ee8
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/libs/vertexai/langchain_google_vertexai/_utils.py b/libs/vertexai/langchain_google_vertexai/_utils.py
@@ -161,9 +161,16 @@ def get_generation_info(
                 if candidate.citation_metadata
                 else None
             ),
+            "usage_metadata": usage_metadata,
         }
-        if usage_metadata:
-            info["usage_metadata"] = usage_metadata
+        try:
+            if candidate.grounding_metadata:
+                info["grounding_metadata"] = proto.Message.to_dict(
+                    candidate.grounding_metadata
+                )
+        except AttributeError:
+            pass
+        info = {k: v for k, v in info.items() if v is not None}
     # https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text-chat#response_body
     else:
         info = dataclasses.asdict(candidate)

diff --git a/libs/vertexai/langchain_google_vertexai/functions_utils.py b/libs/vertexai/langchain_google_vertexai/functions_utils.py
@@ -159,7 +159,7 @@ def _format_to_vertex_tool(
     tool: Union[VertexTool, GapicTool, _VertexToolDict, List[_FunctionDeclarationLike]],
 ) -> GapicTool:
     if isinstance(tool, VertexTool):
-        return tool
+        return tool._raw_tool
     if isinstance(tool, GapicTool):
         return tool
     elif isinstance(tool, (list, dict)):