Merge pull request #138 from svc-design/codex/add-provider-aware-support-for-api-endpoints

This commit is contained in:
shenlan 2025-08-11 00:22:20 +08:00 committed by GitHub
commit 9bd2ff1a62
10 changed files with 119 additions and 91 deletions

View File

@ -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)
}
}

View File

@ -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

View File

@ -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}

View File

@ -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()

View File

@ -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
}

View File

@ -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)

View File

@ -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

View File

@ -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 {

View File

@ -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})

View File

@ -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)
}