ai-workspace-services/accounts
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: support new model config

Browse Source
This commit is contained in:
shenlan 2025-08-10 18:00:17 +08:00
parent 1c8876b69f
commit 4756dfe3e9
6 changed files with 175 additions and 139 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

64
internal/rag/config/config.go
View File

@ -62,25 +62,44 @@ type Sync struct {
} `yaml:"repo"`
}
// Provider defines an LLM provider which can also serve embeddings.
type Provider struct {
Name string `yaml:"name"`
BaseURL string `yaml:"base_url"`
Token string `yaml:"token"`
Models []string `yaml:"models"`
// StringSlice supports unmarshaling from either a single string or a list of strings.
type StringSlice []string
// UnmarshalYAML implements yaml unmarshaling for StringSlice.
func (s *StringSlice) UnmarshalYAML(value *yaml.Node) error {
switch value.Kind {
case yaml.ScalarNode:
var str string
if err := value.Decode(&str); err != nil {
return err
}
*s = []string{str}
case yaml.SequenceNode:
var arr []string
if err := value.Decode(&arr); err != nil {
return err
}
*s = arr
default:
return fmt.Errorf("invalid yaml kind for StringSlice: %v", value.Kind)
}
return nil
}
// ModelCfg describes a model service such as embedder or generator.
type ModelCfg struct {
Provider string `yaml:"provider"`
Models StringSlice `yaml:"models"`
Endpoint string `yaml:"endpoint"`
Token string `yaml:"token"`
}
// EmbeddingCfg describes embedding service settings.
type EmbeddingCfg struct {
Provider string `yaml:"provider"`
BaseURL string `yaml:"base_url"`
Token string `yaml:"token"`
Model string `yaml:"model"`
APIKeyEnv string `yaml:"api_key_env"`
Dimension int `yaml:"dimension"`
RateLimitTPM int `yaml:"rate_limit_tpm"`
MaxBatch int `yaml:"max_batch"`
MaxChars int `yaml:"max_chars"`
MaxBatch int `yaml:"max_batch"`
Dimension int `yaml:"dimension"`
MaxChars int `yaml:"max_chars"`
RateLimitTPM int `yaml:"rate_limit_tpm"`
}
// ChunkingCfg controls how markdown is split into chunks.
@ -94,11 +113,20 @@ type ChunkingCfg struct {
// Config is the root configuration for ingestion.
type Config struct {
Global Global `yaml:"global"`
Sync Sync `yaml:"sync"`
Provider []Provider `yaml:"provider"`
Global Global `yaml:"global"`
Sync Sync `yaml:"sync"`
Models struct {
Embedder ModelCfg `yaml:"embedder"`
Generator ModelCfg `yaml:"generator"`
} `yaml:"models"`
Embedding EmbeddingCfg `yaml:"embedding"`
Chunking ChunkingCfg `yaml:"chunking"`
API struct {
AskAI struct {
Timeout int `yaml:"timeout"`
Retries int `yaml:"retries"`
} `yaml:"askai"`
} `yaml:"api"`
}
// Load reads YAML configuration from the given path.

75
internal/rag/config/runtime.go
View File

@ -1,11 +1,8 @@
package config
import (
"os"
"path/filepath"
"strings"
"gopkg.in/yaml.v3"
)
// RuntimeEmbedding is the resolved embedding configuration used at runtime.
@ -22,38 +19,20 @@ type RuntimeEmbedding struct {
// ResolveEmbedding applies fallback logic to produce runtime embedding settings.
func (c *Config) ResolveEmbedding() RuntimeEmbedding {
e := c.Embedding
var rt RuntimeEmbedding
rt.Provider = e.Provider
rt.Model = e.Model
m := c.Models.Embedder
rt.Provider = m.Provider
if len(m.Models) > 0 {
rt.Model = m.Models[0]
}
rt.BaseURL = strings.TrimRight(m.Endpoint, "/")
rt.APIKey = m.Token
e := c.Embedding
rt.Dimension = e.Dimension
rt.RateLimitTPM = e.RateLimitTPM
rt.MaxBatch = e.MaxBatch
rt.MaxChars = e.MaxChars
// find provider by name
var prov *Provider
for i := range c.Provider {
if c.Provider[i].Name == e.Provider {
prov = &c.Provider[i]
break
}
}
if e.BaseURL != "" {
rt.BaseURL = e.BaseURL
} else if prov != nil {
rt.BaseURL = strings.TrimRight(prov.BaseURL, "/") + "/v1"
}
if e.APIKeyEnv != "" {
rt.APIKey = os.Getenv(e.APIKeyEnv)
} else if e.Token != "" {
rt.APIKey = e.Token
} else if prov != nil {
rt.APIKey = prov.Token
}
return rt
}
@ -84,13 +63,7 @@ type Runtime struct {
VectorDB VectorDB `yaml:"vectordb"`
Datasources []DataSource `yaml:"datasources"`
Proxy string `yaml:"proxy"`
Embedding struct {
Provider string `yaml:"provider"`
BaseURL string `yaml:"base_url"`
Token string `yaml:"token"`
Model string `yaml:"model"`
Dimension int `yaml:"dimension"`
} `yaml:"embedding"`
Embedding RuntimeEmbedding
}
// ServerConfigPath points to the server configuration file.
@ -98,17 +71,18 @@ var ServerConfigPath = filepath.Join("server", "config", "server.yaml")
// LoadServer loads global configuration from ServerConfigPath.
func LoadServer() (*Runtime, error) {
b, err := os.ReadFile(ServerConfigPath)
cfg, err := Load(ServerConfigPath)
if err != nil {
return nil, err
}
var cfg struct {
Global Runtime `yaml:"global"`
rt := &Runtime{
VectorDB: cfg.Global.VectorDB,
Datasources: cfg.Global.Datasources,
Proxy: cfg.Global.Proxy,
}
if err := yaml.Unmarshal(b, &cfg); err != nil {
return nil, err
}
return &cfg.Global, nil
rt.Redis = cfg.Global.Redis
rt.Embedding = cfg.ResolveEmbedding()
return rt, nil
}
// ToConfig converts runtime configuration into service configuration.
@ -121,10 +95,15 @@ func (rt *Runtime) ToConfig() *Config {
c.Global.VectorDB = rt.VectorDB
c.Global.Datasources = rt.Datasources
c.Global.Proxy = rt.Proxy
c.Embedding.Provider = rt.Embedding.Provider
c.Embedding.BaseURL = rt.Embedding.BaseURL
c.Embedding.Token = rt.Embedding.Token
c.Embedding.Model = rt.Embedding.Model
c.Models.Embedder.Provider = rt.Embedding.Provider
c.Models.Embedder.Endpoint = rt.Embedding.BaseURL
c.Models.Embedder.Token = rt.Embedding.APIKey
if rt.Embedding.Model != "" {
c.Models.Embedder.Models = []string{rt.Embedding.Model}
}
c.Embedding.Dimension = rt.Embedding.Dimension
c.Embedding.MaxBatch = rt.Embedding.MaxBatch
c.Embedding.MaxChars = rt.Embedding.MaxChars
c.Embedding.RateLimitTPM = rt.Embedding.RateLimitTPM
return &c
}

21
internal/rag/config/runtime_test.go
View File

@ -19,12 +19,13 @@ func TestVectorDB_DSN(t *testing.T) {
}
func TestResolveEmbedding(t *testing.T) {
cfg := &Config{
Provider: []Provider{{Name: "p1", BaseURL: "https://api.example.com", Token: "tok"}},
Embedding: EmbeddingCfg{Provider: "p1", Model: "m"},
}
cfg := &Config{}
cfg.Models.Embedder.Provider = "p1"
cfg.Models.Embedder.Endpoint = "https://api.example.com"
cfg.Models.Embedder.Token = "tok"
cfg.Models.Embedder.Models = []string{"m"}
e := cfg.ResolveEmbedding()
if e.BaseURL != "https://api.example.com/v1" {
if e.BaseURL != "https://api.example.com" {
t.Fatalf("unexpected base url %q", e.BaseURL)
}
if e.APIKey != "tok" {
@ -49,14 +50,14 @@ func TestResolveChunking(t *testing.T) {
func TestRuntimeToConfigEmbedding(t *testing.T) {
rt := &Runtime{}
rt.Embedding.BaseURL = "http://localhost:8080"
rt.Embedding.Token = "tok"
rt.Embedding.APIKey = "tok"
rt.Embedding.Dimension = 123
cfg := rt.ToConfig()
if cfg.Embedding.BaseURL != "http://localhost:8080" {
t.Fatalf("unexpected base url %q", cfg.Embedding.BaseURL)
if cfg.Models.Embedder.Endpoint != "http://localhost:8080" {
t.Fatalf("unexpected base url %q", cfg.Models.Embedder.Endpoint)
}
if cfg.Embedding.Token != "tok" {
t.Fatalf("unexpected token %q", cfg.Embedding.Token)
if cfg.Models.Embedder.Token != "tok" {
t.Fatalf("unexpected token %q", cfg.Models.Embedder.Token)
}
if cfg.Embedding.Dimension != 123 {
t.Fatalf("unexpected dimension %d", cfg.Embedding.Dimension)

49
server/api/askai.go
View File

@ -47,12 +47,14 @@ func registerAskAIRoutes(r *gin.RouterGroup) {
var ConfigPath = filepath.Join("server", "config", "server.yaml")
type serverConfig struct {
Provider []struct {
Name string `yaml:"name"`
BaseURL string `yaml:"base_url"`
Token string `yaml:"token"`
Models []string `yaml:"models"`
} `yaml:"provider"`
Models struct {
Generator struct {
Provider string `yaml:"provider"`
Models []string `yaml:"models"`
Endpoint string `yaml:"endpoint"`
Token string `yaml:"token"`
} `yaml:"generator"`
} `yaml:"models"`
API struct {
AskAI struct {
Timeout int `yaml:"timeout"` // seconds
@ -74,29 +76,18 @@ func loadConfig() (string, string, string, string, time.Duration, int) {
if err == nil {
var cfg serverConfig
if err := yaml.Unmarshal(data, &cfg); err == nil {
for _, p := range cfg.Provider {
if provider == "" {
provider = p.Name
}
switch p.Name {
case "allama":
if model == "" && len(p.Models) > 0 {
model = p.Models[0]
}
if baseURL == "" {
baseURL = p.BaseURL
}
case "chutes":
if token == "" {
token = p.Token
}
if model == "" && len(p.Models) > 0 {
model = p.Models[0]
}
if baseURL == "" {
baseURL = p.BaseURL
}
}
g := cfg.Models.Generator
if provider == "" {
provider = g.Provider
}
if model == "" && len(g.Models) > 0 {
model = g.Models[0]
}
if baseURL == "" {
baseURL = g.Endpoint
}
if token == "" {
token = g.Token
}
if cfg.API.AskAI.Timeout > 0 {
timeout = time.Duration(cfg.API.AskAI.Timeout) * time.Second

71
server/config/config.go
View File

@ -59,13 +59,6 @@ type Global struct {
VectorDB VectorDB `yaml:"vectordb"`
Datasources []Datasource `yaml:"datasources"`
Proxy string `yaml:"proxy"`
Embedding struct {
Provider string `yaml:"provider"`
BaseURL string `yaml:"base_url"`
Token string `yaml:"token"`
Model string `yaml:"model"`
Dimension int `yaml:"dimension"`
} `yaml:"embedding"`
}
type Sync struct {
@ -74,11 +67,50 @@ type Sync struct {
} `yaml:"repo"`
}
type Provider struct {
Name string `yaml:"name"`
BaseURL string `yaml:"base_url"`
Token string `yaml:"token"`
Models []string `yaml:"models"`
// StringSlice supports unmarshaling from either a single string or a list of strings.
type StringSlice []string
// UnmarshalYAML implements yaml unmarshalling for StringSlice.
func (s *StringSlice) UnmarshalYAML(value *yaml.Node) error {
switch value.Kind {
case yaml.ScalarNode:
var str string
if err := value.Decode(&str); err != nil {
return err
}
*s = []string{str}
case yaml.SequenceNode:
var arr []string
if err := value.Decode(&arr); err != nil {
return err
}
*s = arr
default:
return fmt.Errorf("invalid yaml kind for StringSlice: %v", value.Kind)
}
return nil
}
type ModelCfg struct {
Provider string `yaml:"provider"`
Models StringSlice `yaml:"models"`
Endpoint string `yaml:"endpoint"`
Token string `yaml:"token"`
}
type EmbeddingCfg struct {
MaxBatch int `yaml:"max_batch"`
Dimension int `yaml:"dimension"`
MaxChars int `yaml:"max_chars"`
RateLimitTPM int `yaml:"rate_limit_tpm"`
}
type ChunkingCfg struct {
MaxTokens int `yaml:"max_tokens"`
OverlapTokens int `yaml:"overlap_tokens"`
PreferHeadingSplit bool `yaml:"prefer_heading_split"`
IncludeExts []string `yaml:"include_exts"`
IgnoreDirs []string `yaml:"ignore_dirs"`
}
type API struct {
@ -89,11 +121,16 @@ type API struct {
}
type Config struct {
Log Log `yaml:"log"`
Global Global `yaml:"global"`
Sync Sync `yaml:"sync"`
Provider []Provider `yaml:"provider"`
API API `yaml:"api"`
Log Log `yaml:"log"`
Global Global `yaml:"global"`
Sync Sync `yaml:"sync"`
Models struct {
Embedder ModelCfg `yaml:"embedder"`
Generator ModelCfg `yaml:"generator"`
} `yaml:"models"`
Embedding EmbeddingCfg `yaml:"embedding"`
Chunking ChunkingCfg `yaml:"chunking"`
API API `yaml:"api"`
}
// Load reads the configuration file at the provided path. When path is empty,

34
server/server.go
View File

@ -5,6 +5,7 @@ import (
"os"
"path/filepath"
"runtime"
"strings"
"github.com/gin-gonic/gin"
"gopkg.in/yaml.v3"
@ -12,12 +13,14 @@ import (
// Config represents server configuration loaded from YAML.
type Config struct {
Provider []struct {
Name string `yaml:"name"`
BaseURL string `yaml:"base_url"`
Token string `yaml:"token"`
Models []string `yaml:"models"`
} `yaml:"provider"`
Models struct {
Generator struct {
Provider string `yaml:"provider"`
Models []string `yaml:"models"`
Endpoint string `yaml:"endpoint"`
Token string `yaml:"token"`
} `yaml:"generator"`
} `yaml:"models"`
}
// cfg holds the loaded configuration.
@ -35,20 +38,17 @@ func loadConfig() {
slog.Warn("server config parse", "err", err)
return
}
for _, p := range cfg.Provider {
if p.Name != "chutes" {
continue
g := cfg.Models.Generator
if strings.ToLower(g.Provider) == "chutes" {
if g.Token != "" {
os.Setenv("CHUTES_API_TOKEN", g.Token)
}
if p.Token != "" {
os.Setenv("CHUTES_API_TOKEN", p.Token)
if g.Endpoint != "" {
os.Setenv("CHUTES_API_URL", g.Endpoint)
}
if p.BaseURL != "" {
os.Setenv("CHUTES_API_URL", p.BaseURL)
if len(g.Models) > 0 {
os.Setenv("CHUTES_API_MODEL", g.Models[0])
}
if len(p.Models) > 0 {
os.Setenv("CHUTES_API_MODEL", p.Models[0])
}
break
}
}