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.
This commit is contained in:
@@ -0,0 +1,514 @@
|
||||
// Package aiapi is Track A's HTTP surface: complete, explain, fix, and
|
||||
// optimize, each a thin wrapper around api/ai/router dispatching to
|
||||
// whichever provider.Provider is configured for that operation. Mirrors
|
||||
// queryapi's shape deliberately (same auth wrapper, same request-size
|
||||
// cap, same error-response shape) since this is the same kind of
|
||||
// endpoint -- a JSON-in, JSON-out operation gated by the same RoleViewer
|
||||
// requirement /query uses, nothing AI-specific about the transport.
|
||||
//
|
||||
// What this package does NOT do: execute a query. Every operation here
|
||||
// returns text (a suggestion, an explanation, a fix) for the client to
|
||||
// review -- running anything still goes through the unchanged POST
|
||||
// /query endpoint (queryapi.Handler), never through here. See
|
||||
// /docs/phase-7-ai-design.md for why that split is load-bearing, not
|
||||
// incidental.
|
||||
package aiapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sentry/sentry/api/ai/costguard"
|
||||
"github.com/sentry/sentry/api/ai/provider"
|
||||
"github.com/sentry/sentry/api/ai/router"
|
||||
"github.com/sentry/sentry/api/authz"
|
||||
"github.com/sentry/sentry/api/internal/querylang/ir"
|
||||
"github.com/sentry/sentry/api/internal/querylang/planner"
|
||||
)
|
||||
|
||||
// SchemaContextSource resolves the calling tenant's grounding data.
|
||||
// Core wires a tenant-agnostic adapter around one grounding.Service;
|
||||
// enterprise-api wires one around groundingregistry that reads the
|
||||
// tenant from request context -- same "interface in core, tenant-aware
|
||||
// implementation supplied by whoever constructs the handler" shape as
|
||||
// queryapi.AuditLogger and dashboards.PermissionStore.
|
||||
type SchemaContextSource interface {
|
||||
SchemaContext(ctx context.Context) provider.SchemaContext
|
||||
}
|
||||
|
||||
// InteractionLogger records a translate/fix/optimize suggestion's
|
||||
// accept-or-dismiss outcome into the Phase 4 audit trail (task 12) --
|
||||
// same nil-by-default, fail-open shape as queryapi.AuditLogger: a
|
||||
// single-tenant deployment with no enterprise/ configured just doesn't
|
||||
// log these, same as it doesn't log query executions today.
|
||||
// enterprise/internal/audit supplies the real implementation, writing
|
||||
// into the same append-only audit_log table query executions use
|
||||
// (a new event_type, not a new table -- see
|
||||
// metadata/migrations/0036_add_ai_interaction_event_type.sql).
|
||||
//
|
||||
// Deliberately not wired into Complete (ghost-text): that operation
|
||||
// fires on every keystroke pause, and logging each one at the same
|
||||
// weight as a deliberate Fix/Optimize/Translate review would drown the
|
||||
// signal task 12 actually wants (real accept/reject decisions) in
|
||||
// high-frequency noise. Not wired into Explain either -- it produces no
|
||||
// suggestion to accept or reject, so "accepted vs. rejected" doesn't
|
||||
// apply to it. Both are named scope boundaries, not oversights.
|
||||
type InteractionLogger interface {
|
||||
LogInteraction(ctx context.Context, entry InteractionEntry) error
|
||||
}
|
||||
|
||||
type InteractionEntry struct {
|
||||
// Operation is "translate", "fix", or "optimize" -- the three flows
|
||||
// that produce a suggestion a user explicitly accepts or dismisses.
|
||||
Operation string
|
||||
Input string
|
||||
Output string
|
||||
// Confidence is empty for fix/optimize (provider.Confidence only
|
||||
// applies to Translate/Fix results in a way the frontend surfaces
|
||||
// today -- Optimize's phrasing has no confidence concept).
|
||||
Confidence string
|
||||
// Accepted is false for a dismissed suggestion; Output/FinalQuery
|
||||
// still carry what was offered, since a rejected suggestion is
|
||||
// itself useful signal (task 12: "useful data for improving
|
||||
// grounding/prompting later").
|
||||
Accepted bool
|
||||
// Edited is only meaningful when Accepted -- did the user change
|
||||
// the suggested text before using it. False, not omitted, when
|
||||
// Accepted is false (there's nothing to have edited).
|
||||
Edited bool
|
||||
FinalQuery string
|
||||
}
|
||||
|
||||
// completeTimeout is deliberately tight -- task 5's "low enough latency
|
||||
// to feel responsive" requirement. A slow or hung provider must not
|
||||
// stall the query bar; the frontend's fallback to deterministic
|
||||
// autocomplete (Phase 2/5) kicks in on any error, including a timeout,
|
||||
// so a short timeout here fails fast toward that fallback rather than
|
||||
// making the user wait to find out AI completion isn't going to work
|
||||
// this time.
|
||||
const completeTimeout = 1500 * time.Millisecond
|
||||
|
||||
// Explain/Fix/Optimize are user-initiated (a button press, not
|
||||
// as-you-type), so a more generous budget is the right tradeoff --
|
||||
// correctness/quality over latency here, unlike Complete.
|
||||
const operationTimeout = 15 * time.Second
|
||||
|
||||
type Handler struct {
|
||||
logger *slog.Logger
|
||||
router *router.Router
|
||||
schema SchemaContextSource
|
||||
authz authz.Authorizer
|
||||
interactions InteractionLogger
|
||||
}
|
||||
|
||||
// interactions may be nil -- see InteractionLogger's doc comment.
|
||||
func NewHandler(logger *slog.Logger, r *router.Router, schema SchemaContextSource, authorizer authz.Authorizer, interactions InteractionLogger) *Handler {
|
||||
return &Handler{logger: logger, router: r, schema: schema, authz: authorizer, interactions: interactions}
|
||||
}
|
||||
|
||||
func (h *Handler) RegisterRoutes(mux *http.ServeMux) {
|
||||
mux.HandleFunc("POST /ai/complete", authz.RequireRoleOrService(h.authz, authz.RoleViewer, h.handleComplete))
|
||||
mux.HandleFunc("POST /ai/explain", authz.RequireRoleOrService(h.authz, authz.RoleViewer, h.handleExplain))
|
||||
mux.HandleFunc("POST /ai/fix", authz.RequireRoleOrService(h.authz, authz.RoleViewer, h.handleFix))
|
||||
mux.HandleFunc("POST /ai/optimize", authz.RequireRoleOrService(h.authz, authz.RoleViewer, h.handleOptimize))
|
||||
mux.HandleFunc("POST /ai/translate", authz.RequireRoleOrService(h.authz, authz.RoleViewer, h.handleTranslate))
|
||||
mux.HandleFunc("POST /ai/log-interaction", authz.RequireRoleOrService(h.authz, authz.RoleViewer, h.handleLogInteraction))
|
||||
}
|
||||
|
||||
const maxBodyBytes = 1 << 20 // 1 MiB, same cap queryapi uses -- these bodies are smaller still
|
||||
|
||||
type errorResponse struct {
|
||||
Error string `json:"error"`
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, v any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(v)
|
||||
}
|
||||
|
||||
func writeError(w http.ResponseWriter, status int, msg string) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
_ = json.NewEncoder(w).Encode(errorResponse{Error: msg})
|
||||
}
|
||||
|
||||
func decodeBody(w http.ResponseWriter, r *http.Request, v any) bool {
|
||||
r.Body = http.MaxBytesReader(w, r.Body, maxBodyBytes)
|
||||
if err := json.NewDecoder(r.Body).Decode(v); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid JSON body: "+err.Error())
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// ---- complete ----
|
||||
|
||||
type completeRequest struct {
|
||||
QueryPrefix string `json:"queryPrefix"`
|
||||
Language string `json:"language"`
|
||||
}
|
||||
|
||||
type completeResponse struct {
|
||||
Suggestion string `json:"suggestion"`
|
||||
}
|
||||
|
||||
func (h *Handler) handleComplete(w http.ResponseWriter, r *http.Request) {
|
||||
var req completeRequest
|
||||
if !decodeBody(w, r, &req) {
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(req.QueryPrefix) == "" {
|
||||
writeJSON(w, completeResponse{})
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(r.Context(), completeTimeout)
|
||||
defer cancel()
|
||||
|
||||
result, err := h.router.For(router.OpComplete).Complete(ctx, provider.CompleteRequest{
|
||||
QueryPrefix: req.QueryPrefix,
|
||||
Language: orDefault(req.Language, "spl"),
|
||||
Schema: h.schema.SchemaContext(r.Context()),
|
||||
})
|
||||
if err != nil {
|
||||
// Complete's whole point is to degrade gracefully -- a failed or
|
||||
// slow completion is not worth a scary error response the query
|
||||
// bar has to handle specially. An empty suggestion is exactly
|
||||
// what "no good completion available right now" looks like to
|
||||
// the frontend's fallback logic (task 5).
|
||||
h.logger.Warn("ai complete failed", "error", err)
|
||||
writeJSON(w, completeResponse{})
|
||||
return
|
||||
}
|
||||
writeJSON(w, completeResponse{Suggestion: result.Suggestion})
|
||||
}
|
||||
|
||||
// ---- explain ----
|
||||
|
||||
type explainRequest struct {
|
||||
Query string `json:"query"`
|
||||
Language string `json:"language"`
|
||||
OriginalIntent string `json:"originalIntent"`
|
||||
}
|
||||
|
||||
type explainResponse struct {
|
||||
Explanation string `json:"explanation"`
|
||||
}
|
||||
|
||||
func (h *Handler) handleExplain(w http.ResponseWriter, r *http.Request) {
|
||||
var req explainRequest
|
||||
if !decodeBody(w, r, &req) {
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(req.Query) == "" {
|
||||
writeError(w, http.StatusBadRequest, "query must not be empty")
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(r.Context(), operationTimeout)
|
||||
defer cancel()
|
||||
|
||||
result, err := h.router.For(router.OpExplain).Explain(ctx, provider.ExplainRequest{
|
||||
Query: req.Query,
|
||||
Language: orDefault(req.Language, "spl"),
|
||||
OriginalIntent: req.OriginalIntent,
|
||||
})
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, "explain failed: "+err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, explainResponse{Explanation: result.Explanation})
|
||||
}
|
||||
|
||||
// ---- fix ----
|
||||
|
||||
type fixRequest struct {
|
||||
Query string `json:"query"`
|
||||
Language string `json:"language"`
|
||||
ParseError string `json:"parseError"`
|
||||
ExecutionError string `json:"executionError"`
|
||||
}
|
||||
|
||||
type fixResponse struct {
|
||||
SuggestedQuery string `json:"suggestedQuery"`
|
||||
Explanation string `json:"explanation"`
|
||||
Confidence string `json:"confidence"`
|
||||
// Blocked mirrors costguard's assessment on the *suggested* query --
|
||||
// task 4's stricter AI-track treatment: a reject-level suggestion is
|
||||
// still shown (so the user understands what was tried and why it's
|
||||
// not being offered outright) but the frontend must not present a
|
||||
// plain accept-and-run action for it. CostWarnings is empty unless
|
||||
// Blocked, or the suggestion has a lesser (warn-level) concern worth
|
||||
// surfacing.
|
||||
Blocked bool `json:"blocked"`
|
||||
CostWarnings []string `json:"costWarnings,omitempty"`
|
||||
}
|
||||
|
||||
func (h *Handler) handleFix(w http.ResponseWriter, r *http.Request) {
|
||||
var req fixRequest
|
||||
if !decodeBody(w, r, &req) {
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(req.Query) == "" {
|
||||
writeError(w, http.StatusBadRequest, "query must not be empty")
|
||||
return
|
||||
}
|
||||
if req.ParseError == "" && req.ExecutionError == "" {
|
||||
writeError(w, http.StatusBadRequest, "parseError or executionError must be set")
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(r.Context(), operationTimeout)
|
||||
defer cancel()
|
||||
|
||||
lang := orDefault(req.Language, "spl")
|
||||
result, err := h.router.For(router.OpFix).Fix(ctx, provider.FixRequest{
|
||||
Query: req.Query,
|
||||
Language: lang,
|
||||
ParseError: req.ParseError,
|
||||
ExecutionError: req.ExecutionError,
|
||||
Schema: h.schema.SchemaContext(r.Context()),
|
||||
})
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, "fix failed: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
resp := fixResponse{
|
||||
SuggestedQuery: result.SuggestedQuery,
|
||||
Explanation: result.Explanation,
|
||||
Confidence: string(result.Confidence),
|
||||
}
|
||||
if resp.SuggestedQuery != "" {
|
||||
if plan, err := planner.Compile(resp.SuggestedQuery, planner.Language(lang), time.Now()); err == nil {
|
||||
if assessment := costguard.Assess(plan); assessment.Level != costguard.LevelOK {
|
||||
resp.CostWarnings = assessment.Reasons
|
||||
resp.Blocked = assessment.Level == costguard.LevelReject
|
||||
}
|
||||
}
|
||||
// A suggested query that itself fails to compile is left
|
||||
// unassessed rather than treated as an error -- an unusual
|
||||
// outcome (the model produced something that doesn't parse) the
|
||||
// frontend can still show as a suggestion text, just without a
|
||||
// cost assessment attached to it.
|
||||
}
|
||||
writeJSON(w, resp)
|
||||
}
|
||||
|
||||
// ---- optimize ----
|
||||
|
||||
type optimizeRequest struct {
|
||||
Query string `json:"query"`
|
||||
Language string `json:"language"`
|
||||
}
|
||||
|
||||
type optimizeResponse struct {
|
||||
// Findings is always populated when costguard has anything to say --
|
||||
// rule-based, instant, no model call needed to produce this part.
|
||||
Findings []string `json:"findings"`
|
||||
// Phrased is the AI-phrased version of Findings (task 8: "AI layer
|
||||
// used mainly to phrase the suggestion clearly"). Empty if the
|
||||
// provider is unavailable or fails -- graceful degradation, same as
|
||||
// Complete: the raw Findings are still useful on their own, this is
|
||||
// an enhancement layered on top, not a dependency.
|
||||
Phrased string `json:"phrased"`
|
||||
// SuggestedQuery is a mechanical rewrite, not model-generated --
|
||||
// only populated for the one case this package can safely rewrite
|
||||
// unambiguously (a missing time range; see suggestFix below). Other
|
||||
// findings (an overly large time range, an unfiltered free-text
|
||||
// search) get text-only guidance, honestly, rather than a guessed
|
||||
// rewrite.
|
||||
SuggestedQuery string `json:"suggestedQuery,omitempty"`
|
||||
}
|
||||
|
||||
func (h *Handler) handleOptimize(w http.ResponseWriter, r *http.Request) {
|
||||
var req optimizeRequest
|
||||
if !decodeBody(w, r, &req) {
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(req.Query) == "" {
|
||||
writeError(w, http.StatusBadRequest, "query must not be empty")
|
||||
return
|
||||
}
|
||||
|
||||
lang := orDefault(req.Language, "spl")
|
||||
plan, err := planner.Compile(req.Query, planner.Language(lang), time.Now())
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "query does not compile: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
assessment := costguard.Assess(plan)
|
||||
resp := optimizeResponse{Findings: assessment.Reasons}
|
||||
if assessment.Level == costguard.LevelOK {
|
||||
writeJSON(w, resp)
|
||||
return
|
||||
}
|
||||
|
||||
resp.SuggestedQuery = suggestMechanicalFix(req.Query, plan)
|
||||
|
||||
ctx, cancel := context.WithTimeout(r.Context(), operationTimeout)
|
||||
defer cancel()
|
||||
if result, err := h.router.For(router.OpExplain).Explain(ctx, provider.ExplainRequest{
|
||||
Query: req.Query,
|
||||
Language: lang,
|
||||
RuleFindings: assessment.Reasons,
|
||||
}); err != nil {
|
||||
h.logger.Warn("ai optimize phrasing failed", "error", err)
|
||||
} else {
|
||||
resp.Phrased = result.Explanation
|
||||
}
|
||||
writeJSON(w, resp)
|
||||
}
|
||||
|
||||
// ---- translate (Track B, task 9) ----
|
||||
|
||||
type translateRequest struct {
|
||||
NLQuery string `json:"nlQuery"`
|
||||
}
|
||||
|
||||
type translateResponse struct {
|
||||
Query string `json:"query"`
|
||||
Confidence string `json:"confidence"`
|
||||
LowConfidenceReason string `json:"lowConfidenceReason,omitempty"`
|
||||
// Compiles is false when Query is non-empty but doesn't actually
|
||||
// parse as pipe syntax -- a real, honest outcome (the model
|
||||
// produced something invalid), not folded into "low confidence"
|
||||
// since a model can be confident and still wrong about syntax.
|
||||
// CompileError is set only then.
|
||||
Compiles bool `json:"compiles"`
|
||||
CompileError string `json:"compileError,omitempty"`
|
||||
// Blocked/CostWarnings mirror handleFix's same-named fields exactly
|
||||
// -- task 9's explicit requirement that translation results run
|
||||
// through the shared cost guard "before returning it," same
|
||||
// treatment as an AI-suggested fix gets, not a lesser one.
|
||||
Blocked bool `json:"blocked"`
|
||||
CostWarnings []string `json:"costWarnings,omitempty"`
|
||||
}
|
||||
|
||||
func (h *Handler) handleTranslate(w http.ResponseWriter, r *http.Request) {
|
||||
var req translateRequest
|
||||
if !decodeBody(w, r, &req) {
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(req.NLQuery) == "" {
|
||||
writeError(w, http.StatusBadRequest, "nlQuery must not be empty")
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(r.Context(), operationTimeout)
|
||||
defer cancel()
|
||||
|
||||
result, err := h.router.For(router.OpTranslate).Translate(ctx, provider.TranslateRequest{
|
||||
NLQuery: req.NLQuery,
|
||||
Schema: h.schema.SchemaContext(r.Context()),
|
||||
})
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, "translation failed: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
resp := translateResponse{
|
||||
Query: result.Query,
|
||||
Confidence: string(result.Confidence),
|
||||
LowConfidenceReason: result.LowConfidenceReason,
|
||||
}
|
||||
if resp.Query != "" {
|
||||
// Always pipe syntax -- provider.TranslateResult's own doc
|
||||
// comment requires this (task 9: "prefer this over raw SQL...
|
||||
// narrower, safer surface"), so this always compiles as SPL,
|
||||
// never auto-detected/SQL.
|
||||
plan, compileErr := planner.Compile(resp.Query, planner.SPL, time.Now())
|
||||
if compileErr != nil {
|
||||
resp.Compiles = false
|
||||
resp.CompileError = compileErr.Error()
|
||||
} else {
|
||||
resp.Compiles = true
|
||||
if assessment := costguard.Assess(plan); assessment.Level != costguard.LevelOK {
|
||||
resp.CostWarnings = assessment.Reasons
|
||||
resp.Blocked = assessment.Level == costguard.LevelReject
|
||||
}
|
||||
}
|
||||
}
|
||||
writeJSON(w, resp)
|
||||
}
|
||||
|
||||
// suggestMechanicalFix handles exactly one case: no time bound at all.
|
||||
// Prepending "earliest=-1h " is always syntactically safe (another
|
||||
// AND'd base-search filter term, same as any other) and semantically
|
||||
// the single most common real fix for this specific finding -- not
|
||||
// attempted for any other finding (a too-large span, an unindexed
|
||||
// free-text pattern), which don't have one unambiguous correct rewrite.
|
||||
// Checked against the plan directly (not against costguard's Reasons
|
||||
// text), so this stays correct even if that phrasing changes later.
|
||||
func suggestMechanicalFix(originalQuery string, plan *ir.Plan) string {
|
||||
if plan.RawSQL != "" {
|
||||
return "" // no safe generic rewrite for arbitrary SQL
|
||||
}
|
||||
hasTimeBound := plan.TimeRange != nil && (!plan.TimeRange.From.IsZero() || !plan.TimeRange.To.IsZero())
|
||||
if hasTimeBound {
|
||||
return ""
|
||||
}
|
||||
return "earliest=-1h " + originalQuery
|
||||
}
|
||||
|
||||
func orDefault(s, def string) string {
|
||||
if s == "" {
|
||||
return def
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// ---- log-interaction (task 12) ----
|
||||
|
||||
type logInteractionRequest struct {
|
||||
Operation string `json:"operation"`
|
||||
Input string `json:"input"`
|
||||
Output string `json:"output"`
|
||||
Confidence string `json:"confidence"`
|
||||
Accepted bool `json:"accepted"`
|
||||
Edited bool `json:"edited"`
|
||||
FinalQuery string `json:"finalQuery"`
|
||||
}
|
||||
|
||||
var validInteractionOps = map[string]bool{"translate": true, "fix": true, "optimize": true}
|
||||
|
||||
// handleLogInteraction is called by the frontend at the moment a user
|
||||
// takes a terminal action on a suggestion (accept-and-use or dismiss) --
|
||||
// see InteractionLogger's doc comment for why this is a single
|
||||
// frontend-reported event rather than a backend-correlated
|
||||
// generation-plus-outcome pair. Fail-open, same posture
|
||||
// queryapi.Handler.logAudit uses: a write failure here is logged
|
||||
// server-side and otherwise ignored, never surfaced as an error to a
|
||||
// user who just clicked a button -- audit-trail completeness is a real
|
||||
// requirement, but it shouldn't be able to break the query bar.
|
||||
func (h *Handler) handleLogInteraction(w http.ResponseWriter, r *http.Request) {
|
||||
var req logInteractionRequest
|
||||
if !decodeBody(w, r, &req) {
|
||||
return
|
||||
}
|
||||
if !validInteractionOps[req.Operation] {
|
||||
writeError(w, http.StatusBadRequest, `operation must be one of "translate", "fix", "optimize"`)
|
||||
return
|
||||
}
|
||||
|
||||
if h.interactions != nil {
|
||||
err := h.interactions.LogInteraction(r.Context(), InteractionEntry{
|
||||
Operation: req.Operation,
|
||||
Input: req.Input,
|
||||
Output: req.Output,
|
||||
Confidence: req.Confidence,
|
||||
Accepted: req.Accepted,
|
||||
Edited: req.Edited,
|
||||
FinalQuery: req.FinalQuery,
|
||||
})
|
||||
if err != nil {
|
||||
h.logger.Error("ai interaction audit log write failed", "error", err)
|
||||
}
|
||||
}
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}
|
||||
@@ -0,0 +1,362 @@
|
||||
package aiapi
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/sentry/sentry/api/ai/provider"
|
||||
"github.com/sentry/sentry/api/ai/router"
|
||||
)
|
||||
|
||||
type fakeProvider struct {
|
||||
translateResult provider.TranslateResult
|
||||
completeResult provider.CompleteResult
|
||||
explainResult provider.ExplainResult
|
||||
fixResult provider.FixResult
|
||||
err error
|
||||
|
||||
gotExplainReq provider.ExplainRequest
|
||||
}
|
||||
|
||||
func (f *fakeProvider) Translate(context.Context, provider.TranslateRequest) (provider.TranslateResult, error) {
|
||||
return f.translateResult, f.err
|
||||
}
|
||||
func (f *fakeProvider) Complete(context.Context, provider.CompleteRequest) (provider.CompleteResult, error) {
|
||||
return f.completeResult, f.err
|
||||
}
|
||||
func (f *fakeProvider) Explain(_ context.Context, req provider.ExplainRequest) (provider.ExplainResult, error) {
|
||||
f.gotExplainReq = req
|
||||
return f.explainResult, f.err
|
||||
}
|
||||
func (f *fakeProvider) Fix(context.Context, provider.FixRequest) (provider.FixResult, error) {
|
||||
return f.fixResult, f.err
|
||||
}
|
||||
|
||||
type fakeSchemaSource struct{}
|
||||
|
||||
func (fakeSchemaSource) SchemaContext(context.Context) provider.SchemaContext {
|
||||
return provider.SchemaContext{Services: []string{"api"}}
|
||||
}
|
||||
|
||||
func newTestHandler(p *fakeProvider) *Handler {
|
||||
r := router.New(p)
|
||||
logger := slog.New(slog.NewTextHandler(bytesDiscard{}, nil))
|
||||
return NewHandler(logger, r, fakeSchemaSource{}, nil, nil)
|
||||
}
|
||||
|
||||
type bytesDiscard struct{}
|
||||
|
||||
func (bytesDiscard) Write(p []byte) (int, error) { return len(p), nil }
|
||||
|
||||
func doRequest(t *testing.T, h *Handler, method, path string, body any) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
var buf bytes.Buffer
|
||||
if body != nil {
|
||||
if err := json.NewEncoder(&buf).Encode(body); err != nil {
|
||||
t.Fatalf("encoding request body: %v", err)
|
||||
}
|
||||
}
|
||||
req := httptest.NewRequest(method, path, &buf)
|
||||
mux := http.NewServeMux()
|
||||
h.RegisterRoutes(mux)
|
||||
rec := httptest.NewRecorder()
|
||||
mux.ServeHTTP(rec, req)
|
||||
return rec
|
||||
}
|
||||
|
||||
func TestHandleCompleteReturnsSuggestion(t *testing.T) {
|
||||
p := &fakeProvider{completeResult: provider.CompleteResult{Suggestion: " | stats count"}}
|
||||
h := newTestHandler(p)
|
||||
|
||||
rec := doRequest(t, h, "POST", "/ai/complete", completeRequest{QueryPrefix: "service=api"})
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, body = %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var resp completeResponse
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("decoding response: %v", err)
|
||||
}
|
||||
if resp.Suggestion != " | stats count" {
|
||||
t.Errorf("Suggestion = %q", resp.Suggestion)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleCompleteDegradesGracefullyOnProviderError(t *testing.T) {
|
||||
p := &fakeProvider{err: errors.New("provider down")}
|
||||
h := newTestHandler(p)
|
||||
|
||||
rec := doRequest(t, h, "POST", "/ai/complete", completeRequest{QueryPrefix: "service=api"})
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200 even on provider failure (graceful degradation), body = %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var resp completeResponse
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
if resp.Suggestion != "" {
|
||||
t.Errorf("Suggestion = %q, want empty on provider failure", resp.Suggestion)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleCompleteEmptyPrefixSkipsProviderCall(t *testing.T) {
|
||||
p := &fakeProvider{err: errors.New("should not be called")}
|
||||
h := newTestHandler(p)
|
||||
|
||||
rec := doRequest(t, h, "POST", "/ai/complete", completeRequest{QueryPrefix: " "})
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, body = %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleExplainReturnsExplanation(t *testing.T) {
|
||||
p := &fakeProvider{explainResult: provider.ExplainResult{Explanation: "counts errors per host"}}
|
||||
h := newTestHandler(p)
|
||||
|
||||
rec := doRequest(t, h, "POST", "/ai/explain", explainRequest{Query: "severity=ERROR | stats count by host"})
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, body = %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var resp explainResponse
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
if resp.Explanation == "" {
|
||||
t.Error("expected a non-empty explanation")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleExplainEmptyQueryIsBadRequest(t *testing.T) {
|
||||
h := newTestHandler(&fakeProvider{})
|
||||
rec := doRequest(t, h, "POST", "/ai/explain", explainRequest{Query: ""})
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Errorf("status = %d, want 400", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleExplainProviderErrorIsBadGateway(t *testing.T) {
|
||||
p := &fakeProvider{err: errors.New("model unavailable")}
|
||||
h := newTestHandler(p)
|
||||
rec := doRequest(t, h, "POST", "/ai/explain", explainRequest{Query: "service=api"})
|
||||
if rec.Code != http.StatusBadGateway {
|
||||
t.Errorf("status = %d, want 502", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleFixMissingErrorFieldsIsBadRequest(t *testing.T) {
|
||||
h := newTestHandler(&fakeProvider{})
|
||||
rec := doRequest(t, h, "POST", "/ai/fix", fixRequest{Query: "service=api"})
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Errorf("status = %d, want 400 (neither parseError nor executionError set)", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleFixAssessesSuggestedQueryCost(t *testing.T) {
|
||||
// The provider suggests a fix that aggregates with no time bound --
|
||||
// costguard should reject it (an aggregation gets no implicit row
|
||||
// cap the way a raw-row fetch does), and the handler must mark it
|
||||
// Blocked rather than silently offering it as runnable. Needs a real
|
||||
// leading pipe stage -- bare words with no "|" parse as free-text
|
||||
// search terms, not an aggregation, per the query grammar.
|
||||
p := &fakeProvider{fixResult: provider.FixResult{
|
||||
SuggestedQuery: "service=api | stats count by host",
|
||||
Explanation: "removed the invalid field reference",
|
||||
Confidence: provider.ConfidenceHigh,
|
||||
}}
|
||||
h := newTestHandler(p)
|
||||
|
||||
rec := doRequest(t, h, "POST", "/ai/fix", fixRequest{Query: "bogus_field=1 | stats count by host", ParseError: "unknown field"})
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, body = %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var resp fixResponse
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
if !resp.Blocked {
|
||||
t.Errorf("resp = %+v, want Blocked=true for an unbounded aggregation suggestion", resp)
|
||||
}
|
||||
if len(resp.CostWarnings) == 0 {
|
||||
t.Error("expected non-empty CostWarnings")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleFixBoundedSuggestionIsNotBlocked(t *testing.T) {
|
||||
p := &fakeProvider{fixResult: provider.FixResult{
|
||||
SuggestedQuery: "earliest=-1h | stats count by host",
|
||||
Confidence: provider.ConfidenceHigh,
|
||||
}}
|
||||
h := newTestHandler(p)
|
||||
|
||||
rec := doRequest(t, h, "POST", "/ai/fix", fixRequest{Query: "stats count by host", ParseError: "no time range"})
|
||||
var resp fixResponse
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
if resp.Blocked {
|
||||
t.Errorf("resp = %+v, want Blocked=false for a properly time-bounded suggestion", resp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleOptimizeNoFindingsForBoundedQuery(t *testing.T) {
|
||||
h := newTestHandler(&fakeProvider{})
|
||||
rec := doRequest(t, h, "POST", "/ai/optimize", optimizeRequest{Query: "earliest=-1h severity=ERROR | stats count by host"})
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, body = %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var resp optimizeResponse
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
if len(resp.Findings) != 0 || resp.Phrased != "" {
|
||||
t.Errorf("resp = %+v, want no findings for a bounded query", resp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleOptimizeSuggestsMechanicalFixForMissingTimeRange(t *testing.T) {
|
||||
p := &fakeProvider{explainResult: provider.ExplainResult{Explanation: "add a time range to avoid scanning everything"}}
|
||||
h := newTestHandler(p)
|
||||
|
||||
rec := doRequest(t, h, "POST", "/ai/optimize", optimizeRequest{Query: "stats count by host"})
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, body = %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var resp optimizeResponse
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
if len(resp.Findings) == 0 {
|
||||
t.Error("expected at least one finding for an unbounded aggregation")
|
||||
}
|
||||
if resp.SuggestedQuery != "earliest=-1h stats count by host" {
|
||||
t.Errorf("SuggestedQuery = %q", resp.SuggestedQuery)
|
||||
}
|
||||
if resp.Phrased == "" {
|
||||
t.Error("expected a phrased explanation from the (fake) provider")
|
||||
}
|
||||
if len(p.gotExplainReq.RuleFindings) == 0 {
|
||||
t.Error("expected Explain to have been called with RuleFindings set")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleOptimizeDegradesGracefullyWhenPhraseFails(t *testing.T) {
|
||||
p := &fakeProvider{err: errors.New("model down")}
|
||||
h := newTestHandler(p)
|
||||
|
||||
rec := doRequest(t, h, "POST", "/ai/optimize", optimizeRequest{Query: "stats count by host"})
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200 even when phrasing fails, body = %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var resp optimizeResponse
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
if len(resp.Findings) == 0 {
|
||||
t.Error("Findings should still be populated (rule-based, no model needed) even if phrasing fails")
|
||||
}
|
||||
if resp.Phrased != "" {
|
||||
t.Errorf("Phrased = %q, want empty when the provider fails", resp.Phrased)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleOptimizeInvalidQueryIsBadRequest(t *testing.T) {
|
||||
h := newTestHandler(&fakeProvider{})
|
||||
rec := doRequest(t, h, "POST", "/ai/optimize", optimizeRequest{Query: "| stats"})
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Errorf("status = %d, want 400 for an uncompilable query", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- translate ----
|
||||
|
||||
func TestHandleTranslateEmptyNLQueryIsBadRequest(t *testing.T) {
|
||||
h := newTestHandler(&fakeProvider{})
|
||||
rec := doRequest(t, h, "POST", "/ai/translate", translateRequest{NLQuery: " "})
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Errorf("status = %d, want 400", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleTranslateProviderErrorIsBadGateway(t *testing.T) {
|
||||
p := &fakeProvider{err: errors.New("model unavailable")}
|
||||
h := newTestHandler(p)
|
||||
rec := doRequest(t, h, "POST", "/ai/translate", translateRequest{NLQuery: "errors in the last hour"})
|
||||
if rec.Code != http.StatusBadGateway {
|
||||
t.Errorf("status = %d, want 502", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleTranslateBoundedQueryCompilesCleanly(t *testing.T) {
|
||||
p := &fakeProvider{translateResult: provider.TranslateResult{
|
||||
Query: "earliest=-1h severity=ERROR | stats count by service",
|
||||
Confidence: provider.ConfidenceHigh,
|
||||
}}
|
||||
h := newTestHandler(p)
|
||||
rec := doRequest(t, h, "POST", "/ai/translate", translateRequest{NLQuery: "errors per service in the last hour"})
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, body = %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var resp translateResponse
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
if !resp.Compiles || resp.CompileError != "" {
|
||||
t.Errorf("resp = %+v, want Compiles=true, no CompileError", resp)
|
||||
}
|
||||
if resp.Blocked || len(resp.CostWarnings) != 0 {
|
||||
t.Errorf("resp = %+v, want no cost warnings for a time-bounded query", resp)
|
||||
}
|
||||
if resp.Confidence != string(provider.ConfidenceHigh) {
|
||||
t.Errorf("Confidence = %q", resp.Confidence)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleTranslateUnboundedQueryIsBlocked(t *testing.T) {
|
||||
// The model produced a syntactically valid but unbounded aggregation
|
||||
// -- task 9's explicit requirement that translation results run
|
||||
// through the same cost guard AI-suggested fixes do.
|
||||
p := &fakeProvider{translateResult: provider.TranslateResult{
|
||||
Query: "severity=ERROR | stats count by service",
|
||||
Confidence: provider.ConfidenceHigh,
|
||||
}}
|
||||
h := newTestHandler(p)
|
||||
rec := doRequest(t, h, "POST", "/ai/translate", translateRequest{NLQuery: "errors by service"})
|
||||
var resp translateResponse
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
if !resp.Compiles {
|
||||
t.Fatalf("resp = %+v, want Compiles=true", resp)
|
||||
}
|
||||
if !resp.Blocked || len(resp.CostWarnings) == 0 {
|
||||
t.Errorf("resp = %+v, want Blocked=true with cost warnings for an unbounded aggregation", resp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleTranslateNonCompilingQueryIsHonestlyReported(t *testing.T) {
|
||||
// The model returned something that doesn't actually parse -- a
|
||||
// real, distinct outcome from low confidence (a confident model can
|
||||
// still produce invalid syntax); the handler must say so plainly,
|
||||
// not silently drop it or crash.
|
||||
p := &fakeProvider{translateResult: provider.TranslateResult{
|
||||
Query: "| stats count",
|
||||
Confidence: provider.ConfidenceHigh,
|
||||
}}
|
||||
h := newTestHandler(p)
|
||||
rec := doRequest(t, h, "POST", "/ai/translate", translateRequest{NLQuery: "something odd"})
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200 (a non-compiling suggestion is a reportable outcome, not an HTTP error), body = %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var resp translateResponse
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
if resp.Compiles || resp.CompileError == "" {
|
||||
t.Errorf("resp = %+v, want Compiles=false with a CompileError", resp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleTranslateLowConfidenceCarriesReason(t *testing.T) {
|
||||
p := &fakeProvider{translateResult: provider.TranslateResult{
|
||||
Confidence: provider.ConfidenceLow,
|
||||
LowConfidenceReason: "not sure what 'weird stuff' refers to",
|
||||
}}
|
||||
h := newTestHandler(p)
|
||||
rec := doRequest(t, h, "POST", "/ai/translate", translateRequest{NLQuery: "show me weird stuff"})
|
||||
var resp translateResponse
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
if resp.Query != "" {
|
||||
t.Errorf("Query = %q, want empty for a low-confidence non-answer", resp.Query)
|
||||
}
|
||||
if resp.Confidence != string(provider.ConfidenceLow) || resp.LowConfidenceReason == "" {
|
||||
t.Errorf("resp = %+v, want low confidence with a reason", resp)
|
||||
}
|
||||
if resp.Compiles {
|
||||
t.Errorf("resp = %+v, want Compiles=false when Query is empty", resp)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
// Phase 7 task 13: end-to-end wiring tests distinct from handler_test.go
|
||||
// and ai/provider/ollama's own tests. Those two files each cover one
|
||||
// layer in isolation -- handler_test.go's fakeProvider satisfies
|
||||
// provider.Provider directly, bypassing HTTP/JSON/prompt construction
|
||||
// entirely; ollama_test.go exercises ollama.Client's wire-format parsing
|
||||
// against a stub server, but never through aiapi.Handler's actual HTTP
|
||||
// routes. Neither proves the seam between them actually works: a real
|
||||
// *ollama.Client wired through *router.Router into a real *Handler,
|
||||
// driven by real HTTP requests against the registered routes, with real
|
||||
// planner.Compile/costguard.Assess in the loop.
|
||||
//
|
||||
// No live Ollama or model is needed or used -- mockOllamaServer stands
|
||||
// in for Ollama's real /api/chat wire contract (same technique used for
|
||||
// this phase's live browser verification, see /docs/phase-7-ai-design.md,
|
||||
// just returning a fixed canned response instead of one selected by
|
||||
// inspecting the prompt) with a deterministic canned JSON body, which is
|
||||
// exactly what makes this suite fast and safe to run in CI -- see the
|
||||
// "CI testability" section of the design doc for why testing against a
|
||||
// real model is deliberately kept out of this suite instead.
|
||||
package aiapi
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/sentry/sentry/api/ai/provider/ollama"
|
||||
"github.com/sentry/sentry/api/ai/router"
|
||||
)
|
||||
|
||||
// jsonBody marshals v for use as an http.Post body -- the integration
|
||||
// tests below drive real HTTP requests against a real httptest.Server
|
||||
// (not handler_test.go's doRequest/ResponseRecorder shortcut), since the
|
||||
// point of this file is proving the routes are actually reachable over
|
||||
// real HTTP, not just that Handler's methods dispatch correctly.
|
||||
func jsonBody(t *testing.T, v any) io.Reader {
|
||||
t.Helper()
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
t.Fatalf("marshaling request body: %v", err)
|
||||
}
|
||||
return bytes.NewReader(b)
|
||||
}
|
||||
|
||||
type fakeInteractionLogger struct {
|
||||
entries []InteractionEntry
|
||||
}
|
||||
|
||||
func (f *fakeInteractionLogger) LogInteraction(_ context.Context, entry InteractionEntry) error {
|
||||
f.entries = append(f.entries, entry)
|
||||
return nil
|
||||
}
|
||||
|
||||
// mockOllamaServer returns an httptest.Server that answers any
|
||||
// POST /api/chat with the given assistant message content, matching
|
||||
// Ollama's real response envelope shape byte-for-byte (see
|
||||
// ollama.go's chatResponse).
|
||||
func mockOllamaServer(t *testing.T, content string) *httptest.Server {
|
||||
t.Helper()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/api/chat" {
|
||||
t.Errorf("unexpected request to %s, want /api/chat", r.URL.Path)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"message": map[string]string{"role": "assistant", "content": content},
|
||||
})
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
func newIntegrationHandler(t *testing.T, ollamaContent string) *Handler {
|
||||
t.Helper()
|
||||
mock := mockOllamaServer(t, ollamaContent)
|
||||
client := ollama.New(mock.URL, "test-model")
|
||||
r := router.New(client)
|
||||
logger := slog.New(slog.NewTextHandler(bytesDiscard{}, nil))
|
||||
return NewHandler(logger, r, fakeSchemaSource{}, nil, nil)
|
||||
}
|
||||
|
||||
// TestIntegrationTranslateEndToEnd proves the full path -- HTTP request
|
||||
// in, real ollama.Client HTTP call out to the mock, real JSON parsing,
|
||||
// real planner.Compile, real costguard.Assess, HTTP response out --
|
||||
// works for a query that should pass cleanly (time-bounded, no
|
||||
// aggregation).
|
||||
func TestIntegrationTranslateEndToEnd(t *testing.T) {
|
||||
h := newIntegrationHandler(t, `{"query":"earliest=-1h severity=ERROR","confidence":"high"}`)
|
||||
mux := http.NewServeMux()
|
||||
h.RegisterRoutes(mux)
|
||||
srv := httptest.NewServer(mux)
|
||||
defer srv.Close()
|
||||
|
||||
resp, err := http.Post(srv.URL+"/ai/translate", "application/json",
|
||||
jsonBody(t, map[string]string{"nlQuery": "errors in the last hour"}))
|
||||
if err != nil {
|
||||
t.Fatalf("POST /ai/translate: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200", resp.StatusCode)
|
||||
}
|
||||
|
||||
var got translateResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&got); err != nil {
|
||||
t.Fatalf("decoding response: %v", err)
|
||||
}
|
||||
if got.Query != "earliest=-1h severity=ERROR" {
|
||||
t.Errorf("query = %q", got.Query)
|
||||
}
|
||||
if !got.Compiles {
|
||||
t.Error("compiles = false, want true -- this query is valid pipe syntax")
|
||||
}
|
||||
if got.Blocked {
|
||||
t.Errorf("blocked = true, want false: %v", got.CostWarnings)
|
||||
}
|
||||
if got.Confidence != "high" {
|
||||
t.Errorf("confidence = %q, want high", got.Confidence)
|
||||
}
|
||||
}
|
||||
|
||||
// TestIntegrationTranslateBlockedByCostGuard proves costguard is
|
||||
// actually reached through the full HTTP stack, not just unit-tested
|
||||
// against costguard.Assess in isolation -- an unbounded aggregation
|
||||
// (stats, no time filter) must come back Blocked.
|
||||
func TestIntegrationTranslateBlockedByCostGuard(t *testing.T) {
|
||||
h := newIntegrationHandler(t, `{"query":"severity=ERROR | stats count by service","confidence":"high"}`)
|
||||
mux := http.NewServeMux()
|
||||
h.RegisterRoutes(mux)
|
||||
srv := httptest.NewServer(mux)
|
||||
defer srv.Close()
|
||||
|
||||
resp, err := http.Post(srv.URL+"/ai/translate", "application/json",
|
||||
jsonBody(t, map[string]string{"nlQuery": "error count by service"}))
|
||||
if err != nil {
|
||||
t.Fatalf("POST /ai/translate: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
var got translateResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&got); err != nil {
|
||||
t.Fatalf("decoding response: %v", err)
|
||||
}
|
||||
if !got.Compiles {
|
||||
t.Fatalf("compiles = false, want true: %s", got.CompileError)
|
||||
}
|
||||
if !got.Blocked {
|
||||
t.Error("blocked = false, want true -- unbounded aggregation should be rejected by costguard")
|
||||
}
|
||||
if len(got.CostWarnings) == 0 {
|
||||
t.Error("costWarnings is empty, want at least one reason")
|
||||
}
|
||||
}
|
||||
|
||||
// TestIntegrationFixEndToEnd proves the same seam for /ai/fix.
|
||||
func TestIntegrationFixEndToEnd(t *testing.T) {
|
||||
h := newIntegrationHandler(t, `{"suggested_query":"earliest=-1h severity=ERROR","explanation":"added a time bound","confidence":"high"}`)
|
||||
mux := http.NewServeMux()
|
||||
h.RegisterRoutes(mux)
|
||||
srv := httptest.NewServer(mux)
|
||||
defer srv.Close()
|
||||
|
||||
resp, err := http.Post(srv.URL+"/ai/fix", "application/json", jsonBody(t, map[string]string{
|
||||
"query": "severity=ERROR",
|
||||
"language": "spl",
|
||||
"executionError": "query timed out: no time range specified",
|
||||
}))
|
||||
if err != nil {
|
||||
t.Fatalf("POST /ai/fix: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200", resp.StatusCode)
|
||||
}
|
||||
|
||||
var got fixResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&got); err != nil {
|
||||
t.Fatalf("decoding response: %v", err)
|
||||
}
|
||||
if got.SuggestedQuery != "earliest=-1h severity=ERROR" {
|
||||
t.Errorf("suggestedQuery = %q", got.SuggestedQuery)
|
||||
}
|
||||
if got.Blocked {
|
||||
t.Errorf("blocked = true, want false: %v", got.CostWarnings)
|
||||
}
|
||||
}
|
||||
|
||||
// TestIntegrationCompleteEndToEnd proves the same seam for /ai/complete
|
||||
// (Track A's ghost-text autocomplete).
|
||||
func TestIntegrationCompleteEndToEnd(t *testing.T) {
|
||||
h := newIntegrationHandler(t, `{"suggestion":" severity=ERROR"}`)
|
||||
mux := http.NewServeMux()
|
||||
h.RegisterRoutes(mux)
|
||||
srv := httptest.NewServer(mux)
|
||||
defer srv.Close()
|
||||
|
||||
resp, err := http.Post(srv.URL+"/ai/complete", "application/json", jsonBody(t, map[string]string{
|
||||
"queryPrefix": "service=api ",
|
||||
"language": "spl",
|
||||
}))
|
||||
if err != nil {
|
||||
t.Fatalf("POST /ai/complete: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
var got completeResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&got); err != nil {
|
||||
t.Fatalf("decoding response: %v", err)
|
||||
}
|
||||
if got.Suggestion != " severity=ERROR" {
|
||||
t.Errorf("suggestion = %q", got.Suggestion)
|
||||
}
|
||||
}
|
||||
|
||||
// TestIntegrationLogInteractionEndToEnd proves /ai/log-interaction's
|
||||
// full HTTP decode+validate+dispatch path, using a fake InteractionLogger
|
||||
// (an in-process Go fake is the right seam here, not another mock HTTP
|
||||
// server -- the real implementation is enterprise/internal/audit, which
|
||||
// needs a live Postgres and is covered by that package's own tests).
|
||||
func TestIntegrationLogInteractionEndToEnd(t *testing.T) {
|
||||
logger := &fakeInteractionLogger{}
|
||||
r := router.New(&fakeProvider{})
|
||||
h := NewHandler(slog.New(slog.NewTextHandler(bytesDiscard{}, nil)), r, fakeSchemaSource{}, nil, logger)
|
||||
mux := http.NewServeMux()
|
||||
h.RegisterRoutes(mux)
|
||||
srv := httptest.NewServer(mux)
|
||||
defer srv.Close()
|
||||
|
||||
resp, err := http.Post(srv.URL+"/ai/log-interaction", "application/json", jsonBody(t, map[string]any{
|
||||
"operation": "translate",
|
||||
"input": "errors in the last hour",
|
||||
"output": "earliest=-1h severity=ERROR",
|
||||
"confidence": "high",
|
||||
"accepted": true,
|
||||
"edited": false,
|
||||
"finalQuery": "earliest=-1h severity=ERROR",
|
||||
}))
|
||||
if err != nil {
|
||||
t.Fatalf("POST /ai/log-interaction: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusNoContent {
|
||||
t.Fatalf("status = %d, want 204", resp.StatusCode)
|
||||
}
|
||||
if len(logger.entries) != 1 {
|
||||
t.Fatalf("got %d logged entries, want 1", len(logger.entries))
|
||||
}
|
||||
if logger.entries[0].Operation != "translate" || !logger.entries[0].Accepted {
|
||||
t.Errorf("logged entry = %+v", logger.entries[0])
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
// Package costguard is the shared cost/safety check task 4 asked for --
|
||||
// no such mechanism existed anywhere in Phase 2/3's compiler before this
|
||||
// (confirmed by reading planner.go/sql.go before writing this: plan.TimeRange
|
||||
// can be entirely unset, and nothing downstream rejects that). Built here
|
||||
// as a standalone, pure function operating on the same ir.Plan every
|
||||
// query -- hand-written or AI-generated -- already compiles to, so there
|
||||
// is exactly one cost check, not one per code path.
|
||||
//
|
||||
// This package does not decide what a caller *does* with a Reject-level
|
||||
// Assessment -- see /docs/phase-7-ai-design.md's "Cost/safety guard"
|
||||
// section for how the AI tracks and the existing /query handler each
|
||||
// apply this differently (AI suggestions withhold a Reject-level
|
||||
// suggestion from being offered as directly runnable; the existing
|
||||
// /query handler surfaces the same assessment as a non-blocking warning,
|
||||
// deliberately not a new hard block on hand-written queries this phase
|
||||
// didn't set out to change).
|
||||
package costguard
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sentry/sentry/api/internal/querylang/ir"
|
||||
)
|
||||
|
||||
type Level string
|
||||
|
||||
const (
|
||||
LevelOK Level = "ok"
|
||||
LevelWarn Level = "warn"
|
||||
LevelReject Level = "reject"
|
||||
)
|
||||
|
||||
type Assessment struct {
|
||||
Level Level
|
||||
Reasons []string
|
||||
}
|
||||
|
||||
// maxReasonableSpan and the two below are first-pass heuristic
|
||||
// thresholds, not benchmarked against a production-scale ClickHouse
|
||||
// cluster -- this environment's own data is far smaller than what these
|
||||
// numbers are meant to guard against. Flagged explicitly in
|
||||
// /docs/phase-7-ai-design.md rather than presented as tuned. Revisit
|
||||
// once there's real cluster-size data to check them against.
|
||||
const maxReasonableSpan = 90 * 24 * time.Hour
|
||||
|
||||
// rawSQLTimestampRe is a best-effort, deliberately loose check for
|
||||
// *some* mention of the timestamp column in a raw SQL statement's WHERE
|
||||
// clause -- not a real SQL parser. A false negative here (a query that
|
||||
// does filter by time in a way this regex doesn't recognize) just means
|
||||
// an unnecessary Warn, not a Reject, so being loose-but-safe is the
|
||||
// right failure direction. Raw SQL genuinely can't get the same
|
||||
// structural guarantee the IR-based checks below get, and this package
|
||||
// says so rather than pretending otherwise.
|
||||
var rawSQLTimestampRe = regexp.MustCompile(`(?i)\btimestamp\b\s*[<>=]`)
|
||||
|
||||
// Assess evaluates one compiled plan. Never returns an error -- a plan
|
||||
// that reached this point already parsed successfully; this is a
|
||||
// judgment call about cost, not a correctness check.
|
||||
func Assess(plan *ir.Plan) Assessment {
|
||||
if plan.RawSQL != "" {
|
||||
return assessRawSQL(plan.RawSQL)
|
||||
}
|
||||
return assessIR(plan)
|
||||
}
|
||||
|
||||
func assessRawSQL(sql string) Assessment {
|
||||
if rawSQLTimestampRe.MatchString(sql) {
|
||||
return Assessment{Level: LevelOK}
|
||||
}
|
||||
return Assessment{
|
||||
Level: LevelWarn,
|
||||
Reasons: []string{
|
||||
"no obvious timestamp filter found in this raw SQL -- this is a best-effort text check, not a real parse, so it may be wrong in either direction, but if this query has no time bound it could scan the full table's history",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func assessIR(plan *ir.Plan) Assessment {
|
||||
var reasons []string
|
||||
level := LevelOK
|
||||
|
||||
hasTimeBound := plan.TimeRange != nil && (!plan.TimeRange.From.IsZero() || !plan.TimeRange.To.IsZero())
|
||||
|
||||
if !hasTimeBound {
|
||||
switch {
|
||||
case plan.Aggregation != nil:
|
||||
// Unlike a raw-row query, an aggregation gets no implicit
|
||||
// row cap from the executor regardless of plan.Limit --
|
||||
// see executor/sql.go's buildSQL: the defaultRowLimit
|
||||
// safety net only applies `else if plan.Aggregation ==
|
||||
// nil`. An unbounded aggregation is never merely
|
||||
// "capped but slow" the way a raw-row fetch is.
|
||||
level = LevelReject
|
||||
reasons = append(reasons, "no time range filter, and this query aggregates -- every matching row across the table's entire history must be scanned to compute the aggregate, regardless of how small the output is")
|
||||
default:
|
||||
// A raw-row query with no explicit Limit still gets
|
||||
// executor/sql.go's defaultRowLimit=100 safety net applied
|
||||
// automatically -- it is not actually unbounded output,
|
||||
// just potentially an expensive scan to find those rows
|
||||
// without a time bound to narrow the search. Confirmed by
|
||||
// reading buildSQL directly, not assumed: this is the same
|
||||
// risk level whether plan.Limit is nil or explicitly set,
|
||||
// so both cases share one Warn, not a Reject for one and a
|
||||
// Warn for the other.
|
||||
level = LevelWarn
|
||||
reasons = append(reasons, "no time range filter -- results are capped (explicitly, or by the default 100-row limit), but ClickHouse may still need to scan well beyond that many rows to find them without a time bound to narrow the search")
|
||||
}
|
||||
if len(plan.TextSearch) > 0 {
|
||||
reasons = append(reasons, "the free-text search stage is bounded by the existing 5,000-record Tantivy prefilter cap regardless of time range, which partially limits how bad this is, but doesn't remove the underlying ClickHouse-side cost")
|
||||
}
|
||||
} else if !plan.TimeRange.From.IsZero() && !plan.TimeRange.To.IsZero() {
|
||||
span := plan.TimeRange.To.Sub(plan.TimeRange.From)
|
||||
if span > maxReasonableSpan {
|
||||
level = maxLevel(level, LevelWarn)
|
||||
reasons = append(reasons, "time range spans more than 90 days -- this may be slow depending on data volume")
|
||||
}
|
||||
}
|
||||
|
||||
return Assessment{Level: level, Reasons: reasons}
|
||||
}
|
||||
|
||||
func maxLevel(a, b Level) Level {
|
||||
rank := map[Level]int{LevelOK: 0, LevelWarn: 1, LevelReject: 2}
|
||||
if rank[b] > rank[a] {
|
||||
return b
|
||||
}
|
||||
return a
|
||||
}
|
||||
|
||||
// Summary renders an Assessment as one human-readable line, for
|
||||
// embedding in an AI-suggestion response or a /query warnings entry --
|
||||
// one shared rendering so the two callers don't independently invent
|
||||
// slightly different phrasing for the same underlying reasons.
|
||||
func Summary(a Assessment) string {
|
||||
if a.Level == LevelOK || len(a.Reasons) == 0 {
|
||||
return ""
|
||||
}
|
||||
return strings.Join(a.Reasons, "; ")
|
||||
}
|
||||
@@ -0,0 +1,100 @@
|
||||
package costguard
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/sentry/sentry/api/internal/querylang/ir"
|
||||
)
|
||||
|
||||
// A raw-row (non-aggregation) query with no time range and no explicit
|
||||
// Limit still gets executor/sql.go's defaultRowLimit=100 safety net
|
||||
// applied automatically -- so this is a Warn (a possibly-expensive scan
|
||||
// to find those 100 rows), not a Reject (genuinely unbounded output),
|
||||
// which only an unbounded *aggregation* actually is. See the case
|
||||
// immediately below for that contrast.
|
||||
func TestAssessNoTimeRangeNoLimitRawRowWarns(t *testing.T) {
|
||||
plan := &ir.Plan{Filters: []ir.FilterPredicate{{Field: "service", Op: "=", Value: "api"}}}
|
||||
got := Assess(plan)
|
||||
if got.Level != LevelWarn {
|
||||
t.Errorf("Level = %v, want warn (executor applies a default row limit even with no explicit Limit)", got.Level)
|
||||
}
|
||||
if len(got.Reasons) == 0 {
|
||||
t.Error("expected at least one reason")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssessNoTimeRangeWithAggregationRejects(t *testing.T) {
|
||||
plan := &ir.Plan{
|
||||
TimeRange: &ir.TimeRange{},
|
||||
Aggregation: &ir.Aggregation{Funcs: []ir.AggFunc{{Func: "count", Alias: "count"}}},
|
||||
}
|
||||
got := Assess(plan)
|
||||
if got.Level != LevelReject {
|
||||
t.Errorf("Level = %v, want reject for an unbounded aggregation", got.Level)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssessNoTimeRangeWithLimitWarns(t *testing.T) {
|
||||
plan := &ir.Plan{
|
||||
TimeRange: &ir.TimeRange{},
|
||||
Limit: &ir.Limit{N: 100},
|
||||
}
|
||||
got := Assess(plan)
|
||||
if got.Level != LevelWarn {
|
||||
t.Errorf("Level = %v, want warn (limited, no aggregation)", got.Level)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssessBoundedTimeRangeIsOK(t *testing.T) {
|
||||
now := time.Now()
|
||||
plan := &ir.Plan{
|
||||
TimeRange: &ir.TimeRange{From: now.Add(-1 * time.Hour), To: now},
|
||||
Limit: &ir.Limit{N: 100},
|
||||
}
|
||||
got := Assess(plan)
|
||||
if got.Level != LevelOK {
|
||||
t.Errorf("Level = %v, want ok, reasons: %v", got.Level, got.Reasons)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssessVeryLargeTimeRangeWarns(t *testing.T) {
|
||||
now := time.Now()
|
||||
plan := &ir.Plan{
|
||||
TimeRange: &ir.TimeRange{From: now.Add(-200 * 24 * time.Hour), To: now},
|
||||
Limit: &ir.Limit{N: 100},
|
||||
}
|
||||
got := Assess(plan)
|
||||
if got.Level != LevelWarn {
|
||||
t.Errorf("Level = %v, want warn for a 200-day range", got.Level)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssessRawSQLWithTimestampFilterIsOK(t *testing.T) {
|
||||
plan := &ir.Plan{RawSQL: "SELECT count(*) FROM logs WHERE timestamp > now() - INTERVAL 1 HOUR"}
|
||||
got := Assess(plan)
|
||||
if got.Level != LevelOK {
|
||||
t.Errorf("Level = %v, want ok, reasons: %v", got.Level, got.Reasons)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssessRawSQLWithoutTimestampFilterWarns(t *testing.T) {
|
||||
plan := &ir.Plan{RawSQL: "SELECT service, count(*) FROM logs GROUP BY service"}
|
||||
got := Assess(plan)
|
||||
if got.Level != LevelWarn {
|
||||
t.Errorf("Level = %v, want warn for raw SQL with no detectable time filter", got.Level)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSummaryEmptyForOK(t *testing.T) {
|
||||
if s := Summary(Assessment{Level: LevelOK}); s != "" {
|
||||
t.Errorf("Summary(OK) = %q, want empty", s)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSummaryJoinsReasons(t *testing.T) {
|
||||
a := Assessment{Level: LevelWarn, Reasons: []string{"a", "b"}}
|
||||
if s := Summary(a); s != "a; b" {
|
||||
t.Errorf("Summary = %q, want %q", s, "a; b")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,217 @@
|
||||
// Package ollama implements provider.Provider against a local Ollama
|
||||
// server -- the default, self-hosted primary provider (Phase 7 task 2;
|
||||
// see /docs/phase-7-ai-design.md for why Ollama over vLLM and why
|
||||
// qwen2.5-coder is the recommended model). Same thin-HTTP-client shape
|
||||
// as alerting/internal/queryclient -- net/http + encoding/json, no new
|
||||
// HTTP client dependency, matching this codebase's "boring,
|
||||
// well-understood dependencies" convention.
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/sentry/sentry/api/ai/provider"
|
||||
)
|
||||
|
||||
// Client implements provider.Provider against one Ollama server and one
|
||||
// model. The per-operation routing layer (task 2's "per-operation
|
||||
// provider/model configuration") constructs one Client per distinct
|
||||
// model a deployment configures -- e.g. one for qwen2.5-coder:1.5b
|
||||
// (Complete's fast path) and one for qwen2.5-coder:7b (everything else)
|
||||
// -- rather than this package knowing anything about operation-to-model
|
||||
// routing itself.
|
||||
type Client struct {
|
||||
baseURL string
|
||||
model string
|
||||
http *http.Client
|
||||
}
|
||||
|
||||
// New builds a Client. baseURL defaults to Ollama's standard local
|
||||
// address if empty -- the common case for the primary, self-hosted
|
||||
// deployment target.
|
||||
func New(baseURL, model string) *Client {
|
||||
if baseURL == "" {
|
||||
baseURL = "http://localhost:11434"
|
||||
}
|
||||
return &Client{baseURL: strings.TrimSuffix(baseURL, "/"), model: model, http: &http.Client{}}
|
||||
}
|
||||
|
||||
type chatMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
type chatRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []chatMessage `json:"messages"`
|
||||
Stream bool `json:"stream"`
|
||||
Format string `json:"format,omitempty"`
|
||||
}
|
||||
|
||||
type chatResponse struct {
|
||||
Message chatMessage `json:"message"`
|
||||
}
|
||||
|
||||
// chat calls Ollama's POST /api/chat, non-streaming, and returns the
|
||||
// assistant message content. jsonMode requests Ollama's JSON-constrained
|
||||
// output format -- used by every operation except Explain, which just
|
||||
// wants prose back.
|
||||
func (c *Client) chat(ctx context.Context, system, user string, jsonMode bool) (string, error) {
|
||||
req := chatRequest{
|
||||
Model: c.model,
|
||||
Messages: []chatMessage{
|
||||
{Role: "system", Content: system},
|
||||
{Role: "user", Content: user},
|
||||
},
|
||||
Stream: false,
|
||||
}
|
||||
if jsonMode {
|
||||
req.Format = "json"
|
||||
}
|
||||
|
||||
body, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("ollama: encoding request: %w", err)
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/api/chat", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("ollama: building request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := c.http.Do(httpReq)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("ollama: calling %s: %w", c.baseURL, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
var errBody struct {
|
||||
Error string `json:"error"`
|
||||
}
|
||||
_ = json.NewDecoder(resp.Body).Decode(&errBody)
|
||||
if errBody.Error != "" {
|
||||
return "", fmt.Errorf("ollama: request failed (%d): %s", resp.StatusCode, errBody.Error)
|
||||
}
|
||||
return "", fmt.Errorf("ollama: request failed with status %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var chatResp chatResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&chatResp); err != nil {
|
||||
return "", fmt.Errorf("ollama: decoding response: %w", err)
|
||||
}
|
||||
return chatResp.Message.Content, nil
|
||||
}
|
||||
|
||||
// stripCodeFence handles the common small-model habit of wrapping JSON
|
||||
// output in ```json ... ``` even when explicitly told not to -- a
|
||||
// best-effort cleanup, not a guarantee; a model that returns genuinely
|
||||
// malformed JSON still surfaces as a real decode error to the caller,
|
||||
// which is the correct behavior (better an explicit error than silently
|
||||
// fabricating a result).
|
||||
func stripCodeFence(s string) string {
|
||||
s = strings.TrimSpace(s)
|
||||
if !strings.HasPrefix(s, "```") {
|
||||
return s
|
||||
}
|
||||
s = strings.TrimPrefix(s, "```json")
|
||||
s = strings.TrimPrefix(s, "```")
|
||||
s = strings.TrimSuffix(s, "```")
|
||||
return strings.TrimSpace(s)
|
||||
}
|
||||
|
||||
func parseConfidence(s string) provider.Confidence {
|
||||
switch strings.ToLower(strings.TrimSpace(s)) {
|
||||
case "high":
|
||||
return provider.ConfidenceHigh
|
||||
case "medium":
|
||||
return provider.ConfidenceMedium
|
||||
default:
|
||||
// Unrecognized or missing confidence fails toward caution, not
|
||||
// toward assumed correctness -- an empty/garbled confidence
|
||||
// field from the model is itself a signal something's off.
|
||||
return provider.ConfidenceLow
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) Translate(ctx context.Context, req provider.TranslateRequest) (provider.TranslateResult, error) {
|
||||
raw, err := c.chat(ctx, translateSystemPrompt(req.Schema), req.NLQuery, true)
|
||||
if err != nil {
|
||||
return provider.TranslateResult{}, err
|
||||
}
|
||||
var parsed struct {
|
||||
Query string `json:"query"`
|
||||
Confidence string `json:"confidence"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(stripCodeFence(raw)), &parsed); err != nil {
|
||||
return provider.TranslateResult{}, fmt.Errorf("ollama: parsing translate response: %w", err)
|
||||
}
|
||||
return provider.TranslateResult{
|
||||
Query: parsed.Query,
|
||||
Confidence: parseConfidence(parsed.Confidence),
|
||||
LowConfidenceReason: parsed.Reason,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *Client) Complete(ctx context.Context, req provider.CompleteRequest) (provider.CompleteResult, error) {
|
||||
raw, err := c.chat(ctx, completeSystemPrompt(req.Schema), req.QueryPrefix, true)
|
||||
if err != nil {
|
||||
return provider.CompleteResult{}, err
|
||||
}
|
||||
var parsed struct {
|
||||
Suggestion string `json:"suggestion"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(stripCodeFence(raw)), &parsed); err != nil {
|
||||
return provider.CompleteResult{}, fmt.Errorf("ollama: parsing complete response: %w", err)
|
||||
}
|
||||
return provider.CompleteResult{Suggestion: parsed.Suggestion}, nil
|
||||
}
|
||||
|
||||
func (c *Client) Explain(ctx context.Context, req provider.ExplainRequest) (provider.ExplainResult, error) {
|
||||
user := req.Query
|
||||
switch {
|
||||
case len(req.RuleFindings) > 0:
|
||||
user = fmt.Sprintf("Query: %s\nFindings: %s", req.Query, strings.Join(req.RuleFindings, "; "))
|
||||
case req.OriginalIntent != "":
|
||||
user = fmt.Sprintf("Original request: %q\nGenerated query: %s", req.OriginalIntent, req.Query)
|
||||
}
|
||||
raw, err := c.chat(ctx, explainSystemPrompt(req.OriginalIntent != "", len(req.RuleFindings) > 0), user, false)
|
||||
if err != nil {
|
||||
return provider.ExplainResult{}, err
|
||||
}
|
||||
return provider.ExplainResult{Explanation: strings.TrimSpace(raw)}, nil
|
||||
}
|
||||
|
||||
func (c *Client) Fix(ctx context.Context, req provider.FixRequest) (provider.FixResult, error) {
|
||||
errText := req.ParseError
|
||||
if errText == "" {
|
||||
errText = req.ExecutionError
|
||||
}
|
||||
user := fmt.Sprintf("Query: %s\nError: %s", req.Query, errText)
|
||||
raw, err := c.chat(ctx, fixSystemPrompt(req.Schema), user, true)
|
||||
if err != nil {
|
||||
return provider.FixResult{}, err
|
||||
}
|
||||
var parsed struct {
|
||||
SuggestedQuery string `json:"suggested_query"`
|
||||
Explanation string `json:"explanation"`
|
||||
Confidence string `json:"confidence"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(stripCodeFence(raw)), &parsed); err != nil {
|
||||
return provider.FixResult{}, fmt.Errorf("ollama: parsing fix response: %w", err)
|
||||
}
|
||||
return provider.FixResult{
|
||||
SuggestedQuery: parsed.SuggestedQuery,
|
||||
Explanation: parsed.Explanation,
|
||||
Confidence: parseConfidence(parsed.Confidence),
|
||||
}, nil
|
||||
}
|
||||
|
||||
var _ provider.Provider = (*Client)(nil)
|
||||
@@ -0,0 +1,152 @@
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/sentry/sentry/api/ai/provider"
|
||||
)
|
||||
|
||||
// fakeOllamaServer stands in for a real Ollama server, returning the
|
||||
// given assistant-message content verbatim -- same reasoning
|
||||
// queryclient's tests use httptest against a fake api instead of a real
|
||||
// one: this package's own logic (request shape, response parsing,
|
||||
// JSON-mode contract) is what's under test, not Ollama itself.
|
||||
func fakeOllamaServer(t *testing.T, content string) *httptest.Server {
|
||||
t.Helper()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/api/chat" {
|
||||
t.Errorf("unexpected path %s", r.URL.Path)
|
||||
}
|
||||
var req chatRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
t.Fatalf("decoding request: %v", err)
|
||||
}
|
||||
if len(req.Messages) != 2 || req.Messages[0].Role != "system" || req.Messages[1].Role != "user" {
|
||||
t.Errorf("unexpected messages shape: %+v", req.Messages)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(chatResponse{Message: chatMessage{Role: "assistant", Content: content}})
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
func TestTranslateParsesJSONResponse(t *testing.T) {
|
||||
srv := fakeOllamaServer(t, `{"query": "earliest=-1h severity=ERROR | stats count by service", "confidence": "high", "reason": ""}`)
|
||||
c := New(srv.URL, "test-model")
|
||||
|
||||
got, err := c.Translate(context.Background(), provider.TranslateRequest{NLQuery: "errors in the last hour by service"})
|
||||
if err != nil {
|
||||
t.Fatalf("Translate: %v", err)
|
||||
}
|
||||
if got.Confidence != provider.ConfidenceHigh {
|
||||
t.Errorf("Confidence = %v, want high", got.Confidence)
|
||||
}
|
||||
if got.Query == "" {
|
||||
t.Error("expected a non-empty query")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTranslateHandlesCodeFencedJSON(t *testing.T) {
|
||||
srv := fakeOllamaServer(t, "```json\n{\"query\": \"service=api\", \"confidence\": \"medium\", \"reason\": \"\"}\n```")
|
||||
c := New(srv.URL, "test-model")
|
||||
|
||||
got, err := c.Translate(context.Background(), provider.TranslateRequest{NLQuery: "api logs"})
|
||||
if err != nil {
|
||||
t.Fatalf("Translate with fenced JSON: %v", err)
|
||||
}
|
||||
if got.Query != "service=api" {
|
||||
t.Errorf("Query = %q, want %q", got.Query, "service=api")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTranslateLowConfidenceCarriesReason(t *testing.T) {
|
||||
srv := fakeOllamaServer(t, `{"query": "", "confidence": "low", "reason": "not sure what 'weird stuff' refers to"}`)
|
||||
c := New(srv.URL, "test-model")
|
||||
|
||||
got, err := c.Translate(context.Background(), provider.TranslateRequest{NLQuery: "show me weird stuff"})
|
||||
if err != nil {
|
||||
t.Fatalf("Translate: %v", err)
|
||||
}
|
||||
if got.Confidence != provider.ConfidenceLow || got.LowConfidenceReason == "" {
|
||||
t.Errorf("got = %+v, want low confidence with a reason", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTranslateMalformedJSONIsAnError(t *testing.T) {
|
||||
srv := fakeOllamaServer(t, "not json at all, sorry")
|
||||
c := New(srv.URL, "test-model")
|
||||
|
||||
_, err := c.Translate(context.Background(), provider.TranslateRequest{NLQuery: "anything"})
|
||||
if err == nil {
|
||||
t.Fatal("expected an error for unparseable model output, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompleteReturnsSuggestionOnly(t *testing.T) {
|
||||
srv := fakeOllamaServer(t, `{"suggestion": " | stats count by host"}`)
|
||||
c := New(srv.URL, "test-model")
|
||||
|
||||
got, err := c.Complete(context.Background(), provider.CompleteRequest{QueryPrefix: "service=api", Language: "spl"})
|
||||
if err != nil {
|
||||
t.Fatalf("Complete: %v", err)
|
||||
}
|
||||
if got.Suggestion != " | stats count by host" {
|
||||
t.Errorf("Suggestion = %q", got.Suggestion)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExplainReturnsPlainText(t *testing.T) {
|
||||
srv := fakeOllamaServer(t, "This counts events per host over the last hour.")
|
||||
c := New(srv.URL, "test-model")
|
||||
|
||||
got, err := c.Explain(context.Background(), provider.ExplainRequest{Query: "earliest=-1h | stats count by host"})
|
||||
if err != nil {
|
||||
t.Fatalf("Explain: %v", err)
|
||||
}
|
||||
if got.Explanation == "" {
|
||||
t.Error("expected a non-empty explanation")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFixParsesJSONResponse(t *testing.T) {
|
||||
srv := fakeOllamaServer(t, `{"suggested_query": "earliest=-1h | stats count", "explanation": "added a required time range", "confidence": "high"}`)
|
||||
c := New(srv.URL, "test-model")
|
||||
|
||||
got, err := c.Fix(context.Background(), provider.FixRequest{Query: "stats count", ParseError: "no time range"})
|
||||
if err != nil {
|
||||
t.Fatalf("Fix: %v", err)
|
||||
}
|
||||
if got.SuggestedQuery == "" || got.Explanation == "" {
|
||||
t.Errorf("got = %+v, want both fields populated", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNonOKStatusIsAnError(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_ = json.NewEncoder(w).Encode(map[string]string{"error": "model not found"})
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
c := New(srv.URL, "missing-model")
|
||||
|
||||
_, err := c.Explain(context.Background(), provider.ExplainRequest{Query: "x"})
|
||||
if err == nil {
|
||||
t.Fatal("expected an error for a non-200 response")
|
||||
}
|
||||
if got := err.Error(); !strings.Contains(got, "model not found") {
|
||||
t.Errorf("error = %q, want it to include the server's error message", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultBaseURL(t *testing.T) {
|
||||
c := New("", "m")
|
||||
if c.baseURL != "http://localhost:11434" {
|
||||
t.Errorf("default baseURL = %q", c.baseURL)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/sentry/sentry/api/ai/provider"
|
||||
)
|
||||
|
||||
// grammarReference is a condensed version of
|
||||
// /docs/query-language-reference.md -- every operation's system prompt
|
||||
// includes this so the model is grounded in Sentry's actual pipe syntax,
|
||||
// not whatever generic log-query DSL it may have seen in training.
|
||||
// Trimmed to the parts that matter for generation/explanation (the full
|
||||
// doc's prose and examples aren't needed here); kept in sync with that
|
||||
// doc by hand -- if the grammar changes, this needs updating too, same
|
||||
// as any other place the language is described outside its own parser.
|
||||
const grammarReference = `Sentry query language (pipe syntax):
|
||||
|
||||
<base search> | <stage> | <stage> | ...
|
||||
|
||||
Base search: filter terms and/or free-text search, combined with implicit "and".
|
||||
field=value, field!=value, field>value, field>=value, field<value, field<=value
|
||||
bare word or "quoted phrase" -- free-text search on the message field
|
||||
message:"phrase" -- explicit free-text search
|
||||
earliest=-1h, latest=-5m -- relative time (s/m/h/d/w units), or earliest="2026-08-14T00:00:00Z" (RFC 3339 absolute)
|
||||
"or" only works between free-text terms, never between structured filters
|
||||
|
||||
Pipe stages, in the order they may appear:
|
||||
| where <filter terms> additional filtering, same syntax as base search filters
|
||||
| stats <func>(<field>) as <alias>, ... by <field>, ...
|
||||
functions: count (no field needed), sum, avg, min, max (all need a field)
|
||||
| sort -field, +field, ... "-" descending (default if no sign), "+" ascending
|
||||
| fields field, field, ... choose output columns
|
||||
| head N first N results (default 100)
|
||||
| tail N last N results, chronologically
|
||||
|
||||
Structured columns: timestamp, host, service, severity, message, record_id.
|
||||
Anything else is looked up in per-record attributes (always text; compared
|
||||
numerically when the right-hand side looks like a number).
|
||||
|
||||
Raw SQL (SELECT ...) is also accepted but pipe syntax is strongly preferred
|
||||
for anything generated rather than hand-written -- narrower, safer surface.`
|
||||
|
||||
func renderSchema(s provider.SchemaContext) string {
|
||||
if len(s.Services) == 0 && len(s.Fields) == 0 {
|
||||
return "(no schema grounding data available yet)"
|
||||
}
|
||||
var sb strings.Builder
|
||||
if len(s.Services) > 0 {
|
||||
fmt.Fprintf(&sb, "Known services: %s\n", strings.Join(s.Services, ", "))
|
||||
}
|
||||
if len(s.Fields) > 0 {
|
||||
sb.WriteString("Known fields:\n")
|
||||
for _, f := range s.Fields {
|
||||
if len(f.Examples) > 0 {
|
||||
fmt.Fprintf(&sb, " - %s (examples: %s)\n", f.Name, strings.Join(f.Examples, ", "))
|
||||
} else {
|
||||
fmt.Fprintf(&sb, " - %s\n", f.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
func translateSystemPrompt(schema provider.SchemaContext) string {
|
||||
return fmt.Sprintf(`You translate a plain-English question into a Sentry pipe-syntax query. You never explain, never execute anything, never write raw SQL unless the pipe syntax genuinely cannot express the request.
|
||||
|
||||
%s
|
||||
|
||||
%s
|
||||
|
||||
Respond with ONLY a JSON object, no other text, no markdown fences:
|
||||
{"query": "<the pipe-syntax query>", "confidence": "high"|"medium"|"low", "reason": "<empty unless confidence is low, in which case explain what's ambiguous or unsupported>"}
|
||||
|
||||
If you cannot produce a query you're reasonably confident in, set confidence to "low", leave query empty, and explain why in reason. Never guess with false confidence.`, grammarReference, renderSchema(schema))
|
||||
}
|
||||
|
||||
func completeSystemPrompt(schema provider.SchemaContext) string {
|
||||
return fmt.Sprintf(`You suggest how to continue a partially-typed Sentry query. You are given everything typed so far; respond with ONLY the suggested continuation text (what should appear after the cursor), not the text already typed, not an explanation.
|
||||
|
||||
%s
|
||||
|
||||
%s
|
||||
|
||||
Respond with ONLY a JSON object, no other text, no markdown fences:
|
||||
{"suggestion": "<continuation text, or empty string if you have no good suggestion>"}`, grammarReference, renderSchema(schema))
|
||||
}
|
||||
|
||||
// explainSystemPrompt covers all three contexts provider.ExplainRequest
|
||||
// supports: a plain hand-written-query explanation (both empty), a
|
||||
// post-translation review (hasIntent), or Optimize's "phrase these
|
||||
// findings" mode (hasFindings) -- mutually exclusive in practice, see
|
||||
// ExplainRequest.RuleFindings' doc comment.
|
||||
func explainSystemPrompt(hasIntent, hasFindings bool) string {
|
||||
if hasFindings {
|
||||
return fmt.Sprintf(`A rule-based check already found one or more real issues with a Sentry query's efficiency (e.g. a missing time range). Your only job is to phrase those findings as a short, clear, actionable suggestion for the person who wrote the query -- do not invent additional issues, do not restate the query's own syntax back at them, do not hedge with "might" or "could" about something the check already confirmed. One or two sentences.
|
||||
|
||||
%s`, grammarReference)
|
||||
}
|
||||
|
||||
base := fmt.Sprintf(`You explain what a Sentry query does in plain English, for someone who may not know the query language. Be concise -- two or three sentences, not a line-by-line breakdown unless the query is unusually complex.
|
||||
|
||||
%s`, grammarReference)
|
||||
if hasIntent {
|
||||
base += "\n\nYou are explaining a query that was just generated from a natural-language request. Focus on how the request became this query -- call out any interpretation choices (e.g. how a vague time phrase or field reference was resolved), not just what the query does in isolation."
|
||||
}
|
||||
return base
|
||||
}
|
||||
|
||||
func fixSystemPrompt(schema provider.SchemaContext) string {
|
||||
return fmt.Sprintf(`You fix a broken Sentry query given its error message. Produce a corrected query and a short explanation of what was wrong.
|
||||
|
||||
%s
|
||||
|
||||
%s
|
||||
|
||||
Respond with ONLY a JSON object, no other text, no markdown fences:
|
||||
{"suggested_query": "<corrected query>", "explanation": "<short explanation of what was wrong and what changed>", "confidence": "high"|"medium"|"low"}
|
||||
|
||||
If you cannot determine a fix, set confidence to "low" and suggested_query to an empty string.`, grammarReference, renderSchema(schema))
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
// Package provider defines the model-provider abstraction every AI-assisted
|
||||
// query feature (Phase 7) is built on: translate, complete, explain, fix.
|
||||
// Same narrow-interface pattern as querylang/executor's SQLRunner/
|
||||
// SearchClient -- a small interface a production implementation
|
||||
// (provider/ollama, the default; a cloud adapter, opt-in) and a fake
|
||||
// (for tests) both satisfy, so nothing above this layer needs to know or
|
||||
// care which model actually answered.
|
||||
//
|
||||
// What this package deliberately does NOT do: decide *which* provider or
|
||||
// model answers a given request. That's a routing concern (per-operation
|
||||
// config, per-tenant cloud opt-in) that lives one layer up, once task 3/4
|
||||
// land -- this package only defines the shape every provider must speak.
|
||||
//
|
||||
// Every operation is grounded (SchemaContext) and every result that
|
||||
// produces a query is designed to flow through the unchanged Phase 2
|
||||
// planner.Compile -> executor.Execute path before it ever runs -- this
|
||||
// package returns query *text*, never executes anything itself. See
|
||||
// /docs/phase-7-ai-design.md for the full design this interface was
|
||||
// built against.
|
||||
package provider
|
||||
|
||||
import "context"
|
||||
|
||||
// SchemaContext is the grounding data every operation receives -- known
|
||||
// service names, field names (structured columns plus common attribute
|
||||
// keys), and value examples for enum-like fields (severity, status, and
|
||||
// so on). Sourced from ClickHouse system tables / periodic sampling
|
||||
// (task 3), never hand-maintained, and always scoped to the requesting
|
||||
// tenant's own data -- a provider implementation must never be handed
|
||||
// another tenant's grounding data, the same connection-layer-isolation
|
||||
// discipline Phase 4 applies to query execution itself. This package
|
||||
// doesn't resolve SchemaContext; callers (the schema/metadata service,
|
||||
// task 3) build it and pass it in, so a Provider implementation never
|
||||
// needs ClickHouse access of its own.
|
||||
type SchemaContext struct {
|
||||
Services []string
|
||||
// Fields covers both real columns (timestamp, host, service,
|
||||
// severity, message, record_id) and the common attribute keys seen
|
||||
// in the tenant's own data -- see /docs/query-language-reference.md's
|
||||
// "Field mapping" section for why the distinction mostly doesn't
|
||||
// matter to a query author, and shouldn't need to matter to the model
|
||||
// either.
|
||||
Fields []FieldInfo
|
||||
}
|
||||
|
||||
type FieldInfo struct {
|
||||
Name string
|
||||
// Examples is a short, representative sample of real values seen for
|
||||
// this field -- most useful for enum-like fields (severity, status)
|
||||
// where showing the model the actual vocabulary beats describing it.
|
||||
// Empty for high-cardinality fields (host, message) where examples
|
||||
// wouldn't help and would just spend context budget.
|
||||
Examples []string
|
||||
}
|
||||
|
||||
// Confidence is deliberately a small enum, not a raw float -- a model's
|
||||
// self-reported numeric confidence isn't a calibrated probability, and
|
||||
// pretending it is (via e.g. "reject anything under 0.73") invites false
|
||||
// precision. Three bands are enough to drive real UI behavior (task 10's
|
||||
// "handle low-confidence translation honestly") without pretending to
|
||||
// more precision than a model's self-assessment actually has.
|
||||
type Confidence string
|
||||
|
||||
const (
|
||||
ConfidenceHigh Confidence = "high"
|
||||
ConfidenceMedium Confidence = "medium"
|
||||
ConfidenceLow Confidence = "low"
|
||||
)
|
||||
|
||||
type TranslateRequest struct {
|
||||
NLQuery string
|
||||
Schema SchemaContext
|
||||
}
|
||||
|
||||
type TranslateResult struct {
|
||||
// Query is Phase 2 pipe-syntax, never raw SQL -- the phase brief's
|
||||
// explicit "narrower, safer surface" choice for generation targets.
|
||||
// A provider that can't produce a valid completion should return an
|
||||
// error, not a best-effort raw-SQL fallback.
|
||||
Query string
|
||||
Confidence Confidence
|
||||
// LowConfidenceReason is set (and Query may be empty) when the
|
||||
// provider can't produce a translation it's willing to stand behind
|
||||
// at all -- task 10 wants this said plainly, not papered over with a
|
||||
// guess. Empty when Confidence is High or Medium.
|
||||
LowConfidenceReason string
|
||||
}
|
||||
|
||||
type CompleteRequest struct {
|
||||
// QueryPrefix is everything the user has typed so far, cursor at the
|
||||
// end -- this operation is a full-completion suggestion (ghost text),
|
||||
// not a fill-in-the-middle edit, matching how the query bar's cursor
|
||||
// behaves (Phase 5's QueryEditor.svelte, always append-at-cursor).
|
||||
QueryPrefix string
|
||||
Language string // "spl" or "sql", never "" -- the caller has always already resolved auto-detection by this point
|
||||
Schema SchemaContext
|
||||
}
|
||||
|
||||
type CompleteResult struct {
|
||||
// Suggestion is the suggested continuation only (what ghost-text
|
||||
// should render after the cursor), not QueryPrefix+continuation
|
||||
// restated -- keeps the caller from having to diff its own input
|
||||
// back out of the result.
|
||||
Suggestion string
|
||||
// Empty Suggestion (with no error) is a legitimate response -- "no
|
||||
// good completion here" is not the same failure mode as a timeout or
|
||||
// a down provider, and the caller (task 5's fallback logic) needs to
|
||||
// tell them apart.
|
||||
}
|
||||
|
||||
type ExplainRequest struct {
|
||||
Query string
|
||||
Language string
|
||||
// OriginalIntent, when non-empty, means this Explain call is
|
||||
// reviewing a just-translated query (Track B) rather than an
|
||||
// arbitrary hand-written one (Track A) -- same operation, task 10's
|
||||
// explicit "reuse explain rather than build a separate mechanism"
|
||||
// choice, but the prompt can speak to *how the NL became this query*
|
||||
// instead of only describing the query in isolation.
|
||||
OriginalIntent string
|
||||
// RuleFindings, when non-empty, means this Explain call is task 8's
|
||||
// Optimize suggestion: rule-based detection (costguard) already found
|
||||
// something worth flagging, and the model's only job is phrasing
|
||||
// those specific findings clearly for a user -- not describing what
|
||||
// the query does, not detecting the inefficiency itself. Mutually
|
||||
// exclusive with OriginalIntent in practice (a query is either being
|
||||
// explained, reviewed post-translation, or optimized), but the type
|
||||
// doesn't need to enforce that -- three prompt-shaping contexts for
|
||||
// one operation, matching the same reuse-over-duplication choice
|
||||
// OriginalIntent already made rather than adding a fifth Provider
|
||||
// method for what is still, underneath, "explain something about
|
||||
// this query in plain English."
|
||||
RuleFindings []string
|
||||
}
|
||||
|
||||
type ExplainResult struct {
|
||||
Explanation string
|
||||
}
|
||||
|
||||
type FixRequest struct {
|
||||
Query string
|
||||
Language string
|
||||
// ParseError is set when the query never compiled at all (planner
|
||||
// error text); ExecutionError is set when it compiled but failed at
|
||||
// runtime (executor/ClickHouse error text). Exactly one is set --
|
||||
// the two failure modes want different framing ("this doesn't parse
|
||||
// because..." vs. "this ran but...").
|
||||
ParseError string
|
||||
ExecutionError string
|
||||
Schema SchemaContext
|
||||
}
|
||||
|
||||
type FixResult struct {
|
||||
// SuggestedQuery is the full corrected query text, always shown as a
|
||||
// diff against the original by the caller (task 7's explicit
|
||||
// "never silently applied" requirement) -- this package only
|
||||
// produces the suggestion, the UI owns the diff rendering and the
|
||||
// accept/dismiss decision.
|
||||
SuggestedQuery string
|
||||
// Explanation is a short plain-English note on what was wrong and
|
||||
// what changed -- distinct from Explain's job (describing what a
|
||||
// query *does*), this describes what was *fixed* and why.
|
||||
Explanation string
|
||||
Confidence Confidence
|
||||
}
|
||||
|
||||
// Provider is what every model backend implements: the default
|
||||
// self-hosted Ollama provider, the opt-in cloud adapter, and a fake for
|
||||
// tests. Every method takes a context so a caller can enforce the tight
|
||||
// latency budget Complete needs (task 5) without the interface itself
|
||||
// hard-coding a timeout -- that's a caller concern, since the right
|
||||
// timeout differs by operation (Complete's is much tighter than
|
||||
// Translate's).
|
||||
type Provider interface {
|
||||
Translate(ctx context.Context, req TranslateRequest) (TranslateResult, error)
|
||||
Complete(ctx context.Context, req CompleteRequest) (CompleteResult, error)
|
||||
Explain(ctx context.Context, req ExplainRequest) (ExplainResult, error)
|
||||
Fix(ctx context.Context, req FixRequest) (FixResult, error)
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
// Package router is the per-operation provider/model dispatch layer
|
||||
// /docs/phase-7-ai-design.md's "per-operation provider/model
|
||||
// configuration" section decided to build now rather than defer --
|
||||
// Complete's tight latency budget and Translate/Fix's quality needs are
|
||||
// already in tension under a single-model-for-everything design, not a
|
||||
// hypothetical future conflict.
|
||||
//
|
||||
// Deliberately thin: a lookup table from Operation to whichever
|
||||
// provider.Provider was configured for it, falling back to one default
|
||||
// when an operation has no specific override. This package makes no
|
||||
// decisions about *which* provider is good for an operation -- that's
|
||||
// deployment configuration, resolved once at startup by whoever
|
||||
// constructs a Router (main.go, once the AI HTTP handlers exist to
|
||||
// consume it).
|
||||
package router
|
||||
|
||||
import "github.com/sentry/sentry/api/ai/provider"
|
||||
|
||||
type Operation string
|
||||
|
||||
const (
|
||||
OpTranslate Operation = "translate"
|
||||
OpComplete Operation = "complete"
|
||||
OpExplain Operation = "explain"
|
||||
OpFix Operation = "fix"
|
||||
)
|
||||
|
||||
// Router selects a provider.Provider per Operation. Not itself a
|
||||
// provider.Provider -- callers ask For(op) and then call the operation
|
||||
// they actually need on the result, rather than this type trying to
|
||||
// implement all four methods and dispatch internally, which would just
|
||||
// be an extra layer of indirection for no benefit.
|
||||
type Router struct {
|
||||
byOp map[Operation]provider.Provider
|
||||
fallback provider.Provider
|
||||
}
|
||||
|
||||
// New builds a Router. fallback must not be nil -- every operation
|
||||
// resolves to *some* provider, even a deployment that never calls
|
||||
// SetOperation and just wants one model for everything.
|
||||
func New(fallback provider.Provider) *Router {
|
||||
return &Router{byOp: make(map[Operation]provider.Provider), fallback: fallback}
|
||||
}
|
||||
|
||||
// SetOperation overrides which provider handles op. Call once per
|
||||
// operation that needs a non-default model at startup configuration
|
||||
// time -- not intended to change at runtime (a Router isn't
|
||||
// synchronized for concurrent SetOperation/For calls, matching every
|
||||
// other "assembled once in main.go, read-only after that" config shape
|
||||
// in this codebase, e.g. Handler's fields in queryapi).
|
||||
func (r *Router) SetOperation(op Operation, p provider.Provider) {
|
||||
r.byOp[op] = p
|
||||
}
|
||||
|
||||
// For returns the provider configured for op, or the fallback if none
|
||||
// was set specifically.
|
||||
func (r *Router) For(op Operation) provider.Provider {
|
||||
if p, ok := r.byOp[op]; ok {
|
||||
return p
|
||||
}
|
||||
return r.fallback
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/sentry/sentry/api/ai/provider"
|
||||
)
|
||||
|
||||
// namedFakeProvider lets a test tell which configured provider actually
|
||||
// answered a call -- the thing router.For needs to get right.
|
||||
type namedFakeProvider struct {
|
||||
name string
|
||||
}
|
||||
|
||||
func (f *namedFakeProvider) Translate(context.Context, provider.TranslateRequest) (provider.TranslateResult, error) {
|
||||
return provider.TranslateResult{Query: f.name}, nil
|
||||
}
|
||||
func (f *namedFakeProvider) Complete(context.Context, provider.CompleteRequest) (provider.CompleteResult, error) {
|
||||
return provider.CompleteResult{Suggestion: f.name}, nil
|
||||
}
|
||||
func (f *namedFakeProvider) Explain(context.Context, provider.ExplainRequest) (provider.ExplainResult, error) {
|
||||
return provider.ExplainResult{Explanation: f.name}, nil
|
||||
}
|
||||
func (f *namedFakeProvider) Fix(context.Context, provider.FixRequest) (provider.FixResult, error) {
|
||||
return provider.FixResult{SuggestedQuery: f.name}, nil
|
||||
}
|
||||
|
||||
var _ provider.Provider = (*namedFakeProvider)(nil)
|
||||
|
||||
func TestForReturnsFallbackWhenUnconfigured(t *testing.T) {
|
||||
fallback := &namedFakeProvider{name: "default"}
|
||||
r := New(fallback)
|
||||
|
||||
if got := r.For(OpTranslate); got != fallback {
|
||||
t.Errorf("For(OpTranslate) = %v, want the fallback provider", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetOperationOverridesFallback(t *testing.T) {
|
||||
fallback := &namedFakeProvider{name: "default"}
|
||||
fast := &namedFakeProvider{name: "fast"}
|
||||
r := New(fallback)
|
||||
r.SetOperation(OpComplete, fast)
|
||||
|
||||
if got := r.For(OpComplete); got != fast {
|
||||
t.Errorf("For(OpComplete) = %v, want the fast override", got)
|
||||
}
|
||||
if got := r.For(OpTranslate); got != fallback {
|
||||
t.Errorf("For(OpTranslate) = %v, want unaffected fallback", got)
|
||||
}
|
||||
}
|
||||
@@ -19,6 +19,10 @@ import (
|
||||
"github.com/ClickHouse/clickhouse-go/v2"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
|
||||
"github.com/sentry/sentry/api/ai/aiapi"
|
||||
"github.com/sentry/sentry/api/ai/grounding"
|
||||
"github.com/sentry/sentry/api/ai/provider/ollama"
|
||||
"github.com/sentry/sentry/api/ai/router"
|
||||
"github.com/sentry/sentry/api/authz"
|
||||
"github.com/sentry/sentry/api/dashboards"
|
||||
"github.com/sentry/sentry/api/httpserver"
|
||||
@@ -28,6 +32,13 @@ import (
|
||||
"github.com/sentry/sentry/api/searchclient"
|
||||
)
|
||||
|
||||
// groundingRefreshInterval matches chwriter.Registry/search's
|
||||
// ActiveTenantTracker's own one-minute refresh cadence -- no strong
|
||||
// reason for a different number, and consistency means one interval to
|
||||
// reason about across every "sample something periodically" mechanism
|
||||
// in this codebase, not several slightly different ones.
|
||||
const groundingRefreshInterval = time.Minute
|
||||
|
||||
func main() {
|
||||
logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
|
||||
|
||||
@@ -116,6 +127,32 @@ func main() {
|
||||
queryHandler.RegisterRoutes(mux)
|
||||
dashboardsHandler.RegisterRoutes(mux)
|
||||
|
||||
// AI routes (Phase 7) are only registered at all when OLLAMA_BASE_URL
|
||||
// is set -- an unconfigured deployment gets a plain 404 on /ai/*
|
||||
// rather than every request failing against an unreachable
|
||||
// localhost:11434, matching "no cloud dependency required for the
|
||||
// default deployment" by not forcing a *local* model dependency on a
|
||||
// deployment that doesn't want AI features either.
|
||||
if cfg.AI.OllamaBaseURL != "" {
|
||||
groundingSvc := grounding.New(sqlRunner)
|
||||
groundingSvc.StartRefreshing(ctx, groundingRefreshInterval, func(err error) {
|
||||
logger.Warn("grounding refresh failed", "error", err)
|
||||
})
|
||||
|
||||
defaultProvider := ollama.New(cfg.AI.OllamaBaseURL, cfg.AI.OllamaModel)
|
||||
aiRouter := router.New(defaultProvider)
|
||||
if cfg.AI.OllamaFastModel != "" && cfg.AI.OllamaFastModel != cfg.AI.OllamaModel {
|
||||
aiRouter.SetOperation(router.OpComplete, ollama.New(cfg.AI.OllamaBaseURL, cfg.AI.OllamaFastModel))
|
||||
}
|
||||
|
||||
// nil interaction logger: core has no enterprise/internal/audit
|
||||
// implementation to log translate/fix/optimize interactions
|
||||
// against, same posture as queryHandler's nil audit logger above.
|
||||
aiHandler := aiapi.NewHandler(logger, aiRouter, groundingSvc, authorizer, nil)
|
||||
aiHandler.RegisterRoutes(mux)
|
||||
logger.Info("ai routes enabled", "ollama_base_url", cfg.AI.OllamaBaseURL, "model", cfg.AI.OllamaModel)
|
||||
}
|
||||
|
||||
srv := &http.Server{
|
||||
Addr: cfg.HTTPListenAddr,
|
||||
Handler: httpserver.WithCORS(mux, cfg.CORSAllowedOrigin),
|
||||
|
||||
@@ -17,6 +17,20 @@ type Config struct {
|
||||
QueryTimeout time.Duration
|
||||
CORSAllowedOrigin string
|
||||
EnterpriseAuthURL string
|
||||
AI AIConfig
|
||||
}
|
||||
|
||||
// AIConfig gates Phase 7's AI-assisted query features (Track A/B) --
|
||||
// off unless OllamaBaseURL is set, same "off unless configured"
|
||||
// convention as EnterpriseAuthURL and everything else optional in this
|
||||
// codebase. OllamaFastModel is the per-operation override for
|
||||
// Complete's tight latency budget (/docs/phase-7-ai-design.md's
|
||||
// per-operation provider/model config) -- empty means Complete uses
|
||||
// OllamaModel too, same as every other operation.
|
||||
type AIConfig struct {
|
||||
OllamaBaseURL string
|
||||
OllamaModel string
|
||||
OllamaFastModel string
|
||||
}
|
||||
|
||||
type ClickHouseConfig struct {
|
||||
@@ -65,6 +79,17 @@ func Load() (Config, error) {
|
||||
// base URL (e.g. "http://enterprise-auth:8081") to turn on
|
||||
// real session/service-token enforcement.
|
||||
EnterpriseAuthURL: getenv("ENTERPRISE_AUTH_URL", ""),
|
||||
// Empty OllamaBaseURL means AI features are entirely disabled --
|
||||
// /ai/* routes aren't even registered (see main.go), matching
|
||||
// "no cloud dependency required for the default deployment" and,
|
||||
// by the same reasoning, no *local* model dependency forced on a
|
||||
// deployment that doesn't want one either. Model names default to
|
||||
// the recommendation confirmed in /docs/phase-7-ai-design.md.
|
||||
AI: AIConfig{
|
||||
OllamaBaseURL: getenv("OLLAMA_BASE_URL", ""),
|
||||
OllamaModel: getenv("OLLAMA_MODEL", "qwen2.5-coder:7b"),
|
||||
OllamaFastModel: getenv("OLLAMA_FAST_MODEL", "qwen2.5-coder:1.5b"),
|
||||
},
|
||||
}
|
||||
|
||||
timeoutSec, err := strconv.Atoi(getenv("QUERY_TIMEOUT_SECONDS", "30"))
|
||||
|
||||
+20
-1
@@ -19,6 +19,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sentry/sentry/api/ai/costguard"
|
||||
"github.com/sentry/sentry/api/authz"
|
||||
"github.com/sentry/sentry/api/internal/querylang/planner"
|
||||
"github.com/sentry/sentry/api/querylang/executor"
|
||||
@@ -99,6 +100,20 @@ type queryRequest struct {
|
||||
type queryResponse struct {
|
||||
Columns []string `json:"columns"`
|
||||
Rows [][]any `json:"rows"`
|
||||
// Warnings surfaces costguard's assessment (Phase 7 task 4) for
|
||||
// every query, hand-written or AI-suggested alike -- the same
|
||||
// guard, never a hard block here. AI-suggested queries get a
|
||||
// stricter treatment (a Reject-level assessment withholds the
|
||||
// suggestion entirely, see the ai package) before a query ever
|
||||
// reaches this handler; a hand-written query submitted directly
|
||||
// always runs regardless of what this says, matching every prior
|
||||
// phase's behavior -- this field is informational, not new
|
||||
// enforcement, a deliberate choice recorded in
|
||||
// /docs/phase-7-ai-design.md rather than a retrofit nobody decided
|
||||
// on. Omitted (not an empty array) when there's nothing to say, so
|
||||
// existing callers that don't look for this field see no shape
|
||||
// change at all.
|
||||
Warnings []string `json:"warnings,omitempty"`
|
||||
}
|
||||
|
||||
type errorResponse struct {
|
||||
@@ -149,7 +164,11 @@ func (h *Handler) handleQuery(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
|
||||
h.logAudit(r.Context(), req, len(result.Rows), duration, nil)
|
||||
writeJSON(w, queryResponse{Columns: result.Columns, Rows: result.Rows})
|
||||
resp := queryResponse{Columns: result.Columns, Rows: result.Rows}
|
||||
if assessment := costguard.Assess(plan); assessment.Level != costguard.LevelOK {
|
||||
resp.Warnings = []string{costguard.Summary(assessment)}
|
||||
}
|
||||
writeJSON(w, resp)
|
||||
}
|
||||
|
||||
// logAudit is fail-open by design (see AuditLogger's doc comment): a
|
||||
|
||||
Reference in New Issue
Block a user