Custom LLM Providers

Implement and register a custom LLMProvider to use any model API with ark-operator — step-by-step guide for provider authors.

ark-operator’s agent runtime calls LLMs through the LLMProvider interface. Built-in providers cover Anthropic and OpenAI-compatible APIs. Adding a custom provider lets you integrate any model API — proprietary inference services, on-premise models, or test doubles — without modifying operator code.

The interface

// internal/agent/providers/provider.go
type LLMProvider interface {
    RunTask(
        ctx    context.Context,
        cfg    *config.Config,
        task   queue.Task,
        tools  []mcp.Tool,
        callTool func(context.Context, string, json.RawMessage) (string, error),
        chunkFn func(string),  // nil = no streaming; call with each output chunk if streaming
    ) (string, queue.TokenUsage, error)
}

RunTask is called once per task. It must:

Build the prompt from cfg.SystemPrompt and task.Prompt
Call your model API (with retries if appropriate)
Run the tool-use loop if the model returns tool calls — call callTool, feed results back
Call chunkFn with each output chunk if streaming (pass nil to skip)
Return the final text output, token usage, and any error

Step 1: Create your provider package

mkdir -p internal/agent/providers/myprovider

Implement the interface:

// internal/agent/providers/myprovider/provider.go
package myprovider

import (
    "context"
    "encoding/json"

    "github.com/arkonis-dev/ark-operator/internal/agent/config"
    "github.com/arkonis-dev/ark-operator/internal/agent/mcp"
    "github.com/arkonis-dev/ark-operator/internal/agent/providers"
    "github.com/arkonis-dev/ark-operator/internal/agent/queue"
)

type MyProvider struct {
    // Add your client, base URL, etc. here
}

func (p *MyProvider) RunTask(
    ctx context.Context,
    cfg *config.Config,
    task queue.Task,
    tools []mcp.Tool,
    callTool func(context.Context, string, json.RawMessage) (string, error),
    chunkFn func(string),
) (string, queue.TokenUsage, error) {
    // 1. Build your request
    // 2. Call your API
    // 3. Handle tool calls in a loop
    // 4. Return output, token counts, error
    return output, queue.TokenUsage{InputTokens: inputTok, OutputTokens: outputTok}, nil
}

// Register with the runtime on import
func init() {
    providers.Register("myprovider", func() providers.LLMProvider {
        return &MyProvider{}
    })
}

Step 2: Register the provider

Blank-import your package in the agent runtime entrypoint so the init() function runs:

// runtime/agent/main.go (or cmd/main.go for the operator)
import (
    _ "github.com/arkonis-dev/ark-operator/internal/agent/providers/myprovider"
)

For an external provider (separate module), import from your module path instead.

Step 3: Activate it in agent pods

Set AGENT_PROVIDER to your provider name in the agent pods. Via Helm:

helm upgrade ark-operator arkonis/ark-operator \
  --set agentExtraEnv[0].name=AGENT_PROVIDER,agentExtraEnv[0].value=myprovider

Or set it per-agent via ArkAgent (the AGENT_PROVIDER env var overrides the global setting for that agent’s pods).

Step 4: Test locally

AGENT_PROVIDER=myprovider ark run team.yaml --watch

If the provider name isn’t recognized, ark run exits with:

unknown provider "myprovider" — did you register it with providers.Register?

TokenUsage

Return accurate token counts so budget enforcement and cost tracking work correctly:

return output, queue.TokenUsage{
    InputTokens:  int64(inputTokens),
    OutputTokens: int64(outputTokens),
}, nil

If your API doesn’t report token counts, return zeros — the budget enforcement will not trigger but the cost tracking will show 0.

Tool-use loop pattern

If your model returns tool calls, feed results back in a loop:

for {
    resp, err := callYourAPI(ctx, messages)
    if err != nil {
        return "", queue.TokenUsage{}, err
    }

    if !resp.HasToolCalls() {
        return resp.Text(), resp.Usage(), nil
    }

    for _, tc := range resp.ToolCalls() {
        result, err := callTool(ctx, tc.Name, tc.Input)
        if err != nil {
            result = fmt.Sprintf("error: %s", err)
        }
        messages = append(messages, toolResultMessage(tc.ID, result))
    }
    messages = append(messages, assistantMessage(resp))
}