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

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

@@ -1,12 +1,10 @@
 Metadata-Version: 2.4
 Name: langchain-b12
-Version: 0.1.10
+Version: 0.1.11
 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 → langchain_b12-0.1.11}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "langchain-b12"
-version = "0.1.10"
+version = "0.1.11"
 description = "A reusable collection of tools and implementations for Langchain"
 readme = "README.md"
 authors = [
@@ -9,8 +9,6 @@ authors = [
 requires-python = ">=3.11"
 dependencies = [
     "langchain-core>=0.3.60",
-    "pytest-anyio>=0.0.0",
-    "tenacity>=9.1.2",
 ]
 
 [dependency-groups]

{langchain_b12-0.1.10 → langchain_b12-0.1.11}/src/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,13 @@ 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)
+        return self
+
     @property
     def _llm_type(self) -> str:
         return "vertexai"
@@ -208,64 +210,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 +252,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.11/tests/test_genai.py

@@ -0,0 +1,210 @@
+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 test_max_retries_converts_to_http_retry_options():
+    """Test that max_retries is properly converted to HttpRetryOptions."""
+    client = MagicMock(spec=Client)
+    model = ChatGenAI(client=client, max_retries=5)
+
+    assert model.http_retry_options is not None
+    assert model.http_retry_options.attempts == 5
+
+
+def test_http_retry_options_passed_directly():
+    """Test that http_retry_options can be passed directly."""
+    client = MagicMock(spec=Client)
+    retry_options = types.HttpRetryOptions(
+        attempts=10,
+        initial_delay=2.0,
+        max_delay=30.0,
+    )
+    model = ChatGenAI(client=client, http_retry_options=retry_options)
+
+    assert model.http_retry_options == retry_options
+    assert model.http_retry_options.attempts == 10
+    assert model.http_retry_options.initial_delay == 2.0
+
+
+def test_http_retry_options_overrides_max_retries():
+    """Test that explicit http_retry_options overrides max_retries."""
+    client = MagicMock(spec=Client)
+    retry_options = types.HttpRetryOptions(attempts=7)
+    model = ChatGenAI(client=client, max_retries=3, http_retry_options=retry_options)
+
+    # http_retry_options should take precedence
+    assert model.http_retry_options == retry_options
+    assert model.http_retry_options.attempts == 7
+
+
+def test_retry_options_passed_in_stream_config():
+    """Test that retry options are passed to GenerateContentConfig."""
+    client: Client = MagicMock(spec=Client)
+    client.models.generate_content_stream.return_value = iter(
+        [_make_response_chunk("success")]
+    )
+
+    model = ChatGenAI(client=client, max_retries=5)
+    messages = [HumanMessage(content="foo")]
+    response = model.invoke(messages)
+
+    # Verify the config was called with http_options containing retry_options
+    call_args = client.models.generate_content_stream.call_args
+    config = call_args.kwargs["config"]
+    assert config.http_options is not None
+    assert config.http_options.retry_options is not None
+    assert config.http_options.retry_options.attempts == 5
+
+
+def test_no_retry_options_when_max_retries_none():
+    """Test that no http_retry_options are set when max_retries is None."""
+    client = MagicMock(spec=Client)
+    model = ChatGenAI(client=client, max_retries=None)
+
+    assert model.http_retry_options is None
+
+
+# --- 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"]
+
+
+def test_stream_error_propagates():
+    """Test that errors during streaming are propagated."""
+    client: Client = MagicMock(spec=Client)
+
+    def failing_stream():
+        yield _make_response_chunk("first")
+        raise Exception("Mid-stream error")
+
+    client.models.generate_content_stream.return_value = failing_stream()
+
+    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 before error
+    assert chunks == ["first"]
+
+
+# --- Async streaming tests ---
+
+
+async def _async_iter(items):
+    """Helper to create an async iterator from items."""
+    for item in items:
+        yield item
+
+
+@pytest.mark.asyncio
+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.asyncio
+async def test_astream_error_propagates():
+    """Test that errors during async streaming are propagated."""
+    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"]

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

@@ -252,12 +252,10 @@ wheels = [
 
 [[package]]
 name = "langchain-b12"
-version = "0.1.9"
+version = "0.1.10"
 source = { editable = "." }
 dependencies = [
     { name = "langchain-core" },
-    { name = "pytest-anyio" },
-    { name = "tenacity" },
 ]
 
 [package.dev-dependencies]
@@ -274,11 +272,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" },
-]
+requires-dist = [{ name = "langchain-core", specifier = ">=0.3.60" }]
 
 [package.metadata.requires-dev]
 citations = [
@@ -623,19 +617,6 @@ 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.10/tests/test_genai.py

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