209 lines
7.1 KiB
Go
209 lines
7.1 KiB
Go
// rag.go adds retrieval-augmented search on top of the plain
|
|
// read-a-whole-file skill lookup in skillkit.go. The read_office_skill
|
|
// tool (skillkit.Read) is the right fit when the model already knows
|
|
// which format it needs and wants that skill's full guidance before a
|
|
// multi-step task — but for a narrow, specific question ("can a cell
|
|
// hold a formula", "how do I bold part of a paragraph"), returning an
|
|
// entire 3-4KB skill file just to answer it is the same imprecision
|
|
// Claude Code's own skills avoid by splitting a large skill into a
|
|
// short SKILL.md plus separate reference files loaded on demand.
|
|
// Because every skill here already follows one convention — a "##"
|
|
// heading per topic — each section is naturally a single, coherent
|
|
// chunk, and because the whole corpus is a few dozen short sections
|
|
// fixed at compile time (go:embed, never changes between deploys),
|
|
// embedding it once per process and ranking in memory is simpler and
|
|
// faster than standing up a vector-DB round trip for content this
|
|
// small and this static. search_codebase, by contrast, indexes a
|
|
// workspace's actual source files — large, per-user, and changing —
|
|
// which is exactly the shape services/knowledge's Postgres/pgvector
|
|
// pipeline exists for; that pipeline would be the wrong tool here, not
|
|
// a missing piece.
|
|
package skillkit
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"math"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
)
|
|
|
|
// Embedder is the subset of *vllmkit.Client this package needs —
|
|
// declared as a narrow interface (rather than importing vllmkit
|
|
// directly) so the caller supplies whatever embedding client it
|
|
// already has, and so this package stays testable without a live
|
|
// vLLM instance.
|
|
type Embedder interface {
|
|
Embed(ctx context.Context, text string) ([]float32, error)
|
|
}
|
|
|
|
// Chunk is one semantically coherent section of a skill file — in
|
|
// practice, one "## Heading" and the text under it (or the file's
|
|
// title/intro, for the section before the first heading).
|
|
type Chunk struct {
|
|
Skill string // the skill this section belongs to, e.g. "docx"
|
|
Heading string // the section's heading, e.g. "Editing an existing document" — "Overview" for the lead-in before any "##"
|
|
Text string // the section's full text, heading line included — what gets embedded, and what Search returns
|
|
}
|
|
|
|
// chunkSkill splits one skill's raw markdown body into Chunks, cutting
|
|
// at every "## " line. This mirrors the actual structure every file in
|
|
// skills/*.md already uses, so no separate chunk-size tuning
|
|
// (domain.ChunkText's char budget, needed for arbitrary source code) is
|
|
// necessary here — a skill's own section boundaries are already the
|
|
// right retrieval granularity.
|
|
func chunkSkill(name, body string) []Chunk {
|
|
var (
|
|
chunks []Chunk
|
|
heading string
|
|
cur []string
|
|
)
|
|
flush := func() {
|
|
text := strings.TrimSpace(strings.Join(cur, "\n"))
|
|
if text == "" {
|
|
return
|
|
}
|
|
h := heading
|
|
if h == "" {
|
|
h = "Overview"
|
|
}
|
|
chunks = append(chunks, Chunk{Skill: name, Heading: h, Text: text})
|
|
}
|
|
for _, line := range strings.Split(body, "\n") {
|
|
if strings.HasPrefix(line, "## ") {
|
|
flush()
|
|
heading = strings.TrimPrefix(line, "## ")
|
|
cur = []string{line}
|
|
continue
|
|
}
|
|
cur = append(cur, line)
|
|
}
|
|
flush()
|
|
return chunks
|
|
}
|
|
|
|
// allChunks returns every skill's chunks, in skillkit.Names() order.
|
|
func allChunks() []Chunk {
|
|
var out []Chunk
|
|
for _, name := range Names() {
|
|
out = append(out, chunkSkill(name, registry[name].Body)...)
|
|
}
|
|
return out
|
|
}
|
|
|
|
// Match is one ranked search result.
|
|
type Match struct {
|
|
Chunk
|
|
Score float64 // cosine similarity to the query, higher is more relevant
|
|
}
|
|
|
|
type vectorChunk struct {
|
|
Chunk
|
|
vec []float32
|
|
}
|
|
|
|
// Index is a lazily-built, in-memory semantic index over every skill's
|
|
// chunks. The zero value is not usable — construct with NewIndex. An
|
|
// Index is safe for concurrent use and is meant to be built once and
|
|
// reused for the life of the process (see package doc for why an
|
|
// in-memory index, not a persisted one, is the right call here).
|
|
type Index struct {
|
|
embedder Embedder
|
|
|
|
mu sync.Mutex
|
|
vectors []vectorChunk // nil until a build has succeeded at least once
|
|
}
|
|
|
|
// NewIndex returns an Index that will embed every skill chunk, once,
|
|
// the first time Search is called (not eagerly here) — so a
|
|
// deployment that never calls search_office_skills never pays the
|
|
// embedding cost at all, and the first real call pays it exactly once.
|
|
func NewIndex(embedder Embedder) *Index {
|
|
return &Index{embedder: embedder}
|
|
}
|
|
|
|
// ensureBuilt embeds every chunk if that hasn't already succeeded.
|
|
// Deliberately NOT a sync.Once: if vLLM is unreachable on the very
|
|
// first call, that failure must not be cached forever — the next
|
|
// Search should try again rather than fail for the rest of the
|
|
// process's life over one transient outage. Concurrent callers during
|
|
// a build simply serialize on the mutex rather than building twice.
|
|
func (ix *Index) ensureBuilt(ctx context.Context) error {
|
|
ix.mu.Lock()
|
|
defer ix.mu.Unlock()
|
|
if ix.vectors != nil {
|
|
return nil
|
|
}
|
|
chunks := allChunks()
|
|
vectors := make([]vectorChunk, 0, len(chunks))
|
|
for _, c := range chunks {
|
|
vec, err := ix.embedder.Embed(ctx, c.Text)
|
|
if err != nil {
|
|
return fmt.Errorf("skillkit: embed %s/%s: %w", c.Skill, c.Heading, err)
|
|
}
|
|
vectors = append(vectors, vectorChunk{Chunk: c, vec: vec})
|
|
}
|
|
ix.vectors = vectors
|
|
return nil
|
|
}
|
|
|
|
// Search embeds query and returns the topK chunks whose own embedding
|
|
// is most cosine-similar to it, best first. skillFilter, if non-empty,
|
|
// restricts the search to one skill's chunks (matched the same way
|
|
// Read does: case-insensitive against the skill name). topK<=0 defaults
|
|
// to 4.
|
|
func (ix *Index) Search(ctx context.Context, query string, topK int, skillFilter string) ([]Match, error) {
|
|
if topK <= 0 {
|
|
topK = 4
|
|
}
|
|
if err := ix.ensureBuilt(ctx); err != nil {
|
|
return nil, err
|
|
}
|
|
qvec, err := ix.embedder.Embed(ctx, query)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("skillkit: embed query: %w", err)
|
|
}
|
|
|
|
ix.mu.Lock()
|
|
vectors := ix.vectors
|
|
ix.mu.Unlock()
|
|
|
|
matches := make([]Match, 0, len(vectors))
|
|
for _, v := range vectors {
|
|
if skillFilter != "" && !strings.EqualFold(v.Skill, skillFilter) {
|
|
continue
|
|
}
|
|
matches = append(matches, Match{Chunk: v.Chunk, Score: cosineSimilarity(qvec, v.vec)})
|
|
}
|
|
sort.Slice(matches, func(i, j int) bool { return matches[i].Score > matches[j].Score })
|
|
if len(matches) > topK {
|
|
matches = matches[:topK]
|
|
}
|
|
return matches, nil
|
|
}
|
|
|
|
// cosineSimilarity is the standard cosine of the angle between two
|
|
// vectors, in [-1, 1] for any nonzero inputs (in practice embeddings
|
|
// from one model land in a narrower positive range). Defensively walks
|
|
// only the shorter of the two lengths and returns 0 for a zero vector,
|
|
// rather than panicking or dividing by zero, though in normal operation
|
|
// every vector here comes from the same embedding model and is the
|
|
// same length.
|
|
func cosineSimilarity(a, b []float32) float64 {
|
|
n := len(a)
|
|
if len(b) < n {
|
|
n = len(b)
|
|
}
|
|
var dot, normA, normB float64
|
|
for i := 0; i < n; i++ {
|
|
dot += float64(a[i]) * float64(b[i])
|
|
normA += float64(a[i]) * float64(a[i])
|
|
normB += float64(b[i]) * float64(b[i])
|
|
}
|
|
if normA == 0 || normB == 0 {
|
|
return 0
|
|
}
|
|
return dot / (math.Sqrt(normA) * math.Sqrt(normB))
|
|
}
|