Add TokenCount (#80)

sunshineplan · Jul 10, 2024 · 35be93d · 35be93d
1 parent 43cdb61
commit 35be93d
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 0 deletions.
diff --git a/ai.go b/ai.go
@@ -50,6 +50,13 @@ type ChatStream interface {
 	Close() error
 }
 
+type TokenCount struct {
+	Prompt int
+	Result int
+	Total  int
+}
+
 type ChatResponse interface {
 	Results() []string
+	TokenCount() TokenCount
 }
diff --git a/chatgpt/chatgpt.go b/chatgpt/chatgpt.go
@@ -116,6 +116,22 @@ func (resp *ChatResponse[Response]) Results() (res []string) {
 	return
 }
 
+func (resp *ChatResponse[Response]) TokenCount() (res ai.TokenCount) {
+	switch v := any(resp.resp).(type) {
+	case openai.ChatCompletionResponse:
+		res.Prompt = v.Usage.PromptTokens
+		res.Result = v.Usage.CompletionTokens
+		res.Total = v.Usage.TotalTokens
+	case openai.ChatCompletionStreamResponse:
+		if usage := v.Usage; usage != nil {
+			res.Prompt = usage.PromptTokens
+			res.Result = usage.CompletionTokens
+			res.Total = usage.TotalTokens
+		}
+	}
+	return
+}
+
 func (resp *ChatResponse[Response]) String() string {
 	if res := resp.Results(); len(res) > 0 {
 		return res[0]

diff --git a/gemini/gemini.go b/gemini/gemini.go
@@ -140,6 +140,15 @@ func (resp *ChatResponse) Results() (res []string) {
 	return
 }
 
+func (resp *ChatResponse) TokenCount() (res ai.TokenCount) {
+	if usage := resp.UsageMetadata; usage != nil {
+		res.Prompt = int(usage.PromptTokenCount)
+		res.Result = int(usage.CandidatesTokenCount)
+		res.Total = int(usage.TotalTokenCount)
+	}
+	return
+}
+
 func (resp *ChatResponse) String() string {
 	if res := resp.Results(); len(res) > 0 {
 		return res[0]