TramAI - governed AI workflows for Java and Kotlin

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
Next