langchain-b12 0.1.7__tar.gz → 0.1.9__tar.gz

{langchain_b12-0.1.7 → langchain_b12-0.1.9}/PKG-INFO

@@ -1,10 +1,11 @@
 Metadata-Version: 2.4
 Name: langchain-b12
-Version: 0.1.7
+Version: 0.1.9
 Summary: A reusable collection of tools and implementations for Langchain
 Author-email: Vincent Min <vincent.min@b12-consulting.com>
 Requires-Python: >=3.11
 Requires-Dist: langchain-core>=0.3.60
+Requires-Dist: tenacity>=9.1.2
 Description-Content-Type: text/markdown
 
 # Langchain B12

{langchain_b12-0.1.7 → langchain_b12-0.1.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "langchain-b12"
-version = "0.1.7"
+version = "0.1.9"
 description = "A reusable collection of tools and implementations for Langchain"
 readme = "README.md"
 authors = [
@@ -9,6 +9,7 @@ authors = [
 requires-python = ">=3.11"
 dependencies = [
     "langchain-core>=0.3.60",
+    "tenacity>=9.1.2",
 ]
 
 [dependency-groups]

{langchain_b12-0.1.7 → langchain_b12-0.1.9}/src/langchain_b12/genai/genai.py

@@ -7,10 +7,6 @@ from typing import Any, Literal, cast
 from google import genai
 from google.genai import types
 from google.oauth2 import service_account
-from langchain_b12.genai.genai_utils import (
-    convert_messages_to_contents,
-    parse_response_candidate,
-)
 from langchain_core.callbacks import (
     AsyncCallbackManagerForLLMRun,
     CallbackManagerForLLMRun,
@@ -40,6 +36,18 @@ from langchain_core.utils.function_calling import (
     convert_to_openai_tool,
 )
 from pydantic import BaseModel, ConfigDict, Field
+from tenacity import (
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    stop_never,
+    wait_exponential_jitter,
+)
+
+from langchain_b12.genai.genai_utils import (
+    convert_messages_to_contents,
+    parse_response_candidate,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -75,6 +83,8 @@ class ChatGenAI(BaseChatModel):
     """How many completions to generate for each prompt."""
     seed: int | None = None
     """Random seed for the generation."""
+    max_retries: int | None = Field(default=3)
+    """Maximum number of retries when generation fails. None disables retries."""
     safety_settings: list[types.SafetySetting] | None = None
     """The default safety settings to use for all generations.
 
@@ -173,10 +183,29 @@ class ChatGenAI(BaseChatModel):
         run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
-        stream_iter = self._stream(
-            messages, stop=stop, run_manager=run_manager, **kwargs
+        @retry(
+            reraise=True,
+            stop=stop_after_attempt(self.max_retries + 1)
+            if self.max_retries is not None
+            else stop_never,
+            wait=wait_exponential_jitter(initial=1, max=60),
+            retry=retry_if_exception_type(Exception),
+            before_sleep=lambda retry_state: logger.warning(
+                "ChatGenAI._generate failed (attempt %d/%s). "
+                "Retrying in %.2fs... Error: %s",
+                retry_state.attempt_number,
+                self.max_retries + 1 if self.max_retries is not None else "∞",
+                retry_state.next_action.sleep,
+                retry_state.outcome.exception(),
+            ),
         )
-        return generate_from_stream(stream_iter)
+        def _generate_with_retry() -> ChatResult:
+            stream_iter = self._stream(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            )
+            return generate_from_stream(stream_iter)
+
+        return _generate_with_retry()
 
     async def _agenerate(
         self,
@@ -185,10 +214,29 @@ class ChatGenAI(BaseChatModel):
         run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
-        stream_iter = self._astream(
-            messages, stop=stop, run_manager=run_manager, **kwargs
+        @retry(
+            reraise=True,
+            stop=stop_after_attempt(self.max_retries + 1)
+            if self.max_retries is not None
+            else stop_never,
+            wait=wait_exponential_jitter(initial=1, max=60),
+            retry=retry_if_exception_type(Exception),
+            before_sleep=lambda retry_state: logger.warning(
+                "ChatGenAI._agenerate failed (attempt %d/%s). "
+                "Retrying in %.2fs... Error: %s",
+                retry_state.attempt_number,
+                self.max_retries + 1 if self.max_retries is not None else "∞",
+                retry_state.next_action.sleep,
+                retry_state.outcome.exception(),
+            ),
         )
-        return await agenerate_from_stream(stream_iter)
+        async def _agenerate_with_retry() -> ChatResult:
+            stream_iter = self._astream(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            )
+            return await agenerate_from_stream(stream_iter)
+
+        return await _agenerate_with_retry()
 
     def _stream(
         self,

langchain_b12-0.1.9/tests/test_genai.py

@@ -0,0 +1,124 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+from google.genai import Client, types
+from langchain_b12.genai.genai import ChatGenAI
+from langchain_core.messages import HumanMessage
+
+
+def test_chatgenai():
+    client = MagicMock(spec=Client)
+    model = ChatGenAI(client=client, model="foo", temperature=1)
+    assert model.model_name == "foo"
+    assert model.temperature == 1
+    assert model.client == client
+
+
+def test_chatgenai_invocation():
+    client: Client = MagicMock(spec=Client)
+    client.models.generate_content_stream.return_value = iter(
+        (
+            types.GenerateContentResponse(
+                candidates=[
+                    types.Candidate(
+                        content=types.Content(parts=[types.Part(text="bar")])
+                    ),
+                ]
+            ),
+            types.GenerateContentResponse(
+                candidates=[
+                    types.Candidate(
+                        content=types.Content(parts=[types.Part(text="baz")])
+                    ),
+                ]
+            ),
+        )
+    )
+    model = ChatGenAI(client=client)
+    messages = [HumanMessage(content="foo")]
+    response = model.invoke(messages)
+    method: MagicMock = client.models.generate_content_stream
+    method.assert_called_once()
+    assert response.content == "barbaz"
+
+
+def _make_success_response():
+    """Helper to create a successful streaming response."""
+    return iter(
+        [
+            types.GenerateContentResponse(
+                candidates=[
+                    types.Candidate(
+                        content=types.Content(parts=[types.Part(text="success")])
+                    ),
+                ]
+            ),
+        ]
+    )
+
+
+@patch("langchain_b12.genai.genai.wait_exponential_jitter", return_value=lambda _: 0)
+def test_chatgenai_retry_succeeds_after_failure(mock_wait):
+    """Test that retry logic succeeds after transient failures."""
+    client: Client = MagicMock(spec=Client)
+
+    # First two calls fail, third succeeds
+    client.models.generate_content_stream.side_effect = [
+        Exception("Transient error 1"),
+        Exception("Transient error 2"),
+        _make_success_response(),
+    ]
+
+    model = ChatGenAI(client=client, max_retries=3)
+    messages = [HumanMessage(content="foo")]
+    response = model.invoke(messages)
+
+    assert response.content == "success"
+    assert client.models.generate_content_stream.call_count == 3
+
+
+@patch("langchain_b12.genai.genai.wait_exponential_jitter", return_value=lambda _: 0)
+def test_chatgenai_retry_exhausted_raises(mock_wait):
+    """Test that exception is raised after all retries are exhausted."""
+    client: Client = MagicMock(spec=Client)
+
+    # All calls fail
+    client.models.generate_content_stream.side_effect = Exception("Persistent error")
+
+    model = ChatGenAI(client=client, max_retries=2)
+    messages = [HumanMessage(content="foo")]
+
+    with pytest.raises(Exception, match="Persistent error"):
+        model.invoke(messages)
+
+    # Initial attempt + 2 retries = 3 total calls
+    assert client.models.generate_content_stream.call_count == 3
+
+
+@patch("langchain_b12.genai.genai.wait_exponential_jitter", return_value=lambda _: 0)
+def test_chatgenai_no_retry_when_max_retries_zero(mock_wait):
+    """Test that no retries occur when max_retries=0."""
+    client: Client = MagicMock(spec=Client)
+    client.models.generate_content_stream.side_effect = Exception("Error")
+
+    model = ChatGenAI(client=client, max_retries=0)
+    messages = [HumanMessage(content="foo")]
+
+    with pytest.raises(Exception, match="Error"):
+        model.invoke(messages)
+
+    # Only 1 attempt, no retries
+    assert client.models.generate_content_stream.call_count == 1
+
+
+def test_chatgenai_no_retry_on_success():
+    """Test that no retries occur when first attempt succeeds."""
+    client: Client = MagicMock(spec=Client)
+    client.models.generate_content_stream.return_value = _make_success_response()
+
+    model = ChatGenAI(client=client, max_retries=3)
+    messages = [HumanMessage(content="foo")]
+    response = model.invoke(messages)
+
+    assert response.content == "success"
+    assert client.models.generate_content_stream.call_count == 1