update chat/completions endpoint to allow passing a stream boolean, and handling a streaming response

austin-denoble · austin-denoble · commit c53d16a0d4e1 · 2024-06-13T23:59:34.000-04:00
diff --git a/internal/pkg/assistants/assistant_chat_completions.go b/internal/pkg/assistants/assistant_chat_completions.go
@@ -13,7 +13,7 @@ const (
 	URL_ASSISTANT_CHAT_COMPLETIONS = "/assistant/chat/%s/chat/completions"
 )
 
-func GetAssistantChatCompletions(asstName string, msg string) (*models.ChatCompletionModel, error) {
+func GetAssistantChatCompletions(asstName string, msg string, stream bool) (*models.ChatCompletionModel, error) {
 	outgoingMsg := models.ChatCompletionMessage{
 		Role:    "user",
 		Content: msg,
@@ -30,21 +30,35 @@ func GetAssistantChatCompletions(asstName string, msg string) (*models.ChatCompl
 
 	body := models.ChatCompletionRequest{
 		Messages: chat.Messages,
+		Stream:   stream,
 	}
 
 	assistantDataUrl, err := GetAssistantDataBaseUrl()
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := network.PostAndDecode[models.ChatCompletionRequest, models.ChatCompletionModel](
-		assistantDataUrl,
-		fmt.Sprintf(URL_ASSISTANT_CHAT_COMPLETIONS, asstName),
-		true,
-		body,
-	)
-	if err != nil {
-		return nil, err
+	var resp *models.ChatCompletionModel
+	if !stream {
+		resp, err = network.PostAndDecode[models.ChatCompletionRequest, models.ChatCompletionModel](
+			assistantDataUrl,
+			fmt.Sprintf(URL_ASSISTANT_CHAT_COMPLETIONS, asstName),
+			true,
+			body,
+		)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		resp, err = network.PostAndStreamChatResponse[models.ChatCompletionRequest](
+			assistantDataUrl,
+			fmt.Sprintf(URL_ASSISTANT_CHAT_COMPLETIONS, asstName),
+			true,
+			body,
+		)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	// If the request was successful, update the chat history
diff --git a/internal/pkg/assistants/assistant_file_upload.go b/internal/pkg/assistants/assistant_file_upload.go
@@ -16,7 +16,7 @@ func UploadAssistantFile(name string, filePath string) (*AssistantFileModel, err
 		return nil, err
 	}
 
-	resp, err := network.PostAndDecodeMultipartFormData[AssistantFileModel](
+	resp, err := network.PostMultipartFormDataAndDecode[AssistantFileModel](
 		assistantDataUrl,
 		fmt.Sprintf(URL_ASSISTANT_FILE_UPLOAD, name),
 		true,
diff --git a/internal/pkg/cli/command/assistant/chat.go b/internal/pkg/cli/command/assistant/chat.go
@@ -23,6 +23,7 @@ import (
 type AssistantChatCmdOptions struct {
 	name    string
 	message string
+	stream  bool
 	json    bool
 }
 
@@ -43,19 +44,20 @@ func NewAssistantChatCmd() *cobra.Command {
 				return
 			}
 
-			// If no message is provided drop them into chat
+			// If no message is provided drop them into interactive chat
 			if options.message == "" {
 				startChat(options.name)
 			} else {
 				// If message is provided, send it to the assistant
-				sendMessage(options.name, options.message)
+				sendMessage(options.name, options.message, options.stream)
 			}
 		},
 	}
 
 	cmd.Flags().StringVarP(&options.name, "name", "n", "", "name of the assistant to chat with")
-	cmd.Flags().BoolVar(&options.json, "json", false, "output as JSON")
 	cmd.Flags().StringVarP(&options.message, "message", "m", "", "your message to the assistant")
+	cmd.Flags().BoolVarP(&options.stream, "stream", "s", false, "stream chat message responses")
+	cmd.Flags().BoolVar(&options.json, "json", false, "output as JSON")
 	cmd.MarkFlagRequired("content")
 
 	cmd.AddCommand(NewAssistantChatClearCmd())
@@ -94,7 +96,8 @@ func startChat(asstName string) {
 		checkForChatCommands(text)
 
 		if text != "" {
-			_, err := sendMessage(asstName, text)
+			// Stream here since we're in interactive chat mode
+			_, err := sendMessage(asstName, text, true)
 			if err != nil {
 				pcio.Printf("Error sending message: %s\n", err)
 				continue
@@ -103,27 +106,54 @@ func startChat(asstName string) {
 	}
 }
 
-func sendMessage(asstName string, message string) (*models.ChatCompletionModel, error) {
+func sendMessage(asstName string, message string, stream bool) (*models.ChatCompletionModel, error) {
 	response := &models.ChatCompletionModel{}
 
-	err := style.Spinner("", func() error {
-		chatResponse, err := assistants.GetAssistantChatCompletions(asstName, message)
+	var chatGetter func() error
+	if stream {
+		chatGetter = streamChatResponse(response, asstName, message, stream)
+	} else {
+		chatGetter = getChatResponse(response, asstName, message, stream)
+	}
+
+	err := chatGetter()
+	if err != nil {
+		return nil, err
+	}
+
+	return response, nil
+}
+
+func getChatResponse(resp *models.ChatCompletionModel, asstName string, message string, stream bool) func() error {
+	return func() error {
+		chatResponse, err := assistants.GetAssistantChatCompletions(asstName, message, stream)
 		if err != nil {
 			exit.Error(err)
 		}
 
-		response = chatResponse
+		resp = chatResponse
 
 		for _, choice := range chatResponse.Choices {
 			presenters.PrintAssistantChatResponse(choice.Message.Content)
 		}
 		return nil
-	})
-	if err != nil {
-		return nil, err
 	}
+}
 
-	return response, nil
+func streamChatResponse(resp *models.ChatCompletionModel, asstName string, message string, stream bool) func() error {
+	return func() error {
+		chatResponse, err := assistants.GetAssistantChatCompletions(asstName, message, stream)
+		if err != nil {
+			exit.Error(err)
+		}
+
+		resp = chatResponse
+
+		for _, choice := range chatResponse.Choices {
+			presenters.PrintAssistantChatResponse(choice.Message.Content)
+		}
+		return nil
+	}
 }
 
 func displayChatHistory(asstName string, maxNoMsgs int) {
@@ -137,7 +167,7 @@ func displayChatHistory(asstName string, maxNoMsgs int) {
 	presenters.PrintChatHistory(chat, maxNoMsgs)
 }
 
-// This function checks the input for accepted chat commands
+// Checks the input for accepted chat commands
 func checkForChatCommands(text string) {
 	switch text {
 	case "exit()":
diff --git a/internal/pkg/utils/models/models.go b/internal/pkg/utils/models/models.go
@@ -3,6 +3,7 @@ package models
 import "time"
 
 type ChatCompletionRequest struct {
+	Stream   bool                    `json:"stream"`
 	Messages []ChatCompletionMessage `json:"messages"`
 }
 
@@ -12,6 +13,12 @@ type ChatCompletionModel struct {
 	Model   string        `json:"model"`
 }
 
+type ChoiceModel struct {
+	FinishReason ChatFinishReason      `json:"finish_reason"`
+	Index        int32                 `json:"index"`
+	Message      ChatCompletionMessage `json:"message"`
+}
+
 type ChatCompletionMessage struct {
 	Role    string `json:"role"`
 	Content string `json:"content"`
@@ -33,8 +40,26 @@ const (
 	FunctionCall  ChatFinishReason = "function_call"
 )
 
-type ChoiceModel struct {
+type StreamChatCompletionModel struct {
+	Id      string             `json:"id"`
+	Choices []ChoiceChunkModel `json:"choices"`
+	Model   string             `json:"model"`
+}
+
+type StreamChunk struct {
+	Data StreamChatCompletionModel `json:"data"`
+}
+
+type ChoiceChunkModel struct {
 	FinishReason ChatFinishReason      `json:"finish_reason"`
 	Index        int32                 `json:"index"`
-	Message      ChatCompletionMessage `json:"message"`
+	Delta        ChatCompletionMessage `json:"delta"`
+}
+
+type ContextRefModel struct {
+	Id     string   `json:"id"`
+	Source string   `json:"source"`
+	Text   string   `json:"text"`
+	Score  float64  `json:"score"`
+	Path   []string `json:"path"`
 }
diff --git a/internal/pkg/utils/network/request.go b/internal/pkg/utils/network/request.go
@@ -85,6 +85,45 @@ func decodeResponse[T any](resp *http.Response, target *T) error {
 	return nil
 }
 
+func RequestWithBody[B any](baseUrl string, path string, method string, useApiKey bool, body B) (*http.Response, error) {
+	url := baseUrl + path
+
+	var bodyJson []byte
+	bodyJson, err := json.Marshal(body)
+	if err != nil {
+		log.Error().
+			Err(err).
+			Str("method", method).
+			Str("url", url).
+			Msg("Error marshalling JSON")
+		return nil, pcio.Errorf("error marshalling JSON: %v", err)
+	}
+
+	req, err := buildRequest(method, url, bodyJson)
+	log.Info().
+		Str("method", method).
+		Str("url", url).
+		Msg("Fetching data from dashboard")
+	if err != nil {
+		log.Error().
+			Err(err).
+			Str("url", url).
+			Str("method", method).
+			Msg("Error building request")
+		return nil, pcio.Errorf("error building request: %v", err)
+	}
+
+	resp, err := performRequest(req, useApiKey)
+	if err != nil {
+		log.Error().
+			Err(err).
+			Str("method", method).
+			Str("url", url)
+		return nil, pcio.Errorf("error performing request to %s: %v", url, err)
+	}
+	return resp, nil
+}
+
 func RequestWithBodyAndDecode[B any, R any](baseUrl string, path string, method string, useApiKey bool, body B) (*R, error) {
 	url := baseUrl + path
 
diff --git a/internal/pkg/utils/network/request_post.go b/internal/pkg/utils/network/request_post.go
@@ -1,22 +1,92 @@
 package network
 
 import (
+	"bufio"
 	"bytes"
+	"encoding/json"
+	"fmt"
 	"io"
 	"mime/multipart"
 	"net/http"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"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/pcio"
 )
 
 func PostAndDecode[B any, R any](baseUrl string, path string, useApiKey bool, body B) (*R, error) {
 	return RequestWithBodyAndDecode[B, R](baseUrl, path, http.MethodPost, useApiKey, body)
 }
 
-func PostAndDecodeMultipartFormData[R any](baseUrl string, path string, useApiKey bool, bodyPath string) (*R, error) {
+func PostAndStreamChatResponse[B any](baseUrl string, path string, useApiKey bool, body B) (*models.ChatCompletionModel, error) {
+	resp, err := RequestWithBody[B](baseUrl, path, http.MethodPost, useApiKey, body)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	var completeResponse string
+	var id string
+	var model string
+
+	scanner := bufio.NewScanner(resp.Body)
+	for scanner.Scan() {
+		line := scanner.Text()
+		if strings.HasPrefix(line, "data:") {
+			dataStr := strings.TrimPrefix(line, "data:")
+			dataStr = strings.TrimSpace(dataStr)
+
+			var chunkResp *models.StreamChatCompletionModel
+			if err := json.Unmarshal([]byte(dataStr), &chunkResp); err != nil {
+				return nil, pcio.Errorf("error unmarshaling chunk: %v", err)
+			}
+
+			for _, choice := range chunkResp.Choices {
+				fmt.Print(choice.Delta.Content)
+				os.Stdout.Sync()
+				completeResponse += choice.Delta.Content
+			}
+			id = chunkResp.Id
+			model = chunkResp.Model
+		}
+	}
+
+	completionResp := &models.ChatCompletionModel{
+		Id:    id,
+		Model: model,
+		Choices: []models.ChoiceModel{
+			{
+				FinishReason: "stop",
+				Index:        0,
+				Message: models.ChatCompletionMessage{
+					Content: completeResponse,
+					Role:    "assistant",
+				},
+			},
+		},
+	}
+
+	if err != nil {
+		log.Error().
+			Err(err).
+			Str("method", http.MethodPost).
+			Str("url", baseUrl+path).
+			Str("status", resp.Status).
+			Msg("Error decoding response")
+		return nil, pcio.Errorf("error decoding JSON: %v", err)
+	}
+
+	log.Info().
+		Str("method", http.MethodPost).
+		Str("url", baseUrl+path).
+		Msg("Request completed successfully")
+	return completionResp, nil
+}
+
+func PostMultipartFormDataAndDecode[R any](baseUrl string, path string, useApiKey bool, bodyPath string) (*R, error) {
 	url := baseUrl + path
 
 	var requestBody bytes.Buffer

Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,7 @@ func UploadAssistantFile(name string, filePath string) (*AssistantFileModel, err`
`16`	`16`	`return nil, err`
`17`	`17`	`}`
`18`	`18`
`19`		`- resp, err := network.PostAndDecodeMultipartFormData[AssistantFileModel](`
	`19`	`+ resp, err := network.PostMultipartFormDataAndDecode[AssistantFileModel](`
`20`	`20`	`assistantDataUrl,`
`21`	`21`	`fmt.Sprintf(URL_ASSISTANT_FILE_UPLOAD, name),`
`22`	`22`	`true,`