ai-gateway-go/internal/factcheck/engine.go

package factcheck

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"time"

	platformid "aigateway.local/core/internal/platform/id"
	"github.com/jackc/pgx/v5/pgxpool"
)

// EvidenceHit is a single retrieved knowledge-base excerpt supplied to the
// verifier. The title/content shape mirrors the workbench retriever output.
type EvidenceHit struct {
	DocumentTitle string `json:"document_title"`
	Content       string `json:"content"`
}

// EvidenceRetriever fetches supporting excerpts from a knowledge base. The
// workbench PostgreSQLRetriever satisfies it (through a small adapter), so
// fact-check reuses the exact retrieval path application prompts already use.
type EvidenceRetriever interface {
	Search(context.Context, string, string, int) ([]EvidenceHit, error)
}

// Verifier performs the model-backed fact-check call and returns the raw model
// output text. The engine is responsible for parsing it into claims and a
// verdict. The workbench runtime implements it by routing a non-streaming chat
// completion through the same governed gateway, reusing the caller's own
// credential headers.
type Verifier interface {
	Verify(context.Context, string, string, string, time.Duration) (string, error)
}

// VerifierFunc adapts a function to the Verifier interface.
type VerifierFunc func(context.Context, string, string, string, time.Duration) (string, error)

func (f VerifierFunc) Verify(ctx context.Context, model, system, user string, timeout time.Duration) (string, error) {
	return f(ctx, model, system, user, timeout)
}

// Engine executes fact-check policies against assistant answers and records the
// outcome in fact_check_events. It is deliberately side-effect safe: Check
// never returns an error a caller must propagate — callers treat any failure as
// "fact-check skipped" and keep serving the chat.
type Engine struct {
	pool      *pgxpool.Pool
	retriever EvidenceRetriever
	logger    *slog.Logger
}

func NewEngine(pool *pgxpool.Pool, retriever EvidenceRetriever, logger *slog.Logger) *Engine {
	return &Engine{pool: pool, retriever: retriever, logger: logger}
}

// Check verifies one assistant answer against the configured knowledge bases
// and persists a fact_check_events row. The returned Event has a zero ID when
// fact-checking is not configured or no policy applies; callers should skip
// quietly in that case. scope 形如 department:<uuid>(由调用方从资源部门推导);
// 空 scope 时只应用 global 策略。
func (e *Engine) Check(ctx context.Context, requestID, scope, question, answer string, verifier Verifier) (Event, error) {
	if e == nil || e.retriever == nil || verifier == nil || strings.TrimSpace(answer) == "" {
		return Event{}, nil
	}
	settings, err := e.checkSettings(ctx)
	if err != nil {
		return Event{}, fmt.Errorf("read fact-check settings: %w", err)
	}
	if strings.TrimSpace(settings.Model) == "" {
		return Event{}, nil // not configured; skip without noise
	}
	policy, err := e.enabledPolicy(ctx, scope)
	if err != nil {
		return Event{}, err
	}
	if policy.ID == "" {
		return Event{}, nil
	}
	started := time.Now()
	evidence := e.gatherEvidence(ctx, policy, question)
	if len(evidence) == 0 {
		// Nothing to check against; record an uncertain event so the admin can
		// see retrieval produced no evidence rather than silently passing.
		return e.record(ctx, Event{
			PolicyID: &policy.ID, RequestID: requestID, Model: settings.Model,
			Mode: policy.Mode, Action: policy.Action, Verdict: "uncertain",
			LatencyMS: int(time.Since(started).Milliseconds()), Question: question, Answer: answer,
			Claims: json.RawMessage("[]"), Evidence: evidenceJSON(evidence),
		})
	}
	system, user := buildVerifyPrompt(question, answer, evidence, policy.MaxClaims)
	raw, err := verifier.Verify(ctx, settings.Model, system, user, time.Duration(settings.TimeoutSeconds)*time.Second)
	if err != nil {
		return e.record(ctx, Event{
			PolicyID: &policy.ID, RequestID: requestID, Model: settings.Model,
			Mode: policy.Mode, Action: policy.Action, Verdict: "error",
			LatencyMS: int(time.Since(started).Milliseconds()), Question: question, Answer: answer,
			Claims: json.RawMessage("[]"), Evidence: evidenceJSON(evidence), Error: err.Error(),
		})
	}
	verdict, score, claims := parseVerdict(raw, policy.SupportThreshold)
	return e.record(ctx, Event{
		PolicyID: &policy.ID, RequestID: requestID, Model: settings.Model,
		Mode: policy.Mode, Action: policy.Action, Verdict: verdict, SupportScore: score,
		LatencyMS: int(time.Since(started).Milliseconds()), Question: question, Answer: answer,
		Claims: claimsJSON(claims), Evidence: evidenceJSON(evidence),
	})
}

func (e *Engine) checkSettings(ctx context.Context) (Settings, error) {
	var x Settings
	err := e.pool.QueryRow(ctx, `SELECT f.provider_id::text,coalesce(p.code,''),f.model,f.timeout_seconds,f.updated_at FROM gateway.fact_check_settings f LEFT JOIN gateway.providers p ON p.id=f.provider_id WHERE singleton`).Scan(&x.ProviderID, &x.ProviderCode, &x.Model, &x.TimeoutSeconds, &x.UpdatedAt)
	return x, err
}

// enabledPolicy 选择命中的策略:优先精确匹配调用方 scope(department:<uuid>
// 等),否则回退 global。修复之前 ORDER BY scope LIMIT 1 只取字典序第一条
// 的问题——多策略并存时其余部门的策略被静默忽略或张冠李戴。
func (e *Engine) enabledPolicy(ctx context.Context, scope string) (Policy, error) {
	policy, err := scanPolicy(e.pool.QueryRow(ctx, policySelect+` WHERE enabled AND (scope='global' OR scope=$1) ORDER BY (scope='global'),scope LIMIT 1`, scope))
	if errors.Is(err, ErrNotFound) {
		return Policy{}, nil
	}
	return policy, err
}

// gatherEvidence retrieves up to a bounded number of excerpts across the
// policy's knowledge bases. Retrieval failures never abort the check — a
// broken knowledge base just contributes no evidence.
func (e *Engine) gatherEvidence(ctx context.Context, policy Policy, question string) []EvidenceHit {
	maxHits := policy.MaxClaims * 2
	if maxHits < 4 {
		maxHits = 4
	}
	if maxHits > 20 {
		maxHits = 20
	}
	evidence := []EvidenceHit{}
	for _, kbID := range policy.KnowledgeBaseIDs {
		hits, err := e.retriever.Search(ctx, kbID, question, policy.TopK)
		if err != nil {
			if e.logger != nil {
				e.logger.Warn("fact-check evidence retrieval failed", "knowledge_base_id", kbID, "error", err)
			}
			continue
		}
		for _, hit := range hits {
			if len(evidence) >= maxHits {
				break
			}
			evidence = append(evidence, EvidenceHit{DocumentTitle: hit.DocumentTitle, Content: hit.Content})
		}
		if len(evidence) >= maxHits {
			break
		}
	}
	return evidence
}

func (e *Engine) record(ctx context.Context, event Event) (Event, error) {
	if event.ID == "" {
		id, err := platformid.NewUUID()
		if err != nil {
			return event, err
		}
		event.ID = id
	}
	_, err := e.pool.Exec(ctx, `INSERT INTO gateway.fact_check_events(id,policy_id,request_id,model,mode,action,verdict,support_score,latency_ms,question,answer,claims,evidence,error) VALUES($1,nullif($2,'')::uuid,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12::jsonb,$13::jsonb,$14)`,
		event.ID, ptrString(event.PolicyID), event.RequestID, event.Model, event.Mode, event.Action, event.Verdict, event.SupportScore, event.LatencyMS, event.Question, event.Answer, event.Claims, event.Evidence, event.Error)
	return event, err
}

func ptrString(value *string) string {
	if value == nil {
		return ""
	}
	return *value
}

type verifyClaim struct {
	Claim         string `json:"claim"`
	Verdict       string `json:"verdict"`
	EvidenceIndex []int  `json:"evidence_index"`
}

type verifyResult struct {
	Verdict      string        `json:"verdict"`
	SupportScore float64       `json:"support_score"`
	Claims       []verifyClaim `json:"claims"`
}

func buildVerifyPrompt(question, answer string, evidence []EvidenceHit, maxClaims int) (string, string) {
	if maxClaims < 1 {
		maxClaims = 8
	}
	var buffer bytes.Buffer
	buffer.WriteString("问题：\n")
	buffer.WriteString(question)
	buffer.WriteString("\n\n回答：\n")
	buffer.WriteString(answer)
	buffer.WriteString("\n\n参考资料：\n")
	for i, hit := range evidence {
		fmt.Fprintf(&buffer, "[资料%d] %s\n%s\n\n", i, hit.DocumentTitle, hit.Content)
	}
	system := `你是企业知识库事实核查引擎。你需要逐条判断"回答"中的关键陈述是否能被"参考资料"支持。
输出必须是严格 JSON，不要输出任何其他内容，格式：
{"verdict":"supported|unsupported|uncertain","support_score":0-100,"claims":[{"claim":"...","verdict":"supported|unsupported|uncertain","evidence_index":[0,1]}]}
- verdict：所有关键陈述均被参考资料支持→supported；存在明确被资料否定或资料完全无法支撑的关键陈述→unsupported；资料不足无法判断→uncertain。
- support_score：被支持的陈述占比（0-100）。
- claims：从回答中提取的关键陈述，最多 ` + fmt.Sprint(maxClaims) + ` 条。evidence_index 列出支撑该陈述的资料编号（从0开始）；无支撑填[]。
- 只依据参考资料判断，不要使用你自己的世界知识。
- 回答为空或不包含可核查陈述时，verdict 输出 uncertain，claims 输出 []。`
	return system, buffer.String()
}

func parseVerdict(raw string, threshold int) (string, *int, []verifyClaim) {
	var result verifyResult
	if extracted := extractJSON(raw); json.Unmarshal([]byte(extracted), &result) != nil {
		return "error", nil, []verifyClaim{}
	}
	claims := result.Claims
	if claims == nil {
		claims = []verifyClaim{}
	}
	verdict := strings.ToLower(strings.TrimSpace(result.Verdict))
	var score *int
	if result.SupportScore > 0 {
		value := int(result.SupportScore)
		if value > 100 {
			value = 100
		}
		score = &value
	}
	total := len(claims)
	supported := 0
	for _, c := range claims {
		switch strings.ToLower(strings.TrimSpace(c.Verdict)) {
		case "supported":
			supported++
		case "unsupported", "uncertain":
		default:
			c.Verdict = "uncertain"
		}
	}
	if total > 0 {
		percent := supported * 100 / total
		if score == nil {
			score = &percent
		}
		if verdict != "supported" && verdict != "unsupported" && verdict != "uncertain" {
			switch {
			case percent >= threshold:
				verdict = "supported"
			case supported > 0:
				verdict = "uncertain"
			default:
				verdict = "unsupported"
			}
		}
	} else if verdict != "supported" && verdict != "unsupported" && verdict != "uncertain" {
		verdict = "uncertain"
	}
	return verdict, score, claims
}

// extractJSON returns the text between the first '{' and the last '}', stripping
// markdown code fences models sometimes wrap around their JSON output.
func extractJSON(raw string) string {
	start := strings.Index(raw, "{")
	end := strings.LastIndex(raw, "}")
	if start < 0 || end < start {
		return ""
	}
	return raw[start : end+1]
}

func claimsJSON(claims []verifyClaim) json.RawMessage {
	if claims == nil {
		return json.RawMessage("[]")
	}
	encoded, err := json.Marshal(claims)
	if err != nil {
		return json.RawMessage("[]")
	}
	return encoded
}

func evidenceJSON(evidence []EvidenceHit) json.RawMessage {
	if evidence == nil {
		return json.RawMessage("[]")
	}
	encoded, err := json.Marshal(evidence)
	if err != nil {
		return json.RawMessage("[]")
	}
	return encoded
}