diff --git a/internal/rag/config/config.go b/internal/rag/config/config.go index 7cf91f8..d3369ae 100644 --- a/internal/rag/config/config.go +++ b/internal/rag/config/config.go @@ -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. diff --git a/internal/rag/config/runtime.go b/internal/rag/config/runtime.go index 2976a54..7713dfd 100644 --- a/internal/rag/config/runtime.go +++ b/internal/rag/config/runtime.go @@ -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 } diff --git a/internal/rag/config/runtime_test.go b/internal/rag/config/runtime_test.go index 948b95f..4ab7d66 100644 --- a/internal/rag/config/runtime_test.go +++ b/internal/rag/config/runtime_test.go @@ -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) diff --git a/server/api/askai.go b/server/api/askai.go index 5b4fcc6..cc69bdb 100644 --- a/server/api/askai.go +++ b/server/api/askai.go @@ -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 diff --git a/server/config/config.go b/server/config/config.go index 8e6fb1e..e25e7b5 100644 --- a/server/config/config.go +++ b/server/config/config.go @@ -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, diff --git a/server/server.go b/server/server.go index 7fe9cf1..3389185 100644 --- a/server/server.go +++ b/server/server.go @@ -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 } }