From 525f726265e9969e571126fcbe5494b84e1cd95b Mon Sep 17 00:00:00 2001 From: shenlan Date: Mon, 11 Aug 2025 00:19:18 +0800 Subject: [PATCH] feat: add provider-aware adapters --- client/main.go | 10 +-- example/server/config/server.yaml | 8 +-- internal/rag/config/runtime.go | 6 +- internal/rag/config/runtime_test.go | 6 +- internal/rag/embed/bge.go | 20 +++--- internal/rag/embed/{allama.go => ollama.go} | 34 ++++----- internal/rag/embed/openai.go | 28 ++++---- internal/rag/ingest/ingest.go | 10 +-- internal/rag/service.go | 12 ++-- server/api/askai.go | 76 +++++++++++++-------- 10 files changed, 119 insertions(+), 91 deletions(-) rename internal/rag/embed/{allama.go => ollama.go} (65%) diff --git a/client/main.go b/client/main.go index e0e8a3b..17c1c7e 100644 --- a/client/main.go +++ b/client/main.go @@ -70,13 +70,15 @@ var rootCmd = &cobra.Command{ var embedder embed.Embedder switch embCfg.Provider { - case "allama": - embedder = embed.NewAllama(embCfg.BaseURL, embCfg.Model, embCfg.Dimension) + case "ollama": + embedder = embed.NewOllama(embCfg.Endpoint, embCfg.Model, embCfg.Dimension) + case "chutes": + embedder = embed.NewOpenAI(embCfg.Endpoint, embCfg.APIKey, embCfg.Model, embCfg.Dimension) default: if embCfg.Model != "" { - embedder = embed.NewOpenAI(embCfg.BaseURL, embCfg.APIKey, embCfg.Model, embCfg.Dimension) + embedder = embed.NewOpenAI(embCfg.Endpoint, embCfg.APIKey, embCfg.Model, embCfg.Dimension) } else { - embedder = embed.NewBGE(embCfg.BaseURL, embCfg.APIKey, embCfg.Dimension) + embedder = embed.NewBGE(embCfg.Endpoint, embCfg.APIKey, embCfg.Dimension) } } diff --git a/example/server/config/server.yaml b/example/server/config/server.yaml index f1f4661..5569940 100644 --- a/example/server/config/server.yaml +++ b/example/server/config/server.yaml @@ -19,18 +19,18 @@ sync: proxy: socks5://127.0.0.1:1080 # 仅在同步仓库时使用代理 provider: - - name: allama - base_url: http://localhost:11434 + - name: ollama + endpoint: http://localhost:11434/v1/chat/completions models: - 'gpt-oss:20b' - name: chutes - base_url: https://llm.chutes.ai + endpoint: https://llm.chutes.ai/v1/chat/completions token: "cpk_xxxxxxxxxxxxxxxxxx" models: - 'moonshotai/Kimi-K2-Instruct' embedding: - base_url: https://chutes-baai-bge-m3.chutes.ai/embed + endpoint: https://chutes-baai-bge-m3.chutes.ai/embed/v1/embeddings token: "cpk_xxxxxxxxxxxxxxxxxx" dimension: 0 # 0 = 首次响应自动探测维度 rate_limit_tpm: 120000 diff --git a/internal/rag/config/runtime.go b/internal/rag/config/runtime.go index 7713dfd..682037c 100644 --- a/internal/rag/config/runtime.go +++ b/internal/rag/config/runtime.go @@ -8,7 +8,7 @@ import ( // RuntimeEmbedding is the resolved embedding configuration used at runtime. type RuntimeEmbedding struct { Provider string - BaseURL string + Endpoint string APIKey string Model string Dimension int @@ -25,7 +25,7 @@ func (c *Config) ResolveEmbedding() RuntimeEmbedding { if len(m.Models) > 0 { rt.Model = m.Models[0] } - rt.BaseURL = strings.TrimRight(m.Endpoint, "/") + rt.Endpoint = strings.TrimRight(m.Endpoint, "/") rt.APIKey = m.Token e := c.Embedding @@ -96,7 +96,7 @@ func (rt *Runtime) ToConfig() *Config { c.Global.Datasources = rt.Datasources c.Global.Proxy = rt.Proxy c.Models.Embedder.Provider = rt.Embedding.Provider - c.Models.Embedder.Endpoint = rt.Embedding.BaseURL + c.Models.Embedder.Endpoint = rt.Embedding.Endpoint c.Models.Embedder.Token = rt.Embedding.APIKey if rt.Embedding.Model != "" { c.Models.Embedder.Models = []string{rt.Embedding.Model} diff --git a/internal/rag/config/runtime_test.go b/internal/rag/config/runtime_test.go index 4ab7d66..122be6c 100644 --- a/internal/rag/config/runtime_test.go +++ b/internal/rag/config/runtime_test.go @@ -25,8 +25,8 @@ func TestResolveEmbedding(t *testing.T) { cfg.Models.Embedder.Token = "tok" cfg.Models.Embedder.Models = []string{"m"} e := cfg.ResolveEmbedding() - if e.BaseURL != "https://api.example.com" { - t.Fatalf("unexpected base url %q", e.BaseURL) + if e.Endpoint != "https://api.example.com" { + t.Fatalf("unexpected endpoint %q", e.Endpoint) } if e.APIKey != "tok" { t.Fatalf("unexpected api key %q", e.APIKey) @@ -49,7 +49,7 @@ func TestResolveChunking(t *testing.T) { func TestRuntimeToConfigEmbedding(t *testing.T) { rt := &Runtime{} - rt.Embedding.BaseURL = "http://localhost:8080" + rt.Embedding.Endpoint = "http://localhost:8080" rt.Embedding.APIKey = "tok" rt.Embedding.Dimension = 123 cfg := rt.ToConfig() diff --git a/internal/rag/embed/bge.go b/internal/rag/embed/bge.go index 9115368..d317638 100644 --- a/internal/rag/embed/bge.go +++ b/internal/rag/embed/bge.go @@ -12,19 +12,19 @@ import ( // BGE implements the Embedder interface for a BGE embedding service. type BGE struct { - baseURL string - token string - dim int - client *http.Client + endpoint string + token string + dim int + client *http.Client } // NewBGE returns a new BGE embedder. -func NewBGE(baseURL, token string, dim int) *BGE { +func NewBGE(endpoint, token string, dim int) *BGE { return &BGE{ - baseURL: baseURL, - token: token, - dim: dim, - client: &http.Client{Timeout: 30 * time.Second}, + endpoint: endpoint, + token: token, + dim: dim, + client: &http.Client{Timeout: 30 * time.Second}, } } @@ -37,7 +37,7 @@ func (b *BGE) Embed(ctx context.Context, inputs []string) ([][]float32, int, err for i, text := range inputs { payload := map[string]any{"inputs": text} body, _ := json.Marshal(payload) - req, err := http.NewRequestWithContext(ctx, http.MethodPost, b.baseURL, bytes.NewReader(body)) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, b.endpoint, bytes.NewReader(body)) if err != nil { return nil, 0, err } diff --git a/internal/rag/embed/allama.go b/internal/rag/embed/ollama.go similarity index 65% rename from internal/rag/embed/allama.go rename to internal/rag/embed/ollama.go index 592e415..be027e1 100644 --- a/internal/rag/embed/allama.go +++ b/internal/rag/embed/ollama.go @@ -11,31 +11,31 @@ import ( "time" ) -// Allama implements the Embedder interface using the Allama/Ollama embeddings API. -type Allama struct { - baseURL string - model string - dim int - client *http.Client +// Ollama implements the Embedder interface using the Ollama embeddings API. +type Ollama struct { + endpoint string + model string + dim int + client *http.Client } -// NewAllama creates a new Allama embedder. -func NewAllama(baseURL, model string, dim int) *Allama { - return &Allama{ - baseURL: strings.TrimRight(baseURL, "/"), - model: model, - dim: dim, - client: &http.Client{Timeout: 30 * time.Second}, +// NewOllama creates a new Ollama embedder. +func NewOllama(endpoint, model string, dim int) *Ollama { + return &Ollama{ + endpoint: strings.TrimRight(endpoint, "/"), + model: model, + dim: dim, + client: &http.Client{Timeout: 30 * time.Second}, } } // Dimension returns the embedding dimension if known. -func (a *Allama) Dimension() int { return a.dim } +func (a *Ollama) Dimension() int { return a.dim } -// Embed posts texts to the Allama embeddings endpoint. -func (a *Allama) Embed(ctx context.Context, inputs []string) ([][]float32, int, error) { +// Embed posts texts to the Ollama embeddings endpoint. +func (a *Ollama) Embed(ctx context.Context, inputs []string) ([][]float32, int, error) { vecs := make([][]float32, len(inputs)) - url := a.baseURL + "/api/embeddings" + url := a.endpoint for i, text := range inputs { payload := map[string]any{"model": a.model, "prompt": text} body, _ := json.Marshal(payload) diff --git a/internal/rag/embed/openai.go b/internal/rag/embed/openai.go index af9de34..6162a47 100644 --- a/internal/rag/embed/openai.go +++ b/internal/rag/embed/openai.go @@ -12,21 +12,21 @@ import ( // OpenAI implements the Embedder interface using OpenAI-compatible APIs. type OpenAI struct { - baseURL string - apiKey string - model string - dim int - client *http.Client + endpoint string + apiKey string + model string + dim int + client *http.Client } // NewOpenAI creates a new OpenAI embedder from configuration. -func NewOpenAI(baseURL, apiKey, model string, dim int) *OpenAI { +func NewOpenAI(endpoint, apiKey, model string, dim int) *OpenAI { return &OpenAI{ - baseURL: baseURL, - apiKey: apiKey, - model: model, - dim: dim, - client: &http.Client{Timeout: 30 * time.Second}, + endpoint: endpoint, + apiKey: apiKey, + model: model, + dim: dim, + client: &http.Client{Timeout: 30 * time.Second}, } } @@ -40,12 +40,14 @@ func (o *OpenAI) Embed(ctx context.Context, inputs []string) ([][]float32, int, "input": inputs, } b, _ := json.Marshal(payload) - req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.baseURL+"/embeddings", bytes.NewReader(b)) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.endpoint, bytes.NewReader(b)) if err != nil { return nil, 0, err } req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+o.apiKey) + if o.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+o.apiKey) + } resp, err := o.client.Do(req) if err != nil { return nil, 0, err diff --git a/internal/rag/ingest/ingest.go b/internal/rag/ingest/ingest.go index 0356bbb..4cc111a 100644 --- a/internal/rag/ingest/ingest.go +++ b/internal/rag/ingest/ingest.go @@ -69,13 +69,15 @@ func IngestRepo(ctx context.Context, cfg *cfgpkg.Config, ds cfgpkg.DataSource, o var embedder embed.Embedder switch embCfg.Provider { - case "allama": - embedder = embed.NewAllama(embCfg.BaseURL, embCfg.Model, embCfg.Dimension) + case "ollama": + embedder = embed.NewOllama(embCfg.Endpoint, embCfg.Model, embCfg.Dimension) + case "chutes": + embedder = embed.NewOpenAI(embCfg.Endpoint, embCfg.APIKey, embCfg.Model, embCfg.Dimension) default: if embCfg.Model != "" { - embedder = embed.NewOpenAI(embCfg.BaseURL, embCfg.APIKey, embCfg.Model, embCfg.Dimension) + embedder = embed.NewOpenAI(embCfg.Endpoint, embCfg.APIKey, embCfg.Model, embCfg.Dimension) } else { - embedder = embed.NewBGE(embCfg.BaseURL, embCfg.APIKey, embCfg.Dimension) + embedder = embed.NewBGE(embCfg.Endpoint, embCfg.APIKey, embCfg.Dimension) } } if err := store.EnsureSchema(ctx, conn, embedder.Dimension(), opt.MigrateDim); err != nil { diff --git a/internal/rag/service.go b/internal/rag/service.go index 2cce100..c752cf2 100644 --- a/internal/rag/service.go +++ b/internal/rag/service.go @@ -56,18 +56,20 @@ func (s *Service) Query(ctx context.Context, question string, limit int) ([]Docu return nil, nil } embCfg := s.cfg.ResolveEmbedding() - if embCfg.APIKey == "" || embCfg.BaseURL == "" { + if embCfg.Endpoint == "" { return nil, nil } var emb embed.Embedder switch embCfg.Provider { - case "allama": - emb = embed.NewAllama(embCfg.BaseURL, embCfg.Model, embCfg.Dimension) + case "ollama": + emb = embed.NewOllama(embCfg.Endpoint, embCfg.Model, embCfg.Dimension) + case "chutes": + emb = embed.NewOpenAI(embCfg.Endpoint, embCfg.APIKey, embCfg.Model, embCfg.Dimension) default: if embCfg.Model != "" { - emb = embed.NewOpenAI(embCfg.BaseURL, embCfg.APIKey, embCfg.Model, embCfg.Dimension) + emb = embed.NewOpenAI(embCfg.Endpoint, embCfg.APIKey, embCfg.Model, embCfg.Dimension) } else { - emb = embed.NewBGE(embCfg.BaseURL, embCfg.APIKey, embCfg.Dimension) + emb = embed.NewBGE(embCfg.Endpoint, embCfg.APIKey, embCfg.Dimension) } } vecs, _, err := emb.Embed(ctx, []string{question}) diff --git a/server/api/askai.go b/server/api/askai.go index cc69bdb..d264fa3 100644 --- a/server/api/askai.go +++ b/server/api/askai.go @@ -63,12 +63,12 @@ type serverConfig struct { } `yaml:"api"` } -// loadConfig reads provider, model, URL, timeout and retries from ConfigPath +// loadConfig reads provider, model, endpoint, timeout and retries from ConfigPath // and environment variables. func loadConfig() (string, string, string, string, time.Duration, int) { provider := "" model := os.Getenv("CHUTES_API_MODEL") - baseURL := os.Getenv("CHUTES_API_URL") + endpoint := os.Getenv("CHUTES_API_URL") token := "" timeout := 30 * time.Second retries := 3 @@ -83,8 +83,8 @@ func loadConfig() (string, string, string, string, time.Duration, int) { if model == "" && len(g.Models) > 0 { model = g.Models[0] } - if baseURL == "" { - baseURL = g.Endpoint + if endpoint == "" { + endpoint = g.Endpoint } if token == "" { token = g.Token @@ -104,23 +104,33 @@ func loadConfig() (string, string, string, string, time.Duration, int) { retries = 3 } provider = strings.ToLower(provider) - baseURL = strings.TrimRight(baseURL, "/") - if provider == "allama" { - if baseURL == "" { - baseURL = "http://localhost:11434" + endpoint = strings.TrimRight(endpoint, "/") + switch provider { + case "ollama": + if endpoint == "" { + endpoint = "http://localhost:11434/v1/chat/completions" } if model == "" { model = "gpt-oss:20b" } - return provider, token, model, baseURL + "/api/chat", timeout, retries + return provider, token, model, endpoint, timeout, retries + case "chutes": + if endpoint == "" { + endpoint = "https://llm.chutes.ai/v1/chat/completions" + } + if model == "" { + model = "deepseek-ai/DeepSeek-R1" + } + return provider, token, model, endpoint, timeout, retries + default: + if endpoint == "" { + endpoint = "https://llm.chutes.ai/v1/chat/completions" + } + if model == "" { + model = "deepseek-ai/DeepSeek-R1" + } + return provider, token, model, endpoint, timeout, retries } - if baseURL == "" { - baseURL = "https://llm.chutes.ai" - } - if model == "" { - model = "deepseek-ai/DeepSeek-R1" - } - return "chutes", token, model, baseURL + "/v1/chat/completions", timeout, retries } // callChutes sends the question to the hosted LLM service and returns the reply. @@ -191,12 +201,14 @@ func callChutes(token, model, url string, timeout time.Duration, retries int, qu return "", lastErr } -// callAllama sends the question to a local Allama server. -func callAllama(model, url string, timeout time.Duration, retries int, question string) (string, error) { +// callOllama sends the question to a local Ollama server. +func callOllama(model, url string, timeout time.Duration, retries int, question string) (string, error) { reqBody := map[string]any{ - "model": model, - "messages": []any{map[string]string{"role": "user", "content": question}}, - "stream": false, + "model": model, + "messages": []any{map[string]any{"role": "user", "content": question}}, + "stream": false, + "max_tokens": 1024, + "temperature": 0.7, } data, err := json.Marshal(reqBody) if err != nil { @@ -222,19 +234,25 @@ func callAllama(model, url string, timeout time.Duration, retries int, question continue } if resp.StatusCode != http.StatusOK { - lastErr = fmt.Errorf("allama API error: %s", string(b)) + lastErr = fmt.Errorf("ollama API error: %s", string(b)) continue } var res struct { - Message struct { - Content string `json:"content"` - } `json:"message"` + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` } if err := json.Unmarshal(b, &res); err != nil { lastErr = err continue } - return res.Message.Content, nil + if len(res.Choices) == 0 { + lastErr = errors.New("no choices returned") + continue + } + return res.Choices[0].Message.Content, nil } if lastErr == nil { lastErr = errors.New("request failed") @@ -246,8 +264,10 @@ func callAllama(model, url string, timeout time.Duration, retries int, question func callLLM(question string) (string, error) { provider, token, model, url, timeout, retries := loadConfig() switch provider { - case "allama": - return callAllama(model, url, timeout, retries, question) + case "ollama": + return callOllama(model, url, timeout, retries, question) + case "chutes": + return callChutes(token, model, url, timeout, retries, question) default: return callChutes(token, model, url, timeout, retries, question) }