Tramai Logo
Tramai

Custom Providers

TramAI is intentionally structured so new providers can be added without changing the public service model. The ModelProvider SPI lets you integrate any AI backend — proprietary APIs, internal gateways, community models, or custom inference servers.

What This Covers

  • The ModelProvider SPI and its methods
  • Writing a full provider implementation
  • Provider module boundaries and rules
  • Native structured output support
  • Provider Trust Zones for sovereign mode
  • Registration in standalone and Spring Boot setups

The ModelProvider SPI

All providers implement the ModelProvider interface from tramai-core:

interface ModelProvider {
    suspend fun complete(request: ModelRequest): ModelResponse
    fun providerId(): String
    fun supportsStructuredOutput(): Boolean = false
}
MethodPurpose
completeSend a request to the AI backend and return a response
providerIdUnique string identifier used in model routing and configuration
supportsStructuredOutput(Optional) signal native structured output support to the engine

Full Provider Example

Here is a complete custom provider that wraps a hypothetical internal AI gateway:

class InternalGatewayProvider(
    private val baseUrl: String,
    private val apiKey: String,
    private val httpClient: HttpClient = HttpClient.default(),
) : ModelProvider {

    override fun providerId(): String = "internal-gateway"

    override suspend fun complete(request: ModelRequest): ModelResponse {
        // Build the provider-native HTTP payload from ModelRequest
        val payload = GatewayRequest(
            model = request.model,
            messages = request.messages.map { msg ->
                GatewayMessage(role = msg.role, content = msg.content)
            },
            temperature = request.temperature,
            maxTokens = request.maxTokens,
        )

        // Send the request
        val response: GatewayResponse = try {
            httpClient.post("$baseUrl/v1/chat") {
                header("Authorization", "Bearer $apiKey")
                setBody(payload)
            }.body()
        } catch (e: Exception) {
            throw ProviderException(
                providerId = providerId(),
                message = "Gateway request failed: ${e.message}",
                cause = e,
            )
        }

        // Map the provider-native response to ModelResponse
        return ModelResponse(
            content = response.choices.firstOrNull()?.text ?: "",
            modelUsed = request.model,
            inputTokens = response.usage?.inputTokens ?: 0,
            outputTokens = response.usage?.outputTokens ?: 0,
            finishReason = mapFinishReason(response.choices.firstOrNull()?.finishReason),
        )
    }

    private fun mapFinishReason(raw: String?): FinishReason = when (raw) {
        "stop" -> FinishReason.STOP
        "length" -> FinishReason.LENGTH
        "tool_calls" -> FinishReason.TOOL_CALLS
        else -> FinishReason.OTHER
    }
}

// Provider-native DTOs (internal to the provider module)
data class GatewayRequest(
    val model: String,
    val messages: List<GatewayMessage>,
    val temperature: Double? = null,
    val maxTokens: Int? = null,
)

data class GatewayMessage(
    val role: String,
    val content: String,
)

data class GatewayResponse(
    val choices: List<GatewayChoice>,
    val usage: GatewayUsage? = null,
)

data class GatewayChoice(
    val text: String,
    val finishReason: String? = null,
)

data class GatewayUsage(
    val inputTokens: Int,
    val outputTokens: Int,
)

Registering the Provider

Standalone Builder

val tramai = Tramai {
    provider(
        InternalGatewayProvider(
            baseUrl = System.getenv("GATEWAY_URL"),
            apiKey = System.getenv("GATEWAY_API_KEY"),
        ),
        name = "gateway",
        default = true,
    )
    model("gpt-4o-mini", "gateway")
}

Spring Boot

@Bean
fun gatewayProvider(): ModelProvider = InternalGatewayProvider(
    baseUrl = environment.getProperty("gateway.base-url")!!,
    apiKey = environment.getProperty("gateway.api-key")!!,
)

Provider Implementation Rules

Provider modules are responsible for:

  • Translating ModelRequest into provider-native HTTP payloads
  • Mapping provider responses into ModelResponse
  • Surfacing transport or protocol failures as ProviderException

Do not put in providers:

  • Retry policy for structured output
  • Typed parsing logic
  • Tracing policy
  • Service proxy rules
  • Token budget enforcement
  • Circuit-breaking logic

The engine owns all of those. Your provider should focus on transport and mapping only.

Native Structured Output

If your provider supports structured output natively (like OpenAI's response_format or Anthropic's tool-use trick), override supportsStructuredOutput():

override fun supportsStructuredOutput(): Boolean = true

The engine uses this flag to decide whether to use the provider's native structured output or fall back to schema-in-prompt mode. When true, the engine sends the JSON schema as part of the request payload; the provider must honor it.

Provider Trust Zones (Sovereign Mode)

When running in sovereign mode, each provider is assigned a trust zone through ProviderTrustZone. This controls routing policy based on data residency requirements.

val tramai = Tramai {
    provider(OllamaProvider(baseUrl = "http://localhost:11434"), name = "ollama")
    provider(OpenAiProvider(apiKey = "..."), name = "openai")

    providerZones = mapOf(
        "ollama" to ProviderTrustZone.LOCAL,
        "openai" to ProviderTrustZone.GLOBAL_CLOUD,
    )
}

Available trust zones:

ZoneMeaning
ProviderTrustZone.LOCALSame-machine or local-network provider (Ollama, local vLLM)
ProviderTrustZone.EU_CLOUDEU-hosted cloud provider (GDPR-aligned)
ProviderTrustZone.GLOBAL_CLOUDGlobal cloud provider (standard SaaS)

Zones are enforced by the sovereign policy engine during provider routing. When a deployment mode restricts data to certain zones, routing to providers outside those zones is blocked.

For full sovereign mode setup, see Sovereign Mode.

Module Structure

When creating a standalone provider module, follow this structure:

my-tramai-provider/
├── src/main/kotlin/com/example/
│   ├── MyProvider.kt           # ModelProvider implementation
│   ├── MyProviderConfig.kt     # Configuration data class
│   └── MyProviderFactory.kt    # Optional factory for Spring/standalone
├── build.gradle.kts
└── README.md

Dependency philosophy:

  • tramai-core has zero runtime dependencies beyond Kotlin stdlib
  • Providers depend on tramai-core for shared contracts
  • Providers do NOT depend on tramai-engine, tramai-structured, or tramai-observability
  • Framework adapters (tramai-spring) are thin integration layers, not alternate runtimes

Testing a Custom Provider

Use the testing utilities from tramai-testing to verify your provider in isolation:

class InternalGatewayProviderTest {

    private val provider = InternalGatewayProvider(
        baseUrl = "http://localhost:8080",
        apiKey = "test-key",
    )

    @Test
    fun `should map provider response correctly`() = runTest {
        // Use MockWebServer or similar to simulate the backend
        val response = provider.complete(
            ModelRequest(model = "test-model", messages = listOf(Message.user("hello")))
        )
        assertNotNull(response.content)
    }
}

For more on testing, see Testing TramAI.

Before Adding Large Features

Check whether the feature belongs in TramAI's current scope.

Good fit:

  • New provider module for a public or internal API
  • Better testing helper for provider verification
  • Additional configuration options

Higher-risk fit (update specs first):

  • Agent-style tool orchestration inside the provider
  • Memory and session state management
  • Generated code paths
  • Provider-specific special cases leaking into the engine

Next Steps

Previous
Custom Observers
Next
MCP Integration