Phase 7: AI-assisted query authoring (autocomplete, explain, fix, optimize, NL translation)

Adds a self-hosted (Ollama, qwen2.5-coder) model provider abstraction with a pluggable opt-in cloud adapter, schema grounding, and a shared cost/safety guard every AI-suggested query is assessed against -- compiling to and executing through the same unchanged Phase 2 IR/ compiler and Phase 4 tenant scoping as a hand-written query, no parallel execution path. Track A (built into the query bar): inline ghost-text autocomplete, "Explain this query", "Fix this query" with a diff view, and a rule-based "Optimize" suggestion. Track B: natural-language-to-query translation, always a separate review step from execution, with `sentryctl query --nl` requiring explicit confirmation to run. Every accepted/dismissed translate-fix-optimize interaction is logged into the same append-only audit_log table Phase 4 built. Two real product bugs were found and fixed via live browser verification (a Svelte effect re-running on every keystroke that silently cancelled the ghost-text debounce; a ghost-text widget positioned at document offset 0 instead of the cursor), and a real costguard logic bug (unbounded-aggregation vs. raw-row) was caught by its own test suite. New integration tests wire a real Ollama client through the real HTTP handler against a mock server matching Ollama's wire contract (hack/mock-ollama), keeping model-quality verification out of CI as a disclosed, periodic human-run check instead. See /docs/phase-7-ai-design.md and /docs/phase-7-runbook.md.
2026-08-16 18:06:27 -07:00
parent 661568085e
commit 7d316f92db
37 changed files with 5230 additions and 20 deletions
@@ -0,0 +1,257 @@
+// Package grounding builds provider.SchemaContext from a tenant's own
+// ClickHouse data -- known service names, common attribute keys, and
+// example values for enum-like fields -- sourced by periodic sampling,
+// never hand-maintained (task 3). See /docs/phase-7-ai-design.md's
+// "Schema grounding" section for the embedded-in-prompt-vs-retrieved
+// tradeoff this package's shape is built around.
+//
+// Tenant scoping is structural, not a filter this package applies: a
+// Service wraps exactly one executor.SQLRunner, and that SQLRunner is
+// already tenant-scoped by whoever constructed it (the plain shared
+// runner in a single-tenant deployment, or one specific tenant's
+// chrunner-resolved connection in enterprise-api) -- the same connection-
+// layer isolation discipline Phase 4 established for query execution
+// applies here for free, because grounding queries run through the exact
+// same SQLRunner interface, never a separate admin/shared connection.
+// A multi-tenant deployment needs one Service per active tenant --
+// enterprise/internal/groundingregistry provides that, mirroring
+// enterprise/internal/chwriter.Registry's per-tenant-instance shape.
+package grounding
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/sentry/sentry/api/ai/provider"
+	"github.com/sentry/sentry/api/querylang/executor"
+)
+
+// staticFields are always present regardless of what's actually been
+// ingested -- Phase 0's schema (storage/migrations/0001_create_logs_table.sql)
+// plus record_id (0002). Listed here rather than queried: their existence
+// doesn't depend on sampling, only their *values* do (severity's examples
+// still come from a real query below, in case a deployment's severities
+// diverge from the standard OTel set).
+var staticFields = []string{"timestamp", "host", "service", "severity", "message", "record_id"}
+
+// Tuning constants. First-pass values, not benchmarked against a
+// production-scale cluster -- see the "not yet verified at scale" note
+// in /docs/phase-7-ai-design.md. Deliberately conservative (short lookback,
+// small caps) since grounding data trades completeness for prompt-budget
+// and refresh-query cost, not the other way around.
+const (
+	sampleWindow         = 7 * 24 * time.Hour // how far back sampling queries look
+	maxServices          = 50
+	maxAttributeKeys     = 100                     // how many keys we learn about at all
+	maxEnumCandidateKeys = 15                      // of those, how many get a real example-value query (each is a separate round trip)
+	maxExamplesPerField  = 20                      // a field returning more distinct values than this in the capped query isn't treated as enum-like
+	perFieldQueryLimit   = maxExamplesPerField + 1 // +1 so "more than maxExamplesPerField" is detectable, not just silently truncated
+)
+
+// Service produces provider.SchemaContext for one tenant (or, in a
+// single-tenant deployment, the whole instance) from its own ClickHouse
+// data. Safe for concurrent use: Refresh swaps a snapshot under a mutex,
+// Current reads it under the same lock -- same last-known-good pattern
+// enterprise/internal/chwriter.Registry and search/src/tenants.rs's
+// ActiveTenantTracker already use, so a slow or failing refresh never
+// blocks or blanks a caller mid-request.
+type Service struct {
+	runner executor.SQLRunner
+
+	mu       sync.RWMutex
+	snapshot provider.SchemaContext
+}
+
+func New(runner executor.SQLRunner) *Service {
+	return &Service{runner: runner}
+}
+
+// Current returns the last successfully refreshed SchemaContext --
+// possibly stale, possibly zero-valued if Refresh has never succeeded
+// yet, but never a partial/torn snapshot. Callers (the AI operation
+// handlers, task 5+) should treat a zero-valued Services/Fields as "no
+// grounding data yet available," not an error -- every operation still
+// works with an empty SchemaContext, just less well-grounded, matching
+// this codebase's "absence is a normal state, not a failure" convention
+// (e.g. AuditLogger, getAuthFeatures).
+func (s *Service) Current() provider.SchemaContext {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.snapshot
+}
+
+// SchemaContext implements aiapi.SchemaContextSource directly -- a
+// single-tenant deployment has exactly one Service, so there's no
+// per-request tenant resolution to do here (ctx is unused); it's the
+// same shape as Current, just satisfying the interface aiapi's handlers
+// depend on so main.go can wire *Service in without a separate adapter
+// type. enterprise-api's multi-tenant equivalent (groundingregistry)
+// implements this same interface by actually reading ctx.
+func (s *Service) SchemaContext(context.Context) provider.SchemaContext {
+	return s.Current()
+}
+
+// Refresh runs the sampling queries and swaps the cached snapshot on
+// success. A failed refresh leaves the previous snapshot in place
+// (last-known-good) rather than clearing it -- a transient ClickHouse
+// hiccup shouldn't blank out grounding for every AI request until the
+// next successful refresh.
+func (s *Service) Refresh(ctx context.Context) error {
+	services, err := s.sampleServices(ctx)
+	if err != nil {
+		return fmt.Errorf("grounding: sampling services: %w", err)
+	}
+
+	attrKeys, err := s.sampleAttributeKeys(ctx)
+	if err != nil {
+		return fmt.Errorf("grounding: sampling attribute keys: %w", err)
+	}
+
+	fields := make([]provider.FieldInfo, 0, len(staticFields)+len(attrKeys))
+	for _, name := range staticFields {
+		examples, _ := s.sampleFieldExamples(ctx, name, name == "severity")
+		fields = append(fields, provider.FieldInfo{Name: name, Examples: examples})
+	}
+
+	candidateKeys := attrKeys
+	if len(candidateKeys) > maxEnumCandidateKeys {
+		candidateKeys = candidateKeys[:maxEnumCandidateKeys]
+	}
+	enumExamples := make(map[string][]string, len(candidateKeys))
+	for _, key := range candidateKeys {
+		examples, ok := s.sampleFieldExamples(ctx, key, false)
+		if ok {
+			enumExamples[key] = examples
+		}
+	}
+	for _, key := range attrKeys {
+		fields = append(fields, provider.FieldInfo{Name: key, Examples: enumExamples[key]})
+	}
+
+	s.mu.Lock()
+	s.snapshot = provider.SchemaContext{Services: services, Fields: fields}
+	s.mu.Unlock()
+	return nil
+}
+
+// StartRefreshing runs Refresh once immediately (best-effort -- a failure
+// here just means Current() returns a zero snapshot until the first
+// successful tick, not a fatal startup error, since grounding is an
+// enhancement, not a dependency anything else blocks on) and then on
+// interval until ctx is cancelled. Same shape as chwriter.Registry.
+// StartRefreshing.
+func (s *Service) StartRefreshing(ctx context.Context, interval time.Duration, onError func(error)) {
+	if err := s.Refresh(ctx); err != nil && onError != nil {
+		onError(err)
+	}
+	go func() {
+		ticker := time.NewTicker(interval)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				if err := s.Refresh(ctx); err != nil && onError != nil {
+					onError(err)
+				}
+			}
+		}
+	}()
+}
+
+func (s *Service) sampleServices(ctx context.Context) ([]string, error) {
+	sql := fmt.Sprintf(
+		"SELECT service FROM logs WHERE timestamp > now() - INTERVAL %d SECOND GROUP BY service ORDER BY count() DESC LIMIT %d",
+		int(sampleWindow.Seconds()), maxServices,
+	)
+	res, err := s.runner.RunSQL(ctx, sql)
+	if err != nil {
+		return nil, err
+	}
+	return firstColumnStrings(res), nil
+}
+
+func (s *Service) sampleAttributeKeys(ctx context.Context) ([]string, error) {
+	sql := fmt.Sprintf(
+		"SELECT arrayJoin(mapKeys(attributes)) AS attr_key FROM logs WHERE timestamp > now() - INTERVAL %d SECOND GROUP BY attr_key ORDER BY count() DESC LIMIT %d",
+		int(sampleWindow.Seconds()), maxAttributeKeys,
+	)
+	res, err := s.runner.RunSQL(ctx, sql)
+	if err != nil {
+		return nil, err
+	}
+	return firstColumnStrings(res), nil
+}
+
+// sampleFieldExamples returns up to maxExamplesPerField distinct values
+// for a field, and false if the field turned out not to look enum-like
+// (more distinct values than the cap turned up, or the value column
+// wasn't usable) -- matching FieldInfo's doc comment that a
+// high-cardinality field should carry no examples rather than a
+// truncated, misleading sample. isStructuredColumn distinguishes
+// `severity` (a real column) from an attributes[...] lookup.
+func (s *Service) sampleFieldExamples(ctx context.Context, field string, isStructuredColumn bool) ([]string, bool) {
+	col := "attributes[" + quoteLiteral(field) + "]"
+	if isStructuredColumn {
+		col = "`" + field + "`"
+	}
+	sql := fmt.Sprintf(
+		"SELECT DISTINCT %s AS v FROM logs WHERE timestamp > now() - INTERVAL %d SECOND AND %s != '' LIMIT %d",
+		col, int(sampleWindow.Seconds()), col, perFieldQueryLimit,
+	)
+	res, err := s.runner.RunSQL(ctx, sql)
+	if err != nil {
+		return nil, false
+	}
+	values := firstColumnStrings(res)
+	if len(values) == 0 || len(values) > maxExamplesPerField {
+		return nil, false
+	}
+	return values, true
+}
+
+func firstColumnStrings(res *executor.Result) []string {
+	if res == nil || len(res.Columns) == 0 {
+		return nil
+	}
+	out := make([]string, 0, len(res.Rows))
+	for _, row := range res.Rows {
+		if len(row) == 0 {
+			continue
+		}
+		if s, ok := row[0].(string); ok && s != "" {
+			out = append(out, s)
+		}
+	}
+	return out
+}
+
+// quoteLiteral must stay byte-for-byte in sync with executor/sql.go's
+// unexported function of the same name (ClickHouse SQL string literals
+// use backslash escaping, not SQL-standard doubled quotes -- easy to get
+// wrong by assuming the more common convention, which an earlier draft
+// of this function did). Duplicated rather than exported from executor,
+// since executor's quoteLiteral is deliberately unexported (query-SQL
+// building is that package's own concern) and grounding's use is
+// narrow enough not to justify widening that package's public surface
+// for one helper.
+func quoteLiteral(s string) string {
+	var sb strings.Builder
+	sb.WriteByte('\'')
+	for _, r := range s {
+		switch r {
+		case '\\':
+			sb.WriteString(`\\`)
+		case '\'':
+			sb.WriteString(`\'`)
+		default:
+			sb.WriteRune(r)
+		}
+	}
+	sb.WriteByte('\'')
+	return sb.String()
+}
@@ -0,0 +1,169 @@
+package grounding
+
+import (
+	"context"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/sentry/sentry/api/ai/provider"
+	"github.com/sentry/sentry/api/querylang/executor"
+)
+
+// routingFakeRunner returns a canned result keyed by a substring match
+// against the SQL text -- grounding.Refresh issues several structurally
+// different queries in sequence (services, attribute keys, then one
+// per candidate enum field), unlike executor's tests where a single
+// fixed result/err per call is enough.
+type routingFakeRunner struct {
+	byContains []struct {
+		substr string
+		result *executor.Result
+		err    error
+	}
+	calls int
+}
+
+func (r *routingFakeRunner) on(substr string, result *executor.Result) *routingFakeRunner {
+	r.byContains = append(r.byContains, struct {
+		substr string
+		result *executor.Result
+		err    error
+	}{substr, result, nil})
+	return r
+}
+
+func (r *routingFakeRunner) RunSQL(_ context.Context, sql string) (*executor.Result, error) {
+	r.calls++
+	for _, rule := range r.byContains {
+		if strings.Contains(sql, rule.substr) {
+			if rule.err != nil {
+				return nil, rule.err
+			}
+			return rule.result, nil
+		}
+	}
+	// Unmatched queries (most of the per-field example queries in a
+	// small test fixture) come back empty, same as a field with no data
+	// -- not an error, matching sampleFieldExamples' "unusable -> not
+	// enum-like" treatment.
+	return &executor.Result{Columns: []string{"v"}, Rows: nil}, nil
+}
+
+func strResult(col string, vals ...string) *executor.Result {
+	rows := make([][]any, len(vals))
+	for i, v := range vals {
+		rows[i] = []any{v}
+	}
+	return &executor.Result{Columns: []string{col}, Rows: rows}
+}
+
+func TestRefreshPopulatesServicesAndFields(t *testing.T) {
+	runner := (&routingFakeRunner{}).
+		on("FROM logs WHERE timestamp > now() - INTERVAL 604800 SECOND GROUP BY service", strResult("service", "api", "web", "worker")).
+		on("mapKeys(attributes)", strResult("attr_key", "status", "latency_ms")).
+		on("`severity`", strResult("v", "INFO", "WARN", "ERROR")).
+		on("attributes['status']", strResult("v", "200", "404", "500"))
+
+	svc := New(runner)
+	if err := svc.Refresh(context.Background()); err != nil {
+		t.Fatalf("Refresh: %v", err)
+	}
+
+	got := svc.Current()
+	if len(got.Services) != 3 || got.Services[0] != "api" {
+		t.Errorf("Services = %v, want [api web worker]", got.Services)
+	}
+
+	var severity, status *provider.FieldInfo
+	for i := range got.Fields {
+		f := &got.Fields[i]
+		switch f.Name {
+		case "severity":
+			severity = f
+		case "status":
+			status = f
+		}
+	}
+	if severity == nil || len(severity.Examples) != 3 {
+		t.Errorf("severity field = %+v, want 3 examples", severity)
+	}
+	if status == nil || len(status.Examples) != 3 {
+		t.Errorf("status field = %+v, want 3 examples", status)
+	}
+
+	// Static fields with no configured example rule (host, message,
+	// timestamp, record_id) should still be present, just with no
+	// examples -- Refresh must not drop them.
+	names := make(map[string]bool, len(got.Fields))
+	for _, f := range got.Fields {
+		names[f.Name] = true
+	}
+	for _, want := range []string{"timestamp", "host", "message", "record_id"} {
+		if !names[want] {
+			t.Errorf("static field %q missing from Fields", want)
+		}
+	}
+}
+
+func TestRefreshFailureLeavesPreviousSnapshot(t *testing.T) {
+	good := (&routingFakeRunner{}).
+		on("GROUP BY service", strResult("service", "api"))
+	svc := New(good)
+	if err := svc.Refresh(context.Background()); err != nil {
+		t.Fatalf("first Refresh: %v", err)
+	}
+	first := svc.Current()
+
+	svc.runner = &erroringRunner{}
+	if err := svc.Refresh(context.Background()); err == nil {
+		t.Fatal("expected Refresh to fail with an erroring runner")
+	}
+
+	after := svc.Current()
+	if len(after.Services) != len(first.Services) || after.Services[0] != first.Services[0] {
+		t.Errorf("Current() after a failed Refresh = %+v, want unchanged snapshot %+v", after, first)
+	}
+}
+
+type erroringRunner struct{}
+
+func (erroringRunner) RunSQL(context.Context, string) (*executor.Result, error) {
+	return nil, context.DeadlineExceeded
+}
+
+func TestFieldExampleCapExcludesHighCardinalityFields(t *testing.T) {
+	many := make([]string, maxExamplesPerField+1)
+	for i := range many {
+		many[i] = string(rune('a' + i%26))
+	}
+	runner := (&routingFakeRunner{}).
+		on("GROUP BY service", strResult("service", "api")).
+		on("mapKeys(attributes)", strResult("attr_key", "trace_id")).
+		on("attributes['trace_id']", strResult("v", many...))
+
+	svc := New(runner)
+	if err := svc.Refresh(context.Background()); err != nil {
+		t.Fatalf("Refresh: %v", err)
+	}
+	for _, f := range svc.Current().Fields {
+		if f.Name == "trace_id" && len(f.Examples) != 0 {
+			t.Errorf("high-cardinality field trace_id got %d examples, want 0 (not treated as enum-like)", len(f.Examples))
+		}
+	}
+}
+
+func TestStartRefreshingRunsOnInterval(t *testing.T) {
+	runner := (&routingFakeRunner{}).on("GROUP BY service", strResult("service", "api"))
+	svc := New(runner)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	svc.StartRefreshing(ctx, 10*time.Millisecond, nil)
+
+	// The immediate synchronous refresh should have already happened by
+	// the time StartRefreshing returns.
+	if got := svc.Current().Services; len(got) != 1 {
+		t.Fatalf("Current() immediately after StartRefreshing = %v, want [api]", got)
+	}
+}