From dbad3b8c9bbe33671f18ef86612d2bfd46ad959b Mon Sep 17 00:00:00 2001 From: Leonid Kuligin Date: Thu, 12 Dec 2024 15:00:36 +0100 Subject: [PATCH] enabled tool support for gemini2 (#640) --- .../langchain_google_genai/chat_models.py | 6 +++- .../tests/integration_tests/test_standard.py | 30 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/libs/genai/langchain_google_genai/chat_models.py b/libs/genai/langchain_google_genai/chat_models.py index 45ad1cec..086fa23e 100644 --- a/libs/genai/langchain_google_genai/chat_models.py +++ b/libs/genai/langchain_google_genai/chat_models.py @@ -1374,7 +1374,11 @@ def create_cached_content( @property def _supports_tool_choice(self) -> bool: - return "gemini-1.5-pro" in self.model or "gemini-1.5-flash" in self.model + return ( + "gemini-1.5-pro" in self.model + or "gemini-1.5-flash" in self.model + or "gemini-2" in self.model + ) def _get_tool_name( diff --git a/libs/genai/tests/integration_tests/test_standard.py b/libs/genai/tests/integration_tests/test_standard.py index e8feaee9..a5a7d230 100644 --- a/libs/genai/tests/integration_tests/test_standard.py +++ b/libs/genai/tests/integration_tests/test_standard.py @@ -11,6 +11,36 @@ from langchain_google_genai import ChatGoogleGenerativeAI rate_limiter = InMemoryRateLimiter(requests_per_second=0.25) +rate_limiter_2_0 = InMemoryRateLimiter(requests_per_second=0.1) + + +class TestGeminiAI2Standard(ChatModelIntegrationTests): + @property + def chat_model_class(self) -> Type[BaseChatModel]: + return ChatGoogleGenerativeAI + + @property + def chat_model_params(self) -> dict: + return { + "model": "models/gemini-2.0-flash-exp", + "rate_limiter": rate_limiter_2_0, + } + + @pytest.mark.xfail(reason="with_structured_output with JSON schema not supported.") + async def test_structured_output_async(self, model: BaseChatModel) -> None: + await super().test_structured_output_async(model) + + @pytest.mark.xfail(reason="with_structured_output with JSON schema not supported.") + def test_structured_output(self, model: BaseChatModel) -> None: + super().test_structured_output(model) + + @pytest.mark.xfail(reason="with_structured_output with JSON schema not supported.") + def test_structured_output_pydantic_2_v1(self, model: BaseChatModel) -> None: + super().test_structured_output_pydantic_2_v1(model) + + @pytest.mark.xfail(reason="investigate") + def test_bind_runnables_as_tools(self, model: BaseChatModel) -> None: + super().test_bind_runnables_as_tools(model) class TestGeminiAIStandard(ChatModelIntegrationTests):