mlflow

Documentation

Overview

Package mlflow provides MLflow observability integration for the Go AI SDK. It enables automatic tracking of AI operations including prompts, responses, token usage, latencies, and costs to MLflow Tracking Server.

MLflow Tracing provides experiment tracking and observability for AI applications via OpenTelemetry. When enabled, MLflow records:

• Prompts/messages and generated responses
• Latencies and call hierarchy
• Token usage (when the provider returns it)
• Exceptions and errors

Example usage:

tracker := mlflow.New(mlflow.Config{
    TrackingURI:    "http://localhost:5000",
    ExperimentName: "my-ai-app",
})
defer tracker.Shutdown(context.Background())

// Use with middleware or direct telemetry settings
tracer := tracker.Tracer()

Index

• type Config
• type Tracker
  • func New(cfg Config) (*Tracker, error)
  • func (t *Tracker) ForceFlush(ctx context.Context) error
  • func (t *Tracker) Shutdown(ctx context.Context) error
  • func (t *Tracker) Tracer() trace.Tracer

Types

type Config

type Config struct {
	// TrackingURI is the MLflow tracking server endpoint
	// Example: "http://localhost:5000" or "https://mlflow.example.com"
	TrackingURI string

	// ExperimentName is the name of the MLflow experiment to log to
	// If not provided, uses "default"
	ExperimentName string

	// ExperimentID is the MLflow experiment ID (optional)
	// Takes precedence over ExperimentName if both are provided
	ExperimentID string

	// ServiceName is the name of the service for OpenTelemetry
	// If not provided, uses "go-ai-sdk"
	ServiceName string

	// Insecure controls whether to use insecure HTTP connection
	// Set to true for local development without TLS
	// Default: false (uses HTTPS)
	Insecure bool

	// Headers contains additional headers to send with trace exports
	// Example: map[string]string{"Authorization": "Bearer token"}
	Headers map[string]string
}

Config holds configuration for MLflow observability

type Tracker

type Tracker struct {
	// contains filtered or unexported fields
}

Tracker manages MLflow observability integration

func New

func New(cfg Config) (*Tracker, error)

New creates a new MLflow tracker with the provided configuration

func (*Tracker) ForceFlush

func (t *Tracker) ForceFlush(ctx context.Context) error

ForceFlush forces any pending spans to be exported immediately

func (*Tracker) Shutdown

func (t *Tracker) Shutdown(ctx context.Context) error

Shutdown gracefully shuts down the tracker, flushing any pending spans

func (*Tracker) Tracer

func (t *Tracker) Tracer() trace.Tracer

Tracer returns an OpenTelemetry tracer that can be used with AI SDK telemetry settings