312 lines
7.0 KiB
Go
312 lines
7.0 KiB
Go
package kb
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
"unicode/utf8"
|
|
)
|
|
|
|
const (
|
|
// EmbeddingDim BGE-M3 标准向量维度
|
|
EmbeddingDim = 1024
|
|
|
|
// embedRequestTokenBudget 单批次 Token 上限
|
|
embedRequestTokenBudget = 7000
|
|
// embedMaxBatchItems 单请求文本条数上限
|
|
embedMaxBatchItems = 16
|
|
// embedMaxItemTokens 单条文本截断上限
|
|
embedMaxItemTokens = 6000
|
|
// embedLogBodyMax 错误截断最大字节
|
|
embedLogBodyMax = 2048
|
|
)
|
|
|
|
// Embedder 向量生成客户端
|
|
type Embedder struct {
|
|
Endpoint string
|
|
APIKey string
|
|
Model string
|
|
HTTP *http.Client
|
|
}
|
|
|
|
// NewEmbedder 创建 Embedder 实例
|
|
func NewEmbedder(endpoint, key, model string) *Embedder {
|
|
return &Embedder{
|
|
Endpoint: normalizeAPIEndpoint(endpoint, "/embeddings"),
|
|
APIKey: key,
|
|
Model: model,
|
|
HTTP: &http.Client{Timeout: 60 * time.Second},
|
|
}
|
|
}
|
|
|
|
// Embed 批量生成文本嵌入向量
|
|
func (e *Embedder) Embed(ctx context.Context, texts []string) ([][]float32, error) {
|
|
if e == nil || e.Endpoint == "" {
|
|
return nil, fmt.Errorf("embedder not configured")
|
|
}
|
|
if len(texts) == 0 {
|
|
return nil, nil
|
|
}
|
|
|
|
out := make([][]float32, len(texts))
|
|
for start := 0; start < len(texts); {
|
|
end := nextEmbedBatchEnd(texts, start, embedRequestTokenBudget, embedMaxBatchItems)
|
|
batch := texts[start:end]
|
|
vecs, err := e.embedBatch(ctx, batch)
|
|
if err != nil {
|
|
if isContextLengthError(err) && len(batch) > 1 {
|
|
mid := start + (end-start+1)/2
|
|
if mid <= start {
|
|
mid = start + 1
|
|
}
|
|
if mid >= end {
|
|
mid = end - 1
|
|
}
|
|
left, lerr := e.Embed(ctx, texts[start:mid])
|
|
if lerr != nil {
|
|
return nil, lerr
|
|
}
|
|
right, rerr := e.Embed(ctx, texts[mid:end])
|
|
if rerr != nil {
|
|
return nil, rerr
|
|
}
|
|
copy(out[start:mid], left)
|
|
copy(out[mid:end], right)
|
|
start = end
|
|
continue
|
|
}
|
|
est := estimateTokensBatch(batch)
|
|
return nil, fmt.Errorf("%w (batch=%d est_tokens≈%d)", err, len(batch), est)
|
|
}
|
|
copy(out[start:end], vecs)
|
|
start = end
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// EmbedOne 单条文本向量化辅助方法
|
|
func (e *Embedder) EmbedOne(ctx context.Context, text string) ([]float32, error) {
|
|
vecs, err := e.Embed(ctx, []string{text})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(vecs) == 0 {
|
|
return nil, fmt.Errorf("empty embedding returned")
|
|
}
|
|
return vecs[0], nil
|
|
}
|
|
|
|
func (e *Embedder) embedBatch(ctx context.Context, texts []string) ([][]float32, error) {
|
|
body := map[string]any{"model": e.Model, "input": texts}
|
|
raw, err := e.post(ctx, e.Endpoint, body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var resp struct {
|
|
Data []struct {
|
|
Embedding []float32 `json:"embedding"`
|
|
Index int `json:"index"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(raw, &resp); err != nil {
|
|
return nil, fmt.Errorf("decode embeddings: %s", formatHTTPBodyError(0, raw))
|
|
}
|
|
out := make([][]float32, len(texts))
|
|
for _, d := range resp.Data {
|
|
if d.Index >= 0 && d.Index < len(out) {
|
|
out[d.Index] = padOrTrim(d.Embedding, EmbeddingDim)
|
|
}
|
|
}
|
|
for i, v := range out {
|
|
if v == nil && i < len(resp.Data) {
|
|
out[i] = padOrTrim(resp.Data[i].Embedding, EmbeddingDim)
|
|
}
|
|
if out[i] == nil {
|
|
return nil, fmt.Errorf("missing embedding %d", i)
|
|
}
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (e *Embedder) post(ctx context.Context, url string, body any) ([]byte, error) {
|
|
return postJSON(ctx, e.HTTP, url, e.APIKey, body)
|
|
}
|
|
|
|
func postJSON(ctx context.Context, httpc *http.Client, url, key string, body any) ([]byte, error) {
|
|
b, _ := json.Marshal(body)
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(b))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
if key != "" {
|
|
req.Header.Set("Authorization", "Bearer "+key)
|
|
}
|
|
resp, err := httpc.Do(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer resp.Body.Close()
|
|
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
|
|
if resp.StatusCode >= 300 {
|
|
return nil, fmt.Errorf("http %s: %s", resp.Status, formatHTTPBodyError(resp.StatusCode, raw))
|
|
}
|
|
return raw, nil
|
|
}
|
|
|
|
func formatHTTPBodyError(status int, raw []byte) string {
|
|
msg := extractAPIErrorMessage(raw)
|
|
if msg != "" {
|
|
return msg
|
|
}
|
|
trim := strings.TrimSpace(string(raw))
|
|
if trim == "" {
|
|
if status > 0 {
|
|
return fmt.Sprintf("empty body (status %d)", status)
|
|
}
|
|
return "empty body"
|
|
}
|
|
if !json.Valid(raw) || strings.HasPrefix(trim, "<") {
|
|
preview := truncateUTF8(trim, 200)
|
|
if status > 0 {
|
|
return fmt.Sprintf("non-JSON response (status %d): %s", status, preview)
|
|
}
|
|
return "non-JSON response: " + preview
|
|
}
|
|
return truncateUTF8(trim, embedLogBodyMax)
|
|
}
|
|
|
|
func extractAPIErrorMessage(raw []byte) string {
|
|
var envelope struct {
|
|
Error struct {
|
|
Message string `json:"message"`
|
|
Type string `json:"type"`
|
|
Code any `json:"code"`
|
|
} `json:"error"`
|
|
Message string `json:"message"`
|
|
}
|
|
if err := json.Unmarshal(raw, &envelope); err != nil {
|
|
return ""
|
|
}
|
|
if m := strings.TrimSpace(envelope.Error.Message); m != "" {
|
|
return m
|
|
}
|
|
return strings.TrimSpace(envelope.Message)
|
|
}
|
|
|
|
func isContextLengthError(err error) bool {
|
|
if err == nil {
|
|
return false
|
|
}
|
|
msg := strings.ToLower(err.Error())
|
|
return strings.Contains(msg, "maximum context length") ||
|
|
strings.Contains(msg, "context length") ||
|
|
strings.Contains(msg, "too many tokens") ||
|
|
(strings.Contains(msg, "input tokens") && strings.Contains(msg, "reduce"))
|
|
}
|
|
|
|
func estimateTokens(s string) int {
|
|
s = strings.TrimSpace(s)
|
|
if s == "" {
|
|
return 1
|
|
}
|
|
tokens := 0
|
|
asciiRun := 0
|
|
flushASCII := func() {
|
|
if asciiRun == 0 {
|
|
return
|
|
}
|
|
tokens += (asciiRun + 3) / 4
|
|
asciiRun = 0
|
|
}
|
|
for _, r := range s {
|
|
if r <= 127 {
|
|
asciiRun++
|
|
continue
|
|
}
|
|
flushASCII()
|
|
tokens += 2
|
|
}
|
|
flushASCII()
|
|
if tokens < 1 {
|
|
return 1
|
|
}
|
|
return tokens
|
|
}
|
|
|
|
func estimateTokensBatch(texts []string) int {
|
|
sum := 0
|
|
for _, t := range texts {
|
|
sum += estimateTokens(t)
|
|
}
|
|
return sum
|
|
}
|
|
|
|
func nextEmbedBatchEnd(texts []string, start, budget, maxItems int) int {
|
|
if maxItems <= 0 {
|
|
maxItems = embedMaxBatchItems
|
|
}
|
|
end := start
|
|
used := 0
|
|
for end < len(texts) && end-start < maxItems {
|
|
t := estimateTokens(texts[end])
|
|
if t > embedMaxItemTokens {
|
|
t = embedMaxItemTokens
|
|
}
|
|
if end > start && used+t > budget {
|
|
break
|
|
}
|
|
used += t
|
|
end++
|
|
if used >= budget {
|
|
break
|
|
}
|
|
}
|
|
if end == start {
|
|
end = start + 1
|
|
}
|
|
return end
|
|
}
|
|
|
|
func normalizeAPIEndpoint(endpoint, suffix string) string {
|
|
endpoint = strings.TrimSpace(endpoint)
|
|
if endpoint == "" {
|
|
return ""
|
|
}
|
|
endpoint = strings.TrimRight(endpoint, "/")
|
|
if strings.HasSuffix(endpoint, suffix) {
|
|
return endpoint
|
|
}
|
|
if strings.HasSuffix(endpoint, "/v1") || strings.Contains(endpoint, "/v1/") {
|
|
return endpoint + suffix
|
|
}
|
|
return endpoint + "/v1" + suffix
|
|
}
|
|
|
|
func padOrTrim(v []float32, n int) []float32 {
|
|
if len(v) == n {
|
|
return v
|
|
}
|
|
out := make([]float32, n)
|
|
copy(out, v)
|
|
return out
|
|
}
|
|
|
|
func truncateUTF8(s string, n int) string {
|
|
if n <= 0 || len(s) <= n {
|
|
return s
|
|
}
|
|
for n > 0 && !utf8.RuneStart(s[n]) {
|
|
n--
|
|
}
|
|
if n <= 0 {
|
|
return ""
|
|
}
|
|
return s[:n]
|
|
}
|