#78: Normalize response latency by response token count & renamed Tok…

…enCount to TokenUsage
EinStack · Jan 29, 2024 · 7301bb7 · 7301bb7
1 parent 0b032c7
commit 7301bb7
Show file tree

Hide file tree

Showing 10 changed files with 16 additions and 17 deletions.
diff --git a/docs/docs.go b/docs/docs.go
@@ -650,11 +650,11 @@ const docTemplate = `{
                     }
                 },
                 "tokenCount": {
-                    "$ref": "#/definitions/schemas.TokenCount"
+                    "$ref": "#/definitions/schemas.TokenUsage"
                 }
             }
         },
-        "schemas.TokenCount": {
+        "schemas.TokenUsage": {
             "type": "object",
             "properties": {
                 "promptTokens": {

diff --git a/docs/swagger.json b/docs/swagger.json
@@ -647,11 +647,11 @@
                     }
                 },
                 "tokenCount": {
-                    "$ref": "#/definitions/schemas.TokenCount"
+                    "$ref": "#/definitions/schemas.TokenUsage"
                 }
             }
         },
-        "schemas.TokenCount": {
+        "schemas.TokenUsage": {
             "type": "object",
             "properties": {
                 "promptTokens": {

diff --git a/docs/swagger.yaml b/docs/swagger.yaml
@@ -358,9 +358,9 @@ definitions:
           type: string
         type: object
       tokenCount:
-        $ref: '#/definitions/schemas.TokenCount'
+        $ref: '#/definitions/schemas.TokenUsage'
     type: object
-  schemas.TokenCount:
+  schemas.TokenUsage:
     properties:
       promptTokens:
         type: number

diff --git a/pkg/api/schemas/language.go b/pkg/api/schemas/language.go
@@ -39,10 +39,10 @@ type UnifiedChatResponse struct {
 type ProviderResponse struct {
 	SystemID   map[string]string `json:"responseId,omitempty"`
 	Message    ChatMessage       `json:"message"`
-	TokenCount TokenCount        `json:"tokenCount"`
+	TokenUsage TokenUsage        `json:"tokenCount"`
 }
 
-type TokenCount struct {
+type TokenUsage struct {
 	PromptTokens   float64 `json:"promptTokens"`
 	ResponseTokens float64 `json:"responseTokens"`
 	TotalTokens    float64 `json:"totalTokens"`

diff --git a/pkg/providers/anthropic/chat.go b/pkg/providers/anthropic/chat.go
@@ -178,7 +178,7 @@ func (c *Client) doChatRequest(ctx context.Context, payload *ChatRequest) (*sche
 				Content: anthropicCompletion.Content[0].Text,
 				Name:    "",
 			},
-			TokenCount: schemas.TokenCount{
+			TokenUsage: schemas.TokenCount{
 				PromptTokens:   0, // Anthropic doesn't send prompt tokens
 				ResponseTokens: 0,
 				TotalTokens:    0,

diff --git a/pkg/providers/azureopenai/chat.go b/pkg/providers/azureopenai/chat.go
@@ -190,7 +190,7 @@ func (c *Client) doChatRequest(ctx context.Context, payload *ChatRequest) (*sche
 				Content: openAICompletion.Choices[0].Message.Content,
 				Name:    "",
 			},
-			TokenCount: schemas.TokenCount{
+			TokenUsage: schemas.TokenCount{
 				PromptTokens:   openAICompletion.Usage.PromptTokens,
 				ResponseTokens: openAICompletion.Usage.CompletionTokens,
 				TotalTokens:    openAICompletion.Usage.TotalTokens,

diff --git a/pkg/providers/cohere/chat.go b/pkg/providers/cohere/chat.go
@@ -195,7 +195,7 @@ func (c *Client) doChatRequest(ctx context.Context, payload *ChatRequest) (*sche
 				Content: cohereCompletion.Text,
 				Name:    "",
 			},
-			TokenCount: schemas.TokenCount{
+			TokenUsage: schemas.TokenCount{
 				PromptTokens:   cohereCompletion.TokenCount.PromptTokens,
 				ResponseTokens: cohereCompletion.TokenCount.ResponseTokens,
 				TotalTokens:    cohereCompletion.TokenCount.TotalTokens,

diff --git a/pkg/providers/octoml/chat.go b/pkg/providers/octoml/chat.go
@@ -176,7 +176,7 @@ func (c *Client) doChatRequest(ctx context.Context, payload *ChatRequest) (*sche
 				Content: openAICompletion.Choices[0].Message.Content,
 				Name:    "",
 			},
-			TokenCount: schemas.TokenCount{
+			TokenUsage: schemas.TokenCount{
 				PromptTokens:   openAICompletion.Usage.PromptTokens,
 				ResponseTokens: openAICompletion.Usage.CompletionTokens,
 				TotalTokens:    openAICompletion.Usage.TotalTokens,

diff --git a/pkg/providers/openai/chat.go b/pkg/providers/openai/chat.go
@@ -190,7 +190,7 @@ func (c *Client) doChatRequest(ctx context.Context, payload *ChatRequest) (*sche
 				Content: openAICompletion.Choices[0].Message.Content,
 				Name:    "",
 			},
-			TokenCount: schemas.TokenCount{
+			TokenUsage: schemas.TokenCount{
 				PromptTokens:   openAICompletion.Usage.PromptTokens,
 				ResponseTokens: openAICompletion.Usage.CompletionTokens,
 				TotalTokens:    openAICompletion.Usage.TotalTokens,

diff --git a/pkg/providers/provider.go b/pkg/providers/provider.go
@@ -79,14 +79,13 @@ func (m *LangModel) Weight() int {
 }
 
 func (m *LangModel) Chat(ctx context.Context, request *schemas.UnifiedChatRequest) (*schemas.UnifiedChatResponse, error) {
-	// TODO: we may want to track time-to-first-byte to "normalize" response latency wrt response size
 	startedAt := time.Now()
 	resp, err := m.client.Chat(ctx, request)
 
-	// Do we want to track latency in case of errors as well?
-	m.latency.Add(float64(time.Since(startedAt)))
-
 	if err == nil {
+		// record latency per token to normalize measurements
+		m.latency.Add(float64(time.Since(startedAt)) / resp.ModelResponse.TokenUsage.ResponseTokens)
+
 		// successful response
 		resp.ModelID = m.modelID