ai_scheduler/internal/domain/llm/capability/router.go

49 lines
1.3 KiB
Go

package capability
import (
"ai_scheduler/internal/config"
"ai_scheduler/internal/domain/llm"
"strings"
"time"
)
func Route(cfg *config.Config, ability Ability) (ProviderChoice, llm.Options, error) {
cap, ok := cfg.LLM.Capabilities[string(ability)]
if !ok {
return ProviderChoice{}, llm.Options{}, llm.ErrInvalidCapability
}
prov, ok := cfg.LLM.Providers[cap.Provider]
if !ok {
return ProviderChoice{}, llm.Options{}, llm.ErrProviderNotFound
}
var modelConf config.LLMModle
found := false
for _, m := range prov.Models {
if m.Name == cap.Model || m.ID == cap.Model {
modelConf = m
found = true
break
}
}
if !found {
return ProviderChoice{}, llm.Options{}, llm.ErrModelNotFound
}
to := llm.Options{}
to.Model = modelConf.Name
to.Stream = cap.Parameters.Stream || modelConf.Streaming
if to.Stream && !modelConf.Streaming {
to.Stream = false
}
to.MaxTokens = modelConf.MaxTokens
if cap.Parameters.MaxTokens > 0 && cap.Parameters.MaxTokens <= modelConf.MaxTokens {
to.MaxTokens = cap.Parameters.MaxTokens
}
to.Temperature = float32(cap.Parameters.Temperature)
to.Modalities = append([]string{}, modelConf.Modalities...)
d, _ := time.ParseDuration(strings.TrimSpace(prov.Timeout))
to.Timeout = d
to.Endpoint = prov.Endpoint
choice := ProviderChoice{Provider: cap.Provider, Model: to.Model}
return choice, to, nil
}