langchain-b12 0.1.10__py3-none-any.whl → 0.1.12__py3-none-any.whl

langchain_b12/genai/genai.py

@@ -35,14 +35,7 @@ from langchain_core.tools import BaseTool
 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 pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 from langchain_b12.genai.genai_utils import (
     convert_messages_to_contents,
@@ -84,7 +77,9 @@ 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 retries indefinitely."""
+    """Maximum number of retries. Prefer `http_retry_options`, but this is kept for compatibility."""
+    http_retry_options: types.HttpRetryOptions | None = Field(default=None)
+    """HTTP retry options for API requests. If not set, max_retries will be used to create default options."""
     safety_settings: list[types.SafetySetting] | None = None
     """The default safety settings to use for all generations.
 
@@ -107,6 +102,19 @@ class ChatGenAI(BaseChatModel):
         arbitrary_types_allowed=True,
     )
 
+    @model_validator(mode="after")
+    def _setup_retry_options(self) -> "ChatGenAI":
+        """Convert max_retries to http_retry_options if not explicitly set."""
+        if self.http_retry_options is None and self.max_retries is not None:
+            self.http_retry_options = types.HttpRetryOptions(
+                attempts=self.max_retries,
+                initial_delay=1.0,
+                max_delay=60.0,
+                exp_base=2.0,
+                jitter=0.1,
+            )
+        return self
+
     @property
     def _llm_type(self) -> str:
         return "vertexai"
@@ -208,64 +216,32 @@ class ChatGenAI(BaseChatModel):
         **kwargs: Any,
     ) -> Iterator[ChatGenerationChunk]:
         system_message, contents = self._prepare_request(messages=messages)
-
-        @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(),
-            ),
+        http_options = (
+            types.HttpOptions(retry_options=self.http_retry_options)
+            if self.http_retry_options
+            else 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,
+        response_iter = self.client.models.generate_content_stream(
+            model=self.model_name,
+            contents=contents,
+            config=types.GenerateContentConfig(
+                http_options=http_options,
+                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,
                 ),
-            )
-            # 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)
+                **kwargs,
+            ),
+        )
+        total_lc_usage = None
         for response_chunk in response_iter:
             chunk, total_lc_usage = self._gemini_chunk_to_generation_chunk(
                 response_chunk, prev_total_usage=total_lc_usage
@@ -282,65 +258,33 @@ class ChatGenAI(BaseChatModel):
         **kwargs: Any,
     ) -> AsyncIterator[ChatGenerationChunk]:
         system_message, contents = self._prepare_request(messages=messages)
-
-        @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(),
-            ),
+        http_options = (
+            types.HttpOptions(retry_options=self.http_retry_options)
+            if self.http_retry_options
+            else None
         )
-        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,
+        response_iter = self.client.aio.models.generate_content_stream(
+            model=self.model_name,
+            contents=contents,
+            config=types.GenerateContentConfig(
+                http_options=http_options,
+                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,
                 ),
-            )
-            # 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:
+                **kwargs,
+            ),
+        )
+        total_lc_usage = None
+        async for response_chunk in await response_iter:
             chunk, total_lc_usage = self._gemini_chunk_to_generation_chunk(
                 response_chunk, prev_total_usage=total_lc_usage
             )

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

@@ -1,12 +1,10 @@
 Metadata-Version: 2.4
 Name: langchain-b12
-Version: 0.1.10
+Version: 0.1.12
 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.10.dist-info → langchain_b12-0.1.12.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=JoivVvUBl-mvRpl9UTC_Q6-8N4DkEHK1f7-bqI_V6Y4,20786
+langchain_b12/genai/genai.py,sha256=niECJjU4lGcA0YYU7pGfL4oDZoiHYSdKWXyvB3gEBe4,18230
 langchain_b12/genai/genai_utils.py,sha256=tA6UiJURK25-11vtaX4768UV47jDCYwVKIIWydD4Egw,10736
-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.12.dist-info/METADATA,sha256=c5A1cSulh5NeVxvBn8lkfU4peFVU7ei5eNScNDTkYds,1205
+langchain_b12-0.1.12.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+langchain_b12-0.1.12.dist-info/RECORD,,