diff --git a/cmd/config.go b/cmd/config.go index 5272efc..95cb85d 100644 --- 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 index a736634..da525cb 100644 --- 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 index 0cb7a01..5838c39 100644 --- 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 index 08ac7d3..116e470 100644 --- 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) -} \ No newline at end of file +}