Repository: sashabaranov/go-openai
Pull request #223 (2023-04-03 11:47:40)

to fix v1.5.8 Usage return 
In the pr on [215](https://github.com/sashabaranov/go-openai/pull/215). I found that I used the wrong way to add Usage. when `stream = true` will not  return `Usage` filed.
![usage](https://user-images.githubusercontent.com/596076/229359992-b9018087-17c9-46ae-b1da-5e984aecf912.png)
`Usage` should be calculated by the user himself.

sashabaranov: @nasa1024 hey, you meant to remove this line in this PR? (see diff)
nasa1024: here is the right way to caculator `usage`

```bash
go get github.com/pkoukk/tiktoken-go
```
here is example to get `usage`

```go
package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"

	"github.com/pkoukk/tiktoken-go"
	openai "github.com/sashabaranov/go-openai"
)

func main() {
	chat()
	stream()
}

func chat() {
	client := openai.NewClient("")
	resp, err := client.CreateChatCompletion(
		context.Background(),
		openai.ChatCompletionRequest{
			Model:     openai.GPT3Dot5Turbo,
			MaxTokens: 20,
			Messages: []openai.ChatCompletionMessage{
				{
					Role:    openai.ChatMessageRoleAssistant,
					Content: "你好，我叫小梦",
				},
				{
					Role:    openai.ChatMessageRoleUser,
					Content: "什么是算法稳定币？",
				},
			},
		},
	)

	if err != nil {
		fmt.Printf("ChatCompletion error: %v\n", err)
		return
	}

	fmt.Printf("response: %#v\n", resp)
}

func stream() {
	c := openai.NewClient("")
	ctx := context.Background()
	var response openai.ChatCompletionResponse
	var completionText string

	req := openai.ChatCompletionRequest{
		Model:     openai.GPT3Dot5Turbo,
		MaxTokens: 20,
		Messages: []openai.ChatCompletionMessage{
			{
				Role:    openai.ChatMessageRoleAssistant,
				Content: "你好，我叫小梦",
			},
			{
				Role:    openai.ChatMessageRoleUser,
				Content: "什么是算法稳定币？",
			},
		},
		Stream: true,
	}
	stream, err := c.CreateChatCompletionStream(ctx, req)
	if err != nil {
		fmt.Printf("ChatCompletionStream error: %v\n", err)
		return
	}
	defer stream.Close()

	fmt.Printf("Stream response: ")
	for {
		response, err := stream.Recv()
		if errors.Is(err, io.EOF) {
			fmt.Println("\nStream finished")
			break
		}

		if err != nil {
			fmt.Printf("\nStream error: %v\n", err)
			return
		}

		completionText += response.Choices[0].Delta.Content
		fmt.Printf(response.Choices[0].Delta.Content)
	}

	str, err := json.Marshal(req)
	if err != nil {
		err = fmt.Errorf(": %v", err)
		return
	}
	fmt.Printf("Request: %s", string(str))

	response.Usage.PromptTokens = NumTokensFromMessages(req.Messages, openai.GPT3Dot5Turbo)
	response.Usage.CompletionTokens = getTokenByModel(completionText, openai.GPT3Dot5Turbo)
	response.Usage.TotalTokens = response.Usage.PromptTokens + response.Usage.CompletionTokens
	fmt.Printf("response: %#v\n", response)
}

// getTokenByEncoding
func getTokenByModel(text string, model string) (num_tokens int) {

	tkm, err := tiktoken.EncodingForModel(model)
	if err != nil {
		err = fmt.Errorf(": %v", err)
		return
	}

	token := tkm.Encode(text, nil, nil)

	return len(token)
}

// NumTokensFromMessages
func NumTokensFromMessages(messages []openai.ChatCompletionMessage, model string) (num_tokens int) {
	tkm, err := tiktoken.EncodingForModel(model)
	if err != nil {
		err = fmt.Errorf("EncodingForModel: %v", err)
		fmt.Println(err)
		return
	}

	var tokens_per_message int
	var tokens_per_name int
	if model == "gpt-3.5-turbo-0301" || model == "gpt-3.5-turbo" {
		tokens_per_message = 4
		tokens_per_name = -1
	} else if model == "gpt-4-0314" || model == "gpt-4" {
		tokens_per_message = 3
		tokens_per_name = 1
	} else {
		fmt.Println("Warning: model not found. Using cl100k_base encoding.")
		tokens_per_message = 3
		tokens_per_name = 1
	}

	for _, message := range messages {
		num_tokens += tokens_per_message
		num_tokens += len(tkm.Encode(message.Content, nil, nil))
		if message.Name != "" {
			num_tokens += tokens_per_name
		}
	}
	num_tokens += 3
	return num_tokens
}

```

I'm sorry for increase in your workload and the error of the commit history 
nasa1024: @collinvandyck
there is have prompt token no equal 

```bash
response: openai.ChatCompletionResponse{ID:"chatcmpl-71EycXmMs6ofZi4IdbUwEaYP1T8iR", Object:"chat.completion", Created:1680530254, Model:"gpt-3.5-turbo-0301", Choices:[]openai.ChatCompletionChoice{openai.ChatCompletionChoice{Index:0, Message:openai.ChatCompletionMessage{Role:"assistant", Content:"算法稳定币是一种数字货币，与其他稳定", Name:""}, FinishReason:"length"}}, Usage:openai.Usage{PromptTokens:35, CompletionTokens:20, TotalTokens:55}}
Stream response: 算法稳定币是一种基于算法的数字货币，
Stream finished
Request: {"model":"gpt-3.5-turbo","messages":[{"role":"assistant","content":"你好，我叫小梦"},{"role":"user","content":"什么是算法稳定币？"}],"max_tokens":20,"stream":true}response: openai.ChatCompletionResponse{ID:"", Object:"", Created:0, Model:"", Choices:[]openai.ChatCompletionChoice(nil), Usage:openai.Usage{PromptTokens:33, CompletionTokens:20, TotalTokens:53}}
```
maybe we need some discuss more
collinvandyck: @nasa1024 I believe that @sashabaranov is saying that because usage does not come back in the actual HTTP response for stream completions, that adding `omitempty` is not really the right approach, since it will always be empty. I think it should be removed altogether.

Also thanks for the call out on how to calculate tokens manually. This seems like a cool thing to have in go-openai -- the Recv code could possibly use this and just calculate the usage for you based on the model that generated the stream response. 
codecov[bot]: ## [Codecov](https://codecov.io/gh/sashabaranov/go-openai/pull/223?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#223](https://codecov.io/gh/sashabaranov/go-openai/pull/223?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (bad64ef) into [master](https://codecov.io/gh/sashabaranov/go-openai/commit/bee06561748e30a5286e56e1ef906d8e1075c6ec?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (bee0656) will **not change** coverage.
> The diff coverage is `n/a`.

```diff
@@           Coverage Diff           @@
##           master     #223   +/-   ##
=======================================
  Coverage   72.51%   72.51%           
=======================================
  Files          21       21           
  Lines         593      593           
=======================================
  Hits          430      430           
  Misses        124      124           
  Partials       39       39           
```


| [Impacted Files](https://codecov.io/gh/sashabaranov/go-openai/pull/223?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | |
|---|---|---|
| [chat\_stream.go](https://codecov.io/gh/sashabaranov/go-openai/pull/223?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y2hhdF9zdHJlYW0uZ28=) | `88.00% <ø> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)


Patch:
diff --git a/api.go b/api.go
index 00d6d35..2c978bc 100644
--- a/api.go
+++ b/api.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"strings"
 )
 
 // Client is OpenAI GPT-3 API client.
@@ -39,7 +40,13 @@ func NewOrgClient(authToken, org string) *Client {
 
 func (c *Client) sendRequest(req *http.Request, v interface{}) error {
 	req.Header.Set("Accept", "application/json; charset=utf-8")
-	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.config.authToken))
+	// Azure API Key authentication
+	if c.config.APIType == APITypeAzure {
+		req.Header.Set(AzureAPIKeyHeader, c.config.authToken)
+	} else {
+		// OpenAI or Azure AD authentication
+		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.config.authToken))
+	}
 
 	// Check whether Content-Type is already set, Upload Files API requires
 	// Content-Type == multipart/form-data
@@ -83,6 +90,15 @@ func (c *Client) sendRequest(req *http.Request, v interface{}) error {
 }
 
 func (c *Client) fullURL(suffix string) string {
+	// /openai/deployments/{engine}/chat/completions?api-version={api_version}
+	if c.config.APIType == APITypeAzure || c.config.APIType == APITypeAzureAD {
+		baseURL := c.config.BaseURL
+		baseURL = strings.TrimRight(baseURL, "/")
+		return fmt.Sprintf("%s/%s/%s/%s%s?api-version=%s",
+			baseURL, azureAPIPrefix, azureDeploymentsPrefix, c.config.Engine, suffix, c.config.APIVersion)
+	}
+
+	// c.config.APIType == APITypeOpenAI || c.config.APIType == ""
 	return fmt.Sprintf("%s%s", c.config.BaseURL, suffix)
 }
 
@@ -100,7 +116,14 @@ func (c *Client) newStreamRequest(
 	req.Header.Set("Accept", "text/event-stream")
 	req.Header.Set("Cache-Control", "no-cache")
 	req.Header.Set("Connection", "keep-alive")
-	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.config.authToken))
 
+	// https://learn.microsoft.com/en-us/azure/cognitive-services/openai/reference#authentication
+	// Azure API Key authentication
+	if c.config.APIType == APITypeAzure {
+		req.Header.Set(AzureAPIKeyHeader, c.config.authToken)
+	} else {
+		// OpenAI or Azure AD authentication
+		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.config.authToken))
+	}
 	return req, nil
 }
diff --git a/chat_stream.go b/chat_stream.go
index 009e1b1..8211292 100644
--- a/chat_stream.go
+++ b/chat_stream.go
@@ -21,7 +21,6 @@ type ChatCompletionStreamResponse struct {
 	Created int64                        `json:"created"`
 	Model   string                       `json:"model"`
 	Choices []ChatCompletionStreamChoice `json:"choices"`
-	Usage   Usage                        `json:"usage"`
 }
 
 // ChatCompletionStream
diff --git a/config.go b/config.go
index e09c256..52e1efc 100644
--- a/config.go
+++ b/config.go
@@ -5,28 +5,61 @@ import (
 )
 
 const (
-	apiURLv1                       = "https://api.openai.com/v1"
+	openaiAPIURLv1                 = "https://api.openai.com/v1"
 	defaultEmptyMessagesLimit uint = 300
+
+	azureAPIPrefix         = "openai"
+	azureDeploymentsPrefix = "deployments"
+)
+
+type APIType string
+
+const (
+	APITypeOpenAI  APIType = "OPEN_AI"
+	APITypeAzure   APIType = "AZURE"
+	APITypeAzureAD APIType = "AZURE_AD"
 )
 
+const AzureAPIKeyHeader = "api-key"
+
 // ClientConfig is a configuration of a client.
 type ClientConfig struct {
 	authToken string
 
-	HTTPClient *http.Client
+	BaseURL    string
+	OrgID      string
+	APIType    APIType
+	APIVersion string // required when APIType is APITypeAzure or APITypeAzureAD
+	Engine     string // required when APIType is APITypeAzure or APITypeAzureAD
 
-	BaseURL string
-	OrgID   string
+	HTTPClient *http.Client
 
 	EmptyMessagesLimit uint
 }
 
 func DefaultConfig(authToken string) ClientConfig {
 	return ClientConfig{
+		authToken: authToken,
+		BaseURL:   openaiAPIURLv1,
+		APIType:   APITypeOpenAI,
+		OrgID:     "",
+
 		HTTPClient: &http.Client{},
-		BaseURL:    apiURLv1,
+
+		EmptyMessagesLimit: defaultEmptyMessagesLimit,
+	}
+}
+
+func DefaultAzureConfig(apiKey, baseURL, engine string) ClientConfig {
+	return ClientConfig{
+		authToken:  apiKey,
+		BaseURL:    baseURL,
 		OrgID:      "",
-		authToken:  authToken,
+		APIType:    APITypeAzure,
+		APIVersion: "2023-03-15-preview",
+		Engine:     engine,
+
+		HTTPClient: &http.Client{},
 
 		EmptyMessagesLimit: defaultEmptyMessagesLimit,
 	}


Test patch:
diff --git a/api_internal_test.go b/api_internal_test.go
new file mode 100644
index 0000000..83dcafc
--- /dev/null
+++ b/api_internal_test.go
@@ -0,0 +1,133 @@
+package openai
+
+import (
+	"context"
+	"testing"
+)
+
+func TestOpenAIFullURL(t *testing.T) {
+	cases := []struct {
+		Name   string
+		Suffix string
+		Expect string
+	}{
+		{
+			"ChatCompletionsURL",
+			"/chat/completions",
+			"https://api.openai.com/v1/chat/completions",
+		},
+		{
+			"CompletionsURL",
+			"/completions",
+			"https://api.openai.com/v1/completions",
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.Name, func(t *testing.T) {
+			az := DefaultConfig("dummy")
+			cli := NewClientWithConfig(az)
+			actual := cli.fullURL(c.Suffix)
+			if actual != c.Expect {
+				t.Errorf("Expected %s, got %s", c.Expect, actual)
+			}
+			t.Logf("Full URL: %s", actual)
+		})
+	}
+}
+
+func TestRequestAuthHeader(t *testing.T) {
+	cases := []struct {
+		Name      string
+		APIType   APIType
+		HeaderKey string
+		Token     string
+		Expect    string
+	}{
+		{
+			"OpenAIDefault",
+			"",
+			"Authorization",
+			"dummy-token-openai",
+			"Bearer dummy-token-openai",
+		},
+		{
+			"OpenAI",
+			APITypeOpenAI,
+			"Authorization",
+			"dummy-token-openai",
+			"Bearer dummy-token-openai",
+		},
+		{
+			"AzureAD",
+			APITypeAzureAD,
+			"Authorization",
+			"dummy-token-azure",
+			"Bearer dummy-token-azure",
+		},
+		{
+			"Azure",
+			APITypeAzure,
+			AzureAPIKeyHeader,
+			"dummy-api-key-here",
+			"dummy-api-key-here",
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.Name, func(t *testing.T) {
+			az := DefaultConfig(c.Token)
+			az.APIType = c.APIType
+
+			cli := NewClientWithConfig(az)
+			req, err := cli.newStreamRequest(context.Background(), "POST", "/chat/completions", nil)
+			if err != nil {
+				t.Errorf("Failed to create request: %v", err)
+			}
+			actual := req.Header.Get(c.HeaderKey)
+			if actual != c.Expect {
+				t.Errorf("Expected %s, got %s", c.Expect, actual)
+			}
+			t.Logf("%s: %s", c.HeaderKey, actual)
+		})
+	}
+}
+
+func TestAzureFullURL(t *testing.T) {
+	cases := []struct {
+		Name    string
+		BaseURL string
+		Engine  string
+		Expect  string
+	}{
+		{
+			"AzureBaseURLWithSlashAutoStrip",
+			"https://httpbin.org/",
+			"chatgpt-demo",
+			"https://httpbin.org/" +
+				"openai/deployments/chatgpt-demo" +
+				"/chat/completions?api-version=2023-03-15-preview",
+		},
+		{
+			"AzureBaseURLWithoutSlashOK",
+			"https://httpbin.org",
+			"chatgpt-demo",
+			"https://httpbin.org/" +
+				"openai/deployments/chatgpt-demo" +
+				"/chat/completions?api-version=2023-03-15-preview",
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.Name, func(t *testing.T) {
+			az := DefaultAzureConfig("dummy", c.BaseURL, c.Engine)
+			cli := NewClientWithConfig(az)
+			// /openai/deployments/{engine}/chat/completions?api-version={api_version}
+			actual := cli.fullURL("/chat/completions")
+			if actual != c.Expect {
+				t.Errorf("Expected %s, got %s", c.Expect, actual)
+			}
+			t.Logf("Full URL: %s", actual)
+		})
+	}
+}