Skip to content

Commit 75e26df

Browse files
Siddhant-K-codeona-agent
Siddhant-K-code
and
ona-agent
authored
feat(embedding): add Ollama tests and multi-provider support (#83)
- Add 15 unit tests for the Ollama embedding client (mock HTTP server) - Add --embedding-provider flag to api, serve, and memory commands - Use the embedding registry instead of hardcoding OpenAI in all commands - Accept ollama and cohere as valid providers in config validation - Add base_url field to EmbeddingConfig for custom endpoints Closes #25 Co-authored-by: Ona <no-reply@ona.com>
1 parent 4e70441 commit 75e26df

5 files changed

Lines changed: 380 additions & 51 deletions

File tree

‎cmd/api.go‎

Lines changed: 35 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,10 @@ import (
1313

1414
distillcache "github.com/Siddhant-K-code/distill/pkg/cache"
1515
"github.com/Siddhant-K-code/distill/pkg/contextlab"
16-
"github.com/Siddhant-K-code/distill/pkg/embedding/openai"
16+
"github.com/Siddhant-K-code/distill/pkg/embedding"
17+
_ "github.com/Siddhant-K-code/distill/pkg/embedding/cohere"
18+
_ "github.com/Siddhant-K-code/distill/pkg/embedding/ollama"
19+
_ "github.com/Siddhant-K-code/distill/pkg/embedding/openai"
1720
"github.com/Siddhant-K-code/distill/pkg/metrics"
1821
"github.com/Siddhant-K-code/distill/pkg/sse"
1922
"github.com/Siddhant-K-code/distill/pkg/telemetry"
@@ -44,7 +47,9 @@ func init() {
4447
apiCmd.Flags().IntP("port", "p", 8080, "HTTP server port")
4548
apiCmd.Flags().String("host", "0.0.0.0", "HTTP server host")
4649
apiCmd.Flags().String("openai-key", "", "OpenAI API key for embeddings (or use OPENAI_API_KEY)")
47-
apiCmd.Flags().String("embedding-model", "text-embedding-3-small", "OpenAI embedding model")
50+
apiCmd.Flags().String("embedding-provider", "openai", "Embedding provider (openai, ollama, cohere)")
51+
apiCmd.Flags().String("embedding-model", "text-embedding-3-small", "Embedding model name")
52+
apiCmd.Flags().String("embedding-base-url", "", "Embedding provider base URL (e.g. http://localhost:11434 for Ollama)")
4853
apiCmd.Flags().String("api-keys", "", "Comma-separated list of valid API keys (or use DISTILL_API_KEYS)")
4954
apiCmd.Flags().Bool("memory", false, "Enable persistent memory store")
5055
apiCmd.Flags().Bool("session", false, "Enable session management")
@@ -53,7 +58,9 @@ func init() {
5358
// Bind to viper for config file support
5459
_ = viper.BindPFlag("server.port", apiCmd.Flags().Lookup("port"))
5560
_ = viper.BindPFlag("server.host", apiCmd.Flags().Lookup("host"))
61+
_ = viper.BindPFlag("embedding.provider", apiCmd.Flags().Lookup("embedding-provider"))
5662
_ = viper.BindPFlag("embedding.model", apiCmd.Flags().Lookup("embedding-model"))
63+
_ = viper.BindPFlag("embedding.base_url", apiCmd.Flags().Lookup("embedding-base-url"))
5764
}
5865

5966
// DedupeRequest is the JSON request body for /v1/dedupe.
@@ -117,7 +124,7 @@ type DedupeStats struct {
117124

118125
// APIServer holds the API server state.
119126
type APIServer struct {
120-
embedder *openai.Client
127+
embedder embedding.Provider
121128
validKeys map[string]bool
122129
hasAuth bool
123130
metrics *metrics.Metrics
@@ -132,6 +139,12 @@ func runAPI(cmd *cobra.Command, args []string) error {
132139
embeddingModel := viper.GetString("embedding.model")
133140
apiKeysStr, _ := cmd.Flags().GetString("api-keys")
134141

142+
embeddingProvider := viper.GetString("embedding.provider")
143+
embeddingBaseURL, _ := cmd.Flags().GetString("embedding-base-url")
144+
if embeddingBaseURL == "" {
145+
embeddingBaseURL = viper.GetString("embedding.base_url")
146+
}
147+
135148
// Resolve from environment
136149
if openaiKey == "" {
137150
openaiKey = os.Getenv("OPENAI_API_KEY")
@@ -151,13 +164,26 @@ func runAPI(cmd *cobra.Command, args []string) error {
151164
}
152165
}
153166

154-
// Create embedding provider if OpenAI key is provided
155-
var embedder *openai.Client
156-
if openaiKey != "" {
167+
// Create embedding provider via registry
168+
var embedder embedding.Provider
169+
needsAPIKey := embeddingProvider == "" || embeddingProvider == "openai" || embeddingProvider == "cohere"
170+
if needsAPIKey && openaiKey == "" {
171+
// No API key and cloud provider selected — embeddings disabled
172+
} else {
173+
apiKey := openaiKey
174+
if embeddingProvider == "cohere" && apiKey == "" {
175+
apiKey = os.Getenv("COHERE_API_KEY")
176+
}
177+
if embeddingProvider == "" {
178+
embeddingProvider = "openai"
179+
}
157180
var err error
158-
embedder, err = openai.NewClient(openai.Config{
159-
APIKey: openaiKey,
160-
Model: embeddingModel,
181+
embedder, err = embedding.NewProvider(embedding.ProviderConfig{
182+
Type: embedding.ProviderType(embeddingProvider),
183+
APIKey: apiKey,
184+
Model: embeddingModel,
185+
BaseURL: embeddingBaseURL,
186+
CacheSize: -1, // caching handled at a higher layer
161187
})
162188
if err != nil {
163189
return fmt.Errorf("failed to create embedding provider: %w", err)

‎cmd/memory.go‎

Lines changed: 61 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,10 @@ import (
66
"fmt"
77
"os"
88

9-
"github.com/Siddhant-K-code/distill/pkg/embedding/openai"
9+
"github.com/Siddhant-K-code/distill/pkg/embedding"
10+
_ "github.com/Siddhant-K-code/distill/pkg/embedding/cohere"
11+
_ "github.com/Siddhant-K-code/distill/pkg/embedding/ollama"
12+
_ "github.com/Siddhant-K-code/distill/pkg/embedding/openai"
1013
"github.com/Siddhant-K-code/distill/pkg/memory"
1114
"github.com/spf13/cobra"
1215
"github.com/spf13/viper"
@@ -67,15 +70,17 @@ func init() {
6770
memoryStoreCmd.Flags().String("source", "", "Source of the memory (e.g., code_review, docs)")
6871
memoryStoreCmd.Flags().StringSlice("tags", nil, "Tags for the memory")
6972
memoryStoreCmd.Flags().String("session-id", "", "Session ID")
70-
memoryStoreCmd.Flags().String("openai-key", "", "OpenAI API key for embeddings (or OPENAI_API_KEY)")
73+
memoryStoreCmd.Flags().String("openai-key", "", "API key for embeddings (or OPENAI_API_KEY / COHERE_API_KEY)")
74+
memoryStoreCmd.Flags().String("embedding-provider", "", "Embedding provider (openai, ollama, cohere)")
7175

7276
// Recall flags
7377
memoryRecallCmd.Flags().String("query", "", "Query text")
7478
memoryRecallCmd.Flags().StringSlice("tags", nil, "Filter by tags")
7579
memoryRecallCmd.Flags().Int("max-results", 10, "Maximum results to return")
7680
memoryRecallCmd.Flags().Int("max-tokens", 0, "Maximum token budget (0 = unlimited)")
7781
memoryRecallCmd.Flags().Float64("recency-weight", 0.3, "Weight for recency vs relevance (0-1)")
78-
memoryRecallCmd.Flags().String("openai-key", "", "OpenAI API key for embeddings (or OPENAI_API_KEY)")
82+
memoryRecallCmd.Flags().String("openai-key", "", "API key for embeddings (or OPENAI_API_KEY / COHERE_API_KEY)")
83+
memoryRecallCmd.Flags().String("embedding-provider", "", "Embedding provider (openai, ollama, cohere)")
7984

8085
// Forget flags
8186
memoryForgetCmd.Flags().StringSlice("tags", nil, "Remove memories with these tags")
@@ -92,6 +97,47 @@ func openMemoryStore(cmd *cobra.Command) (*memory.SQLiteStore, error) {
9297
return memory.NewSQLiteStore(dbPath, cfg)
9398
}
9499

100+
// createEmbedder builds an embedding.Provider from CLI flags and config.
101+
func createEmbedder(cmd *cobra.Command) (embedding.Provider, error) {
102+
apiKey, _ := cmd.Flags().GetString("openai-key")
103+
if apiKey == "" {
104+
apiKey = os.Getenv("OPENAI_API_KEY")
105+
}
106+
107+
providerName, _ := cmd.Flags().GetString("embedding-provider")
108+
if providerName == "" {
109+
providerName = viper.GetString("embedding.provider")
110+
}
111+
if providerName == "" {
112+
providerName = "openai"
113+
}
114+
115+
// Ollama doesn't need an API key
116+
needsKey := providerName == "openai" || providerName == "cohere"
117+
if needsKey && apiKey == "" {
118+
if providerName == "cohere" {
119+
apiKey = os.Getenv("COHERE_API_KEY")
120+
}
121+
if apiKey == "" {
122+
return nil, nil // no key available, skip embedding
123+
}
124+
}
125+
126+
model := viper.GetString("embedding.model")
127+
if model == "" {
128+
model = "text-embedding-3-small"
129+
}
130+
baseURL := viper.GetString("embedding.base_url")
131+
132+
return embedding.NewProvider(embedding.ProviderConfig{
133+
Type: embedding.ProviderType(providerName),
134+
APIKey: apiKey,
135+
Model: model,
136+
BaseURL: baseURL,
137+
CacheSize: -1,
138+
})
139+
}
140+
95141
func runMemoryStore(cmd *cobra.Command, args []string) error {
96142
text, _ := cmd.Flags().GetString("text")
97143
if text == "" {
@@ -101,10 +147,6 @@ func runMemoryStore(cmd *cobra.Command, args []string) error {
101147
source, _ := cmd.Flags().GetString("source")
102148
tags, _ := cmd.Flags().GetStringSlice("tags")
103149
sessionID, _ := cmd.Flags().GetString("session-id")
104-
openaiKey, _ := cmd.Flags().GetString("openai-key")
105-
if openaiKey == "" {
106-
openaiKey = os.Getenv("OPENAI_API_KEY")
107-
}
108150

109151
store, err := openMemoryStore(cmd)
110152
if err != nil {
@@ -118,16 +160,12 @@ func runMemoryStore(cmd *cobra.Command, args []string) error {
118160
Tags: tags,
119161
}
120162

121-
// Generate embedding if OpenAI key is available
122-
if openaiKey != "" {
123-
model := viper.GetString("embedding.model")
124-
if model == "" {
125-
model = "text-embedding-3-small"
126-
}
127-
embedder, err := openai.NewClient(openai.Config{APIKey: openaiKey, Model: model})
128-
if err != nil {
129-
return fmt.Errorf("create embedder: %w", err)
130-
}
163+
// Generate embedding if a provider is available
164+
embedder, err := createEmbedder(cmd)
165+
if err != nil {
166+
return fmt.Errorf("create embedder: %w", err)
167+
}
168+
if embedder != nil {
131169
emb, err := embedder.Embed(context.Background(), text)
132170
if err != nil {
133171
return fmt.Errorf("embed text: %w", err)
@@ -158,10 +196,6 @@ func runMemoryRecall(cmd *cobra.Command, args []string) error {
158196
maxResults, _ := cmd.Flags().GetInt("max-results")
159197
maxTokens, _ := cmd.Flags().GetInt("max-tokens")
160198
recencyWeight, _ := cmd.Flags().GetFloat64("recency-weight")
161-
openaiKey, _ := cmd.Flags().GetString("openai-key")
162-
if openaiKey == "" {
163-
openaiKey = os.Getenv("OPENAI_API_KEY")
164-
}
165199

166200
store, err := openMemoryStore(cmd)
167201
if err != nil {
@@ -177,16 +211,12 @@ func runMemoryRecall(cmd *cobra.Command, args []string) error {
177211
RecencyWeight: recencyWeight,
178212
}
179213

180-
// Generate query embedding if OpenAI key is available
181-
if openaiKey != "" {
182-
model := viper.GetString("embedding.model")
183-
if model == "" {
184-
model = "text-embedding-3-small"
185-
}
186-
embedder, err := openai.NewClient(openai.Config{APIKey: openaiKey, Model: model})
187-
if err != nil {
188-
return fmt.Errorf("create embedder: %w", err)
189-
}
214+
// Generate query embedding if a provider is available
215+
embedder, err := createEmbedder(cmd)
216+
if err != nil {
217+
return fmt.Errorf("create embedder: %w", err)
218+
}
219+
if embedder != nil {
190220
emb, err := embedder.Embed(context.Background(), query)
191221
if err != nil {
192222
return fmt.Errorf("embed query: %w", err)

‎cmd/serve.go‎

Lines changed: 34 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,10 @@ import (
1111
"time"
1212

1313
"github.com/Siddhant-K-code/distill/pkg/contextlab"
14-
"github.com/Siddhant-K-code/distill/pkg/embedding/openai"
14+
"github.com/Siddhant-K-code/distill/pkg/embedding"
15+
_ "github.com/Siddhant-K-code/distill/pkg/embedding/cohere"
16+
_ "github.com/Siddhant-K-code/distill/pkg/embedding/ollama"
17+
_ "github.com/Siddhant-K-code/distill/pkg/embedding/openai"
1518
"github.com/Siddhant-K-code/distill/pkg/metrics"
1619
"github.com/Siddhant-K-code/distill/pkg/retriever"
1720
"github.com/Siddhant-K-code/distill/pkg/telemetry"
@@ -53,8 +56,10 @@ func init() {
5356
serveCmd.Flags().StringP("namespace", "n", "", "Default namespace")
5457

5558
// Embedding settings
56-
serveCmd.Flags().String("openai-key", "", "OpenAI API key for embeddings (or use OPENAI_API_KEY)")
57-
serveCmd.Flags().String("embedding-model", "text-embedding-3-small", "OpenAI embedding model")
59+
serveCmd.Flags().String("openai-key", "", "API key for embeddings (or use OPENAI_API_KEY / COHERE_API_KEY)")
60+
serveCmd.Flags().String("embedding-provider", "openai", "Embedding provider (openai, ollama, cohere)")
61+
serveCmd.Flags().String("embedding-model", "text-embedding-3-small", "Embedding model name")
62+
serveCmd.Flags().String("embedding-base-url", "", "Embedding provider base URL (e.g. http://localhost:11434 for Ollama)")
5863

5964
// ContextLab settings
6065
serveCmd.Flags().Int("over-fetch-k", 50, "Number of chunks to over-fetch")
@@ -69,7 +74,9 @@ func init() {
6974
_ = viper.BindPFlag("retriever.backend", serveCmd.Flags().Lookup("backend"))
7075
_ = viper.BindPFlag("retriever.index", serveCmd.Flags().Lookup("index"))
7176
_ = viper.BindPFlag("retriever.namespace", serveCmd.Flags().Lookup("namespace"))
77+
_ = viper.BindPFlag("embedding.provider", serveCmd.Flags().Lookup("embedding-provider"))
7278
_ = viper.BindPFlag("embedding.model", serveCmd.Flags().Lookup("embedding-model"))
79+
_ = viper.BindPFlag("embedding.base_url", serveCmd.Flags().Lookup("embedding-base-url"))
7380
_ = viper.BindPFlag("retriever.top_k", serveCmd.Flags().Lookup("over-fetch-k"))
7481
_ = viper.BindPFlag("retriever.target_k", serveCmd.Flags().Lookup("target-k"))
7582
_ = viper.BindPFlag("dedup.threshold", serveCmd.Flags().Lookup("threshold"))
@@ -204,12 +211,31 @@ func runServe(cmd *cobra.Command, args []string) error {
204211
}
205212
defer func() { _ = ret.Close() }()
206213

207-
// Create embedding provider if OpenAI key is provided
214+
// Create embedding provider via registry
215+
embeddingProvider := viper.GetString("embedding.provider")
216+
embeddingBaseURL, _ := cmd.Flags().GetString("embedding-base-url")
217+
if embeddingBaseURL == "" {
218+
embeddingBaseURL = viper.GetString("embedding.base_url")
219+
}
220+
208221
var embedder retriever.EmbeddingProvider
209-
if openaiKey != "" {
210-
embedder, err = openai.NewClient(openai.Config{
211-
APIKey: openaiKey,
212-
Model: embeddingModel,
222+
needsAPIKey := embeddingProvider == "" || embeddingProvider == "openai" || embeddingProvider == "cohere"
223+
if needsAPIKey && openaiKey == "" {
224+
// No API key and cloud provider selected — embeddings disabled
225+
} else {
226+
apiKeyForEmbed := openaiKey
227+
if embeddingProvider == "cohere" && apiKeyForEmbed == "" {
228+
apiKeyForEmbed = os.Getenv("COHERE_API_KEY")
229+
}
230+
if embeddingProvider == "" {
231+
embeddingProvider = "openai"
232+
}
233+
embedder, err = embedding.NewProvider(embedding.ProviderConfig{
234+
Type: embedding.ProviderType(embeddingProvider),
235+
APIKey: apiKeyForEmbed,
236+
Model: embeddingModel,
237+
BaseURL: embeddingBaseURL,
238+
CacheSize: -1,
213239
})
214240
if err != nil {
215241
return fmt.Errorf("failed to create embedding provider: %w", err)

‎pkg/config/config.go‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ type ServerConfig struct {
3535
type EmbeddingConfig struct {
3636
Provider string `mapstructure:"provider"`
3737
Model string `mapstructure:"model"`
38+
BaseURL string `mapstructure:"base_url"`
3839
BatchSize int `mapstructure:"batch_size"`
3940
}
4041

@@ -166,9 +167,9 @@ func Validate(cfg *Config) error {
166167
}
167168

168169
// Embedding validation
169-
validProviders := map[string]bool{"openai": true, "": true}
170+
validProviders := map[string]bool{"openai": true, "ollama": true, "cohere": true, "": true}
170171
if !validProviders[cfg.Embedding.Provider] {
171-
errs = append(errs, fmt.Sprintf("embedding.provider: unsupported provider %q (supported: openai)", cfg.Embedding.Provider))
172+
errs = append(errs, fmt.Sprintf("embedding.provider: unsupported provider %q (supported: openai, ollama, cohere)", cfg.Embedding.Provider))
172173
}
173174
if cfg.Embedding.BatchSize < 0 {
174175
errs = append(errs, "embedding.batch_size: must be non-negative")
@@ -252,6 +253,7 @@ func interpolateConfig(cfg *Config) {
252253
cfg.Server.Host = InterpolateEnv(cfg.Server.Host)
253254
cfg.Embedding.Provider = InterpolateEnv(cfg.Embedding.Provider)
254255
cfg.Embedding.Model = InterpolateEnv(cfg.Embedding.Model)
256+
cfg.Embedding.BaseURL = InterpolateEnv(cfg.Embedding.BaseURL)
255257
cfg.Dedup.Method = InterpolateEnv(cfg.Dedup.Method)
256258
cfg.Dedup.Linkage = InterpolateEnv(cfg.Dedup.Linkage)
257259
cfg.Retriever.Backend = InterpolateEnv(cfg.Retriever.Backend)
@@ -281,9 +283,10 @@ server:
281283
write_timeout: 60s
282284
283285
embedding:
284-
provider: openai
286+
provider: openai # openai, ollama, or cohere
285287
model: text-embedding-3-small
286288
batch_size: 100
289+
# base_url: "" # override API endpoint (e.g. http://localhost:11434 for Ollama)
287290
288291
dedup:
289292
threshold: 0.15

0 commit comments

Comments
 (0)