# Custom Providers

O `arandu` usa protocolos Python para injeção de dependência. Você pode usar qualquer backend de LLM ou embedding implementando duas interfaces simples — sem herança necessária.

## Os Protocolos

O SDK define dois protocolos em `arandu.protocols`:

### LLMProvider

Seu LLM provider deve implementar um único método `complete`:

```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: ...
```

| Parâmetro | Descrição |
|-----------|-----------|
| `messages` | Lista de dicts de mensagem com chaves `"role"` e `"content"` (formato OpenAI) |
| `temperature` | Temperatura de sampling (0 = determinístico) |
| `response_format` | Especificação de formato opcional (ex: `{"type": "json_object"}` para modo JSON) |
| `max_tokens` | Máximo opcional de tokens para a resposta |
| **Retorna** | O texto de resposta do assistente como string |

!!! important "Suporte a modo JSON"
    O pipeline de memória depende fortemente de respostas em modo JSON (`response_format={"type": "json_object"}`).
    Seu provider deve suportar isso — nativamente ou parseando a resposta.

### EmbeddingProvider

Seu embedding provider deve implementar dois métodos:

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

| Método | Descrição |
|--------|-----------|
| `embed(texts)` | Gera embeddings para um batch de textos. Retorna um vetor por input. |
| `embed_one(text)` | Gera embedding para um único texto. Retorna `None` se o texto for vazio/inválido. |

!!! note "Dimensões de embedding"
    O `embedding_dimensions` padrão em `MemoryConfig` é 1536 (OpenAI `text-embedding-3-small`).
    Se seu provider usa dimensões diferentes, defina `MemoryConfig(embedding_dimensions=...)` de acordo.

---

## Exemplo: Provider Anthropic

Aqui está um exemplo completo implementando ambos os protocolos usando o SDK da Anthropic:

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


class AnthropicLLMProvider:
    """LLM provider usando a API Claude da Anthropic."""

    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:
        # Converter mensagens formato OpenAI para formato Anthropic
        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"],
                })

        # Adicionar instrução JSON se modo json solicitado
        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 "Providers separados"
    Você pode usar providers diferentes para LLM e embeddings. Por exemplo,
    use Anthropic para completions e OpenAI para embeddings:

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

---

## Exemplo: Provider de Modelo Local

Para rodar com modelos locais (ex: via Ollama):

```python
import httpx


class OllamaProvider:
    """LLM + Embedding provider usando um servidor Ollama local."""

    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 "Dimensões de embedding"
    Quando usar modelos locais, verifique as dimensões de embedding e configure de acordo:

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

---

## Testando Seu Provider

Você pode verificar se seu provider funciona com o sistema de memória antes de ir para produção:

```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:
        # Testar 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")

        # Testar 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())
```

## Requisitos Importantes

Ao implementar um custom provider, tenha estes requisitos em mente:

1. **Modo JSON** — O pipeline envia `response_format={"type": "json_object"}` frequentemente. Seu provider deve retornar JSON válido quando isso é definido.

2. **Async** — Ambos os protocolos são async (`async def`). Se o SDK do seu backend é síncrono, encapsule chamadas com `asyncio.to_thread()`.

3. **Tratamento de vazio/erro** — `embed_one` deve retornar `None` para input vazio, não lançar exceção. `embed` deve retornar `[]` para input vazio.

4. **Timeout** — Considere adicionar timeouts ao seu provider. O SDK define timeouts do seu lado via `MemoryConfig`, mas timeouts no nível do provider adicionam uma camada extra de segurança.

5. **Dimensões de embedding** — Defina `MemoryConfig(embedding_dimensions=N)` para corresponder às dimensões de output do seu provider. Dimensões incompatíveis causam erros no pgvector.
