Refactor Query function to include token usage in response

README.md

@@ -36,12 +36,12 @@ Azure, featuring streaming capabilities and extensive configuration options.
 * **Streaming mode**: Real-time interaction with the GPT model.
 * **Query mode**: Single input-output interactions with the GPT model.
 * **Interactive mode**: The interactive mode allows for a more conversational experience with the model. Exit
-  interactive mode by simply typing 'exit'.
+  interactive mode by simply typing 'exit'. Prints the token usage when combined with query mode.
 * **Thread-based context management**: Enjoy seamless conversations with the GPT model with individualized context for
   each thread, much like your experience on the OpenAI website. Each unique thread has its own history, ensuring
   relevant and coherent responses across different chat instances.
 * **Sliding window history**: To stay within token limits, the chat history automatically trims while still preserving
-  the necessary context. The size of this window can be adjusted through the `context-window` setting. 
+  the necessary context. The size of this window can be adjusted through the `context-window` setting.
 * **Custom context from any source**: You can provide the GPT model with a custom context during conversation. This
   context can be piped in from any source, such as local files, standard input, or even another program. This
   flexibility allows the model to adapt to a wide range of conversational scenarios.
@@ -49,7 +49,7 @@ Azure, featuring streaming capabilities and extensive configuration options.
 * **Thread listing**: Display a list of active threads using the `--list-threads` flag.
 * **Advanced configuration options**: The CLI supports a layered configuration system where settings can be specified
   through default values, a `config.yaml` file, and environment variables. For quick adjustments,
-  various `--set-<value>` flags are provided. To verify your current settings, use the `--config` or `-c` flag. 
+  various `--set-<value>` flags are provided. To verify your current settings, use the `--config` or `-c` flag.
 * **Availability Note**: This CLI supports both gpt-4 and gpt-3.5-turbo models. However, the specific ChatGPT model used
   on chat.openai.com may not be available via the OpenAI API.
 

client/client.go

@@ -104,37 +104,35 @@ func (c *Client) ProvideContext(context string) {
 	c.History = append(c.History, messages...)
 }
 
-// Query sends a query to the API and returns the response as a string.
-// It takes an input string as a parameter and returns a string containing
-// the API response or an error if there's any issue during the process.
-// The method creates a request body with the input and then makes an API
-// call using the Post method. If the response is not empty, it decodes the
-// response JSON and returns the content of the first choice.
-func (c *Client) Query(input string) (string, error) {
+// Query sends a query to the API, returning the response as a string along with the token usage.
+// It takes an input string, constructs a request body, and makes a POST API call.
+// Returns the API response string, the number of tokens used, and an error if any issues occur.
+// If the response contains choices, it decodes the JSON and returns the content of the first choice.
+func (c *Client) Query(input string) (string, int, error) {
 	c.prepareQuery(input)
 
 	body, err := c.createBody(false)
 	if err != nil {
-		return "", err
+		return "", 0, err
 	}
 
 	raw, err := c.caller.Post(c.getEndpoint(c.Config.CompletionsPath), body, false)
 	if err != nil {
-		return "", err
+		return "", 0, err
 	}
 
 	var response types.CompletionsResponse
 	if err := c.processResponse(raw, &response); err != nil {
-		return "", err
+		return "", 0, err
 	}
 
 	if len(response.Choices) == 0 {
-		return "", errors.New("no responses returned")
+		return "", response.Usage.TotalTokens, errors.New("no responses returned")
 	}
 
 	c.updateHistory(response.Choices[0].Message.Content)
 
-	return response.Choices[0].Message.Content, nil
+	return response.Choices[0].Message.Content, response.Usage.TotalTokens, nil
 }
 
 // Stream sends a query to the API and processes the response as a stream.

client/client_test.go

@@ -153,7 +153,7 @@ func testClient(t *testing.T, when spec.G, it spec.S) {
 				Expect(err).NotTo(HaveOccurred())
 				mockCaller.EXPECT().Post(subject.Config.URL+subject.Config.CompletionsPath, body, false).Return(respBytes, tt.postError)
 
-				_, err = subject.Query(query)
+				_, _, err = subject.Query(query)
 				Expect(err).To(HaveOccurred())
 				Expect(err.Error()).To(ContainSubstring(tt.expectedError))
 			})
@@ -161,7 +161,10 @@ func testClient(t *testing.T, when spec.G, it spec.S) {
 
 		when("a valid http response is received", func() {
 			testValidHTTPResponse := func(subject *client.Client, history []types.Message, expectedBody []byte, omitHistory bool) {
-				const answer = "content"
+				const (
+					answer = "content"
+					tokens = 789
+				)
 
 				choice := types.Choice{
 					Message: types.Message{
@@ -177,6 +180,11 @@ func testClient(t *testing.T, when spec.G, it spec.S) {
 					Created: 0,
 					Model:   subject.Config.Model,
 					Choices: []types.Choice{choice},
+					Usage: types.Usage{
+						PromptTokens:     123,
+						CompletionTokens: 456,
+						TotalTokens:      tokens,
+					},
 				}
 
 				respBytes, err := json.Marshal(response)
@@ -194,9 +202,10 @@ func testClient(t *testing.T, when spec.G, it spec.S) {
 					}))
 				}
 
-				result, err := subject.Query(query)
+				result, usage, err := subject.Query(query)
 				Expect(err).NotTo(HaveOccurred())
 				Expect(result).To(Equal(answer))
+				Expect(usage).To(Equal(tokens))
 			}
 			it("uses the values specified by the configuration instead of the default values", func() {
 				const (

cmd/chatgpt/main.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"strconv"
 	"strings"
 	"time"
 
@@ -202,22 +203,23 @@ func run(cmd *cobra.Command, args []string) error {
 
 	if interactiveMode {
 		fmt.Printf("Entering interactive mode. Type 'exit' and press Enter or press Ctrl+C to quit.\n\n")
-
-		// Initialize readline with an empty prompt first, as we'll set it dynamically.
 		rl, err := readline.New("")
 		if err != nil {
 			return err
 		}
 		defer rl.Close()
 
-		prompt := func(qNum int) string {
-			return fmt.Sprintf("[%s] Q%d: ", time.Now().Format("2006-01-02 15:04:05"), qNum)
+		prompt := func(counter string) string {
+			return fmt.Sprintf("[%s] [%s]: ", time.Now().Format("2006-01-02 15:04:05"), counter)
 		}
 
-		qNum := 1
+		qNum, usage := 1, 0
 		for {
-			// Set and update the readline prompt dynamically
-			rl.SetPrompt(prompt(qNum))
+			if queryMode {
+				rl.SetPrompt(prompt(strconv.Itoa(usage)))
+			} else {
+				rl.SetPrompt(prompt(fmt.Sprintf("Q%d", qNum)))
+			}
 
 			line, err := rl.Readline()
 			if err == readline.ErrInterrupt || err == io.EOF {
@@ -225,16 +227,29 @@ func run(cmd *cobra.Command, args []string) error {
 				break
 			}
 
-			if line == "exit" {
+			if line == "exit" || line == "/q" {
 				fmt.Println("Bye!")
+				if queryMode {
+					fmt.Printf("Total tokens used: %d\n", usage)
+				}
 				break
 			}
 
-			if err := client.Stream(line); err != nil {
-				fmt.Println("Error:", err)
+			if queryMode {
+				result, qUsage, err := client.Query(line)
+				if err != nil {
+					fmt.Println("Error:", err)
+				} else {
+					fmt.Printf("%s\n\n", result)
+					usage += qUsage
+				}
 			} else {
-				fmt.Println()
-				qNum++
+				if err := client.Stream(line); err != nil {
+					fmt.Println("Error:", err)
+				} else {
+					fmt.Println()
+					qNum++
+				}
 			}
 		}
 	} else {
@@ -242,7 +257,7 @@ func run(cmd *cobra.Command, args []string) error {
 			return errors.New("you must specify your query")
 		}
 		if queryMode {
-			result, err := client.Query(strings.Join(args, " "))
+			result, _, err := client.Query(strings.Join(args, " "))
 			if err != nil {
 				return err
 			}

types/completions.go

@@ -17,18 +17,20 @@ type Message struct {
 }
 
 type CompletionsResponse struct {
-	ID      string `json:"id"`
-	Object  string `json:"object"`
-	Created int    `json:"created"`
-	Model   string `json:"model"`
-	Usage   struct {
-		PromptTokens     int `json:"prompt_tokens"`
-		CompletionTokens int `json:"completion_tokens"`
-		TotalTokens      int `json:"total_tokens"`
-	} `json:"usage"`
+	ID      string   `json:"id"`
+	Object  string   `json:"object"`
+	Created int      `json:"created"`
+	Model   string   `json:"model"`
+	Usage   Usage    `json:"usage"`
 	Choices []Choice `json:"choices"`
 }
 
+type Usage struct {
+	PromptTokens     int `json:"prompt_tokens"`
+	CompletionTokens int `json:"completion_tokens"`
+	TotalTokens      int `json:"total_tokens"`
+}
+
 type Choice struct {
 	Message      Message `json:"message"`
 	FinishReason string  `json:"finish_reason"`