feat(embedding): 嵌入(n-gram 哈希向量 + 关键词 FTS + OpenAI 兼容 API 钩子 + Manager 选择)

This commit is contained in:
Aska committed 2026-09-06 04:37:10 +08:00
1 parent 70fd915242
commit 51b2bdfba2
9 files changed
+1125

No files matched your search

+236
View File
@@ -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,
}
}
+196
View File
@@ -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")
}
}
+34
View File
@@ -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
}
+133
View File
@@ -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)
})
}
+117
View File
@@ -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)
}
}
+43
View File
@@ -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)
}
}
+147
View File
@@ -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
}
}
+155
View File
@@ -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)
}