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

2026-09-06 04:37:10 +08:00
1 parent 70fd915242
commit 51b2bdfba2
@@ -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")
+	}
+}