A8065384/latticelm
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve test coverage

Browse Source
This commit is contained in:
Anibal Angulo
2026-03-05 18:07:33 +00:00
parent 1e0bb0be8c
commit ccb8267813
8 changed files with 7550 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
internal/conversation/sql_store.go
View File

@@ -148,7 +148,20 @@ func (s *SQLStore) Size() int {
}
func (s *SQLStore) cleanup() {
ticker := time.NewTicker(1 * time.Minute)
// Calculate cleanup interval as 10% of TTL, with sensible bounds
interval := s.ttl / 10
// Cap maximum interval at 1 minute for production
if interval > 1*time.Minute {
interval = 1 * time.Minute
}
// Allow small intervals for testing (as low as 10ms)
if interval < 10*time.Millisecond {
interval = 10 * time.Millisecond
}
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {

83
internal/observability/provider_wrapper.go
View File

@@ -132,48 +132,53 @@ func (p *InstrumentedProvider) GenerateStream(ctx context.Context, messages []ap
defer close(outChan)
defer close(outErrChan)
// Helper function to record final metrics
recordMetrics := func() {
duration := time.Since(start).Seconds()
status := "success"
if streamErr != nil {
status = "error"
if p.tracer != nil {
span := trace.SpanFromContext(ctx)
span.RecordError(streamErr)
span.SetStatus(codes.Error, streamErr.Error())
}
} else {
if p.tracer != nil {
span := trace.SpanFromContext(ctx)
span.SetAttributes(
attribute.Int64("provider.input_tokens", totalInputTokens),
attribute.Int64("provider.output_tokens", totalOutputTokens),
attribute.Int64("provider.chunk_count", chunkCount),
attribute.Float64("provider.ttfb_seconds", ttfb.Seconds()),
)
span.SetStatus(codes.Ok, "")
}
// Record token metrics
if p.registry != nil && (totalInputTokens > 0 || totalOutputTokens > 0) {
providerTokensTotal.WithLabelValues(p.base.Name(), req.Model, "input").Add(float64(totalInputTokens))
providerTokensTotal.WithLabelValues(p.base.Name(), req.Model, "output").Add(float64(totalOutputTokens))
}
}
// Record stream metrics
if p.registry != nil {
providerRequestsTotal.WithLabelValues(p.base.Name(), req.Model, "generate_stream", status).Inc()
providerStreamDuration.WithLabelValues(p.base.Name(), req.Model).Observe(duration)
providerStreamChunks.WithLabelValues(p.base.Name(), req.Model).Add(float64(chunkCount))
if ttfb > 0 {
providerStreamTTFB.WithLabelValues(p.base.Name(), req.Model).Observe(ttfb.Seconds())
}
}
}
for {
select {
case delta, ok := <-baseChan:
if !ok {
// Stream finished - record final metrics
duration := time.Since(start).Seconds()
status := "success"
if streamErr != nil {
status = "error"
if p.tracer != nil {
span := trace.SpanFromContext(ctx)
span.RecordError(streamErr)
span.SetStatus(codes.Error, streamErr.Error())
}
} else {
if p.tracer != nil {
span := trace.SpanFromContext(ctx)
span.SetAttributes(
attribute.Int64("provider.input_tokens", totalInputTokens),
attribute.Int64("provider.output_tokens", totalOutputTokens),
attribute.Int64("provider.chunk_count", chunkCount),
attribute.Float64("provider.ttfb_seconds", ttfb.Seconds()),
)
span.SetStatus(codes.Ok, "")
}
// Record token metrics
if p.registry != nil && (totalInputTokens > 0 || totalOutputTokens > 0) {
providerTokensTotal.WithLabelValues(p.base.Name(), req.Model, "input").Add(float64(totalInputTokens))
providerTokensTotal.WithLabelValues(p.base.Name(), req.Model, "output").Add(float64(totalOutputTokens))
}
}
// Record stream metrics
if p.registry != nil {
providerRequestsTotal.WithLabelValues(p.base.Name(), req.Model, "generate_stream", status).Inc()
providerStreamDuration.WithLabelValues(p.base.Name(), req.Model).Observe(duration)
providerStreamChunks.WithLabelValues(p.base.Name(), req.Model).Add(float64(chunkCount))
if ttfb > 0 {
providerStreamTTFB.WithLabelValues(p.base.Name(), req.Model).Observe(ttfb.Seconds())
}
}
recordMetrics()
return
}
@@ -198,8 +203,10 @@ func (p *InstrumentedProvider) GenerateStream(ctx context.Context, messages []ap
if ok && err != nil {
streamErr = err
outErrChan <- err
recordMetrics()
return
}
return
// If error channel closed without error, continue draining baseChan
}
}
}()

2
internal/observability/testing.go
View File

@@ -10,7 +10,7 @@ import (
"go.opentelemetry.io/otel/sdk/resource"
sdktrace "go.opentelemetry.io/otel/sdk/trace"
"go.opentelemetry.io/otel/sdk/trace/tracetest"
semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
semconv "go.opentelemetry.io/otel/semconv/v1.24.0"
)
// NewTestRegistry creates a new isolated Prometheus registry for testing

13
internal/observability/tracing.go
View File

@@ -17,19 +17,14 @@ import (
// InitTracer initializes the OpenTelemetry tracer provider.
func InitTracer(cfg config.TracingConfig) (*sdktrace.TracerProvider, error) {
// Create resource with service information
res, err := resource.Merge(
resource.Default(),
resource.NewWithAttributes(
semconv.SchemaURL,
semconv.ServiceName(cfg.ServiceName),
),
// Use NewSchemaless to avoid schema version conflicts
res := resource.NewSchemaless(
semconv.ServiceName(cfg.ServiceName),
)
if err != nil {
return nil, fmt.Errorf("failed to create resource: %w", err)
}
// Create exporter
var exporter sdktrace.SpanExporter
var err error
switch cfg.Exporter.Type {
case "otlp":
exporter, err = createOTLPExporter(cfg.Exporter)