128 lines
4.0 KiB
Go
128 lines
4.0 KiB
Go
package embeddingprovider
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"math"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
|
|
"data-terminal-cloud/internal/developmentembedding"
|
|
"data-terminal-cloud/internal/outbox"
|
|
)
|
|
|
|
const Dimensions = developmentembedding.Dimensions
|
|
|
|
type Result struct {
|
|
Vector []float64
|
|
ProviderRef string
|
|
}
|
|
|
|
type Provider interface {
|
|
Embed(context.Context, string) (Result, error)
|
|
}
|
|
|
|
type Local struct{}
|
|
|
|
func (Local) Embed(_ context.Context, content string) (Result, error) {
|
|
if strings.TrimSpace(content) == "" {
|
|
return Result{}, errors.New("embedding content is required")
|
|
}
|
|
return Result{Vector: developmentembedding.Vector(content), ProviderRef: "local-deterministic-hash-v1"}, nil
|
|
}
|
|
|
|
type Ark struct {
|
|
endpoint string
|
|
apiKey string
|
|
model string
|
|
client *http.Client
|
|
}
|
|
|
|
func NewArk(baseURL, apiKey, model string, timeout time.Duration) (*Ark, error) {
|
|
baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
|
|
apiKey = strings.TrimSpace(apiKey)
|
|
model = strings.TrimSpace(model)
|
|
parsed, err := url.Parse(baseURL)
|
|
if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
|
return nil, errors.New("Ark base URL must be an absolute HTTP(S) URL")
|
|
}
|
|
if apiKey == "" || model == "" || timeout <= 0 {
|
|
return nil, errors.New("Ark API key, embedding model and positive timeout are required")
|
|
}
|
|
return &Ark{endpoint: baseURL + "/embeddings", apiKey: apiKey, model: model, client: &http.Client{Timeout: timeout}}, nil
|
|
}
|
|
|
|
func (a *Ark) Embed(ctx context.Context, content string) (Result, error) {
|
|
content = strings.TrimSpace(content)
|
|
if content == "" {
|
|
return Result{}, errors.New("embedding content is required")
|
|
}
|
|
body, err := json.Marshal(map[string]any{
|
|
"model": a.model, "input": content, "encoding_format": "float", "dimensions": Dimensions,
|
|
})
|
|
if err != nil {
|
|
return Result{}, err
|
|
}
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, a.endpoint, bytes.NewReader(body))
|
|
if err != nil {
|
|
return Result{}, fmt.Errorf("build Ark embedding request: %w", err)
|
|
}
|
|
req.Header.Set("Authorization", "Bearer "+a.apiKey)
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("User-Agent", "voice-platform-memory-embedding/1.0")
|
|
resp, err := a.client.Do(req)
|
|
if err != nil {
|
|
return Result{}, outbox.Retryable(fmt.Errorf("call Ark embedding model: %w", err))
|
|
}
|
|
defer resp.Body.Close()
|
|
responseBody, readErr := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
|
|
if readErr != nil {
|
|
return Result{}, outbox.Retryable(fmt.Errorf("read Ark embedding response: %w", readErr))
|
|
}
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
requestErr := fmt.Errorf("Ark embedding model returned HTTP %d", resp.StatusCode)
|
|
if resp.StatusCode == http.StatusRequestTimeout || resp.StatusCode == http.StatusTooEarly || resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 {
|
|
return Result{}, outbox.Retryable(requestErr)
|
|
}
|
|
return Result{}, requestErr
|
|
}
|
|
var response struct {
|
|
Data []struct {
|
|
Index int `json:"index"`
|
|
Embedding []float64 `json:"embedding"`
|
|
} `json:"data"`
|
|
}
|
|
if err = json.Unmarshal(responseBody, &response); err != nil || len(response.Data) != 1 || response.Data[0].Index != 0 {
|
|
return Result{}, outbox.Retryable(errors.New("Ark embedding response has no usable vector"))
|
|
}
|
|
if err = Validate(response.Data[0].Embedding); err != nil {
|
|
return Result{}, outbox.Retryable(err)
|
|
}
|
|
return Result{Vector: response.Data[0].Embedding, ProviderRef: "ark:" + a.model}, nil
|
|
}
|
|
|
|
func Validate(vector []float64) error {
|
|
if len(vector) != Dimensions {
|
|
return fmt.Errorf("embedding dimensions=%d want=%d", len(vector), Dimensions)
|
|
}
|
|
for _, value := range vector {
|
|
if math.IsNaN(value) || math.IsInf(value, 0) {
|
|
return errors.New("embedding contains a non-finite value")
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func Literal(vector []float64) (string, error) {
|
|
if err := Validate(vector); err != nil {
|
|
return "", err
|
|
}
|
|
return developmentembedding.Literal(vector), nil
|
|
}
|