feat(executor): add support for Codex image generation tool usage tracking
- Introduced `publishCodexImageToolUsage` to report image generation tool metrics. - Updated executor logic to handle image generation tool events and defaults. - Added parsing logic for `image_gen` tool usage details in `helps/usage_helpers.go`. - Updated `UsageReporter` for additional model-specific usage publishing. - Refactored usage detail normalizations. Closes: #3063
This commit is contained in:
@@ -30,8 +30,9 @@ import (
|
|||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
codexUserAgent = "codex-tui/0.118.0 (Mac OS 26.3.1; arm64) iTerm.app/3.6.9 (codex-tui; 0.118.0)"
|
codexUserAgent = "codex-tui/0.118.0 (Mac OS 26.3.1; arm64) iTerm.app/3.6.9 (codex-tui; 0.118.0)"
|
||||||
codexOriginator = "codex-tui"
|
codexOriginator = "codex-tui"
|
||||||
|
codexDefaultImageToolModel = "gpt-image-2"
|
||||||
)
|
)
|
||||||
|
|
||||||
var dataTag = []byte("data:")
|
var dataTag = []byte("data:")
|
||||||
@@ -263,6 +264,7 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re
|
|||||||
if detail, ok := helps.ParseCodexUsage(eventData); ok {
|
if detail, ok := helps.ParseCodexUsage(eventData); ok {
|
||||||
reporter.Publish(ctx, detail)
|
reporter.Publish(ctx, detail)
|
||||||
}
|
}
|
||||||
|
publishCodexImageToolUsage(ctx, reporter, body, eventData)
|
||||||
|
|
||||||
completedData := eventData
|
completedData := eventData
|
||||||
outputResult := gjson.GetBytes(completedData, "response.output")
|
outputResult := gjson.GetBytes(completedData, "response.output")
|
||||||
@@ -496,6 +498,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au
|
|||||||
if detail, ok := helps.ParseCodexUsage(data); ok {
|
if detail, ok := helps.ParseCodexUsage(data); ok {
|
||||||
reporter.Publish(ctx, detail)
|
reporter.Publish(ctx, detail)
|
||||||
}
|
}
|
||||||
|
publishCodexImageToolUsage(ctx, reporter, body, data)
|
||||||
data = patchCodexCompletedOutput(data, outputItemsByIndex, outputItemsFallback)
|
data = patchCodexCompletedOutput(data, outputItemsByIndex, outputItemsFallback)
|
||||||
translatedLine = append([]byte("data: "), data...)
|
translatedLine = append([]byte("data: "), data...)
|
||||||
}
|
}
|
||||||
@@ -859,6 +862,31 @@ func ensureImageGenerationTool(body []byte, baseModel string, auth *cliproxyauth
|
|||||||
return body
|
return body
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func publishCodexImageToolUsage(ctx context.Context, reporter *helps.UsageReporter, body []byte, completedData []byte) {
|
||||||
|
detail, ok := helps.ParseCodexImageToolUsage(completedData)
|
||||||
|
if !ok {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
reporter.EnsurePublished(ctx)
|
||||||
|
reporter.PublishAdditionalModel(ctx, codexImageGenerationToolModel(body), detail)
|
||||||
|
}
|
||||||
|
|
||||||
|
func codexImageGenerationToolModel(body []byte) string {
|
||||||
|
tools := gjson.GetBytes(body, "tools")
|
||||||
|
if tools.IsArray() {
|
||||||
|
for _, tool := range tools.Array() {
|
||||||
|
if tool.Get("type").String() != "image_generation" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if model := strings.TrimSpace(tool.Get("model").String()); model != "" {
|
||||||
|
return model
|
||||||
|
}
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return codexDefaultImageToolModel
|
||||||
|
}
|
||||||
|
|
||||||
func isCodexModelCapacityError(errorBody []byte) bool {
|
func isCodexModelCapacityError(errorBody []byte) bool {
|
||||||
if len(errorBody) == 0 {
|
if len(errorBody) == 0 {
|
||||||
return false
|
return false
|
||||||
|
|||||||
@@ -48,6 +48,18 @@ func (r *UsageReporter) Publish(ctx context.Context, detail usage.Detail) {
|
|||||||
r.publishWithOutcome(ctx, detail, false)
|
r.publishWithOutcome(ctx, detail, false)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (r *UsageReporter) PublishAdditionalModel(ctx context.Context, model string, detail usage.Detail) {
|
||||||
|
if r == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
model = strings.TrimSpace(model)
|
||||||
|
if model == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
detail = normalizeUsageDetailTotal(detail)
|
||||||
|
usage.PublishRecord(ctx, r.buildRecordForModel(model, detail, false))
|
||||||
|
}
|
||||||
|
|
||||||
func (r *UsageReporter) PublishFailure(ctx context.Context) {
|
func (r *UsageReporter) PublishFailure(ctx context.Context) {
|
||||||
r.publishWithOutcome(ctx, usage.Detail{}, true)
|
r.publishWithOutcome(ctx, usage.Detail{}, true)
|
||||||
}
|
}
|
||||||
@@ -65,15 +77,20 @@ func (r *UsageReporter) publishWithOutcome(ctx context.Context, detail usage.Det
|
|||||||
if r == nil {
|
if r == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
detail = normalizeUsageDetailTotal(detail)
|
||||||
|
r.once.Do(func() {
|
||||||
|
usage.PublishRecord(ctx, r.buildRecord(detail, failed))
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func normalizeUsageDetailTotal(detail usage.Detail) usage.Detail {
|
||||||
if detail.TotalTokens == 0 {
|
if detail.TotalTokens == 0 {
|
||||||
total := detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
|
total := detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
|
||||||
if total > 0 {
|
if total > 0 {
|
||||||
detail.TotalTokens = total
|
detail.TotalTokens = total
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
r.once.Do(func() {
|
return detail
|
||||||
usage.PublishRecord(ctx, r.buildRecord(detail, failed))
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// ensurePublished guarantees that a usage record is emitted exactly once.
|
// ensurePublished guarantees that a usage record is emitted exactly once.
|
||||||
@@ -93,9 +110,16 @@ func (r *UsageReporter) buildRecord(detail usage.Detail, failed bool) usage.Reco
|
|||||||
if r == nil {
|
if r == nil {
|
||||||
return usage.Record{Detail: detail, Failed: failed}
|
return usage.Record{Detail: detail, Failed: failed}
|
||||||
}
|
}
|
||||||
|
return r.buildRecordForModel(r.model, detail, failed)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, failed bool) usage.Record {
|
||||||
|
if r == nil {
|
||||||
|
return usage.Record{Model: model, Detail: detail, Failed: failed}
|
||||||
|
}
|
||||||
return usage.Record{
|
return usage.Record{
|
||||||
Provider: r.provider,
|
Provider: r.provider,
|
||||||
Model: r.model,
|
Model: model,
|
||||||
Source: r.source,
|
Source: r.source,
|
||||||
APIKey: r.apiKey,
|
APIKey: r.apiKey,
|
||||||
AuthID: r.authID,
|
AuthID: r.authID,
|
||||||
@@ -201,18 +225,15 @@ func ParseCodexUsage(data []byte) (usage.Detail, bool) {
|
|||||||
if !usageNode.Exists() {
|
if !usageNode.Exists() {
|
||||||
return usage.Detail{}, false
|
return usage.Detail{}, false
|
||||||
}
|
}
|
||||||
detail := usage.Detail{
|
return parseOpenAIStyleUsageNode(usageNode), true
|
||||||
InputTokens: usageNode.Get("input_tokens").Int(),
|
}
|
||||||
OutputTokens: usageNode.Get("output_tokens").Int(),
|
|
||||||
TotalTokens: usageNode.Get("total_tokens").Int(),
|
func ParseCodexImageToolUsage(data []byte) (usage.Detail, bool) {
|
||||||
|
usageNode := gjson.ParseBytes(data).Get("response.tool_usage.image_gen")
|
||||||
|
if !usageNode.Exists() || !usageNode.IsObject() {
|
||||||
|
return usage.Detail{}, false
|
||||||
}
|
}
|
||||||
if cached := usageNode.Get("input_tokens_details.cached_tokens"); cached.Exists() {
|
return parseOpenAIStyleUsageNode(usageNode), true
|
||||||
detail.CachedTokens = cached.Int()
|
|
||||||
}
|
|
||||||
if reasoning := usageNode.Get("output_tokens_details.reasoning_tokens"); reasoning.Exists() {
|
|
||||||
detail.ReasoningTokens = reasoning.Int()
|
|
||||||
}
|
|
||||||
return detail, true
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func ParseOpenAIUsage(data []byte) usage.Detail {
|
func ParseOpenAIUsage(data []byte) usage.Detail {
|
||||||
@@ -220,6 +241,10 @@ func ParseOpenAIUsage(data []byte) usage.Detail {
|
|||||||
if !usageNode.Exists() {
|
if !usageNode.Exists() {
|
||||||
return usage.Detail{}
|
return usage.Detail{}
|
||||||
}
|
}
|
||||||
|
return parseOpenAIStyleUsageNode(usageNode)
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseOpenAIStyleUsageNode(usageNode gjson.Result) usage.Detail {
|
||||||
inputNode := usageNode.Get("prompt_tokens")
|
inputNode := usageNode.Get("prompt_tokens")
|
||||||
if !inputNode.Exists() {
|
if !inputNode.Exists() {
|
||||||
inputNode = usageNode.Get("input_tokens")
|
inputNode = usageNode.Get("input_tokens")
|
||||||
|
|||||||
Reference in New Issue
Block a user