update staging URL with new naming schema, add helpers for staging vs…

… production targeting
pinecone-io · Jun 11, 2024 · 326b2d7 · 326b2d7
1 parent ecdd185
commit 326b2d7
Show file tree

Hide file tree

Showing 11 changed files with 132 additions and 42 deletions.
diff --git a/internal/pkg/assistants/assistant_chat_completions.go b/internal/pkg/assistants/assistant_chat_completions.go
@@ -3,16 +3,26 @@ package assistants
 import (
 	"fmt"
 
+	"github.com/pinecone-io/cli/internal/pkg/utils/configuration/config"
 	"github.com/pinecone-io/cli/internal/pkg/utils/configuration/state"
 	"github.com/pinecone-io/cli/internal/pkg/utils/log"
 	"github.com/pinecone-io/cli/internal/pkg/utils/models"
 	"github.com/pinecone-io/cli/internal/pkg/utils/network"
 )
 
 const (
-	URL_ASSISTANT_CHAT_COMPLETIONS = "/knowledge/chat/%s/chat/completions"
+	URL_ASSISTANT_CHAT_COMPLETIONS         = "/knowledge/chat/%s/chat/completions"
+	URL_ASSISTANT_CHAT_COMPLETIONS_STAGING = "/assistant/chat/%s/chat/completions"
 )
 
+func getAssistantChatCompletionsUrl() string {
+	if config.Environment.Get() == "production" {
+		return URL_ASSISTANT_CHAT_COMPLETIONS
+	} else {
+		return URL_ASSISTANT_CHAT_COMPLETIONS_STAGING
+	}
+}
+
 func GetAssistantChatCompletions(kmName string, msg string) (*models.ChatCompletionModel, error) {
 	outgoingMsg := models.ChatCompletionMessage{
 		Role:    "user",
@@ -32,14 +42,14 @@ func GetAssistantChatCompletions(kmName string, msg string) (*models.ChatComplet
 		Messages: chat.Messages,
 	}
 
-	knowledgeDataUrl, err := GetKnowledgeDataBaseUrl()
+	assistantDataUrl, err := GetAssistantDataBaseUrl()
 	if err != nil {
 		return nil, err
 	}
 
 	resp, err := network.PostAndDecode[models.ChatCompletionRequest, models.ChatCompletionModel](
-		knowledgeDataUrl,
-		fmt.Sprintf(URL_ASSISTANT_CHAT_COMPLETIONS, kmName),
+		assistantDataUrl,
+		fmt.Sprintf(getAssistantChatCompletionsUrl(), kmName),
 		true,
 		body,
 	)

diff --git a/internal/pkg/assistants/assistant_file_delete.go b/internal/pkg/assistants/assistant_file_delete.go
@@ -3,25 +3,35 @@ package assistants
 import (
 	"net/http"
 
+	"github.com/pinecone-io/cli/internal/pkg/utils/configuration/config"
 	"github.com/pinecone-io/cli/internal/pkg/utils/network"
 	"github.com/pinecone-io/cli/internal/pkg/utils/pcio"
 )
 
 const (
-	URL_DELETE_ASSISTANT_FILE = "/knowledge/files/%s/%s"
+	URL_DELETE_ASSISTANT_FILE         = "/knowledge/files/%s/%s"
+	URL_DELETE_ASSISTANT_FILE_STAGING = "/assistant/files/%s/%s"
 )
 
+func getDeleteAssistantFileUrl() string {
+	if config.Environment.Get() == "production" {
+		return URL_DELETE_ASSISTANT_FILE
+	} else {
+		return URL_DELETE_ASSISTANT_FILE_STAGING
+	}
+}
+
 type DeleteAssistantFileResponse string
 
 func DeleteKnowledgeFile(kmName string, fileId string) (*DeleteAssistantFileResponse, error) {
-	assistantDataUrl, err := GetKnowledgeDataBaseUrl()
+	assistantDataUrl, err := GetAssistantDataBaseUrl()
 	if err != nil {
 		return nil, err
 	}
 
 	resp, err := network.RequestWithoutBodyAndDecode[DeleteAssistantFileResponse](
 		assistantDataUrl,
-		pcio.Sprintf(URL_DELETE_ASSISTANT_FILE, kmName, fileId),
+		pcio.Sprintf(getDeleteAssistantFileUrl(), kmName, fileId),
 		http.MethodDelete,
 		true,
 	)

diff --git a/internal/pkg/assistants/assistant_file_get.go b/internal/pkg/assistants/assistant_file_get.go
@@ -3,24 +3,34 @@ package assistants
 import (
 	"fmt"
 
+	"github.com/pinecone-io/cli/internal/pkg/utils/configuration/config"
 	"github.com/pinecone-io/cli/internal/pkg/utils/exit"
 	"github.com/pinecone-io/cli/internal/pkg/utils/log"
 	"github.com/pinecone-io/cli/internal/pkg/utils/network"
 )
 
 const (
-	URL_DESCRIBE_ASSISTANT_FILE = "/knowledge/files/%s/%s"
+	URL_DESCRIBE_ASSISTANT_FILE         = "/knowledge/files/%s/%s"
+	URL_DESCRIBE_ASSISTANT_FILE_STAGING = "/assistant/files/%s/%s"
 )
 
+func getDescribeAssistantFileUrl() string {
+	if config.Environment.Get() == "production" {
+		return URL_DESCRIBE_ASSISTANT_FILE
+	} else {
+		return URL_DESCRIBE_ASSISTANT_FILE_STAGING
+	}
+}
+
 func DescribeAssistantFile(name string, fileId string) (*AssistantFileModel, error) {
-	assistantDataUrl, err := GetKnowledgeDataBaseUrl()
+	assistantDataUrl, err := GetAssistantDataBaseUrl()
 	if err != nil {
 		return nil, err
 	}
 
 	resp, err := network.GetAndDecode[AssistantFileModel](
 		assistantDataUrl,
-		fmt.Sprintf(URL_DESCRIBE_ASSISTANT_FILE, name, fileId),
+		fmt.Sprintf(getDescribeAssistantFileUrl(), name, fileId),
 		true,
 	)
 	if err != nil {

diff --git a/internal/pkg/assistants/assistant_file_upload.go b/internal/pkg/assistants/assistant_file_upload.go
@@ -3,22 +3,32 @@ package assistants
 import (
 	"fmt"
 
+	"github.com/pinecone-io/cli/internal/pkg/utils/configuration/config"
 	"github.com/pinecone-io/cli/internal/pkg/utils/network"
 )
 
 const (
-	URL_ASSISTANT_FILE_UPLOAD = "/knowledge/files/%s"
+	URL_ASSISTANT_FILE_UPLOAD         = "/knowledge/files/%s"
+	URL_ASSISTANT_FILE_UPLOAD_STAGING = "/assistant/files/%s"
 )
 
+func getAssistantFileUploadUrl() string {
+	if config.Environment.Get() == "production" {
+		return URL_ASSISTANT_FILE_UPLOAD
+	} else {
+		return URL_ASSISTANT_FILE_UPLOAD_STAGING
+	}
+}
+
 func UploadAssistantFile(name string, filePath string) (*AssistantFileModel, error) {
-	assistantDataUrl, err := GetKnowledgeDataBaseUrl()
+	assistantDataUrl, err := GetAssistantDataBaseUrl()
 	if err != nil {
 		return nil, err
 	}
 
 	resp, err := network.PostAndDecodeMultipartFormData[AssistantFileModel](
 		assistantDataUrl,
-		fmt.Sprintf(URL_ASSISTANT_FILE_UPLOAD, name),
+		fmt.Sprintf(getAssistantFileUploadUrl(), name),
 		true,
 		filePath,
 	)

diff --git a/internal/pkg/assistants/assistant_files_list.go b/internal/pkg/assistants/assistant_files_list.go
@@ -3,14 +3,24 @@ package assistants
 import (
 	"fmt"
 
+	"github.com/pinecone-io/cli/internal/pkg/utils/configuration/config"
 	"github.com/pinecone-io/cli/internal/pkg/utils/log"
 	"github.com/pinecone-io/cli/internal/pkg/utils/network"
 )
 
 const (
-	URL_LIST_ASSISTANT_FILES = "/knowledge/files/%s"
+	URL_LIST_ASSISTANT_FILES         = "/knowledge/files/%s"
+	URL_LIST_ASSISTANT_FILES_STAGING = "/assistant/files/%s"
 )
 
+func getListAssistantFilesUrl() string {
+	if config.Environment.Get() == "production" {
+		return URL_LIST_ASSISTANT_FILES
+	} else {
+		return URL_LIST_ASSISTANT_FILES_STAGING
+	}
+}
+
 type AssistantFileModel struct {
 	Name      string                   `json:"name"`
 	Id        string                   `json:"id"`
@@ -33,14 +43,14 @@ type ListAssistantFilesResponse struct {
 }
 
 func ListAssistantFiles(name string) (*ListAssistantFilesResponse, error) {
-	assistantDataUrl, err := GetKnowledgeDataBaseUrl()
+	assistantDataUrl, err := GetAssistantDataBaseUrl()
 	if err != nil {
 		return nil, err
 	}
 
 	resp, err := network.GetAndDecode[ListAssistantFilesResponse](
 		assistantDataUrl,
-		fmt.Sprintf(URL_LIST_ASSISTANT_FILES, name),
+		fmt.Sprintf(getListAssistantFilesUrl(), name),
 		true,
 	)
 	if err != nil {

diff --git a/internal/pkg/assistants/constants.go b/internal/pkg/assistants/constants.go
@@ -5,18 +5,18 @@ import (
 	"github.com/pinecone-io/cli/internal/pkg/utils/environment"
 )
 
-func GetKnowledgeDataBaseUrl() (string, error) {
+func GetAssistantDataBaseUrl() (string, error) {
 	connectionConfigs, err := environment.GetEnvConfig(config.Environment.Get())
 	if err != nil {
 		return "", err
 	}
-	return connectionConfigs.KnowledgeDataPlaneUrl, nil
+	return connectionConfigs.AssistantDataPlaneUrl, nil
 }
 
-func GetKnowledgeControlBaseUrl() (string, error) {
+func GetAssistantControlBaseUrl() (string, error) {
 	connectionConfigs, err := environment.GetEnvConfig(config.Environment.Get())
 	if err != nil {
 		return "", err
 	}
-	return connectionConfigs.KnowledgeControlPlaneUrl, nil
+	return connectionConfigs.AssistantControlPlaneUrl, nil
 }
diff --git a/internal/pkg/assistants/create.go b/internal/pkg/assistants/create.go
@@ -1,13 +1,23 @@
 package assistants
 
 import (
+	"github.com/pinecone-io/cli/internal/pkg/utils/configuration/config"
 	"github.com/pinecone-io/cli/internal/pkg/utils/network"
 )
 
 const (
-	URL_CREATE_ASSISTANT = "/knowledge/models"
+	URL_CREATE_ASSISTANT         = "/knowledge/models"
+	URL_CREATE_ASSISTANT_STAGING = "/assistant/assistants"
 )
 
+func getCreateAssistantUrl() string {
+	if config.Environment.Get() == "production" {
+		return URL_CREATE_ASSISTANT
+	} else {
+		return URL_CREATE_ASSISTANT_STAGING
+	}
+}
+
 type CreateAssistantRequest struct {
 	Name     string                 `json:"name"`
 	Metadata map[string]interface{} `json:"metadata"`
@@ -18,14 +28,14 @@ func CreateAssistant(name string) (*AssistantModel, error) {
 		Name: name,
 	}
 
-	knowledgeControlUrl, err := GetKnowledgeControlBaseUrl()
+	assistantControlUrl, err := GetAssistantControlBaseUrl()
 	if err != nil {
 		return nil, err
 	}
 
 	resp, err := network.PostAndDecode[CreateAssistantRequest, AssistantModel](
-		knowledgeControlUrl,
-		URL_CREATE_ASSISTANT,
+		assistantControlUrl,
+		getCreateAssistantUrl(),
 		true,
 		body,
 	)

diff --git a/internal/pkg/assistants/delete.go b/internal/pkg/assistants/delete.go
@@ -3,28 +3,38 @@ package assistants
 import (
 	"net/http"
 
+	"github.com/pinecone-io/cli/internal/pkg/utils/configuration/config"
 	"github.com/pinecone-io/cli/internal/pkg/utils/network"
 	"github.com/pinecone-io/cli/internal/pkg/utils/pcio"
 )
 
 const (
-	URL_DELETE_ASSISTANT = "/knowledge/models/%s"
+	URL_DELETE_ASSISTANT         = "/knowledge/models/%s"
+	URL_DELETE_ASSISTANT_STAGING = "/assistant/assistants/%s"
 )
 
+func getDeleteAssistantUrl() string {
+	if config.Environment.Get() == "production" {
+		return URL_DELETE_ASSISTANT
+	} else {
+		return URL_DELETE_ASSISTANT_STAGING
+	}
+}
+
 type DeleteKnowledgeModelResponse struct {
 	Success bool `json:"success"`
 }
 
 func DeleteAssistant(name string) (*DeleteKnowledgeModelResponse, error) {
 
-	knowledgeControlUrl, err := GetKnowledgeControlBaseUrl()
+	assistantControlUrl, err := GetAssistantControlBaseUrl()
 	if err != nil {
 		return nil, err
 	}
 
 	resp, err := network.RequestWithoutBodyAndDecode[DeleteKnowledgeModelResponse](
-		knowledgeControlUrl,
-		pcio.Sprintf(URL_DELETE_ASSISTANT, name),
+		assistantControlUrl,
+		pcio.Sprintf(getDeleteAssistantUrl(), name),
 		http.MethodDelete,
 		true,
 	)

diff --git a/internal/pkg/assistants/get.go b/internal/pkg/assistants/get.go
@@ -3,23 +3,33 @@ package assistants
 import (
 	"fmt"
 
+	"github.com/pinecone-io/cli/internal/pkg/utils/configuration/config"
 	"github.com/pinecone-io/cli/internal/pkg/utils/log"
 	"github.com/pinecone-io/cli/internal/pkg/utils/network"
 )
 
 const (
-	URL_DESCRIBE_ASSISTANT = "/knowledge/models/%s"
+	URL_DESCRIBE_ASSISTANT         = "/knowledge/models/%s"
+	URL_DESCRIBE_ASSISTANT_STAGING = "/assistant/assistants/%s"
 )
 
+func getDescribeAssistantUrl() string {
+	if config.Environment.Get() == "production" {
+		return URL_DESCRIBE_ASSISTANT
+	} else {
+		return URL_DESCRIBE_ASSISTANT_STAGING
+	}
+}
+
 func DescribeAssistant(name string) (*AssistantModel, error) {
-	knowledgeControlUrl, err := GetKnowledgeControlBaseUrl()
+	assistantControlUrl, err := GetAssistantControlBaseUrl()
 	if err != nil {
 		return nil, err
 	}
 
 	resp, err := network.GetAndDecode[AssistantModel](
-		knowledgeControlUrl,
-		fmt.Sprintf(URL_DESCRIBE_ASSISTANT, name),
+		assistantControlUrl,
+		fmt.Sprintf(getDescribeAssistantUrl(), name),
 		true,
 	)
 	if err != nil {

diff --git a/internal/pkg/assistants/list.go b/internal/pkg/assistants/list.go
@@ -3,14 +3,24 @@ package assistants
 import (
 	"encoding/json"
 
+	"github.com/pinecone-io/cli/internal/pkg/utils/configuration/config"
 	"github.com/pinecone-io/cli/internal/pkg/utils/log"
 	"github.com/pinecone-io/cli/internal/pkg/utils/network"
 )
 
 const (
-	URL_LIST_ASSISTANTS = "/knowledge/models"
+	URL_LIST_ASSISTANTS         = "/knowledge/models"
+	URL_LIST_ASSISTANTS_STAGING = "/assistant/assistants"
 )
 
+func getListAssistantsUrl() string {
+	if config.Environment.Get() == "production" {
+		return URL_LIST_ASSISTANTS
+	} else {
+		return URL_LIST_ASSISTANTS_STAGING
+	}
+}
+
 type AssistantModel struct {
 	Name      string               `json:"name"`
 	Metadata  AssistantMetadata    `json:"metadata"`
@@ -44,14 +54,14 @@ type ListAssistantsResponse struct {
 }
 
 func ListAssistants() (*ListAssistantsResponse, error) {
-	knowledgeControlUrl, err := GetKnowledgeControlBaseUrl()
+	assistantControlUrl, err := GetAssistantControlBaseUrl()
 	if err != nil {
 		return nil, err
 	}
 
 	resp, err := network.GetAndDecode[ListAssistantsResponse](
-		knowledgeControlUrl,
-		URL_LIST_ASSISTANTS,
+		assistantControlUrl,
+		getListAssistantsUrl(),
 		true,
 	)
 	if err != nil {