mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-09-06 08:02:22 +08:00
fix: complete xAI image usage reporting
This commit is contained in:
@@ -481,8 +481,12 @@ func xaiBuildSSEFrame(eventName string, data []byte) []byte {
|
||||
}
|
||||
|
||||
func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, endpointPath string) (resp cliproxyexecutor.Response, err error) {
|
||||
model := gjson.GetBytes(req.Payload, "model").String()
|
||||
model := strings.TrimSpace(gjson.GetBytes(req.Payload, "model").String())
|
||||
if model == "" {
|
||||
model = strings.TrimSpace(req.Model)
|
||||
}
|
||||
reporter := helps.NewExecutorUsageReporter(ctx, e, model, auth)
|
||||
defer reporter.TrackFailure(ctx, &err)
|
||||
|
||||
token, baseURL := xaiCreds(auth)
|
||||
if baseURL == "" {
|
||||
@@ -502,10 +506,10 @@ func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth
|
||||
e.recordXAIRequest(ctx, auth, url, httpReq.Header.Clone(), req.Payload)
|
||||
|
||||
httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
|
||||
httpClient = reporter.TrackHTTPClient(httpClient)
|
||||
httpResp, err := httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
helps.RecordAPIResponseError(ctx, e.cfg, err)
|
||||
reporter.PublishFailure(ctx, err)
|
||||
return resp, err
|
||||
}
|
||||
defer func() {
|
||||
@@ -518,7 +522,6 @@ func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth
|
||||
data, err := io.ReadAll(httpResp.Body)
|
||||
if err != nil {
|
||||
helps.RecordAPIResponseError(ctx, e.cfg, err)
|
||||
reporter.PublishFailure(ctx, err)
|
||||
return resp, err
|
||||
}
|
||||
helps.AppendAPIResponseChunk(ctx, e.cfg, data)
|
||||
@@ -526,7 +529,6 @@ func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
|
||||
err = xaiStatusErr(httpResp.StatusCode, data)
|
||||
reporter.PublishFailure(ctx, err)
|
||||
return resp, err
|
||||
}
|
||||
|
||||
|
||||
@@ -2088,6 +2088,9 @@ func TestXAIExecutorExecuteImagesUsesImagesEndpointAndPublishesUsage(t *testing.
|
||||
if record.Detail != (usage.Detail{}) {
|
||||
t.Fatalf("detail = %+v, want zero token usage", record.Detail)
|
||||
}
|
||||
if record.TTFT <= 0 {
|
||||
t.Fatalf("ttft = %v, want positive duration", record.TTFT)
|
||||
}
|
||||
assertNoAdditionalXAIUsageRecord(t, plugin.records)
|
||||
}
|
||||
|
||||
@@ -2140,6 +2143,44 @@ func TestXAIExecutorExecuteImagesPublishesFailureUsage(t *testing.T) {
|
||||
assertNoAdditionalXAIUsageRecord(t, plugin.records)
|
||||
}
|
||||
|
||||
func TestXAIExecutorExecuteImagesPublishesRequestBuildFailureUsage(t *testing.T) {
|
||||
const requestedModel = "grok-imagine-image-fallback"
|
||||
|
||||
plugin := &captureXAIUsagePlugin{
|
||||
model: requestedModel,
|
||||
records: make(chan usage.Record, 2),
|
||||
}
|
||||
usage.RegisterPlugin(plugin)
|
||||
|
||||
exec := NewXAIExecutor(&config.Config{})
|
||||
auth := &cliproxyauth.Auth{
|
||||
Provider: "xai",
|
||||
Attributes: map[string]string{"base_url": "://invalid"},
|
||||
}
|
||||
|
||||
_, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{
|
||||
Model: requestedModel,
|
||||
Payload: []byte(`{"prompt":"draw"}`),
|
||||
}, cliproxyexecutor.Options{
|
||||
SourceFormat: sdktranslator.FromString("openai-image"),
|
||||
Metadata: map[string]any{
|
||||
cliproxyexecutor.RequestPathMetadataKey: "/v1/images/generations",
|
||||
},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("Execute() error = nil, want non-nil")
|
||||
}
|
||||
|
||||
record := waitForXAIUsageRecord(t, plugin.records)
|
||||
if record.Model != requestedModel {
|
||||
t.Fatalf("model = %q, want %q", record.Model, requestedModel)
|
||||
}
|
||||
if !record.Failed {
|
||||
t.Fatal("failed = false, want true")
|
||||
}
|
||||
assertNoAdditionalXAIUsageRecord(t, plugin.records)
|
||||
}
|
||||
|
||||
type captureXAIUsagePlugin struct {
|
||||
model string
|
||||
records chan usage.Record
|
||||
|
||||
Reference in New Issue
Block a user