feat: update sonar models #2 (#3)

sgaunet · Sep 23, 2024 · 0bb3599 · 0bb3599
1 parent 2cc00ce
commit 0bb3599
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 26 deletions.
diff --git a/perplexity.go b/perplexity.go
@@ -70,7 +70,7 @@ func NewClient(apiKey string) *Client {
 		httpClient:  &http.Client{},
 		httpTimeout: DefautTimeout,
 	}
-	s.SetModuleLlama3SonarSmall32kOnline()
+	s.SetModuleLlama31SonarSmall128kOnline()
 	return s
 }
 
@@ -84,24 +84,24 @@ func (s *Client) GetModel() string {
 	return s.model
 }
 
-// SetModuleLlama3SonarSmall32kChat sets the model to llama-3-sonar-small-32k-chat.
-func (s *Client) SetModuleLlama3SonarSmall32kChat() {
-	s.setModel("llama-3-sonar-small-32k-chat")
+// SetModuleLlama31SonarSmall128kChat sets the model to llama-3.1-sonar-small-128k-chat.
+func (s *Client) SetModuleLlama31SonarSmall128kChat() {
+	s.setModel("llama-3.1-sonar-small-128k-chat")
 }
 
-// SetModuleLlama3SonarSmall32kOnline sets the model to llama-3-sonar-small-32k-online.
-func (s *Client) SetModuleLlama3SonarSmall32kOnline() {
-	s.setModel("llama-3-sonar-small-32k-online")
+// SetModuleLlama31SonarLarge128kChat sets the model to llama-3.1-sonar-large-128k-chat.
+func (s *Client) SetModuleLlama31SonarLarge128kChat() {
+	s.setModel("llama-3.1-sonar-large-128k-chat")
 }
 
-// SetModuleLlama3SonarLarge32kChat sets the model to llama-3-sonar-large-32k-chat.
-func (s *Client) SetModuleLlama3SonarLarge32kChat() {
-	s.setModel("llama-3-sonar-large-32k-chat")
+// SetModuleLlama31SonarSmall128kOnline sets the model to llama-3.1-sonar-small-128k-online.
+func (s *Client) SetModuleLlama31SonarSmall128kOnline() {
+	s.setModel("llama-3.1-sonar-small-128k-online")
 }
 
-// SetModuleLlama3SonarLarge32kOnline sets the model to llama-3-sonar-large-32k-online.
-func (s *Client) SetModuleLlama3SonarLarge32kOnline() {
-	s.setModel("llama-3-sonar-large-32k-online")
+// SetModuleLlama31SonarLarge128kChat sets the model to llama-3.1-sonar-large-128k-online.
+func (s *Client) SetModuleLlama31SonarLarge128kOnline() {
+	s.setModel("llama-3.1-sonar-large-128k-online")
 }
 
 // SetModuleLlama3SonarLarge32kInstruct sets the model to llama-3-sonar-large-32k-instruct.

diff --git a/perplexity_test.go b/perplexity_test.go
@@ -46,7 +46,7 @@ func TestGetCompletion(t *testing.T) {
 				defer r.Body.Close()
 				b, err := io.ReadAll(r.Body)
 				assert.Nil(t, err)
-				assert.Equal(t, string(b), `{"messages":[{"role":"user","content":"What's the capital of France?"}],"model":"llama-3-sonar-small-32k-online"}`)
+				assert.Equal(t, string(b), `{"messages":[{"role":"user","content":"What's the capital of France?"}],"model":"llama-3.1-sonar-small-128k-online"}`)
 				w.Header().Add("Content-Type", "application/json")
 				fmt.Fprintln(w, "{}")
 			}))
@@ -78,21 +78,21 @@ func TestGetCompletion(t *testing.T) {
 
 func TestSetModels(t *testing.T) {
 	r := perplexity.NewClient("apikey")
-	t.Run("set model llama-3-sonar-small-32k-chat", func(t *testing.T) {
-		r.SetModuleLlama3SonarSmall32kChat()
-		assert.Equal(t, r.GetModel(), "llama-3-sonar-small-32k-chat")
+	t.Run("set model llama-3.1-sonar-small-128k-chat", func(t *testing.T) {
+		r.SetModuleLlama31SonarSmall128kChat()
+		assert.Equal(t, r.GetModel(), "llama-3.1-sonar-small-128k-chat")
 	})
-	t.Run("set model llama-3-sonar-small-32k-online", func(t *testing.T) {
-		r.SetModuleLlama3SonarSmall32kOnline()
-		assert.Equal(t, r.GetModel(), "llama-3-sonar-small-32k-online")
+	t.Run("set model llama-3.1-sonar-large-128k-chat", func(t *testing.T) {
+		r.SetModuleLlama31SonarLarge128kChat()
+		assert.Equal(t, r.GetModel(), "llama-3.1-sonar-large-128k-chat")
 	})
-	t.Run("set model llama-3-sonar-large-32k-chat", func(t *testing.T) {
-		r.SetModuleLlama3SonarLarge32kChat()
-		assert.Equal(t, r.GetModel(), "llama-3-sonar-large-32k-chat")
+	t.Run("set model llama-3.1-sonar-small-128k-online", func(t *testing.T) {
+		r.SetModuleLlama31SonarSmall128kOnline()
+		assert.Equal(t, r.GetModel(), "llama-3.1-sonar-small-128k-online")
 	})
-	t.Run("set model llama-3-sonar-large-32k-online", func(t *testing.T) {
-		r.SetModuleLlama3SonarLarge32kOnline()
-		assert.Equal(t, r.GetModel(), "llama-3-sonar-large-32k-online")
+	t.Run("set model llama-3.1-sonar-large-128k-online", func(t *testing.T) {
+		r.SetModuleLlama31SonarLarge128kOnline()
+		assert.Equal(t, r.GetModel(), "llama-3.1-sonar-large-128k-online")
 	})
 	t.Run("set model llama-3-8b-instruct", func(t *testing.T) {
 		r.SetModuleLlama38bInstruct()