fix(adapters): make the declared LLMClient interface sufficient to run

docs/SUMMARY.md

@@ -30,6 +30,7 @@
 * [Cost Estimation & Budgets](guides/cost-control.md)
 * [Caching](guides/caching.md)
 * [Multi-Provider Routing](guides/routing.md)
+* [Custom LLM Clients](guides/custom-clients.md)
 
 ## Providers
 

docs/guides/custom-clients.md

@@ -0,0 +1,148 @@
+# Custom LLM Clients
+
+Ondine talks to providers through LiteLLM, which covers most cases. When it
+doesn't, supply your own client and the pipeline uses it everywhere — batching,
+concurrency, retries, cost tracking, streaming and checkpointing all work
+unchanged.
+
+Reach for this when you need to:
+
+- route through an **internal gateway** that adds auth, quota or audit logging
+- call an API that is **not OpenAI-shaped** and has no LiteLLM route
+- add a **caching or rate-limiting layer** of your own
+- **stub the provider in tests** without patching internals
+
+## The interface
+
+Implement three methods:
+
+```python
+from decimal import Decimal
+from typing import Any
+
+from ondine.adapters.llm_client import LLMClient
+from ondine.core.models import LLMResponse
+from ondine.core.specifications import LLMSpec
+
+
+class GatewayClient(LLMClient):
+    def invoke(self, prompt: str, **kwargs: Any) -> LLMResponse:
+        text = my_gateway.complete(prompt)          # your call here
+        return LLMResponse(
+            text=text,
+            tokens_in=self.estimate_tokens(prompt),
+            tokens_out=self.estimate_tokens(text),
+            model=self.model,
+            cost=self.calculate_cost(...),
+            latency_ms=...,
+        )
+
+    def structured_invoke(self, prompt: str, output_cls: Any, **kwargs: Any):
+        # Called when the pipeline uses with_structured_output().
+        # Ask your API for JSON and validate into output_cls, or raise
+        # NotImplementedError if it has no structured mode.
+        ...
+
+    def estimate_tokens(self, text: str) -> int:
+        return len(text) // 4
+```
+
+Then inject it:
+
+```python
+pipeline = (
+    PipelineBuilder.create()
+    .from_dataframe(df, input_columns=["text"], output_columns=["result"])
+    .with_prompt("Summarize: {text}")
+    .with_custom_llm_client(GatewayClient(LLMSpec(provider="openai", model="internal-v2")))
+    .build()
+)
+```
+
+`with_llm()` is not needed — the spec you pass to your client carries the model
+name, pricing and limits. If you call both, the custom client wins.
+
+A runnable version is in
+[`examples/13_custom_client.py`](https://github.com/ptimizeroracle/ondine/blob/main/examples/13_custom_client.py).
+
+## Async: the method that actually runs
+
+The default executor is async, so **`ainvoke()` is what a run calls**, not
+`invoke()`. The base class delegates it to `invoke()` for you, which is correct
+but serialises requests — the sync call blocks the event loop.
+
+For real concurrency, override it:
+
+```python
+    async def ainvoke(self, prompt: str, **kwargs: Any) -> LLMResponse:
+        text = await my_async_gateway.complete(prompt)
+        return LLMResponse(...)
+```
+
+Same for `structured_invoke_async()`.
+
+## Optional lifecycle hooks
+
+If your client holds a session, a connection pool or a subprocess, override
+these. They default to doing nothing.
+
+```python
+    async def start(self) -> None:
+        self._session = aiohttp.ClientSession()
+
+    async def stop(self) -> None:
+        await self._session.close()
+```
+
+`start()` runs before the first request of a run and `stop()` after the last.
+
+## What you get for free
+
+Your client is called through the full pipeline, so everything else keeps
+working:
+
+| | |
+|---|---|
+| **Batching** | `with_batch_size()` groups rows into one prompt before reaching you |
+| **Concurrency** | `with_concurrency()` runs many `ainvoke()` calls at once |
+| **Retries** | Transient failures retry per your error policy |
+| **Cost tracking** | Whatever you report in `LLMResponse.cost` is summed and budgeted |
+| **Checkpointing** | Responses are cached for resume, same as any provider |
+
+## One instance, shared
+
+The pipeline uses **the object you passed**, not a copy — including in
+streamed runs, where each chunk is its own sub-pipeline, and in the automatic
+retry pass. State your client holds (a session, a token, a counter, a shared
+rate limiter) is shared across all of them.
+
+That matters if your client is stateful: a connection pool is opened once, and
+a rate limiter genuinely limits the whole run rather than per chunk.
+
+## Testing with a stub
+
+The same mechanism makes provider-free tests straightforward:
+
+```python
+class StubClient(LLMClient):
+    def __init__(self, spec, answer="ok"):
+        super().__init__(spec)
+        self.answer = answer
+        self.prompts: list[str] = []
+
+    def invoke(self, prompt, **kwargs):
+        self.prompts.append(prompt)
+        return LLMResponse(
+            text=self.answer, tokens_in=1, tokens_out=1, model="stub",
+            cost=Decimal("0"), latency_ms=0.0,
+        )
+
+    def structured_invoke(self, prompt, output_cls, **kwargs):
+        return self.invoke(prompt)
+
+    def estimate_tokens(self, text):
+        return 1
+```
+
+Assert on `stub.prompts` to check what your pipeline actually sent — no network,
+no keys, no patching of ondine internals.

examples/13_custom_client.py

@@ -81,11 +81,44 @@ def invoke(self, prompt: str, **kwargs: Any) -> LLMResponse:
             latency_ms=latency_ms,
         )
 
+    def structured_invoke(
+        self, prompt: str, output_cls: Any, **kwargs: Any
+    ) -> LLMResponse:
+        """Invoke and return output conforming to *output_cls*.
+
+        Required by the LLMClient interface — a client that cannot be
+        instantiated without it. Ondine calls this when a pipeline sets
+        with_structured_output(); if your API has no structured mode, ask it
+        for JSON and validate here, or raise NotImplementedError so callers
+        learn immediately rather than at row 40,000.
+        """
+        # A real implementation would parse the response text into output_cls
+        # and re-raise on validation failure. The demo response is not JSON,
+        # so we simply hand the text back.
+        return self.invoke(prompt, **kwargs)
+
+    async def ainvoke(self, prompt: str, **kwargs: Any) -> LLMResponse:
+        """Async variant — this is the one the pipeline actually calls.
+
+        The default executor is async, so leaving this as a sync-only client
+        means every row goes through the base class. Implement it with your
+        API's async client for real concurrency; delegating to invoke() as
+        below is correct but serialises requests.
+        """
+        return self.invoke(prompt, **kwargs)
+
+    async def structured_invoke_async(
+        self, prompt: str, output_cls: Any, **kwargs: Any
+    ) -> LLMResponse:
+        """Async variant of structured_invoke."""
+        return self.structured_invoke(prompt, output_cls, **kwargs)
+
     def estimate_tokens(self, text: str) -> int:
         """
         Estimate token count.
 
-        Replace with your custom tokenization logic.
+        Used for cost estimation before a run, so it should match your
+        provider's tokenizer as closely as you care about the estimate.
         """
         # Simple word count for demo
         return len(text.split())

ondine/adapters/llm_client.py

@@ -101,6 +101,46 @@ def estimate_tokens(self, text: str) -> int:
         """
         pass
 
+    async def ainvoke(self, prompt: str, **kwargs: Any) -> LLMResponse:
+        """Async variant of :meth:`invoke`. Delegates to it by default.
+
+        The default executor is async, so this — not ``invoke`` — is what a
+        run actually calls. Without a default, a client that implemented every
+        abstract method still had all its rows skipped, because the async path
+        found nothing to call.
+
+        Delegating is correct but serialises requests: the sync call blocks
+        the event loop. Override with your provider's async client to get real
+        concurrency.
+        """
+        return self.invoke(prompt, **kwargs)
+
+    async def structured_invoke_async(
+        self, prompt: str, output_cls: Any, **kwargs: Any
+    ) -> Any:
+        """Async variant of :meth:`structured_invoke`. Delegates by default."""
+        return self.structured_invoke(prompt, output_cls, **kwargs)
+
+    async def start(self) -> None:
+        """Open any resources the client needs for a run.
+
+        Defaults to doing nothing, because most clients have nothing to open.
+        The invocation stage calls this before the first request and
+        :meth:`stop` after the last, so a client holding a session, a pool or
+        a subprocess overrides both.
+
+        Not abstract on purpose. The stage has always called it
+        unconditionally, so every custom client that did not happen to define
+        it crashed at the first row with ``AttributeError: object has no
+        attribute 'start'`` — a requirement the interface never stated and the
+        type checker could not see, since the call carried a
+        ``type: ignore[attr-defined]``. A no-op default removes the trap
+        rather than documenting it.
+        """
+
+    async def stop(self) -> None:
+        """Release whatever :meth:`start` opened. No-op by default."""
+
     def batch_invoke(self, prompts: list[str], **kwargs: Any) -> list[LLMResponse]:
         """
         Invoke LLM with multiple prompts.

ondine/stages/llm_invocation_stage.py

@@ -159,7 +159,7 @@ async def _process_async(
     ) -> list[ResponseBatch]:
         """Internal async implementation — shared by process() and process_async()."""
         # Initialize global connection pool
-        await self.llm_client.start()  # type: ignore[attr-defined]
+        await self.llm_client.start()
 
         try:
             # Setup deployment tracker from Router config
@@ -205,7 +205,7 @@ async def _process_async(
             return response_batches
         finally:
             # Cleanup global connection pool
-            await self.llm_client.stop()  # type: ignore[attr-defined]
+            await self.llm_client.stop()
 
     def _get_router_model_list(self) -> list[dict] | None:
         """Get model list from Router if available."""

tests/unit/test_pipeline_builder_custom_client.py

@@ -365,3 +365,57 @@ def test_specifications_stay_serializable(self):
         dumped = pipeline.specifications.model_dump(mode="json")
 
         assert "custom_llm_client" not in dumped.get("metadata", {})
+
+
+class TestMinimalClientInterface:
+    """A client implementing only the declared interface must run.
+
+    LLMClient declares three abstract methods, but the invocation stage also
+    calls start() and stop() — which lived on UnifiedLiteLLMClient only, and
+    were invoked through a `type: ignore[attr-defined]`. So a client that
+    implemented everything the interface asked for still died at the first row
+    with AttributeError, and neither the ABC nor mypy could say why.
+
+    This pins the contract: implement what is abstract, and it runs.
+    """
+
+    class _BareMinimum(LLMClient):
+        """Exactly the three abstract methods — nothing else."""
+
+        def invoke(self, prompt: str, **kwargs: Any) -> LLMResponse:
+            return LLMResponse(
+                text="minimal",
+                tokens_in=1,
+                tokens_out=1,
+                model="minimal",
+                cost=Decimal("0"),
+                latency_ms=1.0,
+                metadata={},
+            )
+
+        def structured_invoke(self, prompt, output_cls, **kwargs):
+            return self.invoke(prompt, **kwargs)
+
+        def estimate_tokens(self, text: str) -> int:
+            return max(1, len(text) // 4)
+
+    def test_client_with_only_abstract_methods_runs(self):
+        import pandas as pd
+
+        client = self._BareMinimum(LLMSpec(provider="openai", model="not-a-model"))
+
+        result = (
+            PipelineBuilder.create()
+            .from_dataframe(
+                pd.DataFrame({"t": ["a", "b"]}),
+                input_columns=["t"],
+                output_columns=["out"],
+            )
+            .with_prompt("Echo: {t}")
+            .with_custom_llm_client(client)
+            .with_batch_size(1)
+            .build()
+            .execute()
+        )
+
+        assert list(result.to_pandas()["out"]) == ["minimal"] * 2