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

langchain_b12/genai/genai.py

@@ -84,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.
 
@@ -183,29 +183,10 @@ class ChatGenAI(BaseChatModel):
         run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
-        @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(),
-            ),
+        stream_iter = self._stream(
+            messages, stop=stop, run_manager=run_manager, **kwargs
         )
-        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()
+        return generate_from_stream(stream_iter)
 
     async def _agenerate(
         self,
@@ -214,6 +195,20 @@ class ChatGenAI(BaseChatModel):
         run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
+        stream_iter = self._astream(
+            messages, stop=stop, run_manager=run_manager, **kwargs
+        )
+        return await agenerate_from_stream(stream_iter)
+
+    def _stream(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> Iterator[ChatGenerationChunk]:
+        system_message, contents = self._prepare_request(messages=messages)
+
         @retry(
             reraise=True,
             stop=stop_after_attempt(self.max_retries + 1)
@@ -222,7 +217,7 @@ class ChatGenAI(BaseChatModel):
             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). "
+                "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 "∞",
@@ -230,42 +225,47 @@ class ChatGenAI(BaseChatModel):
                 retry_state.outcome.exception(),
             ),
         )
-        async def _agenerate_with_retry() -> ChatResult:
-            stream_iter = self._astream(
-                messages, stop=stop, run_manager=run_manager, **kwargs
+        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,
+                ),
             )
-            return await agenerate_from_stream(stream_iter)
+            # 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
 
-        return await _agenerate_with_retry()
+        # Retry only covers stream initialization and first chunk
+        first_chunk, response_iter, total_lc_usage = _initiate_stream()
 
-    def _stream(
-        self,
-        messages: list[BaseMessage],
-        stop: list[str] | None = None,
-        run_manager: CallbackManagerForLLMRun | None = None,
-        **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,
-            ),
-        )
-        total_lc_usage = None
+        # 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
@@ -282,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.9.dist-info → langchain_b12-0.1.10.dist-info}/METADATA

@@ -1,10 +1,11 @@
 Metadata-Version: 2.4
 Name: langchain-b12
-Version: 0.1.9
+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-0.1.9.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=zZdGwkkaxobnA1jT07MvWHaIHKeBis4X1zJ8o5KrnHk,18841
+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.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.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,,