Custom completion provider

Coming soon

Here is an example to provide a custom completion provider to your agent.

import asyncio
import os
from typing import Literal

from mistralai.agents import agents
from mistralai.vibe.harness.runtime.completion import (
    CompletionConfigBase,
    CompletionProvider,
)
from mistralai.vibe.harness.runtime.completion.providers.openai import (
    OpenAICompletionProvider,
)


# A custom provider should implement the CompletionProvider protocol.
# Extending OpenAICompletionProvider provides that implementation for OpenAI-compatible APIs.
class CompanyGatewayCompletionProvider(OpenAICompletionProvider):
    ...


class CompanyGatewayConfig(CompletionConfigBase):
    # Required by the SDK. Declaring it registers this class: the value is serialized
    # with the config, and each process looks the class up by it. It must be unique so
    # one serialized value identifies exactly one class.
    type: Literal["company_gateway"] = "company_gateway"

    # Required by CompletionConfigBase. When this config is passed to
    # Agent(completion=...), its value must exactly match Agent.model; Agent
    # construction rejects a mismatch.
    model: str = "company-model"

    # Example gateway-specific field.
    base_url: str = "https://llm.example.com/v1"

    # Example secret reference. This value is the environment variable's name,
    # not the authentication token itself, so serialization does not expose it.
    api_key_env_var: str = "COMPANY_LLM_API_KEY"

    def build_provider(self) -> CompletionProvider:
        # Instantiate custom completion provider here from config data
        return CompanyGatewayCompletionProvider(
            # Read the secret only when the runtime constructs the live provider.
            api_key=os.environ[self.api_key_env_var],
            model=self.model,
            base_url=self.base_url,
        )


# Then use in agent definition
agent = agents.Agent(
    model="company-model",
    completion=CompanyGatewayConfig(),
    instructions="Be concise.",
)

A completion provider is a class responsible for connecting an agent to a model. The SDK has built-in adapters for Mistral, OpenAI, and Anthropic; select one with Agent(provider=...). Use a custom completion provider when the model must go through another gateway or when its request and streaming response formats need custom handling.

For a complete OpenRouter integration that adds attribution headers, provider routing, and a token limit, see the openrouter_completion.py example.

A custom integration has two parts:

  • a CompletionConfigBase subclass: serializable configuration that can cross process and workflow boundaries;
  • a CompletionProvider: the runtime object that turns a provider-neutral completion request into a stream of provider-neutral completion events.

Define the config, then pass an instance through Agent(completion=...). There is no registration call: declaring type on the subclass registers it when the class statement runs.

If your gateway implements OpenAI's Chat Completions protocol, reuse the SDK's request and stream adapter and keep only the gateway-specific configuration in your class:

Importing the module is the contract

Importing the module is the contract

Registration is per process, and the only thing that performs it is importing the module that defines the config class. The process that authors the agent and the process that calls the model are not the same one: on a worker the config travels as data and is rebuilt from its type, so a worker that never imported your module refuses the payload with UnsupportedCompletionProviderError.

Keep the config class in a module your agent package imports at startup — the module that defines the agent is the simplest place.

Configuration and constraints

Configuration and constraints

Everything declared on CompletionConfigBase is serialized across process and workflow boundaries. The api_key_env_var field above demonstrates how to handle a secret without putting it in that serialized data: the config contains only the environment-variable name ("COMPANY_LLM_API_KEY"). build_provider() reads the token from the environment later, in the process that constructs the live provider. The token is passed to the provider but is never stored back on the serialized config. Make that environment variable available to every process that can build the provider, including workflow activity workers.

When using completion=:

  • Agent.model must equal completion.model;
  • do not also pass provider=; completion= replaces it entirely;
  • the custom provider owns its endpoint, credentials, headers, and transport—the SDK does not inject the Mistral completion gateway settings;
  • Agent(reasoning_effort=...) and Agent(temperature=...) are not supported; add provider-specific reasoning and sampling fields to your config and handle them in your provider instead;
  • run(..., response_format=...) structured output is not supported with a custom provider.

Tools, sessions, hooks, and the rest of the harness continue to work as long as the provider translates the corresponding messages and tool-call stream correctly.

See also

See also

  • The Agent - where completion= is attached.
  • Run modes - local, self-hosted, and managed execution.