# Custom Providers

`arandu` uses Python protocols for dependency injection. You can use any LLM or embedding backend by implementing two simple interfaces — no inheritance required.

## The Protocols

The SDK defines two protocols in `arandu.protocols`:

### LLMProvider

Your LLM provider must implement a single `complete` method:

```python
class LLMProvider(Protocol):
    async def complete(
        self,
        messages: list[dict],
        temperature: float = 0,
        response_format: dict | None = None,
        max_tokens: int | None = None,
    ) -> str: ...
```

| Parameter | Description |
|-----------|-------------|
| `messages` | List of message dicts with `"role"` and `"content"` keys (OpenAI format) |
| `temperature` | Sampling temperature (0 = deterministic) |
| `response_format` | Optional format spec (e.g., `{"type": "json_object"}` for JSON mode) |
| `max_tokens` | Optional maximum tokens for the response |
| **Returns** | The assistant's response text as a string |

!!! important "JSON mode support"
    The memory pipeline relies heavily on JSON-mode responses (`response_format={"type": "json_object"}`).
    Your provider must support this — either natively or by parsing the response.

### EmbeddingProvider

Your embedding provider must implement two methods:

```python
class EmbeddingProvider(Protocol):
    async def embed(self, texts: list[str]) -> list[list[float]]: ...
    async def embed_one(self, text: str) -> list[float] | None: ...
```

| Method | Description |
|--------|-------------|
| `embed(texts)` | Generate embeddings for a batch of texts. Returns one vector per input. |
| `embed_one(text)` | Generate embedding for a single text. Returns `None` if text is empty/invalid. |

!!! note "Embedding dimensions"
    The default `embedding_dimensions` in `MemoryConfig` is 1536 (OpenAI `text-embedding-3-small`).
    If your provider uses different dimensions, set `MemoryConfig(embedding_dimensions=...)` accordingly.

---

## Example: Anthropic Provider

Here's a complete example implementing both protocols using the Anthropic SDK:

```python
import asyncio
import json
from anthropic import AsyncAnthropic


class AnthropicLLMProvider:
    """LLM provider using Anthropic's Claude API."""

    def __init__(self, api_key: str, model: str = "claude-sonnet-4-20250514") -> None:
        self._client = AsyncAnthropic(api_key=api_key)
        self._model = model

    async def complete(
        self,
        messages: list[dict],
        temperature: float = 0,
        response_format: dict | None = None,
        max_tokens: int | None = None,
    ) -> str:
        # Convert OpenAI-format messages to Anthropic format
        system_msg = ""
        chat_messages = []
        for msg in messages:
            if msg["role"] == "system":
                system_msg = msg["content"]
            else:
                chat_messages.append({
                    "role": msg["role"],
                    "content": msg["content"],
                })

        # Add JSON instruction if json mode requested
        if response_format and response_format.get("type") == "json_object":
            system_msg += "\n\nRespond with valid JSON only. No markdown fences."

        response = await self._client.messages.create(
            model=self._model,
            system=system_msg,
            messages=chat_messages,
            temperature=temperature,
            max_tokens=max_tokens or 4096,
        )

        return response.content[0].text
```

!!! tip "Separate providers"
    You can use different providers for LLM and embeddings. For example,
    use Anthropic for completions and OpenAI for embeddings:

    ```python
    memory = MemoryClient(
        database_url="...",
        llm=AnthropicLLMProvider(api_key="sk-ant-..."),
        embeddings=OpenAIProvider(api_key="sk-..."),  # just for embeddings
    )
    ```

---

## Example: Local Model Provider

For running with local models (e.g., via Ollama):

```python
import httpx


class OllamaProvider:
    """LLM + Embedding provider using a local Ollama server."""

    def __init__(
        self,
        base_url: str = "http://localhost:11434",
        model: str = "llama3.1",
        embedding_model: str = "nomic-embed-text",
    ) -> None:
        self._base_url = base_url
        self._model = model
        self._embedding_model = embedding_model
        self._client = httpx.AsyncClient(timeout=60.0)

    # -- LLMProvider --

    async def complete(
        self,
        messages: list[dict],
        temperature: float = 0,
        response_format: dict | None = None,
        max_tokens: int | None = None,
    ) -> str:
        payload: dict = {
            "model": self._model,
            "messages": messages,
            "stream": False,
            "options": {"temperature": temperature},
        }
        if response_format and response_format.get("type") == "json_object":
            payload["format"] = "json"

        response = await self._client.post(
            f"{self._base_url}/api/chat",
            json=payload,
        )
        response.raise_for_status()
        return response.json()["message"]["content"]

    # -- EmbeddingProvider --

    async def embed(self, texts: list[str]) -> list[list[float]]:
        results = []
        for text in texts:
            if not text.strip():
                continue
            response = await self._client.post(
                f"{self._base_url}/api/embed",
                json={"model": self._embedding_model, "input": text},
            )
            response.raise_for_status()
            results.append(response.json()["embeddings"][0])
        return results

    async def embed_one(self, text: str) -> list[float] | None:
        if not text or not text.strip():
            return None
        results = await self.embed([text])
        return results[0] if results else None
```

!!! warning "Embedding dimensions"
    When using local models, check the embedding dimensions and configure accordingly:

    ```python
    config = MemoryConfig(
        embedding_dimensions=768,  # nomic-embed-text uses 768 dims
    )
    ```

---

## Testing Your Provider

You can verify your provider works with the memory system before going to production:

```python
import asyncio
from arandu import MemoryClient, MemoryConfig


async def test_provider():
    provider = YourProvider(...)
    memory = MemoryClient(
        database_url="postgresql+psycopg://memory:memory@localhost/memory",
        llm=provider,
        embeddings=provider,
    )
    await memory.initialize()

    try:
        # Test write
        result = await memory.write(
            user_id="test",
            message="Testing the provider. My name is Alice and I work at Acme.",
        )
        assert len(result.facts_added) > 0, "No facts extracted — check LLM responses"
        assert len(result.entities_resolved) > 0, "No entities resolved"
        print(f"Write OK: {len(result.facts_added)} facts, {len(result.entities_resolved)} entities")

        # Test retrieve
        context = await memory.retrieve(user_id="test", query="who is Alice?")
        assert len(context.facts) > 0, "No facts retrieved — check embeddings"
        print(f"Retrieve OK: {len(context.facts)} facts found")
        print(f"Context: {context.context}")
    finally:
        await memory.close()


asyncio.run(test_provider())
```

## Key Requirements

When implementing a custom provider, keep these requirements in mind:

1. **JSON mode** — The pipeline sends `response_format={"type": "json_object"}` frequently. Your provider must return valid JSON when this is set.

2. **Async** — Both protocols are async (`async def`). If your backend SDK is synchronous, wrap calls with `asyncio.to_thread()`.

3. **Empty/error handling** — `embed_one` should return `None` for empty input, not raise. `embed` should return `[]` for empty input.

4. **Timeout** — Consider adding timeouts to your provider. The SDK sets timeouts on its side via `MemoryConfig`, but provider-level timeouts add an extra safety layer.

5. **Embedding dimensions** — Set `MemoryConfig(embedding_dimensions=N)` to match your provider's output dimensions. Mismatched dimensions will cause pgvector errors.
