Custom LLM Provider

Heartbit ships with Anthropic and OpenRouter providers, but you can integrate any LLM by implementing the LlmProvider trait.

The LlmProvider Trait

use heartbit::{CompletionRequest, CompletionResponse, Error, OnText};
use std::future::Future;

pub trait LlmProvider: Send + Sync {
    fn complete(
        &self,
        request: CompletionRequest,
    ) -> impl Future<Output = Result<CompletionResponse, Error>> + Send;

    fn stream_complete(
        &self,
        request: CompletionRequest,
        on_text: &OnText,
    ) -> impl Future<Output = Result<CompletionResponse, Error>> + Send {
        // Default: falls back to complete() (no incremental streaming)
        let _ = on_text;
        self.complete(request)
    }

    fn model_name(&self) -> Option<&str> {
        None
    }
}

Only complete() is required. stream_complete() defaults to calling complete() if not overridden. model_name() is used for audit trail events and cost tracking.

Skeleton Implementation

use heartbit::{
    CompletionRequest, CompletionResponse, ContentBlock, Error,
    LlmProvider, StopReason, TokenUsage,
};

pub struct MyProvider {
    api_key: String,
    model: String,
}

impl MyProvider {
    pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Self {
        Self {
            api_key: api_key.into(),
            model: model.into(),
        }
    }
}

impl LlmProvider for MyProvider {
    async fn complete(
        &self,
        request: CompletionRequest,
    ) -> Result<CompletionResponse, Error> {
        // 1. Convert request.messages and request.tools to your API format
        // 2. Make the HTTP call to your LLM backend
        // 3. Parse the response into CompletionResponse

        let response_text = call_my_api(
            &self.api_key,
            &self.model,
            &request,
        ).await.map_err(|e| Error::Agent(e.to_string()))?;

        Ok(CompletionResponse {
            content: vec![ContentBlock::Text { text: response_text }],
            stop_reason: StopReason::EndTurn,
            usage: TokenUsage {
                input_tokens: 0,
                output_tokens: 0,
                cache_creation_input_tokens: 0,
                cache_read_input_tokens: 0,
                reasoning_tokens: 0,
            },
            model: Some(self.model.clone()),
        })
    }

    fn model_name(&self) -> Option<&str> {
        Some(&self.model)
    }
}

Using BoxedProvider for Dynamic Dispatch

LlmProvider uses RPITIT (return-position impl Trait in traits), which is not object-safe. To use providers dynamically (e.g., selecting at runtime), wrap them with BoxedProvider:

use heartbit::BoxedProvider;
use std::sync::Arc;

let provider = MyProvider::new("sk-...", "my-model-v1");
let boxed = Arc::new(BoxedProvider::new(provider));

// Now usable as Arc<BoxedProvider> with AgentRunner::builder
let agent = AgentRunner::builder(boxed).build()?;

BoxedProvider internally uses DynLlmProvider, an object-safe adapter that converts the RPITIT methods to Pin<Box<dyn Future>>. A blanket implementation covers all LlmProvider types automatically.

If you already have a provider behind an Arc, use BoxedProvider::from_arc() to avoid consuming it:

let shared = Arc::new(MyProvider::new("sk-...", "my-model-v1"));
let boxed = BoxedProvider::from_arc(shared.clone());

Wrapping with Retry and Cascade

Heartbit provides composable provider wrappers:

use heartbit::{RetryingProvider, CascadingProvider};

// Add automatic retry with default config (3 retries, 500ms base delay)
let retrying = RetryingProvider::with_defaults(provider);

// Or with custom config:
// let retrying = RetryingProvider::new(provider, RetryConfig { max_retries: 5, ..Default::default() });

These wrappers implement LlmProvider themselves, so they compose naturally.