67 lines
1.9 KiB
Go
67 lines
1.9 KiB
Go
package embedding
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"time"
|
|
)
|
|
|
|
type OpenAIProvider struct {
|
|
apiKey, model string
|
|
dims int
|
|
endpoint string
|
|
client *http.Client
|
|
}
|
|
|
|
func NewOpenAIProvider(apiKey, model string, dims int) *OpenAIProvider {
|
|
if model == "" {
|
|
model = "text-embedding-3-small"
|
|
}
|
|
if dims <= 0 {
|
|
dims = 1536
|
|
}
|
|
return &OpenAIProvider{apiKey: apiKey, model: model, dims: dims, endpoint: "https://api.openai.com/v1/embeddings", client: &http.Client{Timeout: 30 * time.Second}}
|
|
}
|
|
func (p *OpenAIProvider) Dimensions() int { return p.dims }
|
|
func (p *OpenAIProvider) Model() string { return p.model }
|
|
func (p *OpenAIProvider) ProviderName() string { return "openai" }
|
|
func (p *OpenAIProvider) Embed(ctx context.Context, input []string) ([]Vector, error) {
|
|
if p.apiKey == "" {
|
|
return nil, fmt.Errorf("OPENAI_API_KEY is required")
|
|
}
|
|
body, _ := json.Marshal(map[string]any{"model": p.model, "input": input, "dimensions": p.dims})
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.endpoint, bytes.NewReader(body))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req.Header.Set("Authorization", "Bearer "+p.apiKey)
|
|
req.Header.Set("Content-Type", "application/json")
|
|
resp, err := p.client.Do(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return nil, fmt.Errorf("openai embeddings request failed: status %d", resp.StatusCode)
|
|
}
|
|
var parsed struct {
|
|
Data []struct {
|
|
Embedding []float32 `json:"embedding"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil {
|
|
return nil, err
|
|
}
|
|
out := make([]Vector, 0, len(parsed.Data))
|
|
for _, item := range parsed.Data {
|
|
if len(item.Embedding) != p.dims {
|
|
return nil, fmt.Errorf("embedding dimension mismatch")
|
|
}
|
|
out = append(out, Vector(item.Embedding))
|
|
}
|
|
return out, nil
|
|
}
|