Skip to content

Commit 7618d11

Browse files
mrh997mrh997
authored
feat(model): support promptTokenDetails (#401)
1 parent 70d60ff commit 7618d11

File tree

22 files changed

+85
-32
lines changed

22 files changed

+85
-32
lines changed

‎components/document/loader/url/examples/html/main.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ func main() {
3535

3636
addr := "127.0.0.1:18001"
3737

38-
go func() { // nolint: byted_goroutine_recover
38+
go func() {
3939
fmt.Println("Serving directory on http://127.0.0.1:18001")
4040
if err := http.ListenAndServe(addr, nil); err != nil {
4141
fmt.Println("Server failed to start:", err)

‎components/document/transformer/splitter/html/header.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -224,7 +224,7 @@ func extractText(node *html.Node) (string, error) {
224224
orig := node
225225
node = node.FirstChild
226226
for node != nil {
227-
if node == orig { // nolint: byted_address_compare_check
227+
if node == orig {
228228
break
229229
}
230230
if _, ok := ll[node]; !ok {

‎components/model/ark/chat_completion_api.go

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,7 @@ func (cm *completionAPIChatModel) Stream(ctx context.Context, in []*schema.Messa
212212
}
213213

214214
sw.Close()
215-
_ = cm.closeArkStreamReader(stream) // nolint: byted_returned_err_should_do_check
215+
_ = cm.closeArkStreamReader(stream)
216216

217217
}()
218218

@@ -622,7 +622,10 @@ func (cm *completionAPIChatModel) toEinoTokenUsage(usage *model.Usage) *schema.T
622622
return &schema.TokenUsage{
623623
CompletionTokens: usage.CompletionTokens,
624624
PromptTokens: usage.PromptTokens,
625-
TotalTokens: usage.TotalTokens,
625+
PromptTokenDetails: schema.PromptTokenDetails{
626+
CachedTokens: usage.PromptTokensDetails.CachedTokens,
627+
},
628+
TotalTokens: usage.TotalTokens,
626629
}
627630
}
628631

@@ -637,6 +640,9 @@ func (cm *completionAPIChatModel) toModelCallbackUsage(respMeta *schema.Response
637640
return &fmodel.TokenUsage{
638641
CompletionTokens: usage.CompletionTokens,
639642
PromptTokens: usage.PromptTokens,
640-
TotalTokens: usage.TotalTokens,
643+
PromptTokenDetails: fmodel.PromptTokenDetails{
644+
CachedTokens: usage.PromptTokenDetails.CachedTokens,
645+
},
646+
TotalTokens: usage.TotalTokens,
641647
}
642648
}

‎components/model/ark/chatmodel.go

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -570,7 +570,10 @@ func (cm *ChatModel) createContextByContextAPI(ctx context.Context, prefix []*sc
570570
return &CacheInfo{
571571
ContextID: resp.ID,
572572
Usage: schema.TokenUsage{
573-
PromptTokens: resp.Usage.PromptTokens,
573+
PromptTokens: resp.Usage.PromptTokens,
574+
PromptTokenDetails: schema.PromptTokenDetails{
575+
CachedTokens: resp.Usage.PromptTokensDetails.CachedTokens,
576+
},
574577
CompletionTokens: resp.Usage.CompletionTokens,
575578
TotalTokens: resp.Usage.TotalTokens,
576579
},

‎components/model/ark/go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ go 1.23.0
44

55
require (
66
github.com/bytedance/mockey v1.2.14
7-
github.com/cloudwego/eino v0.3.47
7+
github.com/cloudwego/eino v0.4.3
88
github.com/getkin/kin-openapi v0.118.0
99
github.com/openai/openai-go v1.10.1
1010
github.com/smartystreets/goconvey v1.8.1

‎components/model/ark/go.sum

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@ github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCy
2020
github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w=
2121
github.com/cloudwego/eino v0.3.47 h1:nl1Q1QZhFAyl169M32KZB8vj1Zp6fqeSjVF1lVzUSsw=
2222
github.com/cloudwego/eino v0.3.47/go.mod h1:wUjz990apdsaOraOXdh6CdhVXq8DJsOvLsVlxNTcNfY=
23+
github.com/cloudwego/eino v0.4.3 h1:ihHVHPMA8G2DPWeN+NTuIKr1QBAet9g84buH+LUl1bU=
24+
github.com/cloudwego/eino v0.4.3/go.mod h1:wUjz990apdsaOraOXdh6CdhVXq8DJsOvLsVlxNTcNfY=
2325
github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY=
2426
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
2527
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=

‎components/model/ark/responses_api.go

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -703,15 +703,21 @@ func (cm *responsesAPIChatModel) toOutputMessage(resp *responses.Response) (*sch
703703

704704
func (cm *responsesAPIChatModel) toEinoTokenUsage(usage responses.ResponseUsage) *schema.TokenUsage {
705705
return &schema.TokenUsage{
706-
PromptTokens: int(usage.InputTokens),
706+
PromptTokens: int(usage.InputTokens),
707+
PromptTokenDetails: schema.PromptTokenDetails{
708+
CachedTokens: int(usage.InputTokensDetails.CachedTokens),
709+
},
707710
CompletionTokens: int(usage.OutputTokens),
708711
TotalTokens: int(usage.TotalTokens),
709712
}
710713
}
711714

712715
func (cm *responsesAPIChatModel) toModelTokenUsage(usage responses.ResponseUsage) *model.TokenUsage {
713716
return &model.TokenUsage{
714-
PromptTokens: int(usage.InputTokens),
717+
PromptTokens: int(usage.InputTokens),
718+
PromptTokenDetails: model.PromptTokenDetails{
719+
CachedTokens: int(usage.InputTokensDetails.CachedTokens),
720+
},
715721
CompletionTokens: int(usage.OutputTokens),
716722
TotalTokens: int(usage.TotalTokens),
717723
}

‎components/model/arkbot/chatmodel.go

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,14 @@ import (
2525
"runtime/debug"
2626
"time"
2727

28+
"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
29+
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
30+
autils "github.com/volcengine/volcengine-go-sdk/service/arkruntime/utils"
31+
2832
"github.com/cloudwego/eino/callbacks"
2933
"github.com/cloudwego/eino/components"
3034
fmodel "github.com/cloudwego/eino/components/model"
3135
"github.com/cloudwego/eino/schema"
32-
"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
33-
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
34-
autils "github.com/volcengine/volcengine-go-sdk/service/arkruntime/utils"
3536
)
3637

3738
var _ fmodel.ToolCallingChatModel = (*ChatModel)(nil)
@@ -373,7 +374,7 @@ func (cm *ChatModel) Stream(ctx context.Context, in []*schema.Message, opts ...f
373374
}
374375

375376
sw.Close()
376-
_ = closeArkStreamReader(stream) // nolint: byted_returned_err_should_do_check
377+
_ = closeArkStreamReader(stream)
377378

378379
}()
379380

‎components/model/claude/claude.go

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import (
3131
"github.com/anthropics/anthropic-sdk-go/packages/param"
3232
awsConfig "github.com/aws/aws-sdk-go-v2/config"
3333
"github.com/aws/aws-sdk-go-v2/credentials"
34+
3435
"github.com/cloudwego/eino/components"
3536

3637
"github.com/cloudwego/eino/callbacks"
@@ -529,7 +530,10 @@ func (cm *ChatModel) getCallbackOutput(output *schema.Message) *model.CallbackOu
529530
}
530531
if output.ResponseMeta != nil && output.ResponseMeta.Usage != nil {
531532
result.TokenUsage = &model.TokenUsage{
532-
PromptTokens: output.ResponseMeta.Usage.PromptTokens,
533+
PromptTokens: output.ResponseMeta.Usage.PromptTokens,
534+
PromptTokenDetails: model.PromptTokenDetails{
535+
CachedTokens: output.ResponseMeta.Usage.PromptTokenDetails.CachedTokens,
536+
},
533537
CompletionTokens: output.ResponseMeta.Usage.CompletionTokens,
534538
TotalTokens: output.ResponseMeta.Usage.TotalTokens,
535539
}
@@ -608,14 +612,19 @@ func convSchemaMessage(message *schema.Message) (mp anthropic.MessageParam, err
608612
}
609613

610614
func convOutputMessage(resp *anthropic.Message) (*schema.Message, error) {
615+
promptTokens := int(resp.Usage.InputTokens + resp.Usage.CacheReadInputTokens + resp.Usage.CacheCreationInputTokens)
616+
611617
message := &schema.Message{
612618
Role: schema.Assistant,
613619
ResponseMeta: &schema.ResponseMeta{
614620
FinishReason: string(resp.StopReason),
615621
Usage: &schema.TokenUsage{
616-
PromptTokens: int(resp.Usage.InputTokens),
622+
PromptTokens: promptTokens,
623+
PromptTokenDetails: schema.PromptTokenDetails{
624+
CachedTokens: int(resp.Usage.CacheReadInputTokens),
625+
},
617626
CompletionTokens: int(resp.Usage.OutputTokens),
618-
TotalTokens: int(resp.Usage.InputTokens + resp.Usage.OutputTokens),
627+
TotalTokens: promptTokens + int(resp.Usage.OutputTokens),
619628
},
620629
},
621630
}

‎components/model/claude/go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ require (
77
github.com/aws/aws-sdk-go-v2/config v1.29.1
88
github.com/aws/aws-sdk-go-v2/credentials v1.17.54
99
github.com/bytedance/mockey v1.2.13
10-
github.com/cloudwego/eino v0.3.47
10+
github.com/cloudwego/eino v0.4.3
1111
github.com/getkin/kin-openapi v0.118.0
1212
github.com/stretchr/testify v1.9.0
1313
)

0 commit comments

Comments
 (0)