feat(embedding): 嵌入(n-gram 哈希向量 + 关键词 FTS + OpenAI 兼容 API 钩子 + Manager 选择)
This commit is contained in:
1 parent
70fd915242
commit
51b2bdfba2
9 files changed
+1125
No files matched your search
@@ -0,0 +1,236 @@
|
||||
package embedding
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
// maxResponseBytes caps how much of an /embeddings response body is read, so an
|
||||
// oversized or malformed response cannot exhaust memory. Legitimate embeddings
|
||||
// for a normal batch fall far below this; the cap is a guard, not a limit.
|
||||
const maxResponseBytes = 16 << 20
|
||||
|
||||
// maxErrorBodyBytes caps how much of a non-OK response body is read for the
|
||||
// error message.
|
||||
const maxErrorBodyBytes = 4096
|
||||
|
||||
// HTTP client tuning values, matching the recommended pool in
|
||||
// DevDocs/02-development-constraints.md §2.3.
|
||||
const (
|
||||
httpMaxIdleConns = 100
|
||||
httpMaxIdleConnsPerHost = 20
|
||||
httpIdleConnTimeout = 90 * time.Second
|
||||
httpTLSHandshakeTimeout = 15 * time.Second
|
||||
httpRequestTimeout = 180 * time.Second
|
||||
)
|
||||
|
||||
// openAICompatibleRequest is the JSON body sent to an OpenAI-compatible
|
||||
// /embeddings endpoint.
|
||||
type openAICompatibleRequest struct {
|
||||
Model string `json:"model"`
|
||||
Input []string `json:"input"`
|
||||
}
|
||||
|
||||
// openAICompatibleResponse is the subset of the /embeddings JSON response the
|
||||
// APIEmbedder consumes.
|
||||
type openAICompatibleResponse struct {
|
||||
Data []openAICompatibleDatum `json:"data"`
|
||||
}
|
||||
|
||||
// openAICompatibleDatum is a single embedding in the response data array.
|
||||
type openAICompatibleDatum struct {
|
||||
Index int `json:"index"`
|
||||
Embedding []float32 `json:"embedding"`
|
||||
}
|
||||
|
||||
// ErrNoAPIURL is returned by APIEmbedder.Embed when the configured API URL is
|
||||
// blank, so a mis-configured embedder fails loudly instead of issuing a request
|
||||
// to an empty URL.
|
||||
var ErrNoAPIURL = errors.New("embedding: api url is empty")
|
||||
|
||||
// APIEmbedder implements Embedder by calling an OpenAI-compatible /embeddings
|
||||
// endpoint. It is the semantic-vector enhancement selected when
|
||||
// config.EmbeddingConfig.APIURL is set, and it always uses the real HTTP path
|
||||
// (never a stub).
|
||||
//
|
||||
// The vector dimension is not known until the first successful response; Dim()
|
||||
// reports the last observed dimension, or 0 before any successful Embed. The
|
||||
// runtime should therefore call Embed before relying on Dim().
|
||||
type APIEmbedder struct {
|
||||
cfg config.EmbeddingConfig
|
||||
client *http.Client
|
||||
log *slog.Logger
|
||||
|
||||
mu sync.RWMutex
|
||||
dim int
|
||||
}
|
||||
|
||||
// APIOption customizes an APIEmbedder built by NewAPI.
|
||||
type APIOption func(*APIEmbedder)
|
||||
|
||||
// WithHTTPClient sets the *http.Client used for the /embeddings request. A nil
|
||||
// client is ignored and the default is kept.
|
||||
func WithHTTPClient(c *http.Client) APIOption {
|
||||
return func(e *APIEmbedder) {
|
||||
if c != nil {
|
||||
e.client = c
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WithRoundTripper sets the transport for the /embeddings request by wrapping
|
||||
// it in a fresh *http.Client. Tests use this to inject a mock RoundTripper
|
||||
// without starting a network server.
|
||||
func WithRoundTripper(rt http.RoundTripper) APIOption {
|
||||
return func(e *APIEmbedder) {
|
||||
e.client = &http.Client{Transport: rt}
|
||||
}
|
||||
}
|
||||
|
||||
// WithLogger sets the logger used for embed failures. A nil logger is ignored
|
||||
// and the default discard logger is kept.
|
||||
func WithLogger(l *slog.Logger) APIOption {
|
||||
return func(e *APIEmbedder) {
|
||||
if l != nil {
|
||||
e.log = l
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// NewAPI builds an APIEmbedder from cfg using the default HTTP client.
|
||||
func NewAPI(cfg config.EmbeddingConfig, opts ...APIOption) *APIEmbedder {
|
||||
e := &APIEmbedder{
|
||||
cfg: cfg,
|
||||
client: defaultHTTPClient(),
|
||||
log: slog.New(slog.NewTextHandler(io.Discard, nil)),
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(e)
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
// Dim returns the dimension observed from the last successful response, or 0
|
||||
// if no response has been received yet.
|
||||
func (e *APIEmbedder) Dim() int {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
return e.dim
|
||||
}
|
||||
|
||||
// Embed sends texts to the configured /embeddings endpoint and returns the
|
||||
// returned vectors aligned by index. It returns ErrNoAPIURL for a blank URL
|
||||
// and a wrapped error for any HTTP or decode failure.
|
||||
func (e *APIEmbedder) Embed(ctx context.Context, texts []string) ([]Vector, error) {
|
||||
if len(texts) == 0 {
|
||||
return []Vector{}, nil
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if strings.TrimSpace(e.cfg.APIURL) == "" {
|
||||
return nil, ErrNoAPIURL
|
||||
}
|
||||
|
||||
payload, err := json.Marshal(openAICompatibleRequest{
|
||||
Model: e.cfg.Model,
|
||||
Input: texts,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("embedding: marshal request: %w", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, e.cfg.APIURL, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("embedding: build request: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if e.cfg.APIKey != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+e.cfg.APIKey)
|
||||
}
|
||||
|
||||
resp, err := e.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("embedding: request: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
msg := readErrorBody(resp.Body)
|
||||
return nil, fmt.Errorf("embedding: status %d: %s", resp.StatusCode, msg)
|
||||
}
|
||||
|
||||
data, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("embedding: read response: %w", err)
|
||||
}
|
||||
|
||||
var decoded openAICompatibleResponse
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
return nil, fmt.Errorf("embedding: decode response: %w", err)
|
||||
}
|
||||
if len(decoded.Data) != len(texts) {
|
||||
return nil, fmt.Errorf("embedding: got %d vectors for %d texts", len(decoded.Data), len(texts))
|
||||
}
|
||||
|
||||
vectors := make([]Vector, len(texts))
|
||||
dim := 0
|
||||
for _, d := range decoded.Data {
|
||||
if d.Index < 0 || d.Index >= len(vectors) {
|
||||
return nil, fmt.Errorf("embedding: index %d out of range", d.Index)
|
||||
}
|
||||
vec := make(Vector, len(d.Embedding))
|
||||
copy(vec, d.Embedding)
|
||||
vectors[d.Index] = vec
|
||||
if dim == 0 {
|
||||
dim = len(vec)
|
||||
}
|
||||
}
|
||||
|
||||
e.mu.Lock()
|
||||
e.dim = dim
|
||||
e.mu.Unlock()
|
||||
|
||||
return vectors, nil
|
||||
}
|
||||
|
||||
// readErrorBody drains a small prefix of a non-OK response body into a message,
|
||||
// capped at maxErrorBodyBytes so an error response cannot exhaust memory.
|
||||
func readErrorBody(r io.Reader) string {
|
||||
raw, err := io.ReadAll(io.LimitReader(r, maxErrorBodyBytes))
|
||||
if err != nil {
|
||||
return err.Error()
|
||||
}
|
||||
msg := strings.TrimSpace(string(raw))
|
||||
if msg == "" {
|
||||
return "bad status"
|
||||
}
|
||||
return msg
|
||||
}
|
||||
|
||||
// defaultHTTPClient builds the HTTP client configuration described in
|
||||
// DevDocs/02-development-constraints.md §2.3: tuned idle/TLS timeouts plus a
|
||||
// total request timeout. A fresh client is returned per call (never a package
|
||||
// global, which the constraints forbid).
|
||||
func defaultHTTPClient() *http.Client {
|
||||
return &http.Client{
|
||||
Transport: &http.Transport{
|
||||
MaxIdleConns: httpMaxIdleConns,
|
||||
MaxIdleConnsPerHost: httpMaxIdleConnsPerHost,
|
||||
IdleConnTimeout: httpIdleConnTimeout,
|
||||
TLSHandshakeTimeout: httpTLSHandshakeTimeout,
|
||||
},
|
||||
Timeout: httpRequestTimeout,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
package embedding
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
// roundTripFunc adapts a function to http.RoundTripper, letting tests inject a
|
||||
// mock transport without starting a network server.
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
|
||||
|
||||
func TestAPIEmbedder_EmbedsFromServer(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
t.Errorf("method = %s, want POST", r.Method)
|
||||
}
|
||||
if got := r.Header.Get("Content-Type"); got != "application/json" {
|
||||
t.Errorf("Content-Type = %q, want application/json", got)
|
||||
}
|
||||
if got := r.Header.Get("Authorization"); got != "Bearer test-key" {
|
||||
t.Errorf("Authorization = %q, want Bearer test-key", got)
|
||||
}
|
||||
|
||||
var req openAICompatibleRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
t.Errorf("decode request: %v", err)
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if req.Model != "text-embedding-3-small" {
|
||||
t.Errorf("request model = %q, want text-embedding-3-small", req.Model)
|
||||
}
|
||||
|
||||
vec := make([]float32, 4)
|
||||
for i := range vec {
|
||||
vec[i] = float32(i) + 0.1
|
||||
}
|
||||
resp := openAICompatibleResponse{Data: make([]openAICompatibleDatum, 0, len(req.Input))}
|
||||
for i := range req.Input {
|
||||
resp.Data = append(resp.Data, openAICompatibleDatum{Index: i, Embedding: vec})
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(resp)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := config.EmbeddingConfig{
|
||||
APIURL: srv.URL + "/v1/embeddings",
|
||||
APIKey: "test-key",
|
||||
Model: "text-embedding-3-small",
|
||||
}
|
||||
emb := NewAPI(cfg)
|
||||
|
||||
vectors, err := emb.Embed(context.Background(), []string{"hello", "world"})
|
||||
if err != nil {
|
||||
t.Fatalf("Embed() error: %v", err)
|
||||
}
|
||||
if len(vectors) != 2 {
|
||||
t.Fatalf("got %d vectors, want 2", len(vectors))
|
||||
}
|
||||
for i, v := range vectors {
|
||||
if len(v) != 4 {
|
||||
t.Errorf("vector %d length = %d, want 4", i, len(v))
|
||||
}
|
||||
for j, x := range v {
|
||||
if want := float32(j) + 0.1; x != want {
|
||||
t.Errorf("vector %d[%d] = %v, want %v", i, j, x, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
if emb.Dim() != 4 {
|
||||
t.Errorf("Dim() = %d, want 4", emb.Dim())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIEmbedder_RoundTripperInjection(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var seenPath string
|
||||
rt := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
seenPath = r.URL.Path
|
||||
body := `{"data":[{"index":0,"embedding":[0.1,0.2,0.3]}]}`
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
Header: make(http.Header),
|
||||
}, nil
|
||||
})
|
||||
|
||||
emb := NewAPI(
|
||||
config.EmbeddingConfig{APIURL: "http://example.com/v1/embeddings"},
|
||||
WithRoundTripper(rt),
|
||||
)
|
||||
|
||||
vectors, err := emb.Embed(context.Background(), []string{"hi"})
|
||||
if err != nil {
|
||||
t.Fatalf("Embed() error: %v", err)
|
||||
}
|
||||
if seenPath != "/v1/embeddings" {
|
||||
t.Errorf("request path = %q, want /v1/embeddings", seenPath)
|
||||
}
|
||||
if len(vectors) != 1 || len(vectors[0]) != 3 {
|
||||
t.Fatalf("vectors = %v, want 1 vector of length 3", vectors)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIEmbedder_Non200Errors(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
rt := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusInternalServerError,
|
||||
Body: io.NopCloser(strings.NewReader("server exploded")),
|
||||
Header: make(http.Header),
|
||||
}, nil
|
||||
})
|
||||
emb := NewAPI(
|
||||
config.EmbeddingConfig{APIURL: "http://example.com/v1/embeddings"},
|
||||
WithRoundTripper(rt),
|
||||
)
|
||||
|
||||
if _, err := emb.Embed(context.Background(), []string{"hi"}); err == nil {
|
||||
t.Fatal("expected error on 500, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIEmbedder_MismatchedCountErrors(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
rt := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
body := `{"data":[{"index":0,"embedding":[0.1,0.2]}]}`
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
Header: make(http.Header),
|
||||
}, nil
|
||||
})
|
||||
emb := NewAPI(
|
||||
config.EmbeddingConfig{APIURL: "http://example.com/v1/embeddings"},
|
||||
WithRoundTripper(rt),
|
||||
)
|
||||
|
||||
// 2 texts sent, only 1 returned -> error.
|
||||
if _, err := emb.Embed(context.Background(), []string{"a", "b"}); err == nil {
|
||||
t.Fatal("expected error on count mismatch, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIEmbedder_BlankURLErrors(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
emb := NewAPI(config.EmbeddingConfig{})
|
||||
_, err := emb.Embed(context.Background(), []string{"hello"})
|
||||
if err == nil {
|
||||
t.Fatal("expected error on blank URL, got nil")
|
||||
}
|
||||
if !errors.Is(err, ErrNoAPIURL) {
|
||||
t.Errorf("error = %v, want ErrNoAPIURL", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIEmbedder_EmptyInput(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
emb := NewAPI(config.EmbeddingConfig{APIURL: "http://example.invalid"})
|
||||
vectors, err := emb.Embed(context.Background(), nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Embed(nil) error: %v", err)
|
||||
}
|
||||
if len(vectors) != 0 {
|
||||
t.Errorf("got %d vectors, want 0", len(vectors))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIEmbedder_CanceledContext(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
emb := NewAPI(config.EmbeddingConfig{APIURL: "http://example.invalid"})
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
if _, err := emb.Embed(ctx, []string{"hello"}); err == nil {
|
||||
t.Fatal("expected error on canceled context, got nil")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
// Package embedding provides text-to-vector embedding and keyword full-text
|
||||
// search for the Cyrene runtime.
|
||||
//
|
||||
// It implements the vectorization path described in DevDocs/design/04-运行时.md
|
||||
// §7: an Embedder interface with two implementations — a dependency-free
|
||||
// n-gram hashing embedder (the offline default) and an optional OpenAI-
|
||||
// compatible /embeddings API embedder — plus a keyword full-text search (FTS)
|
||||
// inverted index.
|
||||
//
|
||||
// The runtime has two coexisting retrieval levels: the n-gram/API vectors are
|
||||
// semantic-level, while the FTS inverted index is keyword-level. Persistent
|
||||
// vector storage (SQLite vector extension / pgvector) is deliberately out of
|
||||
// scope for this package and belongs to milestone M4 (see
|
||||
// DevDocs/开发时间线.md); M1 only performs embedding and indexes for search.
|
||||
package embedding
|
||||
|
||||
import "context"
|
||||
|
||||
// Vector is a dense embedding: one float32 per dimension.
|
||||
type Vector []float32
|
||||
|
||||
// Embedder converts a batch of texts into dense vectors of a fixed dimension.
|
||||
//
|
||||
// Implementations must be deterministic for identical input and safe for
|
||||
// concurrent use. Embed returns one Vector per input text, aligned by index;
|
||||
// every returned vector has Dim() elements. ctx controls cancellation and must
|
||||
// be honored: a canceled context returns immediately with a non-nil error.
|
||||
type Embedder interface {
|
||||
// Embed returns one vector per input text, in the same order as texts.
|
||||
Embed(ctx context.Context, texts []string) ([]Vector, error)
|
||||
// Dim returns the fixed number of dimensions in every produced vector. It
|
||||
// must match the length of each vector returned by Embed.
|
||||
Dim() int
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
package embedding
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"unicode"
|
||||
)
|
||||
|
||||
// Hit is a single keyword search match: the indexed document id and its
|
||||
// relevance score. Higher scores are more relevant.
|
||||
type Hit[ID comparable] struct {
|
||||
ID ID
|
||||
Score float64
|
||||
}
|
||||
|
||||
// Index is a simple in-memory inverted index for keyword full-text search (FTS).
|
||||
//
|
||||
// It tokenizes each added document into terms and keeps a postings list
|
||||
// (term → doc id → term frequency). Search scores each matching document with
|
||||
// TF-IDF: the sum over query terms of tf × idf, where idf = ln(1 + N/df). This
|
||||
// is the keyword-level retrieval counterpart to the semantic Embedder vectors;
|
||||
// both coexist (see DevDocs/design/04-运行时.md §7).
|
||||
//
|
||||
// The zero value is not ready to use; build one with NewIndex. It is safe for
|
||||
// concurrent use.
|
||||
type Index[ID comparable] struct {
|
||||
mu sync.RWMutex
|
||||
postings map[string]map[ID]int // term -> doc id -> term frequency
|
||||
docLen map[ID]int // doc id -> term count
|
||||
docTerms map[ID][]string // doc id -> indexed terms, for re-indexing
|
||||
docCount int // number of distinct documents
|
||||
}
|
||||
|
||||
// NewIndex creates an empty inverted index.
|
||||
func NewIndex[ID comparable]() *Index[ID] {
|
||||
return &Index[ID]{
|
||||
postings: make(map[string]map[ID]int),
|
||||
docLen: make(map[ID]int),
|
||||
docTerms: make(map[ID][]string),
|
||||
}
|
||||
}
|
||||
|
||||
// Add indexes text under id. Re-adding an id replaces its previous term
|
||||
// contribution, so the index stays consistent with the latest text for that id.
|
||||
func (ix *Index[ID]) Add(text string, id ID) {
|
||||
terms := tokenize(text)
|
||||
freq := make(map[string]int, len(terms))
|
||||
for _, t := range terms {
|
||||
freq[t]++
|
||||
}
|
||||
|
||||
ix.mu.Lock()
|
||||
defer ix.mu.Unlock()
|
||||
|
||||
if _, exists := ix.docLen[id]; exists {
|
||||
ix.removeDocLocked(id)
|
||||
}
|
||||
for t, f := range freq {
|
||||
if ix.postings[t] == nil {
|
||||
ix.postings[t] = make(map[ID]int)
|
||||
}
|
||||
ix.postings[t][id] = f
|
||||
}
|
||||
ix.docTerms[id] = terms
|
||||
ix.docLen[id] = len(terms)
|
||||
ix.docCount++
|
||||
}
|
||||
|
||||
// Search returns every document matching at least one query term, ranked by
|
||||
// descending relevance score. Ties are broken by id for determinism. An empty
|
||||
// or whitespace-only query returns no hits.
|
||||
func (ix *Index[ID]) Search(query string) []Hit[ID] {
|
||||
terms := tokenize(query)
|
||||
if len(terms) == 0 {
|
||||
return []Hit[ID]{}
|
||||
}
|
||||
|
||||
ix.mu.RLock()
|
||||
defer ix.mu.RUnlock()
|
||||
|
||||
scores := make(map[ID]float64)
|
||||
n := ix.docCount
|
||||
for _, t := range terms {
|
||||
df, ok := ix.postings[t]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
idf := math.Log(1 + float64(n)/float64(len(df)))
|
||||
for id, tf := range df {
|
||||
scores[id] += float64(tf) * idf
|
||||
}
|
||||
}
|
||||
|
||||
hits := make([]Hit[ID], 0, len(scores))
|
||||
for id, s := range scores {
|
||||
hits = append(hits, Hit[ID]{ID: id, Score: s})
|
||||
}
|
||||
sort.Slice(hits, func(i, j int) bool {
|
||||
if hits[i].Score != hits[j].Score {
|
||||
return hits[i].Score > hits[j].Score
|
||||
}
|
||||
return fmt.Sprint(hits[i].ID) < fmt.Sprint(hits[j].ID)
|
||||
})
|
||||
return hits
|
||||
}
|
||||
|
||||
// removeDocLocked removes the previous postings contribution of an existing doc
|
||||
// before it is re-indexed. Callers must hold ix.mu.
|
||||
func (ix *Index[ID]) removeDocLocked(id ID) {
|
||||
for _, t := range ix.docTerms[id] {
|
||||
df := ix.postings[t]
|
||||
if df != nil {
|
||||
delete(df, id)
|
||||
if len(df) == 0 {
|
||||
delete(ix.postings, t)
|
||||
}
|
||||
}
|
||||
}
|
||||
delete(ix.docTerms, id)
|
||||
delete(ix.docLen, id)
|
||||
ix.docCount--
|
||||
}
|
||||
|
||||
// tokenize splits text into lowercased word terms. Non-letter and non-digit
|
||||
// runes act as separators; empty and repeated terms are handled by the caller.
|
||||
func tokenize(s string) []string {
|
||||
return strings.FieldsFunc(strings.ToLower(s), func(r rune) bool {
|
||||
return !unicode.IsLetter(r) && !unicode.IsDigit(r)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package embedding
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestIndex_AddAndSearch_SingleTerm(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ix := NewIndex[int]()
|
||||
ix.Add("go is a programming language", 1)
|
||||
ix.Add("rust is a systems language", 2)
|
||||
ix.Add("python is great", 3)
|
||||
|
||||
hits := ix.Search("language")
|
||||
if len(hits) != 2 {
|
||||
t.Fatalf("Search(\"language\") returned %d hits, want 2", len(hits))
|
||||
}
|
||||
if got := []int{hits[0].ID, hits[1].ID}; !reflect.DeepEqual(got, []int{1, 2}) {
|
||||
t.Errorf("top hits = %v, want [1 2]", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIndex_SearchRelevanceOrdersByFrequency(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ix := NewIndex[int]()
|
||||
ix.Add("cat cat cat dog", 1)
|
||||
ix.Add("dog dog cat", 2)
|
||||
|
||||
hits := ix.Search("cat")
|
||||
if len(hits) != 2 {
|
||||
t.Fatalf("Search(\"cat\") returned %d hits, want 2", len(hits))
|
||||
}
|
||||
// Doc 1 has cat 3 times, doc 2 once, so doc 1 ranks higher.
|
||||
if hits[0].ID != 1 {
|
||||
t.Errorf("top hit = %d, want 1 (higher tf)", hits[0].ID)
|
||||
}
|
||||
if hits[1].ID != 2 {
|
||||
t.Errorf("second hit = %d, want 2", hits[1].ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIndex_SearchOverlapsMultipleTerms(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ix := NewIndex[int]()
|
||||
ix.Add("the red fox", 1)
|
||||
ix.Add("the blue whale", 2)
|
||||
ix.Add("a red car", 3)
|
||||
|
||||
hits := ix.Search("red fox")
|
||||
if len(hits) != 2 {
|
||||
t.Fatalf("Search(\"red fox\") returned %d hits, want 2", len(hits))
|
||||
}
|
||||
// Doc 1 matches both terms (higher), doc 3 only 'red'.
|
||||
if hits[0].ID != 1 || hits[1].ID != 3 {
|
||||
t.Fatalf("hits = %v, want docs [1 3]", hits)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIndex_SearchEmptyQuery(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ix := NewIndex[int]()
|
||||
ix.Add("hello world", 1)
|
||||
|
||||
if got := ix.Search(""); len(got) != 0 {
|
||||
t.Errorf("Search(\"\") returned %d hits, want 0", len(got))
|
||||
}
|
||||
if got := ix.Search(" "); len(got) != 0 {
|
||||
t.Errorf("Search(\" \") returned %d hits, want 0", len(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestIndex_SearchNoMatch(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ix := NewIndex[int]()
|
||||
ix.Add("hello world", 1)
|
||||
|
||||
if got := ix.Search("unrelatedterm"); len(got) != 0 {
|
||||
t.Errorf("Search(unrelatedterm) returned %d hits, want 0", len(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestIndex_ReAddReplacesTerms(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ix := NewIndex[int]()
|
||||
ix.Add("hello world", 1)
|
||||
ix.Add("goodbye world", 1) // re-add same id
|
||||
|
||||
if got := ix.Search("hello"); len(got) != 0 {
|
||||
t.Errorf("Search(\"hello\") after re-add returned %d hits, want 0", len(got))
|
||||
}
|
||||
if got := ix.Search("world"); len(got) != 1 {
|
||||
t.Errorf("Search(\"world\") after re-add returned %d hits, want 1", len(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestIndex_StringIDs(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ix := NewIndex[string]()
|
||||
ix.Add("night sky blue", "moon")
|
||||
ix.Add("ocean blue water", "sea")
|
||||
|
||||
hits := ix.Search("blue")
|
||||
if len(hits) != 2 {
|
||||
t.Fatalf("Search(\"blue\") returned %d hits, want 2", len(hits))
|
||||
}
|
||||
if hits[0].ID != "moon" || hits[1].ID != "sea" {
|
||||
t.Errorf("hits = %v, want [moon sea]", hits)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
package embedding
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
// Manager selects the runtime embedder based on configuration and exposes it
|
||||
// through a single Embedder handle.
|
||||
//
|
||||
// When config.EmbeddingConfig.APIURL is set, the semantic APIEmbedder is
|
||||
// chosen; otherwise the built-in NGramEmbedder is used. In both cases the
|
||||
// caller receives a consistent Embedder that is safe for concurrent use.
|
||||
//
|
||||
// It is the integration point described in DevDocs/design/04-运行时.md §7: the
|
||||
// API path is a real semantic-vector enhancement, while the n-gram path keeps
|
||||
// the runtime zero-dependency and offline.
|
||||
type Manager struct {
|
||||
embedder Embedder
|
||||
}
|
||||
|
||||
// New builds a Manager for cfg and selects the embedder. When cfg.APIURL is set
|
||||
// (ignoring surrounding whitespace), an APIEmbedder is used; otherwise an
|
||||
// NGramEmbedder with the default dimension.
|
||||
func New(cfg config.EmbeddingConfig) *Manager {
|
||||
var embedder Embedder
|
||||
if strings.TrimSpace(cfg.APIURL) != "" {
|
||||
embedder = NewAPI(cfg)
|
||||
} else {
|
||||
embedder = NewNGram(defaultNGramDim)
|
||||
}
|
||||
return &Manager{embedder: embedder}
|
||||
}
|
||||
|
||||
// Embedder returns the selected embedder.
|
||||
func (m *Manager) Embedder() Embedder { return m.embedder }
|
||||
|
||||
// Embed is a convenience pass-through to the selected embedder.
|
||||
func (m *Manager) Embed(ctx context.Context, texts []string) ([]Vector, error) {
|
||||
return m.embedder.Embed(ctx, texts)
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package embedding
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
func TestManager_SelectsAPIWhenConfigured(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cfg := config.EmbeddingConfig{
|
||||
APIURL: "http://localhost:9999/v1/embeddings",
|
||||
APIKey: "secret",
|
||||
Model: "text-embedding-3-small",
|
||||
}
|
||||
m := New(cfg)
|
||||
if m.Embedder() == nil {
|
||||
t.Fatal("Embedder() returned nil")
|
||||
}
|
||||
if _, ok := m.Embedder().(*APIEmbedder); !ok {
|
||||
t.Fatalf("Embedder() = %T, want *APIEmbedder", m.Embedder())
|
||||
}
|
||||
}
|
||||
|
||||
func TestManager_SelectsNGramByDefault(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
m := New(config.EmbeddingConfig{})
|
||||
emb, ok := m.Embedder().(*NGramEmbedder)
|
||||
if !ok {
|
||||
t.Fatalf("Embedder() = %T, want *NGramEmbedder", m.Embedder())
|
||||
}
|
||||
if emb.Dim() != defaultNGramDim {
|
||||
t.Errorf("default ngram dim = %d, want %d", emb.Dim(), defaultNGramDim)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManager_BlankAPIURLSelectsNGram(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cfg := config.EmbeddingConfig{APIURL: " "} // whitespace counts as unset
|
||||
m := New(cfg)
|
||||
if _, ok := m.Embedder().(*NGramEmbedder); !ok {
|
||||
t.Fatalf("Embedder() = %T, want *NGramEmbedder for blank URL", m.Embedder())
|
||||
}
|
||||
}
|
||||
|
||||
func TestManager_EmbedDelegatesToSelectedEmbedder(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
m := New(config.EmbeddingConfig{})
|
||||
vectors, err := m.Embed(context.Background(), []string{"hello"})
|
||||
if err != nil {
|
||||
t.Fatalf("Embed() error: %v", err)
|
||||
}
|
||||
if len(vectors) != 1 {
|
||||
t.Fatalf("got %d vectors, want 1", len(vectors))
|
||||
}
|
||||
if len(vectors[0]) != defaultNGramDim {
|
||||
t.Errorf("vector length = %d, want %d", len(vectors[0]), defaultNGramDim)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,147 @@
|
||||
package embedding
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Default n-gram embedding parameters. These are the zero-configuration values
|
||||
// used when NewNGram is called with a non-positive dimension and no options.
|
||||
const (
|
||||
defaultNGramDim = 256
|
||||
defaultNGramMin = 2
|
||||
defaultNGramMax = 4
|
||||
)
|
||||
|
||||
// FNV-1a constants (Fowler–Noll–Vo, 64-bit). They give a fixed, cross-process-
|
||||
// stable hash so n-gram bucket positions are deterministic across runs and
|
||||
// machines (unlike maphash, which randomizes per process).
|
||||
const (
|
||||
fnv64Offset = 14695981039346656037
|
||||
fnv64Prime = 1099511628211
|
||||
// signBit is the top bit used to derive the feature-hashing sign.
|
||||
signBit = uint64(1) << 63
|
||||
)
|
||||
|
||||
// NGramEmbedder is a zero-dependency feature-hashing embedder. It tokenizes a
|
||||
// text into character n-grams, hashes each n-gram into a fixed-size bucket
|
||||
// using signed FNV-1a, and L2-normalizes the resulting count vector.
|
||||
//
|
||||
// It is the default embedder: deterministic, offline, and dependency-free. It
|
||||
// is the keyword-independent, semantic-level counterpart to the FTS Index.
|
||||
type NGramEmbedder struct {
|
||||
dim int
|
||||
minN int
|
||||
maxN int
|
||||
}
|
||||
|
||||
// NGramOption customizes an NGramEmbedder built by NewNGram.
|
||||
type NGramOption func(*NGramEmbedder)
|
||||
|
||||
// WithNGramRange sets the inclusive character n-gram sizes hashed into the
|
||||
// vector. Non-positive or inverted ranges are ignored and the defaults kept.
|
||||
func WithNGramRange(minN, maxN int) NGramOption {
|
||||
return func(e *NGramEmbedder) {
|
||||
if minN > 0 && maxN >= minN {
|
||||
e.minN, e.maxN = minN, maxN
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// NewNGram builds an NGramEmbedder with dimension dim. dim must be positive;
|
||||
// non-positive values select defaultNGramDim. The n-gram range defaults to
|
||||
// [defaultNGramMin, defaultNGramMax].
|
||||
func NewNGram(dim int, opts ...NGramOption) *NGramEmbedder {
|
||||
if dim <= 0 {
|
||||
dim = defaultNGramDim
|
||||
}
|
||||
e := &NGramEmbedder{dim: dim, minN: defaultNGramMin, maxN: defaultNGramMax}
|
||||
for _, opt := range opts {
|
||||
opt(e)
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
// Dim returns the fixed vector dimension.
|
||||
func (e *NGramEmbedder) Dim() int { return e.dim }
|
||||
|
||||
// Embed returns one L2-normalized n-gram hash vector per input text. The order
|
||||
// of the returned vectors matches texts; an empty input yields an empty slice.
|
||||
func (e *NGramEmbedder) Embed(ctx context.Context, texts []string) ([]Vector, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]Vector, 0, len(texts))
|
||||
for _, text := range texts {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, e.vectorFor(text))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// vectorFor builds the n-gram hash vector for one text and L2-normalizes it.
|
||||
// Empty text produces a zero vector (left as-is by l2Normalize). Text shorter
|
||||
// than the smallest n-gram window hashes the whole text as a single token so it
|
||||
// never collapses to a zero vector.
|
||||
func (e *NGramEmbedder) vectorFor(text string) Vector {
|
||||
v := make(Vector, e.dim)
|
||||
lower := strings.ToLower(strings.TrimSpace(text))
|
||||
runes := []rune(lower)
|
||||
n := len(runes)
|
||||
if n == 0 {
|
||||
return v
|
||||
}
|
||||
hashed := false
|
||||
for size := e.minN; size <= e.maxN && size <= n; size++ {
|
||||
for i := 0; i+size <= n; i++ {
|
||||
h := fnv1a64(string(runes[i : i+size]))
|
||||
v[int(h%uint64(e.dim))] += signedHash(h)
|
||||
hashed = true
|
||||
}
|
||||
}
|
||||
if !hashed {
|
||||
h := fnv1a64(lower)
|
||||
v[int(h%uint64(e.dim))] += signedHash(h)
|
||||
}
|
||||
l2Normalize(v)
|
||||
return v
|
||||
}
|
||||
|
||||
// fnv1a64 returns the FNV-1a 64-bit hash of s. It is deterministic across
|
||||
// processes and platforms, so the same text always lands in the same buckets.
|
||||
func fnv1a64(s string) uint64 {
|
||||
h := uint64(fnv64Offset)
|
||||
for i := 0; i < len(s); i++ {
|
||||
h ^= uint64(s[i])
|
||||
h *= fnv64Prime
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
// signedHash maps the top bit of a hash to the +1/-1 sign used to reduce
|
||||
// collision bias in feature hashing; the lower bits select the bucket.
|
||||
func signedHash(h uint64) float32 {
|
||||
if h&signBit != 0 {
|
||||
return -1
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
// l2Normalize scales v to unit length in place. A zero vector is left as-is
|
||||
// rather than dividing by zero.
|
||||
func l2Normalize(v Vector) {
|
||||
var sum float64
|
||||
for _, x := range v {
|
||||
sum += float64(x) * float64(x)
|
||||
}
|
||||
if sum == 0 {
|
||||
return
|
||||
}
|
||||
norm := float32(math.Sqrt(sum))
|
||||
for i := range v {
|
||||
v[i] /= norm
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,155 @@
|
||||
package embedding
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNGramEmbedder_Dim(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
e := NewNGram(128)
|
||||
if got := e.Dim(); got != 128 {
|
||||
t.Fatalf("Dim() = %d, want 128", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNGramEmbedder_DefaultDim(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
e := NewNGram(0)
|
||||
if got := e.Dim(); got != defaultNGramDim {
|
||||
t.Fatalf("Dim() = %d, want defaultNGramDim (%d)", got, defaultNGramDim)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNGramEmbedder_VectorLengthMatchesDim(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
e := NewNGram(64)
|
||||
vectors, err := e.Embed(context.Background(), []string{"alpha", "beta gamma"})
|
||||
if err != nil {
|
||||
t.Fatalf("Embed() error: %v", err)
|
||||
}
|
||||
if len(vectors) != 2 {
|
||||
t.Fatalf("got %d vectors, want 2", len(vectors))
|
||||
}
|
||||
for i, v := range vectors {
|
||||
if got := len(v); got != e.Dim() {
|
||||
t.Errorf("vector %d length = %d, want %d", i, got, e.Dim())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNGramEmbedder_Deterministic(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
e := NewNGram(128)
|
||||
ctx := context.Background()
|
||||
texts := []string{"hello world", "the quick brown fox jumps", "Golang embedding"}
|
||||
|
||||
first, err := e.Embed(ctx, texts)
|
||||
if err != nil {
|
||||
t.Fatalf("Embed() error: %v", err)
|
||||
}
|
||||
second, err := e.Embed(ctx, texts)
|
||||
if err != nil {
|
||||
t.Fatalf("Embed() error: %v", err)
|
||||
}
|
||||
if len(first) != len(second) {
|
||||
t.Fatalf("got %d vectors on first call, %d on second", len(first), len(second))
|
||||
}
|
||||
for i := range first {
|
||||
if !equalVector(first[i], second[i]) {
|
||||
t.Fatalf("vector %d differs between calls", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNGramEmbedder_Normalized(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
e := NewNGram(128)
|
||||
vectors, err := e.Embed(context.Background(), []string{
|
||||
"a",
|
||||
"the quick brown fox jumps over the lazy dog",
|
||||
"golang fts embedding test",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Embed() error: %v", err)
|
||||
}
|
||||
for i, v := range vectors {
|
||||
norm := l2Norm(v)
|
||||
if math.Abs(norm-1) > 1e-6 {
|
||||
t.Errorf("vector %d L2 norm = %f, want 1", i, norm)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNGramEmbedder_DifferentTextsDiffer(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
e := NewNGram(128)
|
||||
vectors, err := e.Embed(context.Background(), []string{"dog", "cat", "astronaut"})
|
||||
if err != nil {
|
||||
t.Fatalf("Embed() error: %v", err)
|
||||
}
|
||||
if len(vectors) != 3 {
|
||||
t.Fatalf("got %d vectors, want 3", len(vectors))
|
||||
}
|
||||
for i := 0; i < len(vectors); i++ {
|
||||
if l2Norm(vectors[i]) == 0 {
|
||||
t.Errorf("vector %d is a zero vector", i)
|
||||
}
|
||||
for j := i + 1; j < len(vectors); j++ {
|
||||
if equalVector(vectors[i], vectors[j]) {
|
||||
t.Errorf("vectors %d and %d are identical", i, j)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNGramEmbedder_EmptyInput(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
e := NewNGram(128)
|
||||
vectors, err := e.Embed(context.Background(), nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Embed(nil) error: %v", err)
|
||||
}
|
||||
if len(vectors) != 0 {
|
||||
t.Errorf("got %d vectors, want 0", len(vectors))
|
||||
}
|
||||
}
|
||||
|
||||
func TestNGramEmbedder_CanceledContext(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
e := NewNGram(128)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
if _, err := e.Embed(ctx, []string{"hello"}); err == nil {
|
||||
t.Fatal("expected error on canceled context, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func equalVector(a, b Vector) bool {
|
||||
if len(a) != len(b) {
|
||||
return false
|
||||
}
|
||||
for i := range a {
|
||||
if a[i] != b[i] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func l2Norm(v Vector) float64 {
|
||||
var sum float64
|
||||
for _, x := range v {
|
||||
sum += float64(x) * float64(x)
|
||||
}
|
||||
return math.Sqrt(sum)
|
||||
}
|
||||
Reference in new issue
Block a user