fix(memory): rewrite dedup with n-gram keyword overlap + embedder
- Keyword extraction: 2-3 char sliding window + segment splitting - Dedup: keyword overlap (>25%) OR Jaccard (>0.4) triggers merge - SearchByNGrams: pgvector GIN index on keywords array - Embedder: n-gram FNV hash (zero dep), OpenAI API with auto-fallback - GIN index on keywords column for fast overlap search - All tests pass
This commit is contained in:
Aska committed
1 parent
ba42ea2b45
commit
886fcf2216
4 files changed
+371
-113
No files matched your search
@@ -0,0 +1,193 @@
|
||||
// Package memory — Embedder provides vector embeddings for semantic search.
|
||||
// Uses OpenAI-compatible /embeddings API with automatic fallback to n-gram FNV hashing.
|
||||
package memory
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"hash/fnv"
|
||||
"io"
|
||||
"math"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Embedder generates vector embeddings for text.
|
||||
type Embedder struct {
|
||||
apiURL string
|
||||
apiKey string
|
||||
model string
|
||||
client *http.Client
|
||||
|
||||
available bool // true if the API responded successfully
|
||||
dims int // embedding dimensions (1536 for OpenAI, configured per model)
|
||||
}
|
||||
|
||||
// NGramEmbedDim is the dimension of n-gram FNV hash vectors (fallback mode).
|
||||
const NGramEmbedDim = 256
|
||||
|
||||
// NewEmbedder creates an embedder. If apiURL is empty, falls back to n-gram hashing immediately.
|
||||
func NewEmbedder(apiURL, apiKey, model string) *Embedder {
|
||||
e := &Embedder{
|
||||
apiURL: apiURL,
|
||||
apiKey: apiKey,
|
||||
model: model,
|
||||
dims: NGramEmbedDim,
|
||||
client: &http.Client{Timeout: 15 * time.Second},
|
||||
}
|
||||
if apiURL != "" {
|
||||
e.dims = 1536 // Default for text-embedding-3-small
|
||||
if e.checkAvailability() {
|
||||
e.available = true
|
||||
}
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
// IsAvailable returns true if the embedding API is working.
|
||||
func (e *Embedder) IsAvailable() bool { return e.available }
|
||||
|
||||
// Dims returns the embedding vector dimension.
|
||||
func (e *Embedder) Dims() int { return e.dims }
|
||||
|
||||
// checkAvailability tests whether the embedding API endpoint responds correctly.
|
||||
func (e *Embedder) checkAvailability() bool {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
body := map[string]any{
|
||||
"model": e.model,
|
||||
"input": "ping",
|
||||
}
|
||||
jsonBody, _ := json.Marshal(body)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", e.apiURL+"/embeddings", bytes.NewReader(jsonBody))
|
||||
if err != nil { return false }
|
||||
req.Header.Set("Authorization", "Bearer "+e.apiKey)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := e.client.Do(req)
|
||||
if err != nil { return false }
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != 200 { return false }
|
||||
|
||||
var result struct {
|
||||
Data []struct {
|
||||
Embedding []float64 `json:"embedding"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { return false }
|
||||
if len(result.Data) == 0 || len(result.Data[0].Embedding) == 0 { return false }
|
||||
|
||||
// Detect actual dimension
|
||||
e.dims = len(result.Data[0].Embedding)
|
||||
return true
|
||||
}
|
||||
|
||||
// Embed generates a vector embedding for the given text.
|
||||
// Uses the API if available, otherwise falls back to n-gram FNV hashing.
|
||||
func (e *Embedder) Embed(ctx context.Context, text string) ([]float64, error) {
|
||||
if e.available {
|
||||
return e.apiEmbed(ctx, text)
|
||||
}
|
||||
return ngramFNVEmbed(text, e.dims), nil
|
||||
}
|
||||
|
||||
// EmbedBatch generates embeddings for multiple texts.
|
||||
func (e *Embedder) EmbedBatch(ctx context.Context, texts []string) ([][]float64, error) {
|
||||
if e.available {
|
||||
return e.apiEmbedBatch(ctx, texts)
|
||||
}
|
||||
result := make([][]float64, len(texts))
|
||||
for i, t := range texts {
|
||||
result[i] = ngramFNVEmbed(t, e.dims)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (e *Embedder) apiEmbed(ctx context.Context, text string) ([]float64, error) {
|
||||
body := map[string]any{"model": e.model, "input": text}
|
||||
jsonBody, _ := json.Marshal(body)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", e.apiURL+"/embeddings", bytes.NewReader(jsonBody))
|
||||
if err != nil { return nil, err }
|
||||
req.Header.Set("Authorization", "Bearer "+e.apiKey)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := e.client.Do(req)
|
||||
if err != nil { return nil, fmt.Errorf("embed API: %w", err) }
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
msg, _ := io.ReadAll(resp.Body)
|
||||
return nil, fmt.Errorf("embed API %d: %s", resp.StatusCode, string(msg[:min(len(msg), 200)]))
|
||||
}
|
||||
|
||||
var result struct{ Data []struct{ Embedding []float64 `json:"embedding"` } `json:"data"` }
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { return nil, err }
|
||||
if len(result.Data) == 0 { return nil, fmt.Errorf("empty embedding response") }
|
||||
|
||||
return result.Data[0].Embedding, nil
|
||||
}
|
||||
|
||||
func (e *Embedder) apiEmbedBatch(ctx context.Context, texts []string) ([][]float64, error) {
|
||||
body := map[string]any{"model": e.model, "input": texts}
|
||||
jsonBody, _ := json.Marshal(body)
|
||||
req, _ := http.NewRequestWithContext(ctx, "POST", e.apiURL+"/embeddings", bytes.NewReader(jsonBody))
|
||||
req.Header.Set("Authorization", "Bearer "+e.apiKey)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := e.client.Do(req)
|
||||
if err != nil { return nil, err }
|
||||
defer resp.Body.Close()
|
||||
var result struct{ Data []struct{ Embedding []float64 `json:"embedding"` } `json:"data"` }
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { return nil, err }
|
||||
embeddings := make([][]float64, len(result.Data))
|
||||
for i, d := range result.Data {
|
||||
embeddings[i] = d.Embedding
|
||||
}
|
||||
return embeddings, nil
|
||||
}
|
||||
|
||||
// --- N-gram FNV hash embedding (zero-dependency fallback) ---
|
||||
|
||||
// ngramFNVEmbed generates a pseudo-embedding using n-gram hashing.
|
||||
// Not as accurate as real embeddings but works without any external API.
|
||||
// Based on v1's approach — 2-4 character sliding window FNV hashes projected to N dimensions.
|
||||
func ngramFNVEmbed(text string, dims int) []float64 {
|
||||
if dims <= 0 { dims = NGramEmbedDim }
|
||||
runes := []rune(text)
|
||||
vec := make([]float64, dims)
|
||||
|
||||
// Extract n-grams (n=2, 3, 4) and hash to vector positions
|
||||
for n := 2; n <= 4; n++ {
|
||||
for i := 0; i <= len(runes)-n; i++ {
|
||||
gram := string(runes[i : i+n])
|
||||
h := fnv.New64a()
|
||||
h.Write([]byte(gram))
|
||||
idx := h.Sum64() % uint64(dims)
|
||||
vec[idx] += 1.0
|
||||
}
|
||||
}
|
||||
|
||||
// Also include individual characters for short texts
|
||||
if len(runes) < 6 {
|
||||
for _, r := range runes {
|
||||
h := fnv.New64a()
|
||||
h.Write([]byte(string(r)))
|
||||
idx := h.Sum64() % uint64(dims)
|
||||
vec[idx] += 0.5
|
||||
}
|
||||
}
|
||||
|
||||
// L2 normalize
|
||||
var norm float64
|
||||
for _, v := range vec { norm += v * v }
|
||||
norm = math.Sqrt(norm)
|
||||
if norm > 0 {
|
||||
for i := range vec { vec[i] /= norm }
|
||||
}
|
||||
return vec
|
||||
}
|
||||
Reference in new issue
Block a user