106 lines
2.4 KiB
Go
106 lines
2.4 KiB
Go
package embed
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"math"
|
|
"net/http"
|
|
"time"
|
|
)
|
|
|
|
type EmbedClient struct {
|
|
baseURL string
|
|
client *http.Client
|
|
model string
|
|
}
|
|
|
|
type EmbedRequest struct {
|
|
Model string `json:"model"`
|
|
Input []string `json:"input"`
|
|
}
|
|
|
|
type EmbedResponse struct {
|
|
Data []struct {
|
|
Embedding []float32 `json:"embedding"`
|
|
} `json:"data"`
|
|
}
|
|
|
|
func NewEmbedClient(baseURL string) *EmbedClient {
|
|
return &EmbedClient{
|
|
baseURL: baseURL,
|
|
client: &http.Client{Timeout: 30 * time.Second},
|
|
model: "intfloat/e5-small-v2",
|
|
}
|
|
}
|
|
|
|
func (c *EmbedClient) Embed(ctx context.Context, texts []string) ([][]float32, error) {
|
|
if len(texts) == 0 {
|
|
return [][]float32{}, nil
|
|
}
|
|
if len(texts) > 100 {
|
|
return nil, fmt.Errorf("batch too large: %d > 100", len(texts))
|
|
}
|
|
|
|
body := EmbedRequest{
|
|
Model: c.model,
|
|
Input: texts,
|
|
}
|
|
payload, err := json.Marshal(body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("marshal embed request: %w", err)
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, "POST", c.baseURL+"/v1/embeddings", bytes.NewReader(payload))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("embed request: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
resp, err := c.client.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("embed call: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
bodyBytes, _ := io.ReadAll(resp.Body)
|
|
return nil, fmt.Errorf("embed service returned %d: %s", resp.StatusCode, string(bodyBytes))
|
|
}
|
|
|
|
respBody, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("read embed response: %w", err)
|
|
}
|
|
|
|
var er EmbedResponse
|
|
if err := json.Unmarshal(respBody, &er); err != nil {
|
|
return nil, fmt.Errorf("parse embed response: %w", err)
|
|
}
|
|
if len(er.Data) != len(texts) {
|
|
return nil, fmt.Errorf("embed mismatch: expected %d embeddings, got %d", len(texts), len(er.Data))
|
|
}
|
|
|
|
result := make([][]float32, len(er.Data))
|
|
for i, d := range er.Data {
|
|
result[i] = d.Embedding
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
// CosineSimilarity вычисляет косинусную близость между двумя векторами
|
|
func CosineSimilarity(a, b []float32) float32 {
|
|
var dot, normA, normB float64
|
|
for i := range a {
|
|
dot += float64(a[i]) * float64(b[i])
|
|
normA += float64(a[i]) * float64(a[i])
|
|
normB += float64(b[i]) * float64(b[i])
|
|
}
|
|
if normA == 0 || normB == 0 {
|
|
return 0
|
|
}
|
|
return float32(dot / (math.Sqrt(normA) * math.Sqrt(normB)))
|
|
}
|