Add observabilitty and monitoring

internal/observability/init.go

@@ -0,0 +1,98 @@
+package observability
+
+import (
+	"github.com/ajac-zero/latticelm/internal/conversation"
+	"github.com/ajac-zero/latticelm/internal/providers"
+	"github.com/prometheus/client_golang/prometheus"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+)
+
+// ProviderRegistry defines the interface for provider registries.
+// This matches the interface expected by the server.
+type ProviderRegistry interface {
+	Get(name string) (providers.Provider, bool)
+	Models() []struct{ Provider, Model string }
+	ResolveModelID(model string) string
+	Default(model string) (providers.Provider, error)
+}
+
+// WrapProviderRegistry wraps all providers in a registry with observability.
+func WrapProviderRegistry(registry ProviderRegistry, metricsRegistry *prometheus.Registry, tp *sdktrace.TracerProvider) ProviderRegistry {
+	if registry == nil {
+		return nil
+	}
+
+	// We can't directly modify the registry's internal map, so we'll need to
+	// wrap providers as they're retrieved. Instead, create a new instrumented registry.
+	return &InstrumentedRegistry{
+		base:             registry,
+		metrics:          metricsRegistry,
+		tracer:           tp,
+		wrappedProviders: make(map[string]providers.Provider),
+	}
+}
+
+// InstrumentedRegistry wraps a provider registry to return instrumented providers.
+type InstrumentedRegistry struct {
+	base             ProviderRegistry
+	metrics          *prometheus.Registry
+	tracer           *sdktrace.TracerProvider
+	wrappedProviders map[string]providers.Provider
+}
+
+// Get returns an instrumented provider by entry name.
+func (r *InstrumentedRegistry) Get(name string) (providers.Provider, bool) {
+	// Check if we've already wrapped this provider
+	if wrapped, ok := r.wrappedProviders[name]; ok {
+		return wrapped, true
+	}
+
+	// Get the base provider
+	p, ok := r.base.Get(name)
+	if !ok {
+		return nil, false
+	}
+
+	// Wrap it
+	wrapped := NewInstrumentedProvider(p, r.metrics, r.tracer)
+	r.wrappedProviders[name] = wrapped
+	return wrapped, true
+}
+
+// Default returns the instrumented provider for the given model name.
+func (r *InstrumentedRegistry) Default(model string) (providers.Provider, error) {
+	p, err := r.base.Default(model)
+	if err != nil {
+		return nil, err
+	}
+
+	// Check if we've already wrapped this provider
+	name := p.Name()
+	if wrapped, ok := r.wrappedProviders[name]; ok {
+		return wrapped, nil
+	}
+
+	// Wrap it
+	wrapped := NewInstrumentedProvider(p, r.metrics, r.tracer)
+	r.wrappedProviders[name] = wrapped
+	return wrapped, nil
+}
+
+// Models returns the list of configured models and their provider entry names.
+func (r *InstrumentedRegistry) Models() []struct{ Provider, Model string } {
+	return r.base.Models()
+}
+
+// ResolveModelID returns the provider_model_id for a model.
+func (r *InstrumentedRegistry) ResolveModelID(model string) string {
+	return r.base.ResolveModelID(model)
+}
+
+// WrapConversationStore wraps a conversation store with observability.
+func WrapConversationStore(store conversation.Store, backend string, metricsRegistry *prometheus.Registry, tp *sdktrace.TracerProvider) conversation.Store {
+	if store == nil {
+		return nil
+	}
+
+	return NewInstrumentedStore(store, backend, metricsRegistry, tp)
+}