langchain-b12 0.1.8__py3-none-any.whl → 0.1.10__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__)
 
@@ -76,7 +84,7 @@ class ChatGenAI(BaseChatModel):
     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."""
+    """Maximum number of retries when generation fails. None retries indefinitely."""
     safety_settings: list[types.SafetySetting] | None = None
     """The default safety settings to use for all generations.
 
@@ -175,24 +183,10 @@ class ChatGenAI(BaseChatModel):
         run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
-        attempts = 0
-        while True:
-            try:
-                stream_iter = self._stream(
-                    messages, stop=stop, run_manager=run_manager, **kwargs
-                )
-                return generate_from_stream(stream_iter)
-            except Exception as e:  # noqa: BLE001
-                if self.max_retries is None or attempts >= self.max_retries:
-                    raise
-                attempts += 1
-                logger.warning(
-                    "ChatGenAI._generate failed (attempt %d/%d). "
-                    "Retrying... Error: %s",
-                    attempts,
-                    self.max_retries,
-                    e,
-                )
+        stream_iter = self._stream(
+            messages, stop=stop, run_manager=run_manager, **kwargs
+        )
+        return generate_from_stream(stream_iter)
 
     async def _agenerate(
         self,
@@ -201,24 +195,10 @@ class ChatGenAI(BaseChatModel):
         run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
-        attempts = 0
-        while True:
-            try:
-                stream_iter = self._astream(
-                    messages, stop=stop, run_manager=run_manager, **kwargs
-                )
-                return await agenerate_from_stream(stream_iter)
-            except Exception as e:  # noqa: BLE001
-                if self.max_retries is None or attempts >= self.max_retries:
-                    raise
-                attempts += 1
-                logger.warning(
-                    "ChatGenAI._agenerate failed (attempt %d/%d). "
-                    "Retrying... Error: %s",
-                    attempts,
-                    self.max_retries,
-                    e,
-                )
+        stream_iter = self._astream(
+            messages, stop=stop, run_manager=run_manager, **kwargs
+        )
+        return await agenerate_from_stream(stream_iter)
 
     def _stream(
         self,
@@ -228,26 +208,64 @@ class ChatGenAI(BaseChatModel):
         **kwargs: Any,
     ) -> Iterator[ChatGenerationChunk]:
         system_message, contents = self._prepare_request(messages=messages)
-        response_iter = self.client.models.generate_content_stream(
-            model=self.model_name,
-            contents=contents,
-            config=types.GenerateContentConfig(
-                system_instruction=system_message,
-                temperature=self.temperature,
-                top_k=self.top_k,
-                top_p=self.top_p,
-                max_output_tokens=self.max_output_tokens,
-                candidate_count=self.n,
-                stop_sequences=stop or self.stop,
-                safety_settings=self.safety_settings,
-                thinking_config=self.thinking_config,
-                automatic_function_calling=types.AutomaticFunctionCallingConfig(
-                    disable=True,
-                ),
-                **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._stream failed to start (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(),
             ),
         )
-        total_lc_usage = None
+        def _initiate_stream() -> tuple[
+            ChatGenerationChunk,
+            Iterator[types.GenerateContentResponse],
+            UsageMetadata | None,
+        ]:
+            """Initialize stream and fetch first chunk. Retries only apply here."""
+            response_iter = self.client.models.generate_content_stream(
+                model=self.model_name,
+                contents=contents,
+                config=types.GenerateContentConfig(
+                    system_instruction=system_message,
+                    temperature=self.temperature,
+                    top_k=self.top_k,
+                    top_p=self.top_p,
+                    max_output_tokens=self.max_output_tokens,
+                    candidate_count=self.n,
+                    stop_sequences=stop or self.stop,
+                    safety_settings=self.safety_settings,
+                    thinking_config=self.thinking_config,
+                    automatic_function_calling=types.AutomaticFunctionCallingConfig(
+                        disable=True,
+                    ),
+                    **kwargs,
+                ),
+            )
+            # Fetch first chunk to ensure connection is established
+            first_response = next(iter(response_iter))
+            first_chunk, total_usage = self._gemini_chunk_to_generation_chunk(
+                first_response, prev_total_usage=None
+            )
+            return first_chunk, response_iter, total_usage
+
+        # Retry only covers stream initialization and first chunk
+        first_chunk, response_iter, total_lc_usage = _initiate_stream()
+
+        # Yield first chunk
+        if run_manager and isinstance(first_chunk.message.content, str):
+            run_manager.on_llm_new_token(first_chunk.message.content)
+        yield first_chunk
+
+        # Continue streaming without retry (retries during streaming are not well defined)
         for response_chunk in response_iter:
             chunk, total_lc_usage = self._gemini_chunk_to_generation_chunk(
                 response_chunk, prev_total_usage=total_lc_usage
@@ -264,27 +282,65 @@ class ChatGenAI(BaseChatModel):
         **kwargs: Any,
     ) -> AsyncIterator[ChatGenerationChunk]:
         system_message, contents = self._prepare_request(messages=messages)
-        response_iter = self.client.aio.models.generate_content_stream(
-            model=self.model_name,
-            contents=contents,
-            config=types.GenerateContentConfig(
-                system_instruction=system_message,
-                temperature=self.temperature,
-                top_k=self.top_k,
-                top_p=self.top_p,
-                max_output_tokens=self.max_output_tokens,
-                candidate_count=self.n,
-                stop_sequences=stop or self.stop,
-                safety_settings=self.safety_settings,
-                thinking_config=self.thinking_config,
-                automatic_function_calling=types.AutomaticFunctionCallingConfig(
-                    disable=True,
-                ),
-                **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._astream failed to start (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(),
             ),
         )
-        total_lc_usage = None
-        async for response_chunk in await response_iter:
+        async def _initiate_stream() -> tuple[
+            ChatGenerationChunk,
+            AsyncIterator[types.GenerateContentResponse],
+            UsageMetadata | None,
+        ]:
+            """Initialize stream and fetch first chunk. Retries only apply here."""
+            response_iter = await self.client.aio.models.generate_content_stream(
+                model=self.model_name,
+                contents=contents,
+                config=types.GenerateContentConfig(
+                    system_instruction=system_message,
+                    temperature=self.temperature,
+                    top_k=self.top_k,
+                    top_p=self.top_p,
+                    max_output_tokens=self.max_output_tokens,
+                    candidate_count=self.n,
+                    stop_sequences=stop or self.stop,
+                    safety_settings=self.safety_settings,
+                    thinking_config=self.thinking_config,
+                    automatic_function_calling=types.AutomaticFunctionCallingConfig(
+                        disable=True,
+                    ),
+                    **kwargs,
+                ),
+            )
+            # Fetch first chunk to ensure connection is established
+            first_response = await response_iter.__anext__()
+            first_chunk, total_usage = self._gemini_chunk_to_generation_chunk(
+                first_response, prev_total_usage=None
+            )
+            return first_chunk, response_iter, total_usage
+
+        # Retry only covers stream initialization and first chunk
+        first_chunk, response_iter, total_lc_usage = await _initiate_stream()
+
+        # Yield first chunk
+        if run_manager and isinstance(first_chunk.message.content, str):
+            await run_manager.on_llm_new_token(first_chunk.message.content)
+        yield first_chunk
+
+        # Continue streaming without retry (retries during streaming are not well defined)
+        async for response_chunk in response_iter:
             chunk, total_lc_usage = self._gemini_chunk_to_generation_chunk(
                 response_chunk, prev_total_usage=total_lc_usage
             )

{langchain_b12-0.1.8.dist-info → langchain_b12-0.1.10.dist-info}/METADATA

@@ -1,10 +1,12 @@
 Metadata-Version: 2.4
 Name: langchain-b12
-Version: 0.1.8
+Version: 0.1.10
 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: pytest-anyio>=0.0.0
+Requires-Dist: tenacity>=9.1.2
 Description-Content-Type: text/markdown
 
 # Langchain B12

{langchain_b12-0.1.8.dist-info → langchain_b12-0.1.10.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=4Q0j2YsAPLrHhIy_pYXoncb4hqJJkEUpxB3oD3qaECI,18120
+langchain_b12/genai/genai.py,sha256=JoivVvUBl-mvRpl9UTC_Q6-8N4DkEHK1f7-bqI_V6Y4,20786
 langchain_b12/genai/genai_utils.py,sha256=tA6UiJURK25-11vtaX4768UV47jDCYwVKIIWydD4Egw,10736
-langchain_b12-0.1.8.dist-info/METADATA,sha256=0-KZr-PXjE16ar4LpQbdWHX8CrViLBxlfV9uGwE0Qw0,1204
-langchain_b12-0.1.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-langchain_b12-0.1.8.dist-info/RECORD,,
+langchain_b12-0.1.10.dist-info/METADATA,sha256=LcMlsuxt4CO9Q-FeGqR3tx2mhmEhEMAagkWvBTmUtbo,1271
+langchain_b12-0.1.10.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+langchain_b12-0.1.10.dist-info/RECORD,,

{langchain_b12-0.1.8.dist-info → langchain_b12-0.1.10.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