first commit

embed.go

@@ -0,0 +1,96 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+	"time"
+)
+
+const defaultModel = "text-embedding-3-small"
+const defaultBaseURL = "https://api.openai.com/v1"
+
+type embedder struct {
+	apiKey     string
+	model      string
+	baseURL    string
+	httpClient *http.Client
+}
+
+func newEmbedder(apiKey, model, baseURL string) *embedder {
+	if model == "" {
+		model = defaultModel
+	}
+	if baseURL == "" {
+		baseURL = defaultBaseURL
+	}
+	return &embedder{
+		apiKey:  apiKey,
+		model:   model,
+		baseURL: strings.TrimRight(baseURL, "/"),
+		httpClient: &http.Client{
+			Timeout: 30 * time.Second,
+		},
+	}
+}
+
+type embeddingRequest struct {
+	Input          string `json:"input"`
+	Model          string `json:"model"`
+	EncodingFormat string `json:"encoding_format,omitempty"`
+}
+
+type embeddingResponse struct {
+	Data []struct {
+		Embedding []float32 `json:"embedding"`
+	} `json:"data"`
+}
+
+func (e *embedder) embedText(ctx context.Context, text string) ([]float32, error) {
+	reqBody := embeddingRequest{
+		Input: text,
+		Model: e.model,
+	}
+
+	body, err := json.Marshal(reqBody)
+	if err != nil {
+		return nil, fmt.Errorf("marshal embedding request: %w", err)
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, e.baseURL+"/embeddings", bytes.NewReader(body))
+	if err != nil {
+		return nil, fmt.Errorf("create embedding request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", "Bearer "+e.apiKey)
+
+	resp, err := e.httpClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("embedding request: %w", err)
+	}
+	defer resp.Body.Close()
+
+	respBody, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("read embedding response: %w", err)
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("embedding API error %d: %s", resp.StatusCode, string(respBody))
+	}
+
+	var embResp embeddingResponse
+	if err := json.Unmarshal(respBody, &embResp); err != nil {
+		return nil, fmt.Errorf("unmarshal embedding response: %w", err)
+	}
+
+	if len(embResp.Data) == 0 {
+		return nil, fmt.Errorf("no embedding returned")
+	}
+
+	return embResp.Data[0].Embedding, nil
+}