add suggested changes

haiku_rag_slim/haiku/rag/chunkers/docling_serve.py

@@ -56,9 +56,11 @@ class DoclingServeChunker(DocumentChunker):
 
     def __init__(self, config: AppConfig = Config):
         self.config = config
+        ds = config.providers.docling_serve
         self.client = DoclingServeClient(
-            base_urls=config.providers.docling_serve.base_urls,
+            base_urls=ds.base_urls,
-            api_key=config.providers.docling_serve.api_key,
+            api_key=ds.api_key,
+            breaker_config=ds.circuit_breaker,
         )
         self.chunker_type = config.processing.chunker_type
 

haiku_rag_slim/haiku/rag/circuit_breaker.py

@@ -5,7 +5,11 @@ from haiku.rag.config import CircuitBreakerConfig
 
 
 class CircuitBreaker:
-    """Three-state breaker over discover() failures.
+    """Three-state circuit breaker (closed / open with cooldown).
+
+    Shared by the ingester's per-source discover() breaker and the
+    docling-serve per-instance breaker — lives here, below both
+    ``providers`` and ``ingester``, so neither has to depend on the other.
 
     - closed: failures are counted; threshold flips to open.
     - open: probes are blocked until cooldown elapses, then a single probe

haiku_rag_slim/haiku/rag/config/models.py

@@ -218,6 +218,20 @@ class OllamaConfig(BaseModel):
     )
 
 
+class CircuitBreakerConfig(BaseModel):
+    """Three-state breaker (closed -> open -> probe) over consecutive failures.
+    Used by the ingester's per-source discover() breaker and the docling-serve
+    per-instance breaker."""
+
+    failure_threshold: int = Field(
+        default=5, description="Consecutive failures before the breaker opens."
+    )
+    cooldown_s: float = Field(
+        default=600.0,
+        description="How long the breaker stays open before allowing a probe.",
+    )
+
+
 class DoclingServeConfig(BaseModel):
     """docling-serve endpoints. Accepts a single URL or a list — when a list
     is given, the client round-robins jobs across the URLs. Each job's
@@ -227,6 +241,14 @@ class DoclingServeConfig(BaseModel):
 
     base_url: str | list[str] = "http://localhost:5001"
     api_key: str = ""
+    circuit_breaker: CircuitBreakerConfig = Field(
+        default_factory=lambda: CircuitBreakerConfig(
+            failure_threshold=3, cooldown_s=30.0
+        ),
+        description="Per-instance circuit breaker: skip an instance after this "
+        "many consecutive failures for this cooldown. Defaults are tuned faster "
+        "than the discover breaker since instances restart quickly.",
+    )
 
     @property
     def base_urls(self) -> list[str]:
@@ -295,19 +317,6 @@ class RetryPolicyConfig(BaseModel):
     jitter: float = Field(default=0.25, ge=0.0, le=1.0)
 
 
-class CircuitBreakerConfig(BaseModel):
-    """Per-source breaker over discover() failures. Stops the ingester from
-    hammering a source that's persistently failing."""
-
-    failure_threshold: int = Field(
-        default=5, description="Consecutive failures before the breaker opens."
-    )
-    cooldown_s: float = Field(
-        default=600.0,
-        description="How long the breaker stays open before allowing a probe.",
-    )
-
-
 class WorkerConfig(BaseModel):
     worker_count: int = Field(
         default=4,

haiku_rag_slim/haiku/rag/converters/docling_serve.py

@@ -59,9 +59,11 @@ class DoclingServeConverter(DocumentConverter):
             config: Application configuration containing docling-serve settings.
         """
         self.config = config
+        ds = config.providers.docling_serve
         self.client = DoclingServeClient(
-            base_urls=config.providers.docling_serve.base_urls,
+            base_urls=ds.base_urls,
-            api_key=config.providers.docling_serve.api_key,
+            api_key=ds.api_key,
+            breaker_config=ds.circuit_breaker,
         )
 
     @property

haiku_rag_slim/haiku/rag/ingester/pollers/__init__.py

@@ -1,4 +1,4 @@
-from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
+from haiku.rag.circuit_breaker import CircuitBreaker
 from haiku.rag.ingester.pollers.factory import build_source
 from haiku.rag.ingester.pollers.fs import FSPoller
 from haiku.rag.ingester.pollers.manager import PollerManager

haiku_rag_slim/haiku/rag/ingester/pollers/base.py

@@ -5,7 +5,7 @@ from datetime import UTC, datetime
 
 from haiku.rag.config import SourceConfig
 from haiku.rag.ingester.batch import BatchChange, BatchSourceSummary
-from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
+from haiku.rag.circuit_breaker import CircuitBreaker
 from haiku.rag.ingester.queue.models import JobOp, SyncRow
 from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
 from haiku.rag.ingester.sources.base import (

haiku_rag_slim/haiku/rag/ingester/pollers/fs.py

@@ -12,7 +12,7 @@ from haiku.rag.telemetry import logfire
 
 if TYPE_CHECKING:
     from haiku.rag.config import FSSourceConfig
-    from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
+    from haiku.rag.circuit_breaker import CircuitBreaker
     from haiku.rag.ingester.sources.fs import FSSource
 
 logger = logging.getLogger(__name__)

haiku_rag_slim/haiku/rag/ingester/pollers/manager.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING
 from haiku.rag.config import FSSourceConfig, SourceConfig
 from haiku.rag.ingester.batch import BatchManifest
 from haiku.rag.ingester.pollers.base import BasePoller
-from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
+from haiku.rag.circuit_breaker import CircuitBreaker
 from haiku.rag.ingester.pollers.factory import build_source
 from haiku.rag.ingester.pollers.fs import FSPoller
 from haiku.rag.ingester.pollers.periodic import PeriodicPoller

haiku_rag_slim/haiku/rag/ingester/pollers/periodic.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
 from haiku.rag.ingester.pollers.base import BasePoller
 
 if TYPE_CHECKING:
-    from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
+    from haiku.rag.circuit_breaker import CircuitBreaker
 
 
 class PeriodicPoller(BasePoller):

haiku_rag_slim/haiku/rag/ingester/workers/pool.py

@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING
 
 from haiku.rag.config import CircuitBreakerConfig
 from haiku.rag.ingester.exceptions import PermanentError, TransientError
-from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
+from haiku.rag.circuit_breaker import CircuitBreaker
 from haiku.rag.ingester.queue.models import Job, JobOp
 from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
 from haiku.rag.ingester.workers.pipeline import run_job

haiku_rag_slim/haiku/rag/providers/docling_serve.py

@@ -9,6 +9,9 @@ from typing import Any
 
 import httpx
 
+from haiku.rag.circuit_breaker import CircuitBreaker
+from haiku.rag.config import CircuitBreakerConfig
+
 logger = logging.getLogger(__name__)
 
 # HTTP statuses that count against an instance's health: 408/429 are transient
@@ -29,44 +32,6 @@ def _is_instance_failure(exc: BaseException) -> bool:
     return isinstance(exc, httpx.TransportError)
 
 
-class _InstanceBreaker:
-    """Per-instance circuit breaker (closed / open with cooldown).
-
-    Self-contained so the provider layer doesn't depend on the ingester
-    package. ``is_open`` auto-probes after the cooldown elapses: a single
-    request is allowed through, and its success closes the breaker while another
-    failure re-opens it.
-    """
-
-    def __init__(
-        self,
-        failure_threshold: int,
-        cooldown_s: float,
-        now_fn: Callable[[], float],
-    ):
-        self._threshold = failure_threshold
-        self._cooldown_s = cooldown_s
-        self._now = now_fn
-        self._consecutive_failures = 0
-        self._opened_at: float | None = None
-
-    @property
-    def is_open(self) -> bool:
-        if self._opened_at is None:
-            return False
-        # Cooldown elapsed → let the next request probe the instance.
-        return self._now() - self._opened_at < self._cooldown_s
-
-    def record_success(self) -> None:
-        self._consecutive_failures = 0
-        self._opened_at = None
-
-    def record_failure(self) -> None:
-        self._consecutive_failures += 1
-        if self._consecutive_failures >= self._threshold:
-            self._opened_at = self._now()
-
-
 # Process-global registry of round-robin rotators over docling-serve
 # instances, keyed by the sorted tuple of base URLs. Clients are constructed
 # per-job by `get_converter` / `get_chunker`, so the rotation cursor has to
@@ -80,7 +45,7 @@ _instance_rotators: dict[tuple[str, ...], "itertools.cycle[str]"] = {}
 # skipped across subsequent jobs until its cooldown elapses. The first client
 # to touch a URL fixes that breaker's threshold/cooldown/clock; in practice all
 # clients share one config (the docling-serve provider config).
-_instance_breakers: dict[str, "_InstanceBreaker"] = {}
+_instance_breakers: dict[str, "CircuitBreaker"] = {}
 
 
 class DoclingServeClient:
@@ -99,8 +64,7 @@ class DoclingServeClient:
         api_key: str | None = None,
         timeout: float = 300,
         transport: httpx.AsyncBaseTransport | None = None,
-        breaker_failure_threshold: int = 3,
+        breaker_config: CircuitBreakerConfig | None = None,
-        breaker_cooldown_s: float = 30.0,
         now_fn: Callable[[], float] = time.monotonic,
     ):
         urls = [base_urls] if isinstance(base_urls, str) else list(base_urls)
@@ -111,8 +75,12 @@ class DoclingServeClient:
         self.timeout = timeout
         # transport is for testing — production callers leave it None.
         self._transport = transport
-        self._breaker_failure_threshold = breaker_failure_threshold
+        # Per-instance breaker config; defaults are tuned faster than the
-        self._breaker_cooldown_s = breaker_cooldown_s
+        # ingester's discover breaker since docling-serve instances restart
+        # quickly. Overridden from DoclingServeConfig.circuit_breaker.
+        self._breaker_config = breaker_config or CircuitBreakerConfig(
+            failure_threshold=3, cooldown_s=30.0
+        )
         self._now = now_fn
         # setdefault is atomic under the GIL — concurrent constructors with
         # the same instance set will end up sharing one rotator.
@@ -130,31 +98,27 @@ class DoclingServeClient:
         callers should let `_pick_url` round-robin per request."""
         return self.base_urls[0]
 
-    def _breaker_for(self, url: str) -> _InstanceBreaker:
+    def _breaker_for(self, url: str) -> CircuitBreaker:
         breaker = _instance_breakers.get(url)
         if breaker is None:
-            breaker = _InstanceBreaker(
+            breaker = CircuitBreaker(self._breaker_config, now_fn=self._now)
-                self._breaker_failure_threshold,
-                self._breaker_cooldown_s,
-                self._now,
-            )
             _instance_breakers[url] = breaker
         return breaker
 
     def _pick_url(self) -> str:
         """Next instance in the round-robin, skipping any whose circuit breaker
-        is open (an instance that recently crashed / overloaded). Falls back to
+        is open (an instance that recently crashed / overloaded). When every
-        the first instance seen when every breaker is open — a single-instance
+        breaker is open — a single-instance fleet, or a fully-down one — probe
-        fleet, or a fully-down one, still gets a probe rather than no attempt."""
+        one anyway rather than failing with nothing to pick."""
-        fallback: str | None = None
         for _ in range(len(self.base_urls)):
             url = next(self._instance_rotator)
-            if fallback is None:
-                fallback = url
             if not self._breaker_for(url).is_open:
                 return url
-        assert fallback is not None  # base_urls is non-empty (checked in __init__)
+        # All breakers open. Take one more step (the loop consumed a full
-        return fallback
+        # rotation, landing the cursor back at the start) so consecutive
+        # all-open calls rotate across instances instead of pinning the first
+        # one — an all-429 overload shouldn't pile every retry on one node.
+        return next(self._instance_rotator)
 
     def _record_outcome(self, base_url: str, exc: BaseException) -> None:
         """Record a failed request against the instance's breaker, but only when
@@ -170,8 +134,8 @@ class DoclingServeClient:
                 "docling-serve instance %s breaker opened after %d consecutive "
                 "failure(s); skipping it for %.0fs",
                 base_url,
-                self._breaker_failure_threshold,
+                self._breaker_config.failure_threshold,
-                self._breaker_cooldown_s,
+                self._breaker_config.cooldown_s,
             )
 
     def _get_headers(self) -> dict[str, str]:

tests/ingester/test_circuit_breaker.py

@@ -1,5 +1,5 @@
 from haiku.rag.config import CircuitBreakerConfig
-from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
+from haiku.rag.circuit_breaker import CircuitBreaker
 
 
 class _Clock:

tests/ingester/test_pollers.py

@@ -11,7 +11,7 @@ from haiku.rag.config import (
     S3SourceConfig,
     WebDAVSourceConfig,
 )
-from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
+from haiku.rag.circuit_breaker import CircuitBreaker
 from haiku.rag.ingester.pollers.manager import PollerManager
 from haiku.rag.ingester.pollers.periodic import PeriodicPoller
 from haiku.rag.ingester.queue.models import JobOp, JobStatus

tests/test_docling_serve_client.py

@@ -9,6 +9,7 @@ resume mid-rotation, breaking specific-order assertions.
 import httpx
 import pytest
 
+from haiku.rag.config import CircuitBreakerConfig
 from haiku.rag.providers.docling_serve import DoclingServeClient
 
 
@@ -152,7 +153,7 @@ async def test_open_breaker_skips_crashed_instance():
     client = DoclingServeClient(
         base_urls=["http://crash-x:5001", "http://live-x:5001"],
         transport=transport,
-        breaker_failure_threshold=2,
+        breaker_config=CircuitBreakerConfig(failure_threshold=2, cooldown_s=30.0),
     )
 
     await _poll(client)  # picks crash-x → fail (failures=1)
@@ -179,8 +180,7 @@ async def test_breaker_recovers_after_cooldown():
     client = DoclingServeClient(
         base_urls=[flip, "http://spare-y:5001"],
         transport=transport,
-        breaker_failure_threshold=2,
+        breaker_config=CircuitBreakerConfig(failure_threshold=2, cooldown_s=30.0),
-        breaker_cooldown_s=30.0,
         now_fn=lambda: clock[0],
     )
 
@@ -219,7 +219,7 @@ async def test_4xx_does_not_trip_breaker():
     client = DoclingServeClient(
         base_urls=[bad],
         transport=httpx.MockTransport(handler),
-        breaker_failure_threshold=1,
+        breaker_config=CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0),
     )
 
     for _ in range(3):
@@ -243,7 +243,7 @@ async def test_pick_url_falls_back_when_all_breakers_open():
     client = DoclingServeClient(
         base_urls=[lone],
         transport=transport,
-        breaker_failure_threshold=1,
+        breaker_config=CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0),
     )
 
     await _poll(client)  # fails → breaker opens (threshold=1)
@@ -253,6 +253,54 @@ async def test_pick_url_falls_back_when_all_breakers_open():
     assert client._pick_url() == lone
 
 
+@pytest.mark.asyncio
+async def test_all_open_breakers_rotate_instead_of_pinning_one():
+    """When every breaker is open, successive picks must rotate across the
+    fleet rather than pinning the first instance — an all-429 overload
+    shouldn't pile every retry on one node."""
+    urls = [
+        "http://allopen-a:5001",
+        "http://allopen-b:5001",
+        "http://allopen-c:5001",
+    ]
+    transport, _ = _health_transport(
+        {"allopen-a", "allopen-b", "allopen-c"}, "t", {"ok": True}
+    )
+    client = DoclingServeClient(
+        base_urls=urls,
+        transport=transport,
+        breaker_config=CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0),
+    )
+
+    # One failure each opens all three breakers: each _poll picks the next
+    # still-closed instance, fails, and opens it.
+    for _ in range(3):
+        await _poll(client)
+    assert all(client._breaker_for(u).is_open for u in urls)
+
+    picks = [client._pick_url() for _ in range(6)]
+    # All-open fallback rotates a → b → c → a → b → c, not the same URL six times.
+    assert picks == urls + urls
+
+
+def test_config_breaker_reaches_client():
+    """The breaker config set on DoclingServeConfig must reach the client via
+    get_converter / get_chunker (previously hardcoded constructor defaults)."""
+    from haiku.rag.chunkers.docling_serve import DoclingServeChunker
+    from haiku.rag.config import AppConfig
+    from haiku.rag.converters.docling_serve import DoclingServeConverter
+
+    config = AppConfig()
+    config.providers.docling_serve.base_url = "http://cfg-brk:5001"
+    config.providers.docling_serve.circuit_breaker = CircuitBreakerConfig(
+        failure_threshold=9, cooldown_s=123.0
+    )
+
+    for component in (DoclingServeConverter(config), DoclingServeChunker(config)):
+        assert component.client._breaker_config.failure_threshold == 9
+        assert component.client._breaker_config.cooldown_s == 123.0
+
+
 @pytest.mark.asyncio
 async def test_zip_records_instance_failure():
     """submit_and_poll_zip records a crashed instance against its breaker, the
@@ -262,7 +310,7 @@ async def test_zip_records_instance_failure():
     client = DoclingServeClient(
         base_urls=[z],
         transport=transport,
-        breaker_failure_threshold=1,
+        breaker_config=CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0),
     )
 
     with pytest.raises(httpx.ConnectError):