159 lines
3.8 KiB
Go
159 lines
3.8 KiB
Go
package storage
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"math"
|
||
"net/http"
|
||
"os"
|
||
"sync"
|
||
)
|
||
|
||
// Embedder bge-m3 编码客户端。优先使用本地 vLLM（端口 8000），fallback 到模力方舟 API。
|
||
type Embedder struct {
|
||
endpoint string // 本地 vLLM: http://localhost:8000/v1/embeddings
|
||
modelName string // 模型名，默认 bge-m3，可通过 MODEL_NAME 覆盖
|
||
apiKey string // 模力方舟 API key（fallback 用）
|
||
httpClient *http.Client
|
||
cache sync.Map // string → []float32
|
||
dim int
|
||
}
|
||
|
||
// NewEmbedder 创建编码客户端。endpoint 默认为本地 vLLM。
|
||
func NewEmbedder(endpoint string) *Embedder {
|
||
if endpoint == "" {
|
||
endpoint = os.Getenv("VLLM_ENDPOINT")
|
||
if endpoint == "" {
|
||
endpoint = "http://localhost:8000/v1/embeddings"
|
||
}
|
||
}
|
||
return &Embedder{
|
||
endpoint: endpoint,
|
||
modelName: getModelName(),
|
||
apiKey: os.Getenv("MOLIFANG_API_KEY"),
|
||
httpClient: &http.Client{},
|
||
dim: 1024,
|
||
}
|
||
}
|
||
|
||
// getModelName 获取模型名，优先 MODEL_NAME 环境变量，默认为 bge-m3
|
||
func getModelName() string {
|
||
if m := os.Getenv("MODEL_NAME"); m != "" {
|
||
return m
|
||
}
|
||
return "bge-m3"
|
||
}
|
||
func (e *Embedder) Encode(texts []string) ([][]float32, error) {
|
||
// 检查缓存
|
||
result := make([][]float32, len(texts))
|
||
uncached := make([]int, 0)
|
||
for i, t := range texts {
|
||
if v, ok := e.cache.Load(t); ok {
|
||
result[i] = v.([]float32)
|
||
} else {
|
||
uncached = append(uncached, i)
|
||
}
|
||
}
|
||
if len(uncached) == 0 {
|
||
return result, nil
|
||
}
|
||
|
||
// 收集未缓存的文本
|
||
uncachedTexts := make([]string, len(uncached))
|
||
for idx, i := range uncached {
|
||
uncachedTexts[idx] = texts[i]
|
||
}
|
||
|
||
// 优先本地 vLLM
|
||
vectors, err := e.encodeRemote(e.endpoint, uncachedTexts)
|
||
if err != nil {
|
||
// fallback：模力方舟
|
||
if e.apiKey != "" {
|
||
vectors, err = e.encodeRemote("https://ai.gitee.com/v1/embeddings", uncachedTexts)
|
||
}
|
||
if err != nil {
|
||
return nil, fmt.Errorf("encode: all endpoints failed: %w", err)
|
||
}
|
||
}
|
||
|
||
// L2 归一化 + 缓存
|
||
for idx, i := range uncached {
|
||
normalized := l2Normalize(vectors[idx])
|
||
result[i] = normalized
|
||
e.cache.Store(texts[i], normalized)
|
||
}
|
||
return result, nil
|
||
}
|
||
|
||
// EncodeSingle 编码单条文本。
|
||
func (e *Embedder) EncodeSingle(text string) ([]float32, error) {
|
||
vecs, err := e.Encode([]string{text})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return vecs[0], nil
|
||
}
|
||
|
||
// encodeRemote 调用远端 OpenAI 兼容 embeddings API。
|
||
func (e *Embedder) encodeRemote(endpoint string, texts []string) ([][]float32, error) {
|
||
reqBody := map[string]any{
|
||
"model": e.modelName,
|
||
"input": texts,
|
||
}
|
||
body, _ := json.Marshal(reqBody)
|
||
|
||
req, _ := http.NewRequest("POST", endpoint, bytes.NewReader(body))
|
||
req.Header.Set("Content-Type", "application/json")
|
||
if e.apiKey != "" {
|
||
req.Header.Set("Authorization", "Bearer "+e.apiKey)
|
||
}
|
||
|
||
resp, err := e.httpClient.Do(req)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode != http.StatusOK {
|
||
rbody, _ := io.ReadAll(resp.Body)
|
||
return nil, fmt.Errorf("embedding API: status %d: %s", resp.StatusCode, string(rbody))
|
||
}
|
||
|
||
var result struct {
|
||
Data []struct {
|
||
Embedding []float64 `json:"embedding"`
|
||
} `json:"data"`
|
||
}
|
||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||
return nil, fmt.Errorf("decode embedding response: %w", err)
|
||
}
|
||
|
||
vectors := make([][]float32, len(result.Data))
|
||
for i, d := range result.Data {
|
||
vectors[i] = make([]float32, len(d.Embedding))
|
||
for j, v := range d.Embedding {
|
||
vectors[i][j] = float32(v)
|
||
}
|
||
}
|
||
return vectors, nil
|
||
}
|
||
|
||
// l2Normalize L2 归一化。
|
||
func l2Normalize(v []float32) []float32 {
|
||
var sum float64
|
||
for _, x := range v {
|
||
sum += float64(x) * float64(x)
|
||
}
|
||
norm := float32(math.Sqrt(sum))
|
||
if norm == 0 {
|
||
return v
|
||
}
|
||
result := make([]float32, len(v))
|
||
for i, x := range v {
|
||
result[i] = x / norm
|
||
}
|
||
return result
|
||
}
|