MLflow Observability

MLflow Tracing provides automatic observability for applications built with the Go-AI SDK via OpenTelemetry. When enabled, MLflow records prompts, responses, token usage, latencies, and exceptions for comprehensive experiment tracking and debugging.

What is MLflow?

MLflow is an open-source platform for managing the ML lifecycle, including experimentation, reproducibility, deployment, and a central model registry. MLflow Tracing enables tracking of LLM applications with detailed observability.

What Gets Tracked

When MLflow observability is enabled, the following information is automatically captured:

• Prompts and Messages: Input text and conversation history
• Generated Responses: Complete model outputs
• Token Usage: Input, output, and total tokens (with cache details when available)
• Latencies: Request duration and call hierarchy
• Exceptions: Errors and stack traces
• Metadata: Model IDs, parameters, and explicitly included context

Quick Start

Prerequisites

1. Start MLflow Tracking Server:

mlflow server --backend-store-uri sqlite:///mlruns.db --port 5000

For production, use a proper database backend. See MLflow Setup Guide.

2. Install Go-AI SDK:

go get github.com/digitallysavvy/go-ai

Basic Setup

package main

import (
    "context"
    "log"
    "os"

    "github.com/digitallysavvy/go-ai/pkg/ai"
    "github.com/digitallysavvy/go-ai/pkg/observability/mlflow"
    "github.com/digitallysavvy/go-ai/pkg/providers/openai"
    "github.com/digitallysavvy/go-ai/pkg/telemetry"
)

func main() {
    // Create MLflow tracker
    tracker, err := mlflow.New(mlflow.Config{
        TrackingURI:    "http://localhost:5000",
        ExperimentName: "my-ai-app",
        ServiceName:    "openai-service",
        Insecure:       true, // For local development
    })
    if err != nil {
        log.Fatal(err)
    }
    defer tracker.Shutdown(context.Background())

    // Create provider and model
    provider := openai.New(openai.Config{
        APIKey: os.Getenv("OPENAI_API_KEY"),
    })
    model, _ := provider.LanguageModel("gpt-4")

    // Send spans to MLflow: register an integration built with the MLflow tracer
    telemetry.RegisterTelemetryIntegration(
        telemetry.NewLegacyOpenTelemetry(telemetry.LegacyOpenTelemetryOptions{Tracer: tracker.Tracer()}),
    )

    // Configure telemetry
    telemetrySettings := &telemetry.Options{
        IsEnabled:     telemetry.Bool(true),
        RecordInputs:  true,
        RecordOutputs: true,
        FunctionID:    "text-generation",
    }

    // Generate text with telemetry
    result, err := ai.GenerateText(context.Background(), ai.GenerateTextOptions{
        Model:     model,
        Prompt:    "Explain machine learning in simple terms",
        Telemetry: telemetrySettings,
    })
    if err != nil {
        log.Fatal(err)
    }

    log.Printf("Response: %s\n", result.Text)

    // Flush traces to MLflow
    tracker.ForceFlush(context.Background())

    log.Println("✅ Traces sent to MLflow at http://localhost:5000")
}

View Traces

1. Open MLflow UI: http://localhost:5000
2. Navigate to: Experiments > my-ai-app > Traces
3. Click on individual traces to see detailed information

Configuration Options

MLflow Config

type Config struct {
    // TrackingURI is the MLflow tracking server endpoint (required)
    // Examples: "http://localhost:5000", "https://mlflow.example.com"
    TrackingURI string

    // ExperimentName is the MLflow experiment to log to (optional)
    // Default: "default"
    ExperimentName string

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

    // ServiceName for OpenTelemetry (optional)
    // Default: "go-ai-sdk"
    ServiceName string

    // Insecure controls HTTP vs HTTPS (optional)
    // Set to true for local development
    // Default: false (uses HTTPS)
    Insecure bool

    // Headers for additional authentication (optional)
    // Example: map[string]string{"Authorization": "Bearer token"}
    Headers map[string]string
}

Examples

Local Development:

tracker, err := mlflow.New(mlflow.Config{
    TrackingURI:    "http://localhost:5000",
    ExperimentName: "dev-experiment",
    Insecure:       true,
})

Production with HTTPS:

tracker, err := mlflow.New(mlflow.Config{
    TrackingURI:    "https://mlflow.company.com",
    ExperimentID:   "12345",
    ServiceName:    "production-ai-service",
    Headers: map[string]string{
        "Authorization": "Bearer " + os.Getenv("MLFLOW_TOKEN"),
    },
})

Multiple Experiments:

// Experiment 1: Development
devTracker, _ := mlflow.New(mlflow.Config{
    TrackingURI:    "http://localhost:5000",
    ExperimentName: "development",
})

// Experiment 2: Production
prodTracker, _ := mlflow.New(mlflow.Config{
    TrackingURI:    "https://mlflow.company.com",
    ExperimentName: "production",
})

Telemetry Settings

Configure what gets tracked:

telemetrySettings := &telemetry.Options{
    // Enable/disable telemetry
    IsEnabled: telemetry.Bool(true),

    // Control what gets recorded
    RecordInputs:  true,  // Record prompts and messages
    RecordOutputs: true,  // Record generated responses

    // Function identifier for grouping traces
    FunctionID: "chatbot-generation",

    // Explicit context inclusion for telemetry
    IncludeRuntimeContext: map[string]bool{"request_id": true},
}
// The MLflow tracer is configured on the registered integration
// (telemetry.NewLegacyOpenTelemetry), not on telemetry.Options.

Privacy and Security

Control sensitive data recording:

// Disable input/output recording for sensitive data
telemetrySettings := &telemetry.Options{
    IsEnabled: telemetry.Bool(true),
    RecordInputs:  false, // Don't log user prompts
    RecordOutputs: false, // Don't log responses
}
// Token usage and latency still tracked

Advanced Usage

Streaming with MLflow

stream, err := ai.StreamText(context.Background(), ai.StreamTextOptions{
    Model:  model,
    Prompt: "Write a long article about AI",
    // Telemetry automatically tracks streaming
})

for chunk := range stream.Chunks() {
    fmt.Print(chunk.Text)
}

if err := stream.Err(); err != nil {
    log.Printf("stream error: %v", err)
}

// Final token usage is tracked when stream completes

Multi-Step Tool Calling

MLflow automatically tracks tool calling loops:

result, err := ai.GenerateText(context.Background(), ai.GenerateTextOptions{
    Model:  model,
    Prompt: "What's the weather in San Francisco?",
    Tools: []types.Tool{
        weatherTool,
    },
    StopWhen: []ai.StopCondition{ai.IsStepCount(5)},
    // Each tool call step is tracked as a span
})

// MLflow shows:
// - Initial prompt
// - Model requests tool call
// - Tool execution
// - Final response with tool result

Custom Spans

Add custom tracing for your application logic:

import "go.opentelemetry.io/otel/trace"

func processWithTracing(ctx context.Context, tracker *mlflow.Tracker) {
    tracer := tracker.Tracer()
    ctx, span := tracer.Start(ctx, "custom-processing")
    defer span.End()

    // Your processing logic
    result, err := ai.GenerateText(ctx, options)

    // Add custom attributes
    span.SetAttributes(
        attribute.String("result_type", "summary"),
        attribute.Int64("result_length", int64(len(result.Text))),
    )
}

Best Practices

1. Experiment Organization

// Group related runs by experiment
experiments := map[string]string{
    "development":   "dev-experiment",
    "staging":       "staging-experiment",
    "production":    "prod-experiment",
    "feature-test":  "feature-xyz-experiment",
}

env := os.Getenv("ENVIRONMENT")
tracker, _ := mlflow.New(mlflow.Config{
    TrackingURI:    mlflowURI,
    ExperimentName: experiments[env],
})

2. Include Safe Context

telemetrySettings.IncludeRuntimeContext = map[string]bool{
    "request_id": true,
    "feature_flag": true,
    "ab_test_group": true,
}

3. Graceful Shutdown

func main() {
    tracker, _ := mlflow.New(config)

    // Ensure traces are flushed on shutdown
    defer func() {
        ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
        defer cancel()

        if err := tracker.ForceFlush(ctx); err != nil {
            log.Printf("Error flushing traces: %v", err)
        }
        if err := tracker.Shutdown(ctx); err != nil {
            log.Printf("Error shutting down tracker: %v", err)
        }
    }()

    // Your application logic
}

4. Error Tracking

result, err := ai.GenerateText(ctx, options)
if err != nil {
    // Errors are automatically captured in traces
    log.Printf("Generation failed: %v", err)
    return
}

Troubleshooting

Traces Not Appearing

1. Check MLflow server is running:

curl http://localhost:5000/api/2.0/mlflow/experiments/list

2. Verify experiment exists:

mlflow experiments list

3. Check for flush errors:

if err := tracker.ForceFlush(ctx); err != nil {
    log.Printf("Flush error: %v", err)
}

4. Enable debug logging:

import "go.opentelemetry.io/otel"
import "go.opentelemetry.io/otel/sdk/trace"

// Add this before creating tracker
otel.SetErrorHandler(otel.ErrorHandlerFunc(func(err error) {
    log.Printf("OTEL error: %v", err)
}))

Connection Issues

// Ensure correct URL format
tracker, err := mlflow.New(mlflow.Config{
    TrackingURI: "http://localhost:5000", // Include http://
    Insecure:    true,                     // For local HTTP
})
if err != nil {
    log.Fatalf("Failed to create tracker: %v", err)
}

Network Timeouts

// Create context with timeout
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()

// Use context for operations
result, err := ai.GenerateText(ctx, options)

Comparison with TypeScript SDK

The Go implementation provides the same MLflow integration as the TypeScript AI SDK:

TypeScript:

import { init } from 'mlflow-tracing';
import { generateText } from 'ai';

init(); // Auto-configures from environment

const result = await generateText({
  model: openai('gpt-4'),
  prompt: 'Hello',
  experimental_telemetry: { isEnabled: true },
});

Go:

import "github.com/digitallysavvy/go-ai/pkg/observability/mlflow"

tracker, _ := mlflow.New(mlflow.Config{
    TrackingURI:    os.Getenv("MLFLOW_TRACKING_URI"),
    ExperimentName: "my-app",
})
defer tracker.Shutdown(context.Background())

result, _ := ai.GenerateText(context.Background(), ai.GenerateTextOptions{
    Model:  model,
    Prompt: "Hello",
    // Telemetry configured via tracer
})

See Also

• MLflow Official Documentation
• MLflow Tracing Guide
• OpenTelemetry Go SDK
• Telemetry Package Reference
• Example: observability-mlflow