Fix missing client option (#28)

sunshineplan · Mar 29, 2024 · e317833 · e317833
1 parent 299d1a1
commit e317833
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 9 deletions.
diff --git a/ai_test.go b/ai_test.go
@@ -95,14 +95,12 @@ func TestGemini(t *testing.T) {
 		ai.WithAPIKey(apiKey),
 		ai.WithEndpoint(os.Getenv("GEMINI_ENDPOINT")),
 		ai.WithProxy(os.Getenv("GEMINI_PROXY")),
+		ai.WithModel(os.Getenv("CHATGPT_MODEL")),
 	)
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer gemini.Close()
-	if model := os.Getenv("GEMINI_MODEL"); model != "" {
-		gemini.SetModel(model)
-	}
 	if err := testChat(gemini, "Who are you?"); err != nil {
 		t.Error(err)
 	}
@@ -123,14 +121,12 @@ func TestChatGPT(t *testing.T) {
 		ai.WithAPIKey(apiKey),
 		ai.WithEndpoint(os.Getenv("CHATGPT_ENDPOINT")),
 		ai.WithProxy(os.Getenv("CHATGPT_PROXY")),
+		ai.WithModel(os.Getenv("CHATGPT_MODEL")),
 	)
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer chatgpt.Close()
-	if model := os.Getenv("CHATGPT_MODEL"); model != "" {
-		chatgpt.SetModel(model)
-	}
 	if err := testChat(chatgpt, "Who are you?"); err != nil {
 		t.Error(err)
 	}

diff --git a/config.go b/config.go
@@ -45,7 +45,8 @@ func WithAPIKey(apiKey string) ClientOption           { return withAPIKey(apiKey
 func WithEndpoint(endpoint string) ClientOption       { return withEndpoint(endpoint) }
 func WithProxy(proxy string) ClientOption             { return withProxy(proxy) }
 func WithLimit(limit rate.Limit) ClientOption         { return withLimit(limit) }
-func WithModelConfig(config ModelConfig) ClientOption { return withModel(config) }
+func WithModel(model string) ClientOption             { return withModel(model) }
+func WithModelConfig(config ModelConfig) ClientOption { return withModelConfig(config) }
 
 type withAPIKey string
 
@@ -63,6 +64,10 @@ type withLimit rate.Limit
 
 func (w withLimit) Apply(cfg *ClientConfig) { cfg.Limit = (*rate.Limit)(&w) }
 
-type withModel ModelConfig
+type withModel string
 
-func (w withModel) Apply(cfg *ClientConfig) { cfg.ModelConfig = ModelConfig(w) }
+func (w withModel) Apply(cfg *ClientConfig) { cfg.Model = string(w) }
+
+type withModelConfig ModelConfig
+
+func (w withModelConfig) Apply(cfg *ClientConfig) { cfg.ModelConfig = ModelConfig(w) }