Allow setting of org in config

TimeSurgeLabs · Aug 13, 2023 · 3bff8a9 · 3bff8a9
1 parent 0c4863a
commit 3bff8a9
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 4 deletions.
diff --git a/cmd/config.go b/cmd/config.go
@@ -39,7 +39,7 @@ GitHub Token Generation: https://github.com/settings/tokens
 		}
 
 		// if none of the config options are provided, print a warning
-		if apiKey == "" && model == "" && ghToken == "" && userColor == "" && ottoColor == "" {
+		if apiKey == "" && model == "" && ghToken == "" && userColor == "" && ottoColor == "" && organization == "" {
 			log.Warn("No configuration options provided")
 			os.Exit(0)
 		}
@@ -79,6 +79,12 @@ GitHub Token Generation: https://github.com/settings/tokens
 			c.OttoColor = ottoColor
 		}
 
+		// if the organization is provided, set it
+		if organization != "" {
+			fmt.Println("Setting organization...")
+			c.Org = organization
+		}
+
 		// save the config
 		err = c.Save()
 		if err != nil {
@@ -103,4 +109,6 @@ func init() {
 	configCmd.Flags().StringVarP(&userColor, "userColor", "u", "", "User color for configuration")
 	// set otto color
 	configCmd.Flags().StringVarP(&ottoColor, "ottoColor", "o", "", "Otto color for configuration")
+	// set organization
+	configCmd.Flags().StringVarP(&organization, "organization", "g", "", "Organization to use for documentation")
 }
diff --git a/cmd/vars.go b/cmd/vars.go
@@ -57,6 +57,7 @@ var deleteHistory string
 var readOnly bool
 var clearHistory bool
 var repoContext bool
+var organization string
 
 var log = l.NewWithOptions(os.Stderr, l.Options{
 	Level:           l.InfoLevel,

diff --git a/pkg/ai/req.go b/pkg/ai/req.go
@@ -10,7 +10,11 @@ import (
 
 func request(systemMsg, userMsg string, conf *config.Config) (string, error) {
 
-	c := openai.NewClient(conf.APIKey)
+	config := openai.DefaultConfig(conf.APIKey)
+	config.OrgID = conf.Org
+
+	c := openai.NewClientWithConfig(config)
+
 	ctx := context.Background()
 
 	req := openai.ChatCompletionRequest{
@@ -40,7 +44,11 @@ func request(systemMsg, userMsg string, conf *config.Config) (string, error) {
 }
 
 func requestStream(systemMsg, userMsg string, conf *config.Config) (*openai.ChatCompletionStream, error) {
-	c := openai.NewClient(conf.APIKey)
+	config := openai.DefaultConfig(conf.APIKey)
+	config.OrgID = conf.Org
+
+	c := openai.NewClientWithConfig(config)
+
 	ctx := context.Background()
 
 	req := openai.ChatCompletionRequest{

diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -10,6 +10,7 @@ import (
 // Config represents the configuration file
 type Config struct {
 	APIKey    string `json:"api_key"`
+	Org       string `json:"org_id"`
 	Model     string `json:"model"`
 	GHToken   string `json:"gh_token"`
 	Signature string `json:"signature"`
@@ -117,4 +118,4 @@ func Save(c *Config) error {
 	}
 
 	return json.NewEncoder(file).Encode(c)
-}
+}