thread-master/backend/internal/library/embedding/client.go

// Package embedding provides an optional OpenAI embeddings client.
// It is intentionally not wired into patrol/graph flows; relevance uses
// offline lexical intent scoring in knowledge/semantic_fit.go instead.
package embedding

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"math"
	"net/http"
	"strings"
)

const DefaultModel = "text-embedding-3-small"

type Client struct {
	apiKey string
	model  string
	http   *http.Client
}

func NewClient(apiKey string) *Client {
	return &Client{
		apiKey: strings.TrimSpace(apiKey),
		model:  DefaultModel,
		http:   &http.Client{},
	}
}

func (c *Client) Enabled() bool {
	return c != nil && c.apiKey != ""
}

type EmbeddingRequest struct {
	Input []string `json:"input"`
	Model string   `json:"model"`
}

type EmbeddingResponse struct {
	Data []struct {
		Embedding []float32 `json:"embedding"`
	} `json:"data"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error,omitempty"`
	Usage *struct {
		TotalTokens int `json:"total_tokens"`
	} `json:"usage,omitempty"`
}

func (c *Client) Embed(ctx context.Context, input string) ([]float32, error) {
	if !c.Enabled() {
		return nil, fmt.Errorf("embedding client not configured")
	}
	return c.EmbedBatch(ctx, []string{input})
}

func (c *Client) EmbedBatch(ctx context.Context, inputs []string) ([]float32, error) {
	if !c.Enabled() {
		return nil, fmt.Errorf("embedding client not configured")
	}
	if len(inputs) == 0 {
		return nil, fmt.Errorf("no input text")
	}

	req := EmbeddingRequest{
		Input: inputs,
		Model: c.model,
	}
	body, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://api.openai.com/v1/embeddings", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
	httpReq.Header.Set("Content-Type", "application/json")

	resp, err := c.http.Do(httpReq)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		var errResp struct {
			Error struct {
				Message string `json:"message"`
			} `json:"error"`
		}
		if json.Unmarshal(respBody, &errResp) == nil && errResp.Error.Message != "" {
			return nil, fmt.Errorf("openai embedding: %s", errResp.Error.Message)
		}
		return nil, fmt.Errorf("openai embedding: http %d", resp.StatusCode)
	}

	var result EmbeddingResponse
	if err := json.Unmarshal(respBody, &result); err != nil {
		return nil, err
	}

	if result.Error != nil && result.Error.Message != "" {
		return nil, fmt.Errorf("openai embedding: %s", result.Error.Message)
	}

	if len(result.Data) == 0 {
		return nil, fmt.Errorf("openai embedding: no data returned")
	}

	return result.Data[0].Embedding, nil
}

func CosineSimilarity(a, b []float32) float32 {
	if len(a) == 0 || len(b) == 0 || len(a) != len(b) {
		return 0
	}
	var dot, normA, normB float64
	for i := 0; i < len(a); i++ {
		dot += float64(a[i] * b[i])
		normA += float64(a[i] * a[i])
		normB += float64(b[i] * b[i])
	}
	if normA == 0 || normB == 0 {
		return 0
	}
	return float32(dot / (math.Sqrt(normA) * math.Sqrt(normB)))
}