From 092df753a411db57f4ff35744fa5222800e95310 Mon Sep 17 00:00:00 2001
From: jackspirou <jack.spirou@me.com>
Date: Wed, 13 Mar 2024 20:58:59 -0500
Subject: [PATCH] Update Groq model options and documentation

- Add Groq model options via Groq's OpenAI API compatibility
- Update documentation accordingly
- Fix example
- Fix golangci-lint compliants
- fix comment typo
---
 completion.go   | 13 +++++++++++++
 config.go       | 19 ++++++++++++++++++-
 example_test.go | 35 +++++++++++++++++++++++++++++++++++
 test.mp3        |  1 +
 4 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 test.mp3

diff --git a/completion.go b/completion.go
index ab1dbd6c5..f69124595 100644
--- a/completion.go
+++ b/completion.go
@@ -109,6 +109,19 @@ var disabledModelsForEndpoints = map[string]map[string]bool{
 	},
 }
 
+// Groq contains models that work with Groq's OpenAI Compatibility API.
+//
+// Usage Examples: openai.Groq.Mixtral_8x7b, openai.Groq.LLaMA2_70b, openai.Groq.Gemma_7b_IT.
+var Groq = struct {
+	Mixtral8x7b string
+	LLaMA270b   string
+	Gemma7bIT   string
+}{
+	Mixtral8x7b: "mixtral-8x7b-32768",
+	LLaMA270b:   "llama2-70b-4096",
+	Gemma7bIT:   "gemma-7b-it",
+}
+
 func checkEndpointSupportsModel(endpoint, model string) bool {
 	return !disabledModelsForEndpoints[endpoint][model]
 }
diff --git a/config.go b/config.go
index c58b71ec6..5b383c3e7 100644
--- a/config.go
+++ b/config.go
@@ -6,7 +6,9 @@ import (
 )
 
 const (
-	openaiAPIURLv1                 = "https://api.openai.com/v1"
+	openaiAPIURLv1 = "https://api.openai.com/v1"
+	groqAPIURLv1   = "https://api.groq.com/openai/v1"
+
 	defaultEmptyMessagesLimit uint = 300
 
 	azureAPIPrefix         = "openai"
@@ -19,6 +21,7 @@ const (
 	APITypeOpenAI  APIType = "OPEN_AI"
 	APITypeAzure   APIType = "AZURE"
 	APITypeAzureAD APIType = "AZURE_AD"
+	APITypeGroq    APIType = "GROQ"
 )
 
 const AzureAPIKeyHeader = "api-key"
@@ -67,6 +70,20 @@ func DefaultAzureConfig(apiKey, baseURL string) ClientConfig {
 	}
 }
 
+// DefaultGroqConfig takes a Groq auth token and returns a ClientConfig that works with Groq's OpenAI Compatibility API.
+func DefaultGroqConfig(authToken string) ClientConfig {
+	return ClientConfig{
+		authToken: authToken,
+		BaseURL:   groqAPIURLv1,
+		APIType:   APITypeGroq,
+		OrgID:     "",
+
+		HTTPClient: &http.Client{},
+
+		EmptyMessagesLimit: defaultEmptyMessagesLimit,
+	}
+}
+
 func (ClientConfig) String() string {
 	return "<OpenAI API ClientConfig>"
 }
diff --git a/example_test.go b/example_test.go
index de67c57cd..797a4b2d3 100644
--- a/example_test.go
+++ b/example_test.go
@@ -345,3 +345,38 @@ func ExampleAPIError() {
 		}
 	}
 }
+
+// ExampleDefaultGroqConfig demonstrates how to create a new DefaultGroqConfig, create a Groq client,
+// use a hosted Groq model, and create a chat completion.
+func ExampleDefaultGroqConfig() {
+	config := openai.DefaultGroqConfig(os.Getenv("GROQ_API_KEY"))
+	client := openai.NewClientWithConfig(config)
+
+	req := openai.ChatCompletionRequest{
+		Model: openai.Groq.Mixtral8x7b,
+		Messages: []openai.ChatCompletionMessage{
+			{
+				Role:    openai.ChatMessageRoleSystem,
+				Content: "you are a helpful chatbot",
+			},
+		},
+	}
+	fmt.Println("Conversation")
+	fmt.Println("---------------------")
+	fmt.Print("> ")
+	s := bufio.NewScanner(os.Stdin)
+	for s.Scan() {
+		req.Messages = append(req.Messages, openai.ChatCompletionMessage{
+			Role:    openai.ChatMessageRoleUser,
+			Content: s.Text(),
+		})
+		resp, err := client.CreateChatCompletion(context.Background(), req)
+		if err != nil {
+			fmt.Printf("ChatCompletion error: %v\n", err)
+			continue
+		}
+		fmt.Printf("%s\n\n", resp.Choices[0].Message.Content)
+		req.Messages = append(req.Messages, resp.Choices[0].Message)
+		fmt.Print("> ")
+	}
+}
diff --git a/test.mp3 b/test.mp3
new file mode 100644
index 000000000..b6fc4c620
--- /dev/null
+++ b/test.mp3
@@ -0,0 +1 @@
+hello
\ No newline at end of file