Something went wrong. Try again.
(READ ONLY) Margin is an open annotation layer for the internet. Powered by the AT Protocol.
margin.at
extension web atproto comments
Something went wrong. Try again.
4.9 kB · 220 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221package embeddings
import ( "bytes" "encoding/json" "fmt" "io" "net/http" "os" "strings" "sync" "time"
"margin.at/internal/logger")
const ( Model = "text-embedding-3-small" Dimensions = 1536 MaxTokens = 8191 MaxInputChars = 8000 BatchSize = 64 openAIEndpoint = "https://api.openai.com/v1/embeddings")
type Client struct { apiKey string httpClient *http.Client mu sync.Mutex}
type embeddingRequest struct { Model string `json:"model"` Input []string `json:"input"` Dimensions int `json:"dimensions,omitempty"`}
type embeddingResponse struct { Data []embeddingData `json:"data"` Usage struct { TotalTokens int `json:"total_tokens"` } `json:"usage"` Error *struct { Message string `json:"message"` } `json:"error,omitempty"`}
type embeddingData struct { Index int `json:"index"` Embedding []float32 `json:"embedding"`}
func NewClient() *Client { apiKey := os.Getenv("OPENAI_API_KEY") if apiKey == "" { logger.Info("OPENAI_API_KEY not set — embedding generation will be disabled") } return &Client{ apiKey: apiKey, httpClient: &http.Client{ Timeout: 30 * time.Second, }, }}
func (c *Client) IsEnabled() bool { return c.apiKey != ""}
func (c *Client) Embed(text string) ([]float32, error) { results, err := c.EmbedBatch([]string{text}) if err != nil { return nil, err } if len(results) == 0 { return nil, fmt.Errorf("empty embedding response") } return results[0], nil}
func (c *Client) EmbedBatch(texts []string) ([][]float32, error) { if !c.IsEnabled() { return nil, fmt.Errorf("OpenAI API key not configured") }
truncated := make([]string, len(texts)) for i, t := range texts { t = truncateText(t, MaxInputChars) if strings.TrimSpace(t) == "" { t = " " } truncated[i] = t }
results := make([][]float32, len(texts))
for start := 0; start < len(truncated); start += BatchSize { end := start + BatchSize if end > len(truncated) { end = len(truncated) } batch := truncated[start:end]
embeddings, err := c.callAPI(batch) if err != nil { return nil, fmt.Errorf("embedding batch %d-%d failed: %w", start, end, err) }
for _, emb := range embeddings { idx := start + emb.Index if idx < len(results) { results[idx] = emb.Embedding } } }
return results, nil}
func (c *Client) callAPI(inputs []string) ([]embeddingData, error) { reqBody := embeddingRequest{ Model: Model, Input: inputs, }
body, err := json.Marshal(reqBody) if err != nil { return nil, fmt.Errorf("marshal request: %w", err) }
req, err := http.NewRequest("POST", openAIEndpoint, bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("create request: %w", err) } req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", "Bearer "+c.apiKey)
resp, err := c.httpClient.Do(req) if err != nil { return nil, fmt.Errorf("API request: %w", err) } defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body) if err != nil { return nil, fmt.Errorf("read response: %w", err) }
if resp.StatusCode != http.StatusOK { return nil, fmt.Errorf("API returned %d: %s", resp.StatusCode, string(respBody)) }
var result embeddingResponse if err := json.Unmarshal(respBody, &result); err != nil { return nil, fmt.Errorf("unmarshal response: %w", err) }
if result.Error != nil { return nil, fmt.Errorf("API error: %s", result.Error.Message) }
return result.Data, nil}
func truncateText(text string, maxChars int) string { if len(text) <= maxChars { return text } return text[:maxChars]}
func BuildAnnotationText(bodyValue, selectorJSON, targetTitle, tagsJSON *string) string { var parts []string
if selectorJSON != nil && *selectorJSON != "" { var selector struct { Exact string `json:"exact"` Prefix string `json:"prefix"` Suffix string `json:"suffix"` } if err := json.Unmarshal([]byte(*selectorJSON), &selector); err == nil && selector.Exact != "" { parts = append(parts, selector.Exact) } }
if bodyValue != nil && *bodyValue != "" { parts = append(parts, *bodyValue) }
if targetTitle != nil && *targetTitle != "" { parts = append(parts, *targetTitle) }
if tagsJSON != nil && *tagsJSON != "" { var tags []string if err := json.Unmarshal([]byte(*tagsJSON), &tags); err == nil && len(tags) > 0 { parts = append(parts, strings.Join(tags, ", ")) } }
return strings.Join(parts, " | ")}
func BuildDocumentText(title, description, textContent string, tags []string) string { var parts []string
parts = append(parts, title)
if len(tags) > 0 { parts = append(parts, strings.Join(tags, ", ")) }
if textContent != "" { parts = append(parts, textContent) } else if description != "" { parts = append(parts, description) }
return strings.Join(parts, " | ")}