diff --git a/pkg/memory/store/embed_ollama.go b/pkg/memory/store/embed_ollama.go new file mode 100644 index 000000000..e67e5bd2a --- /dev/null +++ b/pkg/memory/store/embed_ollama.go @@ -0,0 +1,134 @@ +package store + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "time" + + "github.com/sipeed/picoclaw/pkg/memory" +) + +const ( + defaultOllamaBase = "http://localhost:11434" + defaultOllamaModel = "nomic-embed-text" +) + +// OllamaEmbedder implements memory.EmbeddingProvider using Ollama's /api/embed endpoint. +type OllamaEmbedder struct { + base string + model string + dims int + client *http.Client +} + +// OllamaEmbedderConfig configures the Ollama embedding provider. +type OllamaEmbedderConfig struct { + Base string // API base URL. Empty uses default (http://localhost:11434). + Model string // Model name. Empty uses "nomic-embed-text". + Dims int // Expected dimensions. Zero auto-detects on first call. +} + +// NewOllamaEmbedder creates an Ollama embedding provider. +func NewOllamaEmbedder(cfg OllamaEmbedderConfig) *OllamaEmbedder { + base := cfg.Base + if base == "" { + base = defaultOllamaBase + } + model := cfg.Model + if model == "" { + model = defaultOllamaModel + } + return &OllamaEmbedder{ + base: base, + model: model, + dims: cfg.Dims, + client: &http.Client{ + Timeout: 60 * time.Second, + }, + } +} + +type ollamaEmbedRequest struct { + Model string `json:"model"` + Input string `json:"input"` +} + +type ollamaEmbedResponse struct { + Embeddings [][]float32 `json:"embeddings"` +} + +func (o *OllamaEmbedder) Embed(ctx context.Context, text string) (memory.Embedding, error) { + vecs, err := o.embedBatch(ctx, []string{text}) + if err != nil { + return nil, err + } + if len(vecs) == 0 { + return nil, fmt.Errorf("ollama returned no embeddings") + } + return vecs[0], nil +} + +func (o *OllamaEmbedder) EmbedBatch(ctx context.Context, texts []string) ([]memory.Embedding, error) { + return o.embedBatch(ctx, texts) +} + +func (o *OllamaEmbedder) embedBatch(ctx context.Context, texts []string) ([]memory.Embedding, error) { + results := make([]memory.Embedding, len(texts)) + for i, text := range texts { + vec, err := o.embedSingle(ctx, text) + if err != nil { + return nil, fmt.Errorf("embed text %d: %w", i, err) + } + results[i] = vec + + if o.dims == 0 && len(vec) > 0 { + o.dims = len(vec) + } + } + return results, nil +} + +func (o *OllamaEmbedder) embedSingle(ctx context.Context, text string) (memory.Embedding, error) { + body, err := json.Marshal(ollamaEmbedRequest{ + Model: o.model, + Input: text, + }) + if err != nil { + return nil, fmt.Errorf("marshal request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.base+"/api/embed", bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + + resp, err := o.client.Do(req) + if err != nil { + return nil, fmt.Errorf("ollama request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + respBody, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("ollama returned %d: %s", resp.StatusCode, string(respBody)) + } + + var result ollamaEmbedResponse + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, fmt.Errorf("decode response: %w", err) + } + + if len(result.Embeddings) == 0 { + return nil, fmt.Errorf("ollama returned empty embeddings array") + } + + return memory.Embedding(result.Embeddings[0]), nil +} + +func (o *OllamaEmbedder) Dimensions() int { return o.dims } +func (o *OllamaEmbedder) Model() string { return o.model } diff --git a/pkg/memory/store/embed_openai.go b/pkg/memory/store/embed_openai.go new file mode 100644 index 000000000..96d59e927 --- /dev/null +++ b/pkg/memory/store/embed_openai.go @@ -0,0 +1,152 @@ +package store + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "time" + + "github.com/sipeed/picoclaw/pkg/memory" +) + +const ( + defaultOpenAIBase = "https://api.openai.com/v1" + defaultOpenAIModel = "text-embedding-3-small" +) + +// OpenAIEmbedder implements memory.EmbeddingProvider using the OpenAI /v1/embeddings API. +// Compatible with OpenAI, Azure OpenAI, OpenRouter, and any OpenAI-compatible endpoint. +type OpenAIEmbedder struct { + base string + model string + apiKey string + dims int + client *http.Client +} + +// OpenAIEmbedderConfig configures the OpenAI embedding provider. +type OpenAIEmbedderConfig struct { + Base string // API base URL. Empty uses "https://api.openai.com/v1". + Model string // Model name. Empty uses "text-embedding-3-small". + APIKey string // Required API key. + Dims int // Expected dimensions. Zero auto-detects on first call. +} + +// NewOpenAIEmbedder creates an OpenAI-compatible embedding provider. +func NewOpenAIEmbedder(cfg OpenAIEmbedderConfig) *OpenAIEmbedder { + base := cfg.Base + if base == "" { + base = defaultOpenAIBase + } + model := cfg.Model + if model == "" { + model = defaultOpenAIModel + } + return &OpenAIEmbedder{ + base: base, + model: model, + apiKey: cfg.APIKey, + dims: cfg.Dims, + client: &http.Client{ + Timeout: 60 * time.Second, + }, + } +} + +type openAIEmbedRequest struct { + Input interface{} `json:"input"` // string or []string + Model string `json:"model"` + EncodingFormat string `json:"encoding_format,omitempty"` +} + +type openAIEmbedResponse struct { + Data []openAIEmbedData `json:"data"` + Usage openAIEmbedUsage `json:"usage"` +} + +type openAIEmbedData struct { + Index int `json:"index"` + Embedding []float32 `json:"embedding"` +} + +type openAIEmbedUsage struct { + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` +} + +func (o *OpenAIEmbedder) Embed(ctx context.Context, text string) (memory.Embedding, error) { + vecs, err := o.call(ctx, []string{text}) + if err != nil { + return nil, err + } + if len(vecs) == 0 { + return nil, fmt.Errorf("openai returned no embeddings") + } + return vecs[0], nil +} + +func (o *OpenAIEmbedder) EmbedBatch(ctx context.Context, texts []string) ([]memory.Embedding, error) { + return o.call(ctx, texts) +} + +func (o *OpenAIEmbedder) call(ctx context.Context, texts []string) ([]memory.Embedding, error) { + var input interface{} + if len(texts) == 1 { + input = texts[0] + } else { + input = texts + } + + body, err := json.Marshal(openAIEmbedRequest{ + Input: input, + Model: o.model, + EncodingFormat: "float", + }) + if err != nil { + return nil, fmt.Errorf("marshal request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.base+"/embeddings", bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + if o.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+o.apiKey) + } + + resp, err := o.client.Do(req) + if err != nil { + return nil, fmt.Errorf("openai request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + respBody, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("openai returned %d: %s", resp.StatusCode, string(respBody)) + } + + var result openAIEmbedResponse + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, fmt.Errorf("decode response: %w", err) + } + + embeddings := make([]memory.Embedding, len(result.Data)) + for _, d := range result.Data { + if d.Index < len(embeddings) { + embeddings[d.Index] = memory.Embedding(d.Embedding) + } + } + + if o.dims == 0 && len(embeddings) > 0 && len(embeddings[0]) > 0 { + o.dims = len(embeddings[0]) + } + + return embeddings, nil +} + +func (o *OpenAIEmbedder) Dimensions() int { return o.dims } +func (o *OpenAIEmbedder) Model() string { return o.model }