99 lines
3.1 KiB
Go
99 lines
3.1 KiB
Go
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)
|
|
}
|