diff --git a/pkg/plugins/openai/cmd/main/getcustomcosts_test.go b/pkg/plugins/openai/cmd/main/getcustomcosts_test.go index 2b111df..4484e9f 100644 --- a/pkg/plugins/openai/cmd/main/getcustomcosts_test.go +++ b/pkg/plugins/openai/cmd/main/getcustomcosts_test.go @@ -34,9 +34,8 @@ func TestGetCustomCosts(t *testing.T) { config: &config, } - windowStart := time.Date(2024, 10, 9, 0, 0, 0, 0, time.UTC) - // query for qty 2 of 1 hour windows - windowEnd := time.Date(2024, 10, 10, 0, 0, 0, 0, time.UTC) + windowEnd := time.Now().UTC().AddDate(0, -1, 0).Truncate(24 * time.Hour) + windowStart := windowEnd.AddDate(0, 0, -1) req := &pb.CustomCostRequest{ Start: timestamppb.New(windowStart), @@ -50,4 +49,207 @@ func TestGetCustomCosts(t *testing.T) { if len(resp) == 0 { t.Fatalf("empty response") } + + for _, r := range resp { + if len(r.Errors) > 0 { + t.Errorf("response has errors: %v", r.Errors) + } + } +} + +func TestNormalizeModelKey(t *testing.T) { + tests := []struct { + input string + want string + }{ + {"GPT-4o mini", "gpt4omini"}, + {"GPT-4o", "gpt4o"}, + {"gpt-4o-mini-2024-07-18", "gpt4omini"}, + {"gpt-4o-2024-08-06", "gpt4o"}, + {"text-embedding-ada-002-v2", "textembeddingada002v2"}, + {"dall-e-3", "dalle3"}, + {"Image models", "imagemodels"}, + {"", ""}, + } + for _, tt := range tests { + got := normalizeModelKey(tt.input) + if got != tt.want { + t.Errorf("normalizeModelKey(%q) = %q, want %q", tt.input, got, tt.want) + } + } +} + +func TestIsDateLike(t *testing.T) { + tests := []struct { + input string + want bool + }{ + {"2024-07-18", true}, + {"2024-08-06", true}, + {"not-a-date", false}, + {"20240718", false}, + {"2024-7-18", false}, + } + for _, tt := range tests { + got := isDateLike(tt.input) + if got != tt.want { + t.Errorf("isDateLike(%q) = %v, want %v", tt.input, got, tt.want) + } + } +} + +func TestBuildTokenMap(t *testing.T) { + model := "gpt-4o-2024-08-06" + projectID := "proj_abc" + + usage := &openaiplugin.CompletionsUsageResponse{ + Data: []openaiplugin.CompletionsTimeBucket{ + { + Results: []openaiplugin.CompletionsUsageResult{ + { + InputTokens: 100, + OutputTokens: 50, + Model: &model, + ProjectID: &projectID, + }, + { + InputTokens: 200, + OutputTokens: 75, + Model: &model, + ProjectID: &projectID, + }, + }, + }, + }, + } + + tokenMap := buildTokenMap(usage) + + key := normalizeModelKey(model) + "|" + projectID + expected := 100 + 50 + 200 + 75 + if tokenMap[key] != expected { + t.Errorf("buildTokenMap key %q = %d, want %d", key, tokenMap[key], expected) + } +} + +func TestBuildTokenMapNil(t *testing.T) { + tokenMap := buildTokenMap(nil) + if len(tokenMap) != 0 { + t.Errorf("expected empty token map for nil input, got %d entries", len(tokenMap)) + } +} + +func TestBuildCustomCosts(t *testing.T) { + lineItem := "GPT-4o" + projectID := "proj_abc" + model := "gpt-4o-2024-08-06" + + costs := &openaiplugin.CostsResponse{ + Data: []openaiplugin.CostsTimeBucket{ + { + Results: []openaiplugin.CostsResult{ + { + Amount: openaiplugin.CostAmount{Value: 0.42, Currency: "usd"}, + LineItem: &lineItem, + ProjectID: &projectID, + }, + }, + }, + }, + } + + usage := &openaiplugin.CompletionsUsageResponse{ + Data: []openaiplugin.CompletionsTimeBucket{ + { + Results: []openaiplugin.CompletionsUsageResult{ + { + InputTokens: 1000, + OutputTokens: 500, + Model: &model, + ProjectID: &projectID, + }, + }, + }, + }, + } + + customCosts := buildCustomCosts(costs, usage) + + if len(customCosts) != 1 { + t.Fatalf("expected 1 custom cost, got %d", len(customCosts)) + } + + cc := customCosts[0] + if cc.BilledCost != float32(0.42) { + t.Errorf("BilledCost = %f, want 0.42", cc.BilledCost) + } + if cc.ResourceName != "GPT-4o" { + t.Errorf("ResourceName = %q, want %q", cc.ResourceName, "GPT-4o") + } + if cc.UsageQuantity != 1500 { + t.Errorf("UsageQuantity = %f, want 1500", cc.UsageQuantity) + } + if cc.ChargeCategory != "Usage" { + t.Errorf("ChargeCategory = %q, want %q", cc.ChargeCategory, "Usage") + } +} + +func TestBuildCustomCostsSkipsZero(t *testing.T) { + lineItem := "GPT-4o" + projectID := "proj_abc" + + costs := &openaiplugin.CostsResponse{ + Data: []openaiplugin.CostsTimeBucket{ + { + Results: []openaiplugin.CostsResult{ + { + Amount: openaiplugin.CostAmount{Value: 0, Currency: "usd"}, + LineItem: &lineItem, + ProjectID: &projectID, + }, + }, + }, + }, + } + + customCosts := buildCustomCosts(costs, nil) + + if len(customCosts) != 0 { + t.Errorf("expected 0 custom costs for zero-value entry, got %d", len(customCosts)) + } +} + +func TestBuildCustomCostsNilInputs(t *testing.T) { + customCosts := buildCustomCosts(nil, nil) + if customCosts != nil { + t.Errorf("expected nil for nil costs input, got %d entries", len(customCosts)) + } +} + +func TestBuildCustomCostsNoTokenMatch(t *testing.T) { + lineItem := "Image models" + projectID := "proj_abc" + + costs := &openaiplugin.CostsResponse{ + Data: []openaiplugin.CostsTimeBucket{ + { + Results: []openaiplugin.CostsResult{ + { + Amount: openaiplugin.CostAmount{Value: 1.50, Currency: "usd"}, + LineItem: &lineItem, + ProjectID: &projectID, + }, + }, + }, + }, + } + + customCosts := buildCustomCosts(costs, nil) + + if len(customCosts) != 1 { + t.Fatalf("expected 1 custom cost, got %d", len(customCosts)) + } + if customCosts[0].UsageQuantity != -1 { + t.Errorf("UsageQuantity = %f, want -1 (no token match)", customCosts[0].UsageQuantity) + } } diff --git a/pkg/plugins/openai/cmd/main/main.go b/pkg/plugins/openai/cmd/main/main.go index 964fb97..4bfa1d4 100644 --- a/pkg/plugins/openai/cmd/main/main.go +++ b/pkg/plugins/openai/cmd/main/main.go @@ -6,8 +6,8 @@ import ( "fmt" "io" "net/http" + "net/url" "os" - "regexp" "strconv" "strings" "time" @@ -26,21 +26,15 @@ import ( "github.com/opencost/opencost/core/pkg/util/timeutil" ) -// handshakeConfigs are used to just do a basic handshake between -// a plugin and host. If the handshake fails, a user friendly error is shown. -// This prevents users from executing bad plugins or executing a plugin -// directory. It is a UX feature, not a security feature. var handshakeConfig = plugin.HandshakeConfig{ ProtocolVersion: 1, MagicCookieKey: "PLUGIN_NAME", MagicCookieValue: "openai", } -const openAIUsageURLFmt = "https://api.openai.com/v1/usage?date=%s" -const openAIBillingURLFmt = "https://api.openai.com/v1/dashboard/billing/usage/export?exclude_project_costs=false&file_format=json&new_endpoint=true&project_id&start_date=%s&end_date=%s" -const openAIAPIDateFormat = "2006-01-02" +const openAICostsURL = "https://api.openai.com/v1/organization/costs" +const openAICompletionsUsageURL = "https://api.openai.com/v1/organization/usage/completions" -// Implementation of CustomCostSource type OpenAICostSource struct { rateLimiter *rate.Limiter config *openaiplugin.OpenAIConfig @@ -65,7 +59,6 @@ func (d *OpenAICostSource) GetCustomCosts(req *pb.CustomCostRequest) []*pb.Custo } for _, target := range targets { - // don't allow future request if target.Start().After(time.Now().UTC()) { log.Debugf("skipping future window %v", target) continue @@ -80,7 +73,6 @@ func (d *OpenAICostSource) GetCustomCosts(req *pb.CustomCostRequest) []*pb.Custo } func main() { - configFile, err := commonconfig.GetConfigFilePath() if err != nil { log.Fatalf("error opening config file: %v", err) @@ -91,14 +83,12 @@ func main() { log.Fatalf("error building OpenAI config: %v", err) } log.SetLogLevel(oaiConfig.LogLevel) - // rate limit to 1 request per second rateLimiter := rate.NewLimiter(0.5, 1) oaiCostSrc := OpenAICostSource{ rateLimiter: rateLimiter, config: oaiConfig, } - // pluginMap is the map of plugins we can dispense. var pluginMap = map[string]plugin.Plugin{ "CustomCostSource": &ocplugin.CustomCostPlugin{Impl: &oaiCostSrc}, } @@ -123,177 +113,249 @@ func boilerplateOpenAICustomCost(win opencost.Window) pb.CustomCostResponse { Costs: []*pb.CustomCost{}, } } + func (d *OpenAICostSource) getOpenAICostsForWindow(window opencost.Window) *pb.CustomCostResponse { ccResp := boilerplateOpenAICustomCost(window) - oaiTokenUsages, err := d.getOpenAITokenUsages(*window.Start()) + costsResp, err := d.getOpenAICosts(*window.Start(), *window.End()) if err != nil { - ccResp.Errors = append(ccResp.Errors, fmt.Sprintf("error getting OpenAI token usages: %v", err)) + ccResp.Errors = append(ccResp.Errors, fmt.Sprintf("error getting OpenAI costs: %v", err)) } - oaiBilling, err := d.getOpenAIBilling(*window.Start(), *window.End()) + usageResp, err := d.getOpenAICompletionsUsage(*window.Start(), *window.End()) if err != nil { - ccResp.Errors = append(ccResp.Errors, fmt.Sprintf("error getting OpenAI billing data: %v", err)) + ccResp.Errors = append(ccResp.Errors, fmt.Sprintf("error getting OpenAI usage: %v", err)) } - customCosts, err := getCustomCostsFromUsageAndBilling(oaiTokenUsages, oaiBilling) - if err != nil { - ccResp.Errors = append(ccResp.Errors, fmt.Sprintf("error converting API responses into custom costs: %v", err)) - } + customCosts := buildCustomCosts(costsResp, usageResp) ccResp.Costs = customCosts return &ccResp } -func getCustomCostsFromUsageAndBilling(usage *openaiplugin.OpenAIUsage, billing *openaiplugin.OpenAIBilling) ([]*pb.CustomCost, error) { - customCosts := []*pb.CustomCost{} - +// buildCustomCosts merges cost data (dollar amounts per line_item/project) with +// usage data (token counts per model/project) into FOCUS-compliant CustomCost entries. +// +// The costs API groups by line_item (e.g. "GPT-4o", "Embeddings") and project_id. +// The usage API groups by model (e.g. "gpt-4o-2024-08-06") and project_id. +// We use the costs response as the primary source of truth for billed amounts, and +// attach token counts where we can match on a normalized model/line_item key. +func buildCustomCosts(costs *openaiplugin.CostsResponse, usage *openaiplugin.CompletionsUsageResponse) []*pb.CustomCost { tokenMap := buildTokenMap(usage) - for _, billingEntry := range billing.Data { - tokenMapKey := strings.ReplaceAll(strings.ToLower(billingEntry.Name), "-", "") - tokenMapKey = strings.ReplaceAll(tokenMapKey, " ", "") - tokenMapKey = strings.ReplaceAll(tokenMapKey, "_", "") - - tokenCount, ok := tokenMap[tokenMapKey] - if !ok { - log.Debugf("no token usage found for %s", billingEntry.Name) - tokenCount = -1 - } + var customCosts []*pb.CustomCost - extendedAttrs := pb.CustomCostExtendedAttributes{ - AccountId: &billingEntry.OrganizationID, - SubAccountId: &billingEntry.ProjectID, - } - customCost := pb.CustomCost{ - BilledCost: float32(billingEntry.CostInMajor), - AccountName: billingEntry.OrganizationName, - ChargeCategory: "Usage", - Description: fmt.Sprintf("OpenAI usage for model %s", billingEntry.Name), - ResourceName: billingEntry.Name, - ResourceType: "AI Model", - Id: uuid.New().String(), - ProviderId: fmt.Sprintf("%s/%s/%s", billingEntry.OrganizationID, billingEntry.ProjectID, billingEntry.Name), - UsageQuantity: float32(tokenCount), - UsageUnit: "tokens - All snapshots, all projects", - ExtendedAttributes: &extendedAttrs, - } + if costs == nil { + return customCosts + } + + for _, bucket := range costs.Data { + for _, result := range bucket.Results { + lineItem := derefStr(result.LineItem, "Unknown") + projectID := derefStr(result.ProjectID, "") + cost := result.Amount.Value + + if cost == 0 { + log.Debugf("skipping zero-cost line item %s", lineItem) + continue + } + + tokenKey := normalizeModelKey(lineItem) + "|" + projectID + tokenCount, hasTokens := tokenMap[tokenKey] + if !hasTokens { + tokenKeyNoProject := normalizeModelKey(lineItem) + "|" + tokenCount, hasTokens = tokenMap[tokenKeyNoProject] + } + + usageQty := float32(-1) + if hasTokens { + usageQty = float32(tokenCount) + } + + extendedAttrs := pb.CustomCostExtendedAttributes{ + SubAccountId: strPtr(projectID), + } + cc := pb.CustomCost{ + BilledCost: float32(cost), + ChargeCategory: "Usage", + Description: fmt.Sprintf("OpenAI usage for %s", lineItem), + ResourceName: lineItem, + ResourceType: "AI Model", + Id: uuid.New().String(), + ProviderId: fmt.Sprintf("%s/%s", projectID, lineItem), + UsageQuantity: usageQty, + UsageUnit: "tokens", + ExtendedAttributes: &extendedAttrs, + } - customCosts = append(customCosts, &customCost) + customCosts = append(customCosts, &cc) + } } - return customCosts, nil + return customCosts } -var snapshotRe = regexp.MustCompile(`-\d{4}-\d{2}-\d{2}|-`) - -func buildTokenMap(usage *openaiplugin.OpenAIUsage) map[string]int { +// buildTokenMap aggregates token counts from the completions usage response, +// keyed by normalizedModel|projectID. +func buildTokenMap(usage *openaiplugin.CompletionsUsageResponse) map[string]int { tokenMap := make(map[string]int) if usage == nil { return tokenMap } - for _, usageData := range usage.Data { - key := snapshotRe.ReplaceAllString(usageData.SnapshotID, "") - key = strings.ToLower(key) - if _, ok := tokenMap[key]; !ok { - tokenMap[key] = 0 + for _, bucket := range usage.Data { + for _, result := range bucket.Results { + model := derefStr(result.Model, "unknown") + projectID := derefStr(result.ProjectID, "") + key := normalizeModelKey(model) + "|" + projectID + tokenMap[key] += result.InputTokens + result.OutputTokens } - - tokenMap[key] += (usageData.NGeneratedTokensTotal + usageData.NContextTokensTotal) } return tokenMap } -func (d *OpenAICostSource) getOpenAIBilling(start time.Time, end time.Time) (*openaiplugin.OpenAIBilling, error) { - client := &http.Client{} - openAIBillingURL := fmt.Sprintf(openAIBillingURLFmt, start.Format(openAIAPIDateFormat), end.Format(openAIAPIDateFormat)) - log.Debugf("fetching OpenAI billing data from %s", openAIBillingURL) - var errReq error - var resp *http.Response - for i := 0; i < 3; i++ { - err := d.rateLimiter.Wait(context.Background()) - if err != nil { - log.Warnf("error waiting for rate limiter: %v", err) - return nil, fmt.Errorf("error waiting for rate limiter: %v", err) +// normalizeModelKey strips date suffixes, dashes, spaces, underscores, and lowercases +// so that "GPT-4o mini" and "gpt-4o-mini-2024-07-18" both map to "gpt4omini". +func normalizeModelKey(s string) string { + s = strings.ToLower(s) + // strip trailing date suffix like "-2024-07-18" (11 chars: dash + YYYY-MM-DD) + if len(s) > 11 { + candidate := s[len(s)-10:] + if s[len(s)-11] == '-' && isDateLike(candidate) { + s = s[:len(s)-11] } - var req *http.Request - req, errReq = http.NewRequest("GET", openAIBillingURL, nil) - if errReq != nil { - log.Warnf("error creating billing export request: %v", errReq) - log.Infof("retrying request after 30s") - time.Sleep(30 * time.Second) - continue + } + s = strings.ReplaceAll(s, "-", "") + s = strings.ReplaceAll(s, " ", "") + s = strings.ReplaceAll(s, "_", "") + return s +} + +func isDateLike(s string) bool { + if len(s) != 10 { + return false + } + for i, c := range s { + if i == 4 || i == 7 { + if c != '-' { + return false + } + } else if c < '0' || c > '9' { + return false } + } + return true +} - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", d.config.APIKey)) +func derefStr(p *string, fallback string) string { + if p != nil { + return *p + } + return fallback +} - resp, errReq = client.Do(req) - if errReq != nil { - log.Warnf("error doing billing export request: %v", errReq) - log.Warnf("retrying requestafter 30s") - time.Sleep(30 * time.Second) - continue +func strPtr(s string) *string { + return &s +} + +// getOpenAICosts fetches cost data from /v1/organization/costs grouped by +// line_item and project_id. Handles pagination. +func (d *OpenAICostSource) getOpenAICosts(start, end time.Time) (*openaiplugin.CostsResponse, error) { + params := url.Values{} + params.Set("start_time", strconv.FormatInt(start.Unix(), 10)) + params.Set("end_time", strconv.FormatInt(end.Unix(), 10)) + params.Set("bucket_width", "1d") + params.Set("limit", "31") + params.Add("group_by", "line_item") + params.Add("group_by", "project_id") + + var allBuckets []openaiplugin.CostsTimeBucket + page := "" + + for { + if page != "" { + params.Set("page", page) } - if resp.StatusCode != http.StatusOK { - bodyBytes, err := io.ReadAll(resp.Body) - bodyString := "" - if err != nil { - log.Warnf("error reading body of non-200 response: %v", err) - } else { - bodyString = string(bodyBytes) - } + fullURL := openAICostsURL + "?" + params.Encode() + log.Debugf("fetching OpenAI costs from %s", fullURL) - errReq = fmt.Errorf("received non-200 response for billing export request: %d", resp.StatusCode) - log.Warnf("got non-200 response for billing export request: %d, body is: %s", resp.StatusCode, bodyString) - log.Infof("retrying request after 30s") - time.Sleep(30 * time.Second) - continue - } else { - errReq = nil + body, err := d.doRequestWithRetry("GET", fullURL) + if err != nil { + return nil, fmt.Errorf("costs request failed: %w", err) } - // request was successful, break out of loop - break - } - if errReq != nil { - return nil, fmt.Errorf("error making request after retries: %v", errReq) - } - var billingData openaiplugin.OpenAIBilling - if err := json.NewDecoder(resp.Body).Decode(&billingData); err != nil { - return nil, fmt.Errorf("error decoding billing export response: %v", err) + var resp openaiplugin.CostsResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("error decoding costs response: %w", err) + } + + allBuckets = append(allBuckets, resp.Data...) + + if !resp.HasMore || resp.NextPage == "" { + break + } + page = resp.NextPage } - resp.Body.Close() - for i := range billingData.Data { - asFloat, err := strconv.ParseFloat(billingData.Data[i].CostInMajorStr, 64) + + return &openaiplugin.CostsResponse{Data: allBuckets}, nil +} + +// getOpenAICompletionsUsage fetches token usage from /v1/organization/usage/completions +// grouped by model and project_id. Handles pagination. +func (d *OpenAICostSource) getOpenAICompletionsUsage(start, end time.Time) (*openaiplugin.CompletionsUsageResponse, error) { + params := url.Values{} + params.Set("start_time", strconv.FormatInt(start.Unix(), 10)) + params.Set("end_time", strconv.FormatInt(end.Unix(), 10)) + params.Set("bucket_width", "1d") + params.Set("limit", "31") + params.Add("group_by", "model") + params.Add("group_by", "project_id") + + var allBuckets []openaiplugin.CompletionsTimeBucket + page := "" + + for { + if page != "" { + params.Set("page", page) + } + + fullURL := openAICompletionsUsageURL + "?" + params.Encode() + log.Debugf("fetching OpenAI completions usage from %s", fullURL) + + body, err := d.doRequestWithRetry("GET", fullURL) if err != nil { - return nil, fmt.Errorf("error parsing cost: %v", err) + return nil, fmt.Errorf("completions usage request failed: %w", err) } - billingData.Data[i].CostInMajor = asFloat + + var resp openaiplugin.CompletionsUsageResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("error decoding completions usage response: %w", err) + } + + allBuckets = append(allBuckets, resp.Data...) + + if !resp.HasMore || resp.NextPage == "" { + break + } + page = resp.NextPage } - return &billingData, nil + return &openaiplugin.CompletionsUsageResponse{Data: allBuckets}, nil } -func (d *OpenAICostSource) getOpenAITokenUsages(targetTime time.Time) (*openaiplugin.OpenAIUsage, error) { - client := &http.Client{} +// doRequestWithRetry performs an authenticated GET with up to 3 retries on failure. +func (d *OpenAICostSource) doRequestWithRetry(method, fullURL string) ([]byte, error) { + client := &http.Client{Timeout: 30 * time.Second} - openAIUsageURL := fmt.Sprintf(openAIUsageURLFmt, targetTime.Format(openAIAPIDateFormat)) - log.Debugf("fetching OpenAI usage data from %s", openAIUsageURL) - var errReq error - var resp *http.Response - for i := 0; i < 3; i++ { - errReq = nil - err := d.rateLimiter.Wait(context.Background()) - if err != nil { - log.Warnf("error waiting for rate limiter: %v", err) - return nil, fmt.Errorf("error waiting for rate limiter: %v", err) + var lastErr error + for attempt := 0; attempt < 3; attempt++ { + if err := d.rateLimiter.Wait(context.Background()); err != nil { + return nil, fmt.Errorf("rate limiter error: %w", err) } - var req *http.Request - req, errReq = http.NewRequest("GET", openAIUsageURL, nil) - if errReq != nil { - log.Warnf("error creating usage request: %v", errReq) - log.Warnf("retrying request after 30s") + + req, err := http.NewRequest(method, fullURL, nil) + if err != nil { + lastErr = fmt.Errorf("error creating request: %w", err) + log.Warnf("%v, retrying in 30s", lastErr) time.Sleep(30 * time.Second) continue } @@ -301,45 +363,34 @@ func (d *OpenAICostSource) getOpenAITokenUsages(targetTime time.Time) (*openaipl req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", d.config.APIKey)) - resp, errReq = client.Do(req) - if errReq != nil { - log.Warnf("error doing token request: %v", errReq) - log.Infof("retrying request after 30s") + resp, err := client.Do(req) + if err != nil { + lastErr = fmt.Errorf("request failed: %w", err) + log.Warnf("%v, retrying in 30s", lastErr) + time.Sleep(30 * time.Second) + continue + } + + body, err := io.ReadAll(resp.Body) + resp.Body.Close() + if err != nil { + lastErr = fmt.Errorf("error reading response body: %w", err) + log.Warnf("%v, retrying in 30s", lastErr) time.Sleep(30 * time.Second) continue } - defer resp.Body.Close() if resp.StatusCode != http.StatusOK { - errReq = fmt.Errorf("received non-200 response for token usage request: %d", resp.StatusCode) - bodyBytes, err := io.ReadAll(resp.Body) - bodyString := "" - if err != nil { - log.Warnf("error reading body of non-200 response: %v", err) - } else { - bodyString = string(bodyBytes) - } - log.Warnf("got non-200 response for token usage request: %d, body is: %s", resp.StatusCode, bodyString) - log.Infof("retrying request after 30s") + lastErr = fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body[:min(len(body), 500)])) + log.Warnf("non-200 response: %v, retrying in 30s", lastErr) time.Sleep(30 * time.Second) continue - } else { - errReq = nil } - // request was successful, break out of loop - break - } - - if errReq != nil { - return nil, fmt.Errorf("error making request after retries: %v", errReq) - } - var usageData openaiplugin.OpenAIUsage - if err := json.NewDecoder(resp.Body).Decode(&usageData); err != nil { - return nil, fmt.Errorf("error decoding response: %v", err) + return body, nil } - return &usageData, nil + return nil, fmt.Errorf("request failed after 3 attempts: %w", lastErr) } func getOpenAIConfig(configFilePath string) (*openaiplugin.OpenAIConfig, error) { diff --git a/pkg/plugins/openai/cmd/validator/main/main.go b/pkg/plugins/openai/cmd/validator/main/main.go index 5e16a7a..dd5e37a 100644 --- a/pkg/plugins/openai/cmd/validator/main/main.go +++ b/pkg/plugins/openai/cmd/validator/main/main.go @@ -4,6 +4,7 @@ import ( "encoding/json" "fmt" "os" + "strings" "time" "github.com/hashicorp/go-multierror" @@ -12,6 +13,14 @@ import ( "google.golang.org/protobuf/encoding/protojson" ) +const maxDailyBilledCostPerLineItem = float32(5) + +// expectedCostSubstrings spot-checks common OpenAI line item families across the window. +// Line item names vary by model version and charge type (e.g. "gpt-4o-mini-2024-07-18, input"). +var expectedCostSubstrings = []string{ + "gpt-4o-mini", +} + // the validator is designed to allow plugin implementors to validate their plugin information // as called by the central test harness. // this avoids having to ask folks to re-implement the test harness over again for each plugin @@ -106,30 +115,36 @@ func validate(respDaily, respHourly []*pb.CustomCostResponse) bool { if cost.GetBilledCost() == 0 { log.Debugf("got zero cost for %v", cost) } - if cost.GetBilledCost() > 2 { - log.Errorf("daily cost returned by plugin openai for %v is greater than 1", cost) + if cost.GetBilledCost() > maxDailyBilledCostPerLineItem { + log.Errorf("daily cost returned by plugin openai for %v is greater than %v", cost, maxDailyBilledCostPerLineItem) return false } } - } + if costSum == 0 { log.Errorf("daily costs returned by openai plugin are zero") return false } - expectedCosts := []string{ - "GPT-4o mini", - "GPT-4o", + + if len(seenCosts) == 0 { + log.Errorf("no cost line items found in openai plugin response") + return false } - for _, cost := range expectedCosts { - if !seenCosts[cost] { - log.Errorf("daily cost %s not found in plugin openai response", cost) + for _, substr := range expectedCostSubstrings { + if !costLineItemSeen(seenCosts, substr) { + log.Errorf("daily cost containing %q not found in plugin openai response", substr) return false } } + if !gpt4oLineItemSeen(seenCosts) { + log.Errorf("daily GPT-4o cost line item not found in plugin openai response") + return false + } + + log.Infof("found %d distinct cost line items: %v", len(seenCosts), seenCosts) - // verify the domain matches the plugin name for _, resp := range respDaily { if resp.Domain != "openai" { log.Errorf("daily domain returned by plugin openai does not match plugin name") @@ -137,13 +152,30 @@ func validate(respDaily, respHourly []*pb.CustomCostResponse) bool { } } - if len(seenCosts) < len(expectedCosts)-1 || len(seenCosts) > len(expectedCosts)+1 { - log.Errorf("daily costs returned by openai plugin are very different than expected") - return false - } return true } +func costLineItemSeen(seen map[string]bool, substr string) bool { + substr = strings.ToLower(substr) + for name := range seen { + if strings.Contains(strings.ToLower(name), substr) { + return true + } + } + return false +} + +// gpt4oLineItemSeen matches GPT-4o line items but not GPT-4o mini variants. +func gpt4oLineItemSeen(seen map[string]bool) bool { + for name := range seen { + lower := strings.ToLower(name) + if strings.Contains(lower, "gpt-4o") && !strings.Contains(lower, "gpt-4o-mini") { + return true + } + } + return false +} + func Unmarshal(data []byte) ([]*pb.CustomCostResponse, error) { var raw []json.RawMessage if err := json.Unmarshal(data, &raw); err != nil { diff --git a/pkg/plugins/openai/openaiplugin/openaibilling.go b/pkg/plugins/openai/openaiplugin/openaibilling.go index e767e9a..c6dd22a 100644 --- a/pkg/plugins/openai/openaiplugin/openaibilling.go +++ b/pkg/plugins/openai/openaiplugin/openaibilling.go @@ -1,22 +1,64 @@ package openaiplugin -// OpenAIBilling represents the structure of the response JSON -type OpenAIBilling struct { - Object string `json:"object"` - Data []BillingData `json:"data"` +import ( + "encoding/json" + "fmt" + "strconv" +) + +// CostsResponse represents GET /v1/organization/costs with group_by=line_item,project_id +type CostsResponse struct { + Object string `json:"object"` + Data []CostsTimeBucket `json:"data"` + HasMore bool `json:"has_more"` + NextPage string `json:"next_page"` +} + +type CostsTimeBucket struct { + Object string `json:"object"` + StartTime int64 `json:"start_time"` + EndTime int64 `json:"end_time"` + Results []CostsResult `json:"results"` +} + +type CostsResult struct { + Object string `json:"object"` + Amount CostAmount `json:"amount"` + LineItem *string `json:"line_item"` + ProjectID *string `json:"project_id"` } -// BillingData represents the individual Billing data entries -type BillingData struct { - Timestamp float64 `json:"timestamp"` - Currency string `json:"currency"` - Name string `json:"name"` - Cost float64 `json:"cost"` - OrganizationID string `json:"organization_id"` - ProjectID string `json:"project_id"` - ProjectName string `json:"project_name"` - OrganizationName string `json:"organization_name"` - CostInMajorStr string `json:"cost_in_major"` - CostInMajor float64 `json:"-"` - Date string `json:"date"` +// CostAmount handles the API returning value as either a number or a string. +type CostAmount struct { + Value float64 + Currency string +} + +func (c *CostAmount) UnmarshalJSON(data []byte) error { + var raw struct { + Value json.RawMessage `json:"value"` + Currency string `json:"currency"` + } + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + c.Currency = raw.Currency + + var f float64 + if err := json.Unmarshal(raw.Value, &f); err == nil { + c.Value = f + return nil + } + + var s string + if err := json.Unmarshal(raw.Value, &s); err == nil { + parsed, err := strconv.ParseFloat(s, 64) + if err != nil { + return fmt.Errorf("cannot parse amount value %q as float: %w", s, err) + } + c.Value = parsed + return nil + } + + return fmt.Errorf("cannot unmarshal amount value: %s", string(raw.Value)) } diff --git a/pkg/plugins/openai/openaiplugin/openaiusage.go b/pkg/plugins/openai/openaiplugin/openaiusage.go index d541a40..ff52dd4 100644 --- a/pkg/plugins/openai/openaiplugin/openaiusage.go +++ b/pkg/plugins/openai/openaiplugin/openaiusage.go @@ -1,26 +1,32 @@ package openaiplugin -type OpenAIUsage struct { - Object string `json:"object"` - Data []UsageData `json:"data"` +// CompletionsUsageResponse represents GET /v1/organization/usage/completions +// with group_by=model,project_id +type CompletionsUsageResponse struct { + Object string `json:"object"` + Data []CompletionsTimeBucket `json:"data"` + HasMore bool `json:"has_more"` + NextPage string `json:"next_page"` } -type UsageData struct { - OrganizationID string `json:"organization_id"` - OrganizationName string `json:"organization_name"` - AggregationTimestamp int `json:"aggregation_timestamp"` - NRequests int `json:"n_requests"` - Operation string `json:"operation"` - SnapshotID string `json:"snapshot_id"` - NContextTokensTotal int `json:"n_context_tokens_total"` - NGeneratedTokensTotal int `json:"n_generated_tokens_total"` - Email *string `json:"email"` - APIKeyID *string `json:"api_key_id"` - APIKeyName *string `json:"api_key_name"` - APIKeyRedacted *string `json:"api_key_redacted"` - APIKeyType *string `json:"api_key_type"` - ProjectID *string `json:"project_id"` - ProjectName *string `json:"project_name"` - RequestType string `json:"request_type"` - NCachedContextTokensTotal int `json:"n_cached_context_tokens_total"` +type CompletionsTimeBucket struct { + Object string `json:"object"` + StartTime int64 `json:"start_time"` + EndTime int64 `json:"end_time"` + Results []CompletionsUsageResult `json:"results"` +} + +type CompletionsUsageResult struct { + Object string `json:"object"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + InputCachedTokens int `json:"input_cached_tokens"` + InputAudioTokens int `json:"input_audio_tokens"` + OutputAudioTokens int `json:"output_audio_tokens"` + NumModelRequests int `json:"num_model_requests"` + ProjectID *string `json:"project_id"` + UserID *string `json:"user_id"` + APIKeyID *string `json:"api_key_id"` + Model *string `json:"model"` + Batch *bool `json:"batch"` }