langchain-b12 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

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.7.dist-info → langchain_b12-0.1.9.dist-info}/METADATA

@@ -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.dist-info → langchain_b12-0.1.9.dist-info}/RECORD

@@ -2,8 +2,8 @@ langchain_b12/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langchain_b12/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langchain_b12/citations/citations.py,sha256=ZQvYayjQXIUaRosJ0qwL3Nc7kC8sBzmaIkE-BOslaVI,12261
 langchain_b12/genai/embeddings.py,sha256=h0Z-5PltDW9q79AjSrLemsz-_QKMB-043XXDvYSRQds,3483
-langchain_b12/genai/genai.py,sha256=Nsbwe0nlMW2p5bDtKRunDzamRcWM62zHqUNDBlbNFSg,16936
+langchain_b12/genai/genai.py,sha256=zZdGwkkaxobnA1jT07MvWHaIHKeBis4X1zJ8o5KrnHk,18841
 langchain_b12/genai/genai_utils.py,sha256=tA6UiJURK25-11vtaX4768UV47jDCYwVKIIWydD4Egw,10736
-langchain_b12-0.1.7.dist-info/METADATA,sha256=_wBqxchZ-bO6Ol5cw2IbRLgB0BgF8fa_-XqjlKLsvVU,1204
-langchain_b12-0.1.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-langchain_b12-0.1.7.dist-info/RECORD,,
+langchain_b12-0.1.9.dist-info/METADATA,sha256=CpmYdZwkej43d-DvTOR8FU65QSHAZ4iiFu7Ale8Gy0I,1235
+langchain_b12-0.1.9.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+langchain_b12-0.1.9.dist-info/RECORD,,

{langchain_b12-0.1.7.dist-info → langchain_b12-0.1.9.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.27.0
+Generator: hatchling 1.28.0
 Root-Is-Purelib: true
 Tag: py3-none-any