244 lines
7.2 KiB
Go
244 lines
7.2 KiB
Go
|
|
package portfolio
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"os"
|
||
|
|
"path/filepath"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/VictorVargas/rony-chat-bot/internal/embed"
|
||
|
|
)
|
||
|
|
|
||
|
|
// stubEmbedder returns a fixed vector per exact text, so a test can decide
|
||
|
|
// which chunk a query should land on without running a model.
|
||
|
|
type stubEmbedder struct {
|
||
|
|
vecs map[string][]float32
|
||
|
|
err error
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *stubEmbedder) Embed(_ context.Context, texts []string) ([][]float32, error) {
|
||
|
|
if s.err != nil {
|
||
|
|
return nil, s.err
|
||
|
|
}
|
||
|
|
out := make([][]float32, len(texts))
|
||
|
|
for i, t := range texts {
|
||
|
|
if v, ok := s.vecs[t]; ok {
|
||
|
|
out[i] = embed.Normalize(v)
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
out[i] = embed.Normalize([]float32{1, 1, 1}) // neutral
|
||
|
|
}
|
||
|
|
return out, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func newTestStore(t *testing.T) *Store {
|
||
|
|
t.Helper()
|
||
|
|
dir := t.TempDir()
|
||
|
|
src := filepath.Join(dir, "projects")
|
||
|
|
if err := os.MkdirAll(src, 0o755); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
docs := map[string]string{
|
||
|
|
"tienda.md": "# Tienda\n\n## Tech stack\n\nNext.js and Stripe for payments.\n",
|
||
|
|
"dash.md": "# Dashboard\n\n## Tech stack\n\nGo backend with WebSockets and D3.\n",
|
||
|
|
}
|
||
|
|
for name, body := range docs {
|
||
|
|
if err := os.WriteFile(filepath.Join(src, name), []byte(body), 0o644); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
store, err := OpenStore(filepath.Join(dir, "t.db"))
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
t.Cleanup(func() { store.Close() })
|
||
|
|
if _, _, err := store.Reindex(context.Background(), SourcesFor(src, ""), DefaultChunkerConfig()); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
return store
|
||
|
|
}
|
||
|
|
|
||
|
|
// The point of the whole feature: a query whose words appear nowhere in the
|
||
|
|
// corpus still finds the right chunk, because the vector matches.
|
||
|
|
func TestHybridSearchFindsChunkKeywordSearchCannot(t *testing.T) {
|
||
|
|
store := newTestStore(t)
|
||
|
|
|
||
|
|
// "¿Cómo se paga?" shares no word with "Next.js and Stripe for payments".
|
||
|
|
query := "¿Cómo se paga?"
|
||
|
|
target := []float32{1, 0, 0}
|
||
|
|
emb := &stubEmbedder{vecs: map[string][]float32{query: target}}
|
||
|
|
|
||
|
|
// Embed the corpus so the Stripe chunk owns the target direction.
|
||
|
|
corpus := &stubEmbedder{vecs: map[string][]float32{}}
|
||
|
|
rows, err := store.db.Query(`SELECT id, project_id, section, content FROM portfolio_chunks`)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
for rows.Next() {
|
||
|
|
var id, project, section, content string
|
||
|
|
if err := rows.Scan(&id, &project, §ion, &content); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
v := []float32{0, 1, 0}
|
||
|
|
if project == "tienda" && section == "Tech stack" {
|
||
|
|
v = target
|
||
|
|
}
|
||
|
|
corpus.vecs[project+" — "+section+"\n"+content] = v
|
||
|
|
}
|
||
|
|
rows.Close()
|
||
|
|
|
||
|
|
if _, err := store.EmbedChunks(context.Background(), corpus); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Keyword search alone finds nothing useful for this query.
|
||
|
|
kw, err := store.Search(context.Background(), query, 5)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
for _, h := range kw {
|
||
|
|
if h.ProjectID == "tienda" && h.Section == "Tech stack" {
|
||
|
|
t.Skip("keyword search already found it; this corpus can't demonstrate the gap")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
hits, err := store.HybridSearch(context.Background(), emb, query, 3)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if len(hits) == 0 {
|
||
|
|
t.Fatal("hybrid search returned nothing")
|
||
|
|
}
|
||
|
|
if hits[0].ProjectID != "tienda" || hits[0].Section != "Tech stack" {
|
||
|
|
t.Errorf("top hit = %s/%s, want tienda/Tech stack", hits[0].ProjectID, hits[0].Section)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// A dead embeddings endpoint must degrade to keyword search, not fail the
|
||
|
|
// visitor's question.
|
||
|
|
func TestHybridSearchFallsBackWhenEmbedderFails(t *testing.T) {
|
||
|
|
store := newTestStore(t)
|
||
|
|
broken := &stubEmbedder{err: errors.New("connection refused")}
|
||
|
|
|
||
|
|
hits, err := store.HybridSearch(context.Background(), broken, "Stripe", 3)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("hybrid search should fall back, got error: %v", err)
|
||
|
|
}
|
||
|
|
if len(hits) == 0 {
|
||
|
|
t.Fatal("fallback returned no keyword hits for a term that is in the corpus")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// With no embedder configured at all, hybrid search is plain keyword search.
|
||
|
|
func TestHybridSearchWithoutEmbedderIsKeywordSearch(t *testing.T) {
|
||
|
|
store := newTestStore(t)
|
||
|
|
hits, err := store.HybridSearch(context.Background(), nil, "Stripe", 3)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if len(hits) == 0 {
|
||
|
|
t.Fatal("expected keyword hits for Stripe")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Vectors from a different embedding model (wrong dimension) must be ignored
|
||
|
|
// rather than compared against and ranked.
|
||
|
|
func TestVectorSearchIgnoresMismatchedDimensions(t *testing.T) {
|
||
|
|
store := newTestStore(t)
|
||
|
|
corpus := &stubEmbedder{} // neutral 3-dim vectors
|
||
|
|
if _, err := store.EmbedChunks(context.Background(), corpus); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
hits, err := store.VectorSearch(context.Background(), []float32{1, 0, 0, 0}, 5)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if len(hits) != 0 {
|
||
|
|
t.Errorf("got %d hits for a 4-dim query against 3-dim rows, want 0", len(hits))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// A chunk's id is derived from file + heading + position, so editing the body
|
||
|
|
// of a section keeps the id. A vector keyed only by id would then rank that
|
||
|
|
// chunk by text that no longer exists — silently, with no error anywhere.
|
||
|
|
// The stored content hash makes such a row invisible instead of wrong.
|
||
|
|
func TestVectorSearchIgnoresVectorsWhoseChunkChanged(t *testing.T) {
|
||
|
|
dir := t.TempDir()
|
||
|
|
src := filepath.Join(dir, "projects")
|
||
|
|
if err := os.MkdirAll(src, 0o755); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
file := filepath.Join(src, "demo.md")
|
||
|
|
write := func(body string) {
|
||
|
|
if err := os.WriteFile(file, []byte(body), 0o644); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
write("# Demo\n\n## Tech stack\n\nThe payment provider is Stripe.\n")
|
||
|
|
|
||
|
|
store, err := OpenStore(filepath.Join(dir, "t.db"))
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
defer store.Close()
|
||
|
|
sources := SourcesFor(src, "")
|
||
|
|
if _, _, err := store.Reindex(context.Background(), sources, DefaultChunkerConfig()); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if _, err := store.EmbedChunks(context.Background(), &stubEmbedder{}); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
hits, err := store.VectorSearch(context.Background(), embed.Normalize([]float32{1, 1, 1}), 5)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if len(hits) == 0 {
|
||
|
|
t.Fatal("expected the freshly embedded chunk to be searchable")
|
||
|
|
}
|
||
|
|
|
||
|
|
// Rewrite the body, keeping every heading — the chunk id is unchanged.
|
||
|
|
write("# Demo\n\n## Tech stack\n\nThe payment provider is PayPal. Stripe was removed.\n")
|
||
|
|
if _, _, err := store.Reindex(context.Background(), sources, DefaultChunkerConfig()); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
// Deliberately do NOT re-embed: this is the "endpoint was down" case.
|
||
|
|
|
||
|
|
var ids []string
|
||
|
|
rows, err := store.db.Query(`SELECT chunk_id FROM portfolio_vectors`)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
for rows.Next() {
|
||
|
|
var id string
|
||
|
|
if err := rows.Scan(&id); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
ids = append(ids, id)
|
||
|
|
}
|
||
|
|
rows.Close()
|
||
|
|
if len(ids) == 0 {
|
||
|
|
t.Fatal("precondition: the stale vector row should still be present")
|
||
|
|
}
|
||
|
|
|
||
|
|
hits, err = store.VectorSearch(context.Background(), embed.Normalize([]float32{1, 1, 1}), 5)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if len(hits) != 0 {
|
||
|
|
t.Errorf("stale vector was used for ranking: %+v", hits)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Re-embedding restores it.
|
||
|
|
if _, err := store.EmbedChunks(context.Background(), &stubEmbedder{}); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
hits, err = store.VectorSearch(context.Background(), embed.Normalize([]float32{1, 1, 1}), 5)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if len(hits) == 0 {
|
||
|
|
t.Error("re-embedding should make the chunk searchable again")
|
||
|
|
}
|
||
|
|
}
|