Merge pull request #138 from svc-design/codex/add-provider-aware-support-for-api-endpoints
This commit is contained in:
commit
9bd2ff1a62
@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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}
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
@ -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)
|
||||
@ -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
|
||||
|
||||
@ -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 {
|
||||
|
||||
@ -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})
|
||||
|
||||
@ -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)
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user