From 51b2bdfba2232dfe1c90646d56b64e52c63c3e18 Mon Sep 17 00:00:00 2001 From: Aska Date: Sun, 6 Sep 2026 04:37:10 +0800 Subject: [PATCH] =?UTF-8?q?feat(embedding):=20=E5=B5=8C=E5=85=A5=EF=BC=88n?= =?UTF-8?q?-gram=20=E5=93=88=E5=B8=8C=E5=90=91=E9=87=8F=20+=20=E5=85=B3?= =?UTF-8?q?=E9=94=AE=E8=AF=8D=20FTS=20+=20OpenAI=20=E5=85=BC=E5=AE=B9=20AP?= =?UTF-8?q?I=20=E9=92=A9=E5=AD=90=20+=20Manager=20=E9=80=89=E6=8B=A9?= =?UTF-8?q?=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/runtime/embedding/api.go | 236 +++++++++++++++++++++ internal/runtime/embedding/api_test.go | 196 +++++++++++++++++ internal/runtime/embedding/embedding.go | 34 +++ internal/runtime/embedding/fts.go | 133 ++++++++++++ internal/runtime/embedding/fts_test.go | 117 ++++++++++ internal/runtime/embedding/manager.go | 43 ++++ internal/runtime/embedding/manager_test.go | 64 ++++++ internal/runtime/embedding/ngram.go | 147 +++++++++++++ internal/runtime/embedding/ngram_test.go | 155 ++++++++++++++ 9 files changed, 1125 insertions(+) create mode 100644 internal/runtime/embedding/api.go create mode 100644 internal/runtime/embedding/api_test.go create mode 100644 internal/runtime/embedding/embedding.go create mode 100644 internal/runtime/embedding/fts.go create mode 100644 internal/runtime/embedding/fts_test.go create mode 100644 internal/runtime/embedding/manager.go create mode 100644 internal/runtime/embedding/manager_test.go create mode 100644 internal/runtime/embedding/ngram.go create mode 100644 internal/runtime/embedding/ngram_test.go diff --git a/internal/runtime/embedding/api.go b/internal/runtime/embedding/api.go new file mode 100644 index 0000000..53e8caf --- /dev/null +++ b/internal/runtime/embedding/api.go @@ -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, + } +} diff --git a/internal/runtime/embedding/api_test.go b/internal/runtime/embedding/api_test.go new file mode 100644 index 0000000..f288a13 --- /dev/null +++ b/internal/runtime/embedding/api_test.go @@ -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") + } +} diff --git a/internal/runtime/embedding/embedding.go b/internal/runtime/embedding/embedding.go new file mode 100644 index 0000000..24f5056 --- /dev/null +++ b/internal/runtime/embedding/embedding.go @@ -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 +} diff --git a/internal/runtime/embedding/fts.go b/internal/runtime/embedding/fts.go new file mode 100644 index 0000000..00b0f63 --- /dev/null +++ b/internal/runtime/embedding/fts.go @@ -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) + }) +} diff --git a/internal/runtime/embedding/fts_test.go b/internal/runtime/embedding/fts_test.go new file mode 100644 index 0000000..c3aed28 --- /dev/null +++ b/internal/runtime/embedding/fts_test.go @@ -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) + } +} diff --git a/internal/runtime/embedding/manager.go b/internal/runtime/embedding/manager.go new file mode 100644 index 0000000..99f1fd2 --- /dev/null +++ b/internal/runtime/embedding/manager.go @@ -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) +} diff --git a/internal/runtime/embedding/manager_test.go b/internal/runtime/embedding/manager_test.go new file mode 100644 index 0000000..e7d3200 --- /dev/null +++ b/internal/runtime/embedding/manager_test.go @@ -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) + } +} diff --git a/internal/runtime/embedding/ngram.go b/internal/runtime/embedding/ngram.go new file mode 100644 index 0000000..552f58c --- /dev/null +++ b/internal/runtime/embedding/ngram.go @@ -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 + } +} diff --git a/internal/runtime/embedding/ngram_test.go b/internal/runtime/embedding/ngram_test.go new file mode 100644 index 0000000..af71465 --- /dev/null +++ b/internal/runtime/embedding/ngram_test.go @@ -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) +}