cb525e6598
All checks were successful
Beta Release / beta (push) Successful in 5m9s
Major additions: - RAG pipeline (indexing, chunking, search) with sidebar upload button - Memory system with CRUD API - Plugins and lessons modules - MCP discovery and MCP server - Advanced skills (auto-create, conditional, improver) - Agent browser/image support, delegate, sessions - File editor with CodeMirror in split panes - Markdown rendering via react-markdown + KaTeX + highlight.js - Raw markdown toggle - PWA manifest + service worker - Extension UI redesign with new design tokens and studio-style chat - Pipeline API for chat streaming - Mobile responsive layout 💘 Generated with Crush Assisted-by: GLM-5.1 via Crush <crush@charm.land>
114 lines
2.4 KiB
Go
114 lines
2.4 KiB
Go
package rag
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
type EmbeddingClient struct {
|
|
apiKey string
|
|
baseURL string
|
|
client *http.Client
|
|
}
|
|
|
|
func NewEmbeddingClient(apiKey, baseURL string) *EmbeddingClient {
|
|
if baseURL == "" {
|
|
baseURL = "https://api.openai.com/v1"
|
|
}
|
|
return &EmbeddingClient{
|
|
apiKey: apiKey,
|
|
baseURL: strings.TrimRight(baseURL, "/"),
|
|
client: &http.Client{Timeout: 30 * time.Second},
|
|
}
|
|
}
|
|
|
|
type embeddingRequest struct {
|
|
Model string `json:"model"`
|
|
Input []string `json:"input"`
|
|
}
|
|
|
|
type embeddingResponse struct {
|
|
Data []struct {
|
|
Embedding []float64 `json:"embedding"`
|
|
Index int `json:"index"`
|
|
} `json:"data"`
|
|
Usage struct {
|
|
TotalTokens int `json:"total_tokens"`
|
|
} `json:"usage"`
|
|
}
|
|
|
|
func (c *EmbeddingClient) Embed(texts []string, model string) ([][]float64, error) {
|
|
if len(texts) == 0 {
|
|
return nil, nil
|
|
}
|
|
if model == "" {
|
|
model = "text-embedding-3-small"
|
|
}
|
|
|
|
body := embeddingRequest{
|
|
Model: model,
|
|
Input: texts,
|
|
}
|
|
|
|
bodyBytes, err := json.Marshal(body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("marshal embedding request: %w", err)
|
|
}
|
|
|
|
url := c.baseURL + "/embeddings"
|
|
req, err := http.NewRequest("POST", url, bytes.NewReader(bodyBytes))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create embedding request: %w", err)
|
|
}
|
|
|
|
req.Header.Set("Content-Type", "application/json")
|
|
if c.apiKey != "" {
|
|
req.Header.Set("Authorization", "Bearer "+c.apiKey)
|
|
}
|
|
|
|
resp, err := c.client.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("send 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("parse embedding response: %w", err)
|
|
}
|
|
|
|
result := make([][]float64, len(texts))
|
|
for _, data := range embResp.Data {
|
|
if data.Index < len(result) {
|
|
result[data.Index] = data.Embedding
|
|
}
|
|
}
|
|
|
|
return result, nil
|
|
}
|
|
|
|
func (c *EmbeddingClient) EmbedSingle(text, model string) ([]float64, error) {
|
|
results, err := c.Embed([]string{text}, model)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(results) == 0 {
|
|
return nil, fmt.Errorf("no embedding returned")
|
|
}
|
|
return results[0], nil
|
|
}
|