diff --git a/libs/partners/google-vertexai/langchain_google_vertexai/_utils.py b/libs/partners/google-vertexai/langchain_google_vertexai/_utils.py index 4f0bdf3723b02..feedd701a6197 100644 --- 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 index f9d2f890cbbab..a9283e2d5c862 100644 --- 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 index a688476bfc7ad..f239727ecac9a 100644 --- 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 index 052cf559f0fda..6bcbb6e5abff3 100644 --- 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()