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

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

@@ -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 → langchain_b12-0.1.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "langchain-b12"
-version = "0.1.9"
+version = "0.1.10"
 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",
+    "pytest-anyio>=0.0.0",
     "tenacity>=9.1.2",
 ]
 

{langchain_b12-0.1.9 → langchain_b12-0.1.10}/src/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.10/tests/test_genai.py

@@ -0,0 +1,279 @@
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from google.genai import Client, types
+from langchain_b12.genai.genai import ChatGenAI
+from langchain_core.messages import HumanMessage
+
+
+def _make_response_chunk(text: str) -> types.GenerateContentResponse:
+    """Helper to create a response chunk."""
+    return types.GenerateContentResponse(
+        candidates=[
+            types.Candidate(content=types.Content(parts=[types.Part(text=text)]))
+        ]
+    )
+
+
+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(
+        (
+            _make_response_chunk("bar"),
+            _make_response_chunk("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_iter():
+    """Helper to create a successful streaming iterator."""
+    return iter([_make_response_chunk("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_iter(),
+    ]
+
+    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_iter()
+
+    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
+
+
+# --- Streaming behavior tests ---
+
+
+def test_stream_yields_chunks_immediately():
+    """Test that stream yields chunks as they arrive, not buffered."""
+    client: Client = MagicMock(spec=Client)
+    chunks_yielded: list[str] = []
+
+    def mock_stream():
+        for text in ["chunk1", "chunk2", "chunk3"]:
+            # Track when chunks are yielded from the source
+            chunks_yielded.append(f"source:{text}")
+            yield _make_response_chunk(text)
+
+    client.models.generate_content_stream.return_value = mock_stream()
+
+    model = ChatGenAI(client=client, max_retries=3)
+    messages = [HumanMessage(content="foo")]
+
+    received: list[str] = []
+    for chunk in model.stream(messages):
+        received.append(chunk.content)
+        # After receiving each chunk, check that source yielded it
+        assert len(received) == len([c for c in chunks_yielded if c.startswith("source:")])
+
+    assert received == ["chunk1", "chunk2", "chunk3"]
+
+
+@patch("langchain_b12.genai.genai.wait_exponential_jitter", return_value=lambda _: 0)
+def test_stream_no_retry_after_first_chunk(mock_wait):
+    """Test that errors after first chunk are NOT retried."""
+    client: Client = MagicMock(spec=Client)
+
+    def failing_after_first():
+        yield _make_response_chunk("first")
+        raise Exception("Mid-stream error")
+
+    client.models.generate_content_stream.return_value = failing_after_first()
+
+    model = ChatGenAI(client=client, max_retries=3)
+    messages = [HumanMessage(content="foo")]
+
+    chunks = []
+    with pytest.raises(Exception, match="Mid-stream error"):
+        for chunk in model.stream(messages):
+            chunks.append(chunk.content)
+
+    # First chunk was received
+    assert chunks == ["first"]
+    # Only one call - no retry after first chunk
+    assert client.models.generate_content_stream.call_count == 1
+
+
+@patch("langchain_b12.genai.genai.wait_exponential_jitter", return_value=lambda _: 0)
+def test_stream_retry_on_first_chunk_failure(mock_wait):
+    """Test that failure on first chunk triggers retry."""
+    client: Client = MagicMock(spec=Client)
+
+    def fail_on_first_next():
+        raise Exception("First chunk error")
+        yield  # Make it a generator
+
+    def success_stream():
+        yield _make_response_chunk("success1")
+        yield _make_response_chunk("success2")
+
+    client.models.generate_content_stream.side_effect = [
+        fail_on_first_next(),
+        success_stream(),
+    ]
+
+    model = ChatGenAI(client=client, max_retries=3)
+    messages = [HumanMessage(content="foo")]
+
+    chunks = [chunk.content for chunk in model.stream(messages)]
+    assert chunks == ["success1", "success2"]
+    assert client.models.generate_content_stream.call_count == 2
+
+
+# --- Async streaming tests ---
+
+
+async def _async_iter(items):
+    """Helper to create an async iterator from items."""
+    for item in items:
+        yield item
+
+
+@pytest.mark.anyio
+async def test_astream_yields_chunks_immediately():
+    """Test that async stream yields chunks as they arrive."""
+    client: Client = MagicMock(spec=Client)
+
+    chunks = [
+        _make_response_chunk("async1"),
+        _make_response_chunk("async2"),
+        _make_response_chunk("async3"),
+    ]
+
+    # generate_content_stream returns a coroutine that resolves to async iterator
+    client.aio.models.generate_content_stream = AsyncMock(
+        return_value=_async_iter(chunks)
+    )
+
+    model = ChatGenAI(client=client, max_retries=3)
+    messages = [HumanMessage(content="foo")]
+
+    received: list[str] = []
+    async for chunk in model.astream(messages):
+        received.append(chunk.content)
+
+    assert received == ["async1", "async2", "async3"]
+
+
+@pytest.mark.anyio
+@patch("langchain_b12.genai.genai.wait_exponential_jitter", return_value=lambda _: 0)
+async def test_astream_no_retry_after_first_chunk(mock_wait):
+    """Test that errors after first chunk are NOT retried in async."""
+    client: Client = MagicMock(spec=Client)
+
+    async def failing_after_first():
+        yield _make_response_chunk("first")
+        raise Exception("Async mid-stream error")
+
+    client.aio.models.generate_content_stream = AsyncMock(
+        return_value=failing_after_first()
+    )
+
+    model = ChatGenAI(client=client, max_retries=3)
+    messages = [HumanMessage(content="foo")]
+
+    chunks = []
+    with pytest.raises(Exception, match="Async mid-stream error"):
+        async for chunk in model.astream(messages):
+            chunks.append(chunk.content)
+
+    assert chunks == ["first"]
+    assert client.aio.models.generate_content_stream.call_count == 1
+
+
+@pytest.mark.anyio
+@patch("langchain_b12.genai.genai.wait_exponential_jitter", return_value=lambda _: 0)
+async def test_astream_retry_succeeds_after_failure(mock_wait):
+    """Test that async retry logic works for initial failures."""
+    client: Client = MagicMock(spec=Client)
+
+    call_count = 0
+
+    async def side_effect_fn(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+            raise Exception("Async transient error")
+        return _async_iter([_make_response_chunk("async_success")])
+
+    client.aio.models.generate_content_stream = AsyncMock(side_effect=side_effect_fn)
+
+    model = ChatGenAI(client=client, max_retries=3)
+    messages = [HumanMessage(content="foo")]
+
+    chunks = []
+    async for chunk in model.astream(messages):
+        chunks.append(chunk.content)
+
+    assert chunks == ["async_success"]
+    assert client.aio.models.generate_content_stream.call_count == 2

{langchain_b12-0.1.9 → langchain_b12-0.1.10}/uv.lock

@@ -252,10 +252,11 @@ wheels = [
 
 [[package]]
 name = "langchain-b12"
-version = "0.1.8"
+version = "0.1.9"
 source = { editable = "." }
 dependencies = [
     { name = "langchain-core" },
+    { name = "pytest-anyio" },
     { name = "tenacity" },
 ]
 
@@ -275,6 +276,7 @@ google = [
 [package.metadata]
 requires-dist = [
     { name = "langchain-core", specifier = ">=0.3.60" },
+    { name = "pytest-anyio", specifier = ">=0.0.0" },
     { name = "tenacity", specifier = ">=9.1.2" },
 ]
 
@@ -621,6 +623,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/29/16/c8a903f4c4dffe7a12843191437d7cd8e32751d5de349d45d3fe69544e87/pytest-8.4.1-py3-none-any.whl", hash = "sha256:539c70ba6fcead8e78eebbf1115e8b589e7565830d7d006a8723f19ac8a0afb7", size = 365474 },
 ]
 
+[[package]]
+name = "pytest-anyio"
+version = "0.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "anyio" },
+    { name = "pytest" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/00/44/a02e5877a671b0940f21a7a0d9704c22097b123ed5cdbcca9cab39f17acc/pytest-anyio-0.0.0.tar.gz", hash = "sha256:b41234e9e9ad7ea1dbfefcc1d6891b23d5ef7c9f07ccf804c13a9cc338571fd3", size = 1560 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c6/25/bd6493ae85d0a281b6a0f248d0fdb1d9aa2b31f18bcd4a8800cf397d8209/pytest_anyio-0.0.0-py2.py3-none-any.whl", hash = "sha256:dc8b5c4741cb16ff90be37fddd585ca943ed12bbeb563de7ace6cd94441d8746", size = 1999 },
+]
+
 [[package]]
 name = "pytest-asyncio"
 version = "1.1.0"

langchain_b12-0.1.9/tests/test_genai.py

@@ -1,124 +0,0 @@
-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