config,llm: validate llm_base_url and inherit model name from config

getsavvyinc · Dec 21, 2024 · f5e3bd0 · f5e3bd0
1 parent a4f4c10
commit f5e3bd0
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 7 deletions.
diff --git a/config/config.go b/config/config.go
@@ -15,9 +15,10 @@ var (
 )
 
 type Config struct {
-	Token         string `json:"token"`
-	OpenAIBaseURL string `json:"openai_base_url"`
-	OpenAIKey     string `json:"openai_key"`
+	Token        string `json:"token"`
+	LLMBaseURL   string `json:"llm_base_url"`
+	LLMModelName string `json:"llm_model_name"`
+	LLMAPIKey    string `json:"llm_api_key"`
 }
 
 func (c *Config) Save() error {
@@ -63,5 +64,9 @@ func (c *Config) validate() error {
 	if c.Token == "" {
 		return ErrMissingToken
 	}
+
+	if c.LLMBaseURL != "" && c.LLMModelName == "" {
+		return errors.New("missing or empty llm_model_name")
+	}
 	return nil
 }
diff --git a/llm/service/custom_llm.go b/llm/service/custom_llm.go
@@ -28,8 +28,8 @@ type customSvc struct {
 }
 
 func newCustomService(cfg *config.Config) Service {
-	baseURL := cfg.OpenAIBaseURL
-	apiKey := cfg.OpenAIKey
+	baseURL := cfg.LLMBaseURL
+	apiKey := cfg.LLMAPIKey
 
 	clientConfig := openai.DefaultConfig(apiKey)
 	clientConfig.BaseURL = baseURL
@@ -38,7 +38,7 @@ func newCustomService(cfg *config.Config) Service {
 
 	return &customSvc{
 		cl:        openaiClient,
-		modelName: "llama3.2:latest",
+		modelName: cfg.LLMModelName,
 	}
 }
 

diff --git a/llm/service/service.go b/llm/service/service.go
@@ -18,7 +18,7 @@ type Service interface {
 }
 
 func New(cfg *config.Config) Service {
-	if cfg.OpenAIBaseURL == "" {
+	if cfg.LLMBaseURL == "" {
 		return newDefaultService(cfg)
 	}
 	return newCustomService(cfg)