RAG 是什么：从知识检索到有据可查的模型回答

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strings"
	"unicode"
)

type Chunk struct {
	ID, Document, Revision, Tenant, Text string
	Readers                             []string
}

type Evidence struct {
	ID, Document, Revision, Text string
}

type Answer struct {
	Text, Citation, Quote string
}

func terms(text string) map[string]bool {
	out := make(map[string]bool)
	for _, word := range strings.FieldsFunc(strings.ToLower(text), func(r rune) bool {
		return !unicode.IsLetter(r) && !unicode.IsDigit(r)
	}) {
		out[word] = true
	}
	return out
}

func retrieve(chunks []Chunk, active map[string]string, tenant, reader, question string, k, maxBytes int) ([]Evidence, error) {
	if tenant == "" || reader == "" || k <= 0 || maxBytes < 2 {
		return nil, errors.New("invalid retrieval configuration")
	}
	query := terms(question)
	if len(query) == 0 {
		return nil, errors.New("empty query")
	}
	type ranked struct {
		evidence Evidence
		score    int
	}
	var candidates []ranked
	seen := make(map[string]bool)
	for _, chunk := range chunks {
		if chunk.Tenant != tenant {
			continue
		}
		allowed := false
		for _, member := range chunk.Readers {
			allowed = allowed || member == reader
		}
		if !allowed {
			continue
		}
		revision, exists := active[chunk.Document]
		if !exists || revision != chunk.Revision {
			continue
		}
		if chunk.ID == "" || chunk.Document == "" || revision == "" || seen[chunk.ID] {
			return nil, errors.New("invalid or duplicate active chunk identity")
		}
		seen[chunk.ID] = true
		score := 0
		documentTerms := terms(chunk.Text)
		for term := range query {
			if documentTerms[term] {
				score++
			}
		}
		if score > 0 {
			candidates = append(candidates, ranked{
				Evidence{chunk.ID, chunk.Document, chunk.Revision, chunk.Text}, score,
			})
		}
	}
	sort.Slice(candidates, func(i, j int) bool {
		if candidates[i].score == candidates[j].score {
			return candidates[i].evidence.ID < candidates[j].evidence.ID
		}
		return candidates[i].score > candidates[j].score
	})
	result := make([]Evidence, 0)
	for _, candidate := range candidates {
		trial := append(append([]Evidence{}, result...), candidate.evidence)
		encoded, err := json.Marshal(trial)
		if err != nil {
			return nil, err
		}
		if len(encoded) > maxBytes {
			continue
		}
		result = trial
		if len(result) == k {
			break
		}
	}
	return result, nil
}

func generate(evidence []Evidence) (Answer, error) {
	if len(evidence) == 0 {
		return Answer{}, errors.New("insufficient evidence")
	}
	first := evidence[0]
	return Answer{first.Text, first.ID, first.Text}, nil
}

func validate(answer Answer, evidence []Evidence) error {
	if strings.TrimSpace(answer.Text) == "" || strings.TrimSpace(answer.Quote) == "" {
		return errors.New("empty answer or quotation")
	}
	for _, item := range evidence {
		if answer.Citation == item.ID {
			if !strings.Contains(item.Text, answer.Quote) {
				return errors.New("quotation not in cited evidence")
			}
			if answer.Text != answer.Quote {
				return errors.New("extractive demo rejects added claims")
			}
			return nil
		}
	}
	return errors.New("citation outside request evidence")
}

func main() {
	chunks := []Chunk{
		{"old", "refund-policy", "v1", "shop-a", "Refund requests must arrive within 30 days.", []string{"alice"}},
		{"current", "refund-policy", "v2", "shop-a", "Refund requests for unopened items must arrive within 14 days.", []string{"alice"}},
		{"private", "refund-exceptions", "v1", "shop-a", "Refund exceptions require finance approval.", []string{"finance"}},
		{"other", "refund-policy", "v2", "shop-b", "Refund requests must arrive within 60 days.", []string{"alice"}},
	}
	evidence, err := retrieve(chunks, map[string]string{
		"refund-policy": "v2", "refund-exceptions": "v1",
	}, "shop-a", "alice", "refund requests", 2, 512)
	if err != nil {
		panic(err)
	}
	answer, err := generate(evidence)
	if err != nil {
		panic(err)
	}
	if err := validate(answer, evidence); err != nil {
		panic(err)
	}
	fmt.Printf("sources=%d citation=%s\n%s\n", len(evidence), answer.Citation, answer.Text)
}

sources=1 citation=current
Refund requests for unopened items must arrive within 14 days.