From 0e1175c1a6a3470e46e040a72e4c22f76b93177e Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 5 Nov 2024 01:56:12 +0530 Subject: [PATCH 01/36] x-pack/metricbeat/module/openai: Add new module --- x-pack/metricbeat/include/list.go | 2 + x-pack/metricbeat/metricbeat.reference.yml | 29 ++ .../metricbeat/module/openai/_meta/config.yml | 28 ++ .../module/openai/_meta/docs.asciidoc | 2 + .../metricbeat/module/openai/_meta/fields.yml | 10 + x-pack/metricbeat/module/openai/doc.go | 2 + x-pack/metricbeat/module/openai/fields.go | 23 ++ .../module/openai/usage/_meta/data.json | 19 ++ .../module/openai/usage/_meta/docs.asciidoc | 1 + .../module/openai/usage/_meta/fields.yml | 10 + .../metricbeat/module/openai/usage/client.go | 44 +++ .../metricbeat/module/openai/usage/config.go | 76 +++++ .../metricbeat/module/openai/usage/helper.go | 19 ++ .../metricbeat/module/openai/usage/schema.go | 82 +++++ .../metricbeat/module/openai/usage/usage.go | 316 ++++++++++++++++++ 15 files changed, 663 insertions(+) create mode 100644 x-pack/metricbeat/module/openai/_meta/config.yml create mode 100644 x-pack/metricbeat/module/openai/_meta/docs.asciidoc create mode 100644 x-pack/metricbeat/module/openai/_meta/fields.yml create mode 100644 x-pack/metricbeat/module/openai/doc.go create mode 100644 x-pack/metricbeat/module/openai/fields.go create mode 100644 x-pack/metricbeat/module/openai/usage/_meta/data.json create mode 100644 x-pack/metricbeat/module/openai/usage/_meta/docs.asciidoc create mode 100644 x-pack/metricbeat/module/openai/usage/_meta/fields.yml create mode 100644 x-pack/metricbeat/module/openai/usage/client.go create mode 100644 x-pack/metricbeat/module/openai/usage/config.go create mode 100644 x-pack/metricbeat/module/openai/usage/helper.go create mode 100644 x-pack/metricbeat/module/openai/usage/schema.go create mode 100644 x-pack/metricbeat/module/openai/usage/usage.go diff --git a/x-pack/metricbeat/include/list.go b/x-pack/metricbeat/include/list.go index 01ce86edf78c..fdf19a870249 100644 --- a/x-pack/metricbeat/include/list.go +++ b/x-pack/metricbeat/include/list.go @@ -53,6 +53,8 @@ import ( _ "github.com/elastic/beats/v7/x-pack/metricbeat/module/mssql" _ "github.com/elastic/beats/v7/x-pack/metricbeat/module/mssql/performance" _ "github.com/elastic/beats/v7/x-pack/metricbeat/module/mssql/transaction_log" + _ "github.com/elastic/beats/v7/x-pack/metricbeat/module/openai" + _ "github.com/elastic/beats/v7/x-pack/metricbeat/module/openai/usage" _ "github.com/elastic/beats/v7/x-pack/metricbeat/module/oracle" _ "github.com/elastic/beats/v7/x-pack/metricbeat/module/oracle/performance" _ "github.com/elastic/beats/v7/x-pack/metricbeat/module/oracle/sysmetric" diff --git a/x-pack/metricbeat/metricbeat.reference.yml b/x-pack/metricbeat/metricbeat.reference.yml index 240acb2cfd6a..48fc2eefa302 100644 --- a/x-pack/metricbeat/metricbeat.reference.yml +++ b/x-pack/metricbeat/metricbeat.reference.yml @@ -1257,6 +1257,35 @@ metricbeat.modules: # Path to server status. Default nginx_status server_status_path: "nginx_status" +#-------------------------------- Openai Module -------------------------------- +- module: openai + metricsets: ["usage"] + enabled: false + period: 1h + + # # Project API Keys - Multiple API keys can be specified for different projects + # api_keys: + # - key: "api_key1" + # - key: "api_key2" + + # # API Configuration + # ## Base URL for the OpenAI usage API endpoint + # api_url: "https://api.openai.com/v1/usage" + # ## Custom headers to be included in API requests + # headers: + # - "k1: v1" + # - "k2: v2" + ## Rate Limiting Configuration + # rate_limit: + # limit: 60 # requests per second + # burst: 5 # burst size + # ## Request timeout duration + # timeout: 30s + + # # Data Collection Configuration + # collection: + # ## Number of days to look back when collecting usage data + # lookback_days: 30 #----------------------------- Openmetrics Module ----------------------------- - module: openmetrics metricsets: ['collector'] diff --git a/x-pack/metricbeat/module/openai/_meta/config.yml b/x-pack/metricbeat/module/openai/_meta/config.yml new file mode 100644 index 000000000000..dadcd35b1bae --- /dev/null +++ b/x-pack/metricbeat/module/openai/_meta/config.yml @@ -0,0 +1,28 @@ +- module: openai + metricsets: ["usage"] + enabled: false + period: 1h + + # # Project API Keys - Multiple API keys can be specified for different projects + # api_keys: + # - key: "api_key1" + # - key: "api_key2" + + # # API Configuration + # ## Base URL for the OpenAI usage API endpoint + # api_url: "https://api.openai.com/v1/usage" + # ## Custom headers to be included in API requests + # headers: + # - "k1: v1" + # - "k2: v2" + ## Rate Limiting Configuration + # rate_limit: + # limit: 60 # requests per second + # burst: 5 # burst size + # ## Request timeout duration + # timeout: 30s + + # # Data Collection Configuration + # collection: + # ## Number of days to look back when collecting usage data + # lookback_days: 30 \ No newline at end of file diff --git a/x-pack/metricbeat/module/openai/_meta/docs.asciidoc b/x-pack/metricbeat/module/openai/_meta/docs.asciidoc new file mode 100644 index 000000000000..744909c7a1ab --- /dev/null +++ b/x-pack/metricbeat/module/openai/_meta/docs.asciidoc @@ -0,0 +1,2 @@ +This is the openai module. + diff --git a/x-pack/metricbeat/module/openai/_meta/fields.yml b/x-pack/metricbeat/module/openai/_meta/fields.yml new file mode 100644 index 000000000000..d514eb010f15 --- /dev/null +++ b/x-pack/metricbeat/module/openai/_meta/fields.yml @@ -0,0 +1,10 @@ +- key: openai + title: "openai" + release: beta + description: > + openai module + fields: + - name: openai + type: group + description: > + fields: diff --git a/x-pack/metricbeat/module/openai/doc.go b/x-pack/metricbeat/module/openai/doc.go new file mode 100644 index 000000000000..b6d7a70b1215 --- /dev/null +++ b/x-pack/metricbeat/module/openai/doc.go @@ -0,0 +1,2 @@ +// Package openai is a Metricbeat module that contains MetricSets. +package openai diff --git a/x-pack/metricbeat/module/openai/fields.go b/x-pack/metricbeat/module/openai/fields.go new file mode 100644 index 000000000000..a23f01a7810d --- /dev/null +++ b/x-pack/metricbeat/module/openai/fields.go @@ -0,0 +1,23 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License; +// you may not use this file except in compliance with the Elastic License. + +// Code generated by beats/dev-tools/cmd/asset/asset.go - DO NOT EDIT. + +package openai + +import ( + "github.com/elastic/beats/v7/libbeat/asset" +) + +func init() { + if err := asset.SetFields("metricbeat", "openai", asset.ModuleFieldsPri, AssetOpenai); err != nil { + panic(err) + } +} + +// AssetOpenai returns asset data. +// This is the base64 encoded zlib format compressed contents of module/openai. +func AssetOpenai() string { + return "eJx8jkHOgyAUhPecYuLeC7D4d/9BXsvUEFEIYFpv31i0UUL6Ld/kfTM9Rq4aPnAWq4Bss6NGVw6dAiIdJVHjxiwKMEz3aEO2ftb4UwD2b0zeLI4KeFg6k/Qn6zHLxFPDRl4DNYbol7BfGtar5+xakgz8Xlu6jXr4QbOqUIvrAecRfMkUHC/ZMWXk+vTRVNmP4o3/Iiyl6h0AAP//O4NsEQ==" +} diff --git a/x-pack/metricbeat/module/openai/usage/_meta/data.json b/x-pack/metricbeat/module/openai/usage/_meta/data.json new file mode 100644 index 000000000000..91e18aa44afe --- /dev/null +++ b/x-pack/metricbeat/module/openai/usage/_meta/data.json @@ -0,0 +1,19 @@ +{ + "@timestamp":"2016-05-23T08:05:34.853Z", + "beat":{ + "hostname":"beathost", + "name":"beathost" + }, + "metricset":{ + "host":"localhost", + "module":"openai", + "name":"usage", + "rtt":44269 + }, + "openai":{ + "usage":{ + "example": "usage" + } + }, + "type":"metricsets" +} diff --git a/x-pack/metricbeat/module/openai/usage/_meta/docs.asciidoc b/x-pack/metricbeat/module/openai/usage/_meta/docs.asciidoc new file mode 100644 index 000000000000..dc88baf82a09 --- /dev/null +++ b/x-pack/metricbeat/module/openai/usage/_meta/docs.asciidoc @@ -0,0 +1 @@ +This is the usage metricset of the module openai. diff --git a/x-pack/metricbeat/module/openai/usage/_meta/fields.yml b/x-pack/metricbeat/module/openai/usage/_meta/fields.yml new file mode 100644 index 000000000000..7de86e521631 --- /dev/null +++ b/x-pack/metricbeat/module/openai/usage/_meta/fields.yml @@ -0,0 +1,10 @@ +- name: usage + type: group + release: beta + description: > + usage + fields: + - name: example + type: keyword + description: > + Example field diff --git a/x-pack/metricbeat/module/openai/usage/client.go b/x-pack/metricbeat/module/openai/usage/client.go new file mode 100644 index 000000000000..4fdbcc0b7753 --- /dev/null +++ b/x-pack/metricbeat/module/openai/usage/client.go @@ -0,0 +1,44 @@ +package usage + +import ( + "context" + "net/http" + "time" + + "github.com/elastic/elastic-agent-libs/logp" + "golang.org/x/time/rate" +) + +// RLHTTPClient implements a rate-limited HTTP client that wraps the standard http.Client +// with a rate limiter to control API request frequency. +type RLHTTPClient struct { + ctx context.Context + client *http.Client + logger *logp.Logger + Ratelimiter *rate.Limiter +} + +// Do executes an HTTP request while respecting rate limits. +// It waits for rate limit token before proceeding with the request. +// Returns the HTTP response and any error encountered. +func (c *RLHTTPClient) Do(req *http.Request) (*http.Response, error) { + c.logger.Warn("Waiting for rate limit token") + err := c.Ratelimiter.Wait(context.TODO()) + if err != nil { + return nil, err + } + c.logger.Warn("Rate limit token acquired") + return c.client.Do(req) +} + +// newClient creates a new rate-limited HTTP client with specified rate limiter and timeout. +func newClient(ctx context.Context, logger *logp.Logger, rl *rate.Limiter, timeout time.Duration) *RLHTTPClient { + var client = http.DefaultClient + client.Timeout = timeout + return &RLHTTPClient{ + ctx: ctx, + client: client, + logger: logger, + Ratelimiter: rl, + } +} diff --git a/x-pack/metricbeat/module/openai/usage/config.go b/x-pack/metricbeat/module/openai/usage/config.go new file mode 100644 index 000000000000..2ca63bf0c10c --- /dev/null +++ b/x-pack/metricbeat/module/openai/usage/config.go @@ -0,0 +1,76 @@ +package usage + +import ( + "fmt" + "time" +) + +type Config struct { + APIKeys []apiKeyConfig `config:"api_keys" validate:"required"` + APIURL string `config:"api_url"` + Headers []string `config:"headers"` + RateLimit *rateLimitConfig `config:"rate_limit"` + Timeout time.Duration `config:"timeout"` + Collection collectionConfig `config:"collection"` +} + +type rateLimitConfig struct { + Limit *int `config:"limit"` + Burst *int `config:"burst"` +} + +type apiKeyConfig struct { + Key string `config:"key"` +} + +type collectionConfig struct { + LookbackDays int `config:"lookback_days"` +} + +func defaultConfig() Config { + return Config{ + APIURL: "https://api.openai.com/v1/usage", + Timeout: 30 * time.Second, + RateLimit: &rateLimitConfig{ + Limit: ptr(60), + Burst: ptr(5), + }, + Collection: collectionConfig{ + LookbackDays: 5, // 5 days + }, + } +} + +func (c *Config) Validate() error { + switch { + case len(c.APIKeys) == 0: + return fmt.Errorf("at least one API key must be configured") + + case c.APIURL == "": + return fmt.Errorf("api_url cannot be empty") + + case c.RateLimit == nil: + return fmt.Errorf("rate_limit must be configured") + + case c.RateLimit.Limit == nil: + return fmt.Errorf("rate_limit.limit must be configured") + + case c.RateLimit.Burst == nil: + return fmt.Errorf("rate_limit.burst must be configured") + + case c.Timeout <= 0: + return fmt.Errorf("timeout must be greater than 0") + + case c.Collection.LookbackDays <= 0: + return fmt.Errorf("lookback_days must be greater than 0") + } + + // API keys validation in a separate loop since it needs iteration + for i, apiKey := range c.APIKeys { + if apiKey.Key == "" { + return fmt.Errorf("API key at position %d cannot be empty", i) + } + } + + return nil +} diff --git a/x-pack/metricbeat/module/openai/usage/helper.go b/x-pack/metricbeat/module/openai/usage/helper.go new file mode 100644 index 000000000000..08c7046c6c8e --- /dev/null +++ b/x-pack/metricbeat/module/openai/usage/helper.go @@ -0,0 +1,19 @@ +package usage + +import "strings" + +func ptr[T any](value T) *T { + return &value +} + +func processHeaders(headers []string) map[string]string { + headersMap := make(map[string]string, len(headers)) + for _, header := range headers { + parts := strings.Split(header, ":") + if len(parts) != 2 { + continue + } + headersMap[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1]) + } + return headersMap +} diff --git a/x-pack/metricbeat/module/openai/usage/schema.go b/x-pack/metricbeat/module/openai/usage/schema.go new file mode 100644 index 000000000000..5b44ef687d8c --- /dev/null +++ b/x-pack/metricbeat/module/openai/usage/schema.go @@ -0,0 +1,82 @@ +package usage + +type UsageResponse struct { + Object string `json:"object"` + Data []UsageData `json:"data"` + FtData []interface{} `json:"ft_data"` + DalleApiData []DalleData `json:"dalle_api_data"` + WhisperApiData []WhisperData `json:"whisper_api_data"` + TtsApiData []TtsData `json:"tts_api_data"` + AssistantCodeInterpreterData []interface{} `json:"assistant_code_interpreter_data"` + RetrievalStorageData []interface{} `json:"retrieval_storage_data"` +} + +type UsageData struct { + OrganizationID string `json:"organization_id"` + OrganizationName string `json:"organization_name"` + AggregationTimestamp int64 `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 DalleData struct { + Timestamp int64 `json:"timestamp"` + NumImages int `json:"num_images"` + NumRequests int `json:"num_requests"` + ImageSize string `json:"image_size"` + Operation string `json:"operation"` + UserID *string `json:"user_id"` + OrganizationID string `json:"organization_id"` + ApiKeyID *string `json:"api_key_id"` + ApiKeyName *string `json:"api_key_name"` + ApiKeyRedacted *string `json:"api_key_redacted"` + ApiKeyType *string `json:"api_key_type"` + OrganizationName string `json:"organization_name"` + ModelID string `json:"model_id"` + ProjectID *string `json:"project_id"` + ProjectName *string `json:"project_name"` +} + +type WhisperData struct { + Timestamp int64 `json:"timestamp"` + ModelID string `json:"model_id"` + NumSeconds int `json:"num_seconds"` + NumRequests int `json:"num_requests"` + UserID *string `json:"user_id"` + OrganizationID string `json:"organization_id"` + ApiKeyID *string `json:"api_key_id"` + ApiKeyName *string `json:"api_key_name"` + ApiKeyRedacted *string `json:"api_key_redacted"` + ApiKeyType *string `json:"api_key_type"` + OrganizationName string `json:"organization_name"` + ProjectID *string `json:"project_id"` + ProjectName *string `json:"project_name"` +} + +type TtsData struct { + Timestamp int64 `json:"timestamp"` + ModelID string `json:"model_id"` + NumCharacters int `json:"num_characters"` + NumRequests int `json:"num_requests"` + UserID *string `json:"user_id"` + OrganizationID string `json:"organization_id"` + ApiKeyID *string `json:"api_key_id"` + ApiKeyName *string `json:"api_key_name"` + ApiKeyRedacted *string `json:"api_key_redacted"` + ApiKeyType *string `json:"api_key_type"` + OrganizationName string `json:"organization_name"` + ProjectID *string `json:"project_id"` + ProjectName *string `json:"project_name"` +} diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go new file mode 100644 index 000000000000..523cb40b719b --- /dev/null +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -0,0 +1,316 @@ +package usage + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "time" + + "github.com/elastic/beats/v7/libbeat/common/cfgwarn" + "github.com/elastic/beats/v7/metricbeat/mb" + "github.com/elastic/elastic-agent-libs/logp" + "github.com/elastic/elastic-agent-libs/mapstr" + "golang.org/x/time/rate" +) + +// init registers the MetricSet with the central registry as soon as the program +// starts. The New function will be called later to instantiate an instance of +// the MetricSet for each host is defined in the module's configuration. After the +// MetricSet has been created then Fetch will begin to be called periodically. +func init() { + mb.Registry.MustAddMetricSet("openai", "usage", New) +} + +// MetricSet holds any configuration or state information. It must implement +// the mb.MetricSet interface. And this is best achieved by embedding +// mb.BaseMetricSet because it implements all of the required mb.MetricSet +// interface methods except for Fetch. +type MetricSet struct { + mb.BaseMetricSet + logger *logp.Logger + config Config + report mb.ReporterV2 +} + +// New creates a new instance of the MetricSet. New is responsible for unpacking +// any MetricSet specific configuration options if there are any. +func New(base mb.BaseMetricSet) (mb.MetricSet, error) { + cfgwarn.Beta("The openai usage metricset is beta.") + + config := defaultConfig() + if err := base.Module().UnpackConfig(&config); err != nil { + return nil, err + } + + if err := config.Validate(); err != nil { + return nil, err + } + + return &MetricSet{ + BaseMetricSet: base, + logger: logp.NewLogger("openai.usage"), + config: config, + }, nil +} + +// Fetch method implements the data gathering and data conversion to the right +// format. It publishes the event which is then forwarded to the output. In case +// of an error set the Error field of mb.Event or simply call report.Error(). +func (m *MetricSet) Fetch(report mb.ReporterV2) error { + httpClient := newClient( + context.TODO(), + m.logger, + rate.NewLimiter( + rate.Every(time.Duration(*m.config.RateLimit.Limit)*time.Second), + *m.config.RateLimit.Burst, + ), + m.config.Timeout, + ) + + m.report = report + + endDate := time.Now().UTC() + startDate := endDate.AddDate(0, 0, -m.config.Collection.LookbackDays) + + return m.fetchDateRange(startDate, endDate, httpClient) +} + +func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLHTTPClient) error { + for d := startDate; !d.After(endDate); d = d.AddDate(0, 0, 1) { + dateStr := d.Format("2006-01-02") + for _, apiKey := range m.config.APIKeys { + if err := m.fetchSingleDay(dateStr, apiKey.Key, httpClient); err != nil { + m.logger.Errorf("Error fetching data for date %s: %v", dateStr, err) + continue + } + } + } + return nil +} + +func (m *MetricSet) fetchSingleDay(dateStr, apiKey string, httpClient *RLHTTPClient) error { + req, err := m.createRequest(dateStr, apiKey) + if err != nil { + return fmt.Errorf("error creating request: %w", err) + } + + resp, err := httpClient.Do(req) + if err != nil { + return fmt.Errorf("error making request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("error response from API: %s", resp.Status) + } + + return m.processResponse(resp, dateStr) +} + +func (m *MetricSet) createRequest(dateStr, apiKey string) (*http.Request, error) { + req, err := http.NewRequest(http.MethodGet, m.config.APIURL, nil) + if err != nil { + return nil, fmt.Errorf("error creating request: %w", err) + } + + q := req.URL.Query() + q.Add("date", dateStr) + req.URL.RawQuery = q.Encode() + + req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", apiKey)) + for key, value := range processHeaders(m.config.Headers) { + req.Header.Add(key, value) + } + + return req, nil +} + +func (m *MetricSet) processResponse(resp *http.Response, dateStr string) error { + var usageResponse UsageResponse + if err := json.NewDecoder(resp.Body).Decode(&usageResponse); err != nil { + return fmt.Errorf("error decoding response: %w", err) + } + + m.logger.Info("Fetched usage metrics for date:", dateStr) + + events := make([]mb.Event, 0, len(usageResponse.Data)) + + m.processUsageData(events, usageResponse.Data) + m.processDalleData(events, usageResponse.DalleApiData) + m.processWhisperData(events, usageResponse.WhisperApiData) + m.processTTSData(events, usageResponse.TtsApiData) + + // Process additional data. + // + // NOTE(shmsr): During testing, could not get the usage data for the following + // and found no documentation, example responses, etc. That's why let's store them + // as it is so that we can use processors later on to process them as needed. + m.processFTData(events, usageResponse.FtData) + m.processAssistantCodeInterpreterData(events, usageResponse.AssistantCodeInterpreterData) + m.processRetrievalStorageData(events, usageResponse.RetrievalStorageData) + + return nil +} + +func (m *MetricSet) processUsageData(events []mb.Event, data []UsageData) { + for _, usage := range data { + event := mb.Event{ + Timestamp: time.Unix(usage.AggregationTimestamp, 0), + MetricSetFields: mapstr.M{ + "data": mapstr.M{ + "organization_id": usage.OrganizationID, + "organization_name": usage.OrganizationName, + "n_requests": usage.NRequests, + "operation": usage.Operation, + "snapshot_id": usage.SnapshotID, + "n_context_tokens_total": usage.NContextTokensTotal, + "n_generated_tokens_total": usage.NGeneratedTokensTotal, + "email": usage.Email, + "api_key_id": usage.ApiKeyID, + "api_key_name": usage.ApiKeyName, + "api_key_redacted": usage.ApiKeyRedacted, + "api_key_type": usage.ApiKeyType, + "project_id": usage.ProjectID, + "project_name": usage.ProjectName, + "request_type": usage.RequestType, + "n_cached_context_tokens_total": usage.NCachedContextTokensTotal, + }, + }, + } + events = append(events, event) + } + m.processEvents(events) +} + +func (m *MetricSet) processDalleData(events []mb.Event, data []DalleData) { + for _, dalle := range data { + event := mb.Event{ + Timestamp: time.Unix(dalle.Timestamp, 0), + MetricSetFields: mapstr.M{ + "dalle": mapstr.M{ + "num_images": dalle.NumImages, + "num_requests": dalle.NumRequests, + "image_size": dalle.ImageSize, + "operation": dalle.Operation, + "user_id": dalle.UserID, + "organization_id": dalle.OrganizationID, + "api_key_id": dalle.ApiKeyID, + "api_key_name": dalle.ApiKeyName, + "api_key_redacted": dalle.ApiKeyRedacted, + "api_key_type": dalle.ApiKeyType, + "organization_name": dalle.OrganizationName, + "model_id": dalle.ModelID, + "project_id": dalle.ProjectID, + "project_name": dalle.ProjectName, + }, + }, + } + events = append(events, event) + } + m.processEvents(events) +} + +func (m *MetricSet) processWhisperData(events []mb.Event, data []WhisperData) { + for _, whisper := range data { + event := mb.Event{ + Timestamp: time.Unix(whisper.Timestamp, 0), + MetricSetFields: mapstr.M{ + "whisper": mapstr.M{ + "model_id": whisper.ModelID, + "num_seconds": whisper.NumSeconds, + "num_requests": whisper.NumRequests, + "user_id": whisper.UserID, + "organization_id": whisper.OrganizationID, + "api_key_id": whisper.ApiKeyID, + "api_key_name": whisper.ApiKeyName, + "api_key_redacted": whisper.ApiKeyRedacted, + "api_key_type": whisper.ApiKeyType, + "organization_name": whisper.OrganizationName, + "project_id": whisper.ProjectID, + "project_name": whisper.ProjectName, + }, + }, + } + events = append(events, event) + } + m.processEvents(events) +} + +func (m *MetricSet) processTTSData(events []mb.Event, data []TtsData) { + for _, tts := range data { + event := mb.Event{ + Timestamp: time.Unix(tts.Timestamp, 0), + MetricSetFields: mapstr.M{ + "tts": mapstr.M{ + "model_id": tts.ModelID, + "num_characters": tts.NumCharacters, + "num_requests": tts.NumRequests, + "user_id": tts.UserID, + "organization_id": tts.OrganizationID, + "api_key_id": tts.ApiKeyID, + "api_key_name": tts.ApiKeyName, + "api_key_redacted": tts.ApiKeyRedacted, + "api_key_type": tts.ApiKeyType, + "organization_name": tts.OrganizationName, + "project_id": tts.ProjectID, + "project_name": tts.ProjectName, + }, + }, + } + events = append(events, event) + } + + m.processEvents(events) +} + +func (m *MetricSet) processFTData(events []mb.Event, data []interface{}) { + for _, ft := range data { + event := mb.Event{ + MetricSetFields: mapstr.M{ + "ft_data": mapstr.M{ + "original": ft, + }, + }, + } + events = append(events, event) + } + m.processEvents(events) +} + +func (m *MetricSet) processAssistantCodeInterpreterData(events []mb.Event, data []interface{}) { + for _, aci := range data { + event := mb.Event{ + MetricSetFields: mapstr.M{ + "assistant_code_interpreter": mapstr.M{ + "original": aci, + }, + }, + } + events = append(events, event) + } + m.processEvents(events) +} + +func (m *MetricSet) processRetrievalStorageData(events []mb.Event, data []interface{}) { + for _, rs := range data { + event := mb.Event{ + MetricSetFields: mapstr.M{ + "retrieval_storage": mapstr.M{ + "original": rs, + }, + }, + } + events = append(events, event) + } + m.processEvents(events) +} + +func (m *MetricSet) processEvents(events []mb.Event) { + if len(events) > 0 { + for i := range events { + m.report.Event(events[i]) + } + } + clear(events) +} From 8725cbb430396bd88de03bee1767e6e4e4679a7c Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 5 Nov 2024 13:20:00 +0530 Subject: [PATCH 02/36] update module --- .../module/openai/usage/persistcache.go | 61 +++++++++++++++++++ .../metricbeat/module/openai/usage/usage.go | 60 +++++++++++++++--- .../metricbeat/modules.d/openai.yml.disabled | 31 ++++++++++ 3 files changed, 142 insertions(+), 10 deletions(-) create mode 100644 x-pack/metricbeat/module/openai/usage/persistcache.go create mode 100644 x-pack/metricbeat/modules.d/openai.yml.disabled diff --git a/x-pack/metricbeat/module/openai/usage/persistcache.go b/x-pack/metricbeat/module/openai/usage/persistcache.go new file mode 100644 index 000000000000..8da3f9071c35 --- /dev/null +++ b/x-pack/metricbeat/module/openai/usage/persistcache.go @@ -0,0 +1,61 @@ +package usage + +import ( + "fmt" + "os" + "path" +) + +// stateStore handles persistence of state markers using the filesystem +type stateStore struct { + Dir string // Base directory for storing state files +} + +// newStateStore creates a new state store instance at the specified path +func newStateStore(path string) (*stateStore, error) { + if err := os.MkdirAll(path, 0o755); err != nil { + return nil, fmt.Errorf("creating state directory: %w", err) + } + return &stateStore{ + Dir: path, + }, nil +} + +// getStatePath builds the full file path for a given state key +func (s *stateStore) getStatePath(name string) string { + return path.Join(s.Dir, name) +} + +// Put creates a state marker file for the given key +func (s *stateStore) Put(key string) error { + filePath := s.getStatePath(key) + f, err := os.Create(filePath) + if err != nil { + return fmt.Errorf("creating state file: %w", err) + } + return f.Close() +} + +// Has checks if a state exists for the given key +func (s *stateStore) Has(key string) bool { + filePath := s.getStatePath(key) + _, err := os.Stat(filePath) + return err == nil +} + +// Remove deletes the state marker file for the given key +func (s *stateStore) Remove(key string) error { + filePath := s.getStatePath(key) + if err := os.Remove(filePath); err != nil && !os.IsNotExist(err) { + return fmt.Errorf("removing state file: %w", err) + } + return nil +} + +// Clear removes all state markers by deleting and recreating the state directory +func (s *stateStore) Clear() error { + if err := os.RemoveAll(s.Dir); err != nil { + return fmt.Errorf("clearing state directory: %w", err) + } + return os.MkdirAll(s.Dir, 0o755) +} diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index 523cb40b719b..61f5b30ae9cc 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -2,6 +2,8 @@ package usage import ( "context" + "crypto/sha256" + "encoding/hex" "encoding/json" "fmt" "net/http" @@ -11,6 +13,7 @@ import ( "github.com/elastic/beats/v7/metricbeat/mb" "github.com/elastic/elastic-agent-libs/logp" "github.com/elastic/elastic-agent-libs/mapstr" + "github.com/elastic/elastic-agent-libs/paths" "golang.org/x/time/rate" ) @@ -28,9 +31,10 @@ func init() { // interface methods except for Fetch. type MetricSet struct { mb.BaseMetricSet - logger *logp.Logger - config Config - report mb.ReporterV2 + logger *logp.Logger + config Config + report mb.ReporterV2 + stateStore *stateStore } // New creates a new instance of the MetricSet. New is responsible for unpacking @@ -47,10 +51,16 @@ func New(base mb.BaseMetricSet) (mb.MetricSet, error) { return nil, err } + st, err := newStateStore(paths.Resolve(paths.Data, base.Name())) + if err != nil { + return nil, fmt.Errorf("creating state store: %w", err) + } + return &MetricSet{ BaseMetricSet: base, logger: logp.NewLogger("openai.usage"), config: config, + stateStore: st, }, nil } @@ -77,14 +87,44 @@ func (m *MetricSet) Fetch(report mb.ReporterV2) error { } func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLHTTPClient) error { - for d := startDate; !d.After(endDate); d = d.AddDate(0, 0, 1) { - dateStr := d.Format("2006-01-02") - for _, apiKey := range m.config.APIKeys { + for _, apiKey := range m.config.APIKeys { + // SHA-256 produces a fixed-length (64 characters) hexadecimal string + // that is safe for filenames across all major platforms. Hex encoding + // ensures that the hash is safe for use in file paths as it uses only + // alphanumeric characters. + // + // Also, SHA-256 is a strong cryptographic hash function that is + // deterministic, meaning that the same input will always produce + // the same output and it is an one-way function, meaning that it is + // computationally infeasible to reverse the hash to obtain the + // original. + hasher := sha256.New() + hasher.Write([]byte(apiKey.Key)) + hashedKey := hex.EncodeToString(hasher.Sum(nil)) + stateKey := "state_" + hashedKey + + // If state exists, only fetch current day + if m.stateStore.Has(stateKey) { + currentDay := endDate.Format("2006-01-02") + if err := m.fetchSingleDay(currentDay, apiKey.Key, httpClient); err != nil { + m.logger.Errorf("Error fetching data for date %s: %v", currentDay, err) + } + continue + } + + // First run for this API key - fetch historical data + for d := startDate; !d.After(endDate); d = d.AddDate(0, 0, 1) { + dateStr := d.Format("2006-01-02") if err := m.fetchSingleDay(dateStr, apiKey.Key, httpClient); err != nil { m.logger.Errorf("Error fetching data for date %s: %v", dateStr, err) continue } } + + // Mark this API key as processed + if err := m.stateStore.Put(stateKey); err != nil { + m.logger.Errorf("Error storing state for API key: %v", err) + } } return nil } @@ -156,11 +196,11 @@ func (m *MetricSet) processResponse(resp *http.Response, dateStr string) error { func (m *MetricSet) processUsageData(events []mb.Event, data []UsageData) { for _, usage := range data { event := mb.Event{ - Timestamp: time.Unix(usage.AggregationTimestamp, 0), MetricSetFields: mapstr.M{ "data": mapstr.M{ "organization_id": usage.OrganizationID, "organization_name": usage.OrganizationName, + "aggregation_timestamp": usage.AggregationTimestamp, "n_requests": usage.NRequests, "operation": usage.Operation, "snapshot_id": usage.SnapshotID, @@ -186,9 +226,9 @@ func (m *MetricSet) processUsageData(events []mb.Event, data []UsageData) { func (m *MetricSet) processDalleData(events []mb.Event, data []DalleData) { for _, dalle := range data { event := mb.Event{ - Timestamp: time.Unix(dalle.Timestamp, 0), MetricSetFields: mapstr.M{ "dalle": mapstr.M{ + "timestamp": dalle.Timestamp, "num_images": dalle.NumImages, "num_requests": dalle.NumRequests, "image_size": dalle.ImageSize, @@ -214,9 +254,9 @@ func (m *MetricSet) processDalleData(events []mb.Event, data []DalleData) { func (m *MetricSet) processWhisperData(events []mb.Event, data []WhisperData) { for _, whisper := range data { event := mb.Event{ - Timestamp: time.Unix(whisper.Timestamp, 0), MetricSetFields: mapstr.M{ "whisper": mapstr.M{ + "timestamp": whisper.Timestamp, "model_id": whisper.ModelID, "num_seconds": whisper.NumSeconds, "num_requests": whisper.NumRequests, @@ -240,9 +280,9 @@ func (m *MetricSet) processWhisperData(events []mb.Event, data []WhisperData) { func (m *MetricSet) processTTSData(events []mb.Event, data []TtsData) { for _, tts := range data { event := mb.Event{ - Timestamp: time.Unix(tts.Timestamp, 0), MetricSetFields: mapstr.M{ "tts": mapstr.M{ + "timestamp": tts.Timestamp, "model_id": tts.ModelID, "num_characters": tts.NumCharacters, "num_requests": tts.NumRequests, diff --git a/x-pack/metricbeat/modules.d/openai.yml.disabled b/x-pack/metricbeat/modules.d/openai.yml.disabled new file mode 100644 index 000000000000..41d604c3ea9c --- /dev/null +++ b/x-pack/metricbeat/modules.d/openai.yml.disabled @@ -0,0 +1,31 @@ +# Module: openai +# Docs: https://www.elastic.co/guide/en/beats/metricbeat/main/metricbeat-module-openai.html + +- module: openai + metricsets: ["usage"] + enabled: false + period: 1h + + # # Project API Keys - Multiple API keys can be specified for different projects + # api_keys: + # - key: "api_key1" + # - key: "api_key2" + + # # API Configuration + # ## Base URL for the OpenAI usage API endpoint + # api_url: "https://api.openai.com/v1/usage" + # ## Custom headers to be included in API requests + # headers: + # - "k1: v1" + # - "k2: v2" + ## Rate Limiting Configuration + # rate_limit: + # limit: 60 # requests per second + # burst: 5 # burst size + # ## Request timeout duration + # timeout: 30s + + # # Data Collection Configuration + # collection: + # ## Number of days to look back when collecting usage data + # lookback_days: 30 \ No newline at end of file From 03995f656a59f88db9102cf7c56160188b83b751 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 5 Nov 2024 17:35:13 +0530 Subject: [PATCH 03/36] update module --- metricbeat/docs/fields.asciidoc | 31 +++++++++ metricbeat/docs/modules/openai.asciidoc | 68 +++++++++++++++++++ metricbeat/docs/modules/openai/usage.asciidoc | 29 ++++++++ metricbeat/docs/modules_list.asciidoc | 3 + .../metricbeat/module/openai/usage/client.go | 3 +- .../module/openai/usage/persistcache.go | 16 ++++- .../metricbeat/module/openai/usage/usage.go | 6 +- 7 files changed, 152 insertions(+), 4 deletions(-) create mode 100644 metricbeat/docs/modules/openai.asciidoc create mode 100644 metricbeat/docs/modules/openai/usage.asciidoc diff --git a/metricbeat/docs/fields.asciidoc b/metricbeat/docs/fields.asciidoc index 09fd5e532450..018c7017dd22 100644 --- a/metricbeat/docs/fields.asciidoc +++ b/metricbeat/docs/fields.asciidoc @@ -67,6 +67,7 @@ grouped in the following categories: * <> * <> * <> +* <> * <> * <> * <> @@ -56654,6 +56655,36 @@ type: long -- +[[exported-fields-openai]] +== openai fields + +openai module + + + +[float] +=== openai + + + + +[float] +=== usage + +usage + + + +*`openai.usage.example`*:: ++ +-- +Example field + + +type: keyword + +-- + [[exported-fields-openmetrics]] == Openmetrics fields diff --git a/metricbeat/docs/modules/openai.asciidoc b/metricbeat/docs/modules/openai.asciidoc new file mode 100644 index 000000000000..f961b1e09fbf --- /dev/null +++ b/metricbeat/docs/modules/openai.asciidoc @@ -0,0 +1,68 @@ +//// +This file is generated! See scripts/mage/docs_collector.go +//// + +:modulename: openai +:edit_url: https://github.com/elastic/beats/edit/main/x-pack/metricbeat/module/openai/_meta/docs.asciidoc + + +[[metricbeat-module-openai]] +[role="xpack"] +== openai module + +beta[] + +This is the openai module. + + + +:edit_url: + +[float] +=== Example configuration + +The openai module supports the standard configuration options that are described +in <>. Here is an example configuration: + +[source,yaml] +---- +metricbeat.modules: +- module: openai + metricsets: ["usage"] + enabled: false + period: 1h + + # # Project API Keys - Multiple API keys can be specified for different projects + # api_keys: + # - key: "api_key1" + # - key: "api_key2" + + # # API Configuration + # ## Base URL for the OpenAI usage API endpoint + # api_url: "https://api.openai.com/v1/usage" + # ## Custom headers to be included in API requests + # headers: + # - "k1: v1" + # - "k2: v2" + ## Rate Limiting Configuration + # rate_limit: + # limit: 60 # requests per second + # burst: 5 # burst size + # ## Request timeout duration + # timeout: 30s + + # # Data Collection Configuration + # collection: + # ## Number of days to look back when collecting usage data + # lookback_days: 30---- + +[float] +=== Metricsets + +The following metricsets are available: + +* <> + +include::openai/usage.asciidoc[] + +:edit_url!: diff --git a/metricbeat/docs/modules/openai/usage.asciidoc b/metricbeat/docs/modules/openai/usage.asciidoc new file mode 100644 index 000000000000..69d0ba313d9c --- /dev/null +++ b/metricbeat/docs/modules/openai/usage.asciidoc @@ -0,0 +1,29 @@ +//// +This file is generated! See scripts/mage/docs_collector.go +//// +:edit_url: https://github.com/elastic/beats/edit/main/x-pack/metricbeat/module/openai/usage/_meta/docs.asciidoc + + +[[metricbeat-metricset-openai-usage]] +[role="xpack"] +=== openai usage metricset + +beta[] + +include::../../../../x-pack/metricbeat/module/openai/usage/_meta/docs.asciidoc[] + + +:edit_url: + +==== Fields + +For a description of each field in the metricset, see the +<> section. + +Here is an example document generated by this metricset: + +[source,json] +---- +include::../../../../x-pack/metricbeat/module/openai/usage/_meta/data.json[] +---- +:edit_url!: \ No newline at end of file diff --git a/metricbeat/docs/modules_list.asciidoc b/metricbeat/docs/modules_list.asciidoc index f68dc8e1e65a..3b66228a75f1 100644 --- a/metricbeat/docs/modules_list.asciidoc +++ b/metricbeat/docs/modules_list.asciidoc @@ -240,6 +240,8 @@ This file is generated! See scripts/mage/docs_collector.go |<> |<> |image:./images/icon-yes.png[Prebuilt dashboards are available] | .1+| .1+| |<> +|<> beta[] |image:./images/icon-no.png[No prebuilt dashboards] | +.1+| .1+| |<> beta[] |<> beta[] |image:./images/icon-no.png[No prebuilt dashboards] | .1+| .1+| |<> beta[] |<> |image:./images/icon-yes.png[Prebuilt dashboards are available] | @@ -381,6 +383,7 @@ include::modules/munin.asciidoc[] include::modules/mysql.asciidoc[] include::modules/nats.asciidoc[] include::modules/nginx.asciidoc[] +include::modules/openai.asciidoc[] include::modules/openmetrics.asciidoc[] include::modules/oracle.asciidoc[] include::modules/panw.asciidoc[] diff --git a/x-pack/metricbeat/module/openai/usage/client.go b/x-pack/metricbeat/module/openai/usage/client.go index 4fdbcc0b7753..b5d1b3e60785 100644 --- a/x-pack/metricbeat/module/openai/usage/client.go +++ b/x-pack/metricbeat/module/openai/usage/client.go @@ -5,8 +5,9 @@ import ( "net/http" "time" - "github.com/elastic/elastic-agent-libs/logp" "golang.org/x/time/rate" + + "github.com/elastic/elastic-agent-libs/logp" ) // RLHTTPClient implements a rate-limited HTTP client that wraps the standard http.Client diff --git a/x-pack/metricbeat/module/openai/usage/persistcache.go b/x-pack/metricbeat/module/openai/usage/persistcache.go index 8da3f9071c35..34bb18a6b53f 100644 --- a/x-pack/metricbeat/module/openai/usage/persistcache.go +++ b/x-pack/metricbeat/module/openai/usage/persistcache.go @@ -4,11 +4,13 @@ import ( "fmt" "os" "path" + "sync" ) // stateStore handles persistence of state markers using the filesystem type stateStore struct { - Dir string // Base directory for storing state files + Dir string // Base directory for storing state files + sync.RWMutex // Protects access to the state store } // newStateStore creates a new state store instance at the specified path @@ -28,6 +30,9 @@ func (s *stateStore) getStatePath(name string) string { // Put creates a state marker file for the given key func (s *stateStore) Put(key string) error { + s.Lock() + defer s.Unlock() + filePath := s.getStatePath(key) f, err := os.Create(filePath) if err != nil { @@ -38,6 +43,9 @@ func (s *stateStore) Put(key string) error { // Has checks if a state exists for the given key func (s *stateStore) Has(key string) bool { + s.RLock() + defer s.RUnlock() + filePath := s.getStatePath(key) _, err := os.Stat(filePath) return err == nil @@ -45,6 +53,9 @@ func (s *stateStore) Has(key string) bool { // Remove deletes the state marker file for the given key func (s *stateStore) Remove(key string) error { + s.Lock() + defer s.Unlock() + filePath := s.getStatePath(key) if err := os.Remove(filePath); err != nil && !os.IsNotExist(err) { return fmt.Errorf("removing state file: %w", err) @@ -54,6 +65,9 @@ func (s *stateStore) Remove(key string) error { // Clear removes all state markers by deleting and recreating the state directory func (s *stateStore) Clear() error { + s.Lock() + defer s.Unlock() + if err := os.RemoveAll(s.Dir); err != nil { return fmt.Errorf("clearing state directory: %w", err) } diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index 61f5b30ae9cc..0988e6353677 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -7,14 +7,16 @@ import ( "encoding/json" "fmt" "net/http" + "path" "time" + "golang.org/x/time/rate" + "github.com/elastic/beats/v7/libbeat/common/cfgwarn" "github.com/elastic/beats/v7/metricbeat/mb" "github.com/elastic/elastic-agent-libs/logp" "github.com/elastic/elastic-agent-libs/mapstr" "github.com/elastic/elastic-agent-libs/paths" - "golang.org/x/time/rate" ) // init registers the MetricSet with the central registry as soon as the program @@ -51,7 +53,7 @@ func New(base mb.BaseMetricSet) (mb.MetricSet, error) { return nil, err } - st, err := newStateStore(paths.Resolve(paths.Data, base.Name())) + st, err := newStateStore(paths.Resolve(paths.Data, path.Join(base.Module().Name(), base.Name()))) if err != nil { return nil, fmt.Errorf("creating state store: %w", err) } From 406c1a4ff61f9ac2509e706a71b440d8645b7894 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 5 Nov 2024 20:43:10 +0530 Subject: [PATCH 04/36] update module --- metricbeat/docs/fields.asciidoc | 594 +++++++++++++++++- x-pack/metricbeat/module/openai/fields.go | 2 +- .../module/openai/usage/_meta/data.json | 47 +- .../module/openai/usage/_meta/fields.yml | 231 ++++++- .../metricbeat/module/openai/usage/client.go | 4 + .../metricbeat/module/openai/usage/config.go | 8 +- .../metricbeat/module/openai/usage/helper.go | 4 + .../module/openai/usage/persistcache.go | 4 + .../metricbeat/module/openai/usage/schema.go | 4 + .../metricbeat/module/openai/usage/usage.go | 13 +- .../openai/usage/usage_integration_test.go | 175 ++++++ 11 files changed, 1058 insertions(+), 28 deletions(-) create mode 100644 x-pack/metricbeat/module/openai/usage/usage_integration_test.go diff --git a/metricbeat/docs/fields.asciidoc b/metricbeat/docs/fields.asciidoc index 018c7017dd22..1705b79fe707 100644 --- a/metricbeat/docs/fields.asciidoc +++ b/metricbeat/docs/fields.asciidoc @@ -56671,20 +56671,608 @@ openai module [float] === usage -usage +OpenAI API usage metrics and statistics -*`openai.usage.example`*:: +[float] +=== data + +General usage data metrics + + + +*`openai.usage.data.organization_id`*:: + -- -Example field +Organization identifier + +type: keyword + +-- + +*`openai.usage.data.organization_name`*:: ++ +-- +Organization name + +type: keyword + +-- + +*`openai.usage.data.aggregation_timestamp`*:: ++ +-- +Timestamp of data aggregation + +type: date + +-- + +*`openai.usage.data.n_requests`*:: ++ +-- +Number of requests made + +type: long + +-- + +*`openai.usage.data.operation`*:: ++ +-- +Operation type + +type: keyword + +-- + +*`openai.usage.data.snapshot_id`*:: ++ +-- +Snapshot identifier + +type: keyword + +-- + +*`openai.usage.data.n_context_tokens_total`*:: ++ +-- +Total number of context tokens used + +type: long + +-- + +*`openai.usage.data.n_generated_tokens_total`*:: ++ +-- +Total number of generated tokens + +type: long + +-- + +*`openai.usage.data.n_cached_context_tokens_total`*:: ++ +-- +Total number of cached context tokens + +type: long + +-- + +*`openai.usage.data.email`*:: ++ +-- +User email + +type: keyword + +-- + +*`openai.usage.data.api_key_id`*:: ++ +-- +API key identifier + +type: keyword + +-- + +*`openai.usage.data.api_key_name`*:: ++ +-- +API key name + +type: keyword + +-- + +*`openai.usage.data.api_key_redacted`*:: ++ +-- +Redacted API key + +type: keyword + +-- + +*`openai.usage.data.api_key_type`*:: ++ +-- +Type of API key + +type: keyword + +-- + +*`openai.usage.data.project_id`*:: ++ +-- +Project identifier + +type: keyword + +-- + +*`openai.usage.data.project_name`*:: ++ +-- +Project name + +type: keyword + +-- + +*`openai.usage.data.request_type`*:: ++ +-- +Type of request + +type: keyword + +-- + +[float] +=== dalle + +DALL-E API usage metrics + + + +*`openai.usage.dalle.timestamp`*:: ++ +-- +Timestamp of request + +type: date + +-- + +*`openai.usage.dalle.num_images`*:: ++ +-- +Number of images generated + +type: long + +-- + +*`openai.usage.dalle.num_requests`*:: ++ +-- +Number of requests + +type: long + +-- + +*`openai.usage.dalle.image_size`*:: ++ +-- +Size of generated images + +type: keyword + +-- + +*`openai.usage.dalle.operation`*:: ++ +-- +Operation type + +type: keyword + +-- + +*`openai.usage.dalle.user_id`*:: ++ +-- +User identifier + +type: keyword + +-- + +*`openai.usage.dalle.organization_id`*:: ++ +-- +Organization identifier + +type: keyword + +-- + +*`openai.usage.dalle.api_key_id`*:: ++ +-- +API key identifier + +type: keyword + +-- + +*`openai.usage.dalle.api_key_name`*:: ++ +-- +API key name + +type: keyword + +-- + +*`openai.usage.dalle.api_key_redacted`*:: ++ +-- +Redacted API key + +type: keyword + +-- + +*`openai.usage.dalle.api_key_type`*:: ++ +-- +Type of API key + +type: keyword + +-- + +*`openai.usage.dalle.organization_name`*:: ++ +-- +Organization name + +type: keyword + +-- + +*`openai.usage.dalle.model_id`*:: ++ +-- +Model identifier + +type: keyword + +-- + +*`openai.usage.dalle.project_id`*:: ++ +-- +Project identifier + +type: keyword + +-- + +*`openai.usage.dalle.project_name`*:: ++ +-- +Project name +type: keyword + +-- + +[float] +=== whisper + +Whisper API usage metrics + + + +*`openai.usage.whisper.timestamp`*:: ++ +-- +Timestamp of request + +type: date + +-- + +*`openai.usage.whisper.model_id`*:: ++ +-- +Model identifier + +type: keyword + +-- + +*`openai.usage.whisper.num_seconds`*:: ++ +-- +Number of seconds processed + +type: long + +-- + +*`openai.usage.whisper.num_requests`*:: ++ +-- +Number of requests + +type: long + +-- + +*`openai.usage.whisper.user_id`*:: ++ +-- +User identifier + +type: keyword + +-- + +*`openai.usage.whisper.organization_id`*:: ++ +-- +Organization identifier + +type: keyword + +-- + +*`openai.usage.whisper.api_key_id`*:: ++ +-- +API key identifier + +type: keyword + +-- + +*`openai.usage.whisper.api_key_name`*:: ++ +-- +API key name + +type: keyword + +-- + +*`openai.usage.whisper.api_key_redacted`*:: ++ +-- +Redacted API key + +type: keyword + +-- + +*`openai.usage.whisper.api_key_type`*:: ++ +-- +Type of API key + +type: keyword + +-- + +*`openai.usage.whisper.organization_name`*:: ++ +-- +Organization name + +type: keyword + +-- + +*`openai.usage.whisper.project_id`*:: ++ +-- +Project identifier + +type: keyword + +-- + +*`openai.usage.whisper.project_name`*:: ++ +-- +Project name + +type: keyword + +-- + +[float] +=== tts + +Text-to-Speech API usage metrics + + + +*`openai.usage.tts.timestamp`*:: ++ +-- +Timestamp of request + +type: date + +-- + +*`openai.usage.tts.model_id`*:: ++ +-- +Model identifier type: keyword -- +*`openai.usage.tts.num_characters`*:: ++ +-- +Number of characters processed + +type: long + +-- + +*`openai.usage.tts.num_requests`*:: ++ +-- +Number of requests + +type: long + +-- + +*`openai.usage.tts.user_id`*:: ++ +-- +User identifier + +type: keyword + +-- + +*`openai.usage.tts.organization_id`*:: ++ +-- +Organization identifier + +type: keyword + +-- + +*`openai.usage.tts.api_key_id`*:: ++ +-- +API key identifier + +type: keyword + +-- + +*`openai.usage.tts.api_key_name`*:: ++ +-- +API key name + +type: keyword + +-- + +*`openai.usage.tts.api_key_redacted`*:: ++ +-- +Redacted API key + +type: keyword + +-- + +*`openai.usage.tts.api_key_type`*:: ++ +-- +Type of API key + +type: keyword + +-- + +*`openai.usage.tts.organization_name`*:: ++ +-- +Organization name + +type: keyword + +-- + +*`openai.usage.tts.project_id`*:: ++ +-- +Project identifier + +type: keyword + +-- + +*`openai.usage.tts.project_name`*:: ++ +-- +Project name + +type: keyword + +-- + +[float] +=== ft_data + +Fine-tuning data metrics + + + +*`openai.usage.ft_data.original`*:: ++ +-- +Raw fine-tuning data + +type: object + +-- + +[float] +=== assistant_code_interpreter + +Assistant Code Interpreter usage metrics + + + +*`openai.usage.assistant_code_interpreter.original`*:: ++ +-- +Raw assistant code interpreter data + +type: object + +-- + +[float] +=== retrieval_storage + +Retrieval storage usage metrics + + + +*`openai.usage.retrieval_storage.original`*:: ++ +-- +Raw retrieval storage data + +type: object + +-- + [[exported-fields-openmetrics]] == Openmetrics fields diff --git a/x-pack/metricbeat/module/openai/fields.go b/x-pack/metricbeat/module/openai/fields.go index a23f01a7810d..fa62e586ec53 100644 --- a/x-pack/metricbeat/module/openai/fields.go +++ b/x-pack/metricbeat/module/openai/fields.go @@ -19,5 +19,5 @@ func init() { // AssetOpenai returns asset data. // This is the base64 encoded zlib format compressed contents of module/openai. func AssetOpenai() string { - return "eJx8jkHOgyAUhPecYuLeC7D4d/9BXsvUEFEIYFpv31i0UUL6Ld/kfTM9Rq4aPnAWq4Bss6NGVw6dAiIdJVHjxiwKMEz3aEO2ftb4UwD2b0zeLI4KeFg6k/Qn6zHLxFPDRl4DNYbol7BfGtar5+xakgz8Xlu6jXr4QbOqUIvrAecRfMkUHC/ZMWXk+vTRVNmP4o3/Iiyl6h0AAP//O4NsEQ==" + return "eJzsms9uozwUxfd5iqvu8wJZfFL0/VOlzrRqM5olcvEN8QRsxr5Mmz79yAEnQIBA4qbtFC8hnN/h+uJjRKawxs0MVIqSiQkACYpxBlf5gasJgMYYmcEZPCKxCQBHE2qRklByBn9NAKC4GhLFsxgnAEuBMTez7bkpSJZgiWAHbVKcQaRVlhZHGlSrOmWtzLAId0eb5OyoG3ejEZWP2xTl/Brmd9c5AxIkLUIDTHIwxEgYEqEpXVO3WLbJWYXb7vSIKzv+R4maxYUtq+y81X7aZKgyDzpiUrwwCwoEP/id87jGzZPSTecrTm9LciA4ShJLgbof3h7yaqBR0KFZFGmMcjKJBA2xpD4Lezxn1OStwl44FVDLfE5KjFYfMtD4M0ND9anbw2Mlo2Pwr1nyiNqSnRwkjLffvkpRNxsbWnEntL2slWckS81KkYcmeyik+jSYDEIlCZ8pILVGaQJSxOKzKr2wCiB39S4AkAMgM3h4B3s70fbJJeSvZ2iHKCx1FYeFK+QXqNGWUytVqy9MmGjn9+ySbwZ1i9JuCUhFsMaNh5a0EbHGTZ+OdFAfq53Ddi90BVAjZyE1NOdA6H2h4+hHwY3rwkDoYpOibaRjzFSrHxj6WGXucqU+U+qgPqbUYTuntFjivVa20Nz9vGUPE8d14smbmH/mNzfTfw+3VwO3MK8Q3vVa1JEySwKRsAh9ZXYutl+5O8mvtF9oZW7NBUa8nN1rD+IFqwnVUsW32qVkBrWHtWObPUO3v5fffY/5dwr0XeffG75QJYpj7KGXvlidIbH7Z2V9Y+4+rYRJD1ydnLzfc7mPGL0X7zKbuAZDJbmvwC3UbCuFaDpfF98g7ccIPI08RmAH+JNE4CdKJDpYQk5OowU+05TU9CFFDFdjKPUMpXDFtH0Kta9o2AuO0TRG0xhNYzQ1Mt99NC0p8Pmt9T8hcUqZFDI670uriITs+K6iHu39NZzOTwTDHgP2BMua8e6qMWOEISYpCBXHQEhCnWokf2+dc0eAvxVHuN4Tzsr7yxd2VyqwpYJSqXrUWdubxF8sDgwpXf3nwlnlvXfCUAh/sKrqA//bYv4OAAD//+T56FM=" } diff --git a/x-pack/metricbeat/module/openai/usage/_meta/data.json b/x-pack/metricbeat/module/openai/usage/_meta/data.json index 91e18aa44afe..7a6355842f0e 100644 --- a/x-pack/metricbeat/module/openai/usage/_meta/data.json +++ b/x-pack/metricbeat/module/openai/usage/_meta/data.json @@ -1,19 +1,38 @@ { - "@timestamp":"2016-05-23T08:05:34.853Z", - "beat":{ - "hostname":"beathost", - "name":"beathost" + "@timestamp": "2017-10-12T08:05:34.853Z", + "event": { + "dataset": "openai.usage", + "duration": 115000, + "module": "openai" }, - "metricset":{ - "host":"localhost", - "module":"openai", - "name":"usage", - "rtt":44269 + "metricset": { + "name": "usage", + "period": 10000 }, - "openai":{ - "usage":{ - "example": "usage" + "openai": { + "usage": { + "data": { + "aggregation_timestamp": 1730696460, + "api_key_id": null, + "api_key_name": null, + "api_key_redacted": null, + "api_key_type": null, + "email": null, + "n_cached_context_tokens_total": 0, + "n_context_tokens_total": 118, + "n_generated_tokens_total": 35, + "n_requests": 1, + "operation": "completion-realtime", + "organization_id": "org-dummy", + "organization_name": "Personal", + "project_id": null, + "project_name": null, + "request_type": "", + "snapshot_id": "gpt-4o-realtime-preview-2024-10-01" + } } }, - "type":"metricsets" -} + "service": { + "type": "openai" + } +} \ No newline at end of file diff --git a/x-pack/metricbeat/module/openai/usage/_meta/fields.yml b/x-pack/metricbeat/module/openai/usage/_meta/fields.yml index 7de86e521631..203dacaced0d 100644 --- a/x-pack/metricbeat/module/openai/usage/_meta/fields.yml +++ b/x-pack/metricbeat/module/openai/usage/_meta/fields.yml @@ -2,9 +2,232 @@ type: group release: beta description: > - usage + OpenAI API usage metrics and statistics fields: - - name: example - type: keyword + - name: data + type: group description: > - Example field + General usage data metrics + fields: + - name: organization_id + type: keyword + description: Organization identifier + - name: organization_name + type: keyword + description: Organization name + - name: aggregation_timestamp + type: date + description: Timestamp of data aggregation + - name: n_requests + type: long + description: Number of requests made + - name: operation + type: keyword + description: Operation type + - name: snapshot_id + type: keyword + description: Snapshot identifier + - name: n_context_tokens_total + type: long + description: Total number of context tokens used + - name: n_generated_tokens_total + type: long + description: Total number of generated tokens + - name: n_cached_context_tokens_total + type: long + description: Total number of cached context tokens + - name: email + type: keyword + description: User email + - name: api_key_id + type: keyword + description: API key identifier + - name: api_key_name + type: keyword + description: API key name + - name: api_key_redacted + type: keyword + description: Redacted API key + - name: api_key_type + type: keyword + description: Type of API key + - name: project_id + type: keyword + description: Project identifier + - name: project_name + type: keyword + description: Project name + - name: request_type + type: keyword + description: Type of request + + - name: dalle + type: group + description: > + DALL-E API usage metrics + fields: + - name: timestamp + type: date + description: Timestamp of request + - name: num_images + type: long + description: Number of images generated + - name: num_requests + type: long + description: Number of requests + - name: image_size + type: keyword + description: Size of generated images + - name: operation + type: keyword + description: Operation type + - name: user_id + type: keyword + description: User identifier + - name: organization_id + type: keyword + description: Organization identifier + - name: api_key_id + type: keyword + description: API key identifier + - name: api_key_name + type: keyword + description: API key name + - name: api_key_redacted + type: keyword + description: Redacted API key + - name: api_key_type + type: keyword + description: Type of API key + - name: organization_name + type: keyword + description: Organization name + - name: model_id + type: keyword + description: Model identifier + - name: project_id + type: keyword + description: Project identifier + - name: project_name + type: keyword + description: Project name + + - name: whisper + type: group + description: > + Whisper API usage metrics + fields: + - name: timestamp + type: date + description: Timestamp of request + - name: model_id + type: keyword + description: Model identifier + - name: num_seconds + type: long + description: Number of seconds processed + - name: num_requests + type: long + description: Number of requests + - name: user_id + type: keyword + description: User identifier + - name: organization_id + type: keyword + description: Organization identifier + - name: api_key_id + type: keyword + description: API key identifier + - name: api_key_name + type: keyword + description: API key name + - name: api_key_redacted + type: keyword + description: Redacted API key + - name: api_key_type + type: keyword + description: Type of API key + - name: organization_name + type: keyword + description: Organization name + - name: project_id + type: keyword + description: Project identifier + - name: project_name + type: keyword + description: Project name + + - name: tts + type: group + description: > + Text-to-Speech API usage metrics + fields: + - name: timestamp + type: date + description: Timestamp of request + - name: model_id + type: keyword + description: Model identifier + - name: num_characters + type: long + description: Number of characters processed + - name: num_requests + type: long + description: Number of requests + - name: user_id + type: keyword + description: User identifier + - name: organization_id + type: keyword + description: Organization identifier + - name: api_key_id + type: keyword + description: API key identifier + - name: api_key_name + type: keyword + description: API key name + - name: api_key_redacted + type: keyword + description: Redacted API key + - name: api_key_type + type: keyword + description: Type of API key + - name: organization_name + type: keyword + description: Organization name + - name: project_id + type: keyword + description: Project identifier + - name: project_name + type: keyword + description: Project name + + - name: ft_data + type: group + description: > + Fine-tuning data metrics + fields: + - name: original + type: object + object_type: keyword + description: Raw fine-tuning data + + - name: assistant_code_interpreter + type: group + description: > + Assistant Code Interpreter usage metrics + fields: + - name: original + type: object + object_type: keyword + description: Raw assistant code interpreter data + + - name: retrieval_storage + type: group + description: > + Retrieval storage usage metrics + fields: + - name: original + type: object + object_type: keyword + description: Raw retrieval storage data diff --git a/x-pack/metricbeat/module/openai/usage/client.go b/x-pack/metricbeat/module/openai/usage/client.go index b5d1b3e60785..c8f24aa8a50c 100644 --- a/x-pack/metricbeat/module/openai/usage/client.go +++ b/x-pack/metricbeat/module/openai/usage/client.go @@ -1,3 +1,7 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License; +// you may not use this file except in compliance with the Elastic License. + package usage import ( diff --git a/x-pack/metricbeat/module/openai/usage/config.go b/x-pack/metricbeat/module/openai/usage/config.go index 2ca63bf0c10c..5908bdd57f92 100644 --- a/x-pack/metricbeat/module/openai/usage/config.go +++ b/x-pack/metricbeat/module/openai/usage/config.go @@ -1,3 +1,7 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License; +// you may not use this file except in compliance with the Elastic License. + package usage import ( @@ -61,8 +65,8 @@ func (c *Config) Validate() error { case c.Timeout <= 0: return fmt.Errorf("timeout must be greater than 0") - case c.Collection.LookbackDays <= 0: - return fmt.Errorf("lookback_days must be greater than 0") + case c.Collection.LookbackDays < 0: + return fmt.Errorf("lookback_days must be >= 0") } // API keys validation in a separate loop since it needs iteration diff --git a/x-pack/metricbeat/module/openai/usage/helper.go b/x-pack/metricbeat/module/openai/usage/helper.go index 08c7046c6c8e..f2d3adadd01e 100644 --- a/x-pack/metricbeat/module/openai/usage/helper.go +++ b/x-pack/metricbeat/module/openai/usage/helper.go @@ -1,3 +1,7 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License; +// you may not use this file except in compliance with the Elastic License. + package usage import "strings" diff --git a/x-pack/metricbeat/module/openai/usage/persistcache.go b/x-pack/metricbeat/module/openai/usage/persistcache.go index 34bb18a6b53f..9b7d725461a5 100644 --- a/x-pack/metricbeat/module/openai/usage/persistcache.go +++ b/x-pack/metricbeat/module/openai/usage/persistcache.go @@ -1,3 +1,7 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License; +// you may not use this file except in compliance with the Elastic License. + package usage import ( diff --git a/x-pack/metricbeat/module/openai/usage/schema.go b/x-pack/metricbeat/module/openai/usage/schema.go index 5b44ef687d8c..668d5b03370c 100644 --- a/x-pack/metricbeat/module/openai/usage/schema.go +++ b/x-pack/metricbeat/module/openai/usage/schema.go @@ -1,3 +1,7 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License; +// you may not use this file except in compliance with the Elastic License. + package usage type UsageResponse struct { diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index 0988e6353677..92319af0c557 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -1,3 +1,7 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License; +// you may not use this file except in compliance with the Elastic License. + package usage import ( @@ -196,13 +200,14 @@ func (m *MetricSet) processResponse(resp *http.Response, dateStr string) error { } func (m *MetricSet) processUsageData(events []mb.Event, data []UsageData) { + for _, usage := range data { event := mb.Event{ MetricSetFields: mapstr.M{ "data": mapstr.M{ "organization_id": usage.OrganizationID, "organization_name": usage.OrganizationName, - "aggregation_timestamp": usage.AggregationTimestamp, + "aggregation_timestamp": time.Unix(usage.AggregationTimestamp, 0), "n_requests": usage.NRequests, "operation": usage.Operation, "snapshot_id": usage.SnapshotID, @@ -230,7 +235,7 @@ func (m *MetricSet) processDalleData(events []mb.Event, data []DalleData) { event := mb.Event{ MetricSetFields: mapstr.M{ "dalle": mapstr.M{ - "timestamp": dalle.Timestamp, + "timestamp": time.Unix(dalle.Timestamp, 0), "num_images": dalle.NumImages, "num_requests": dalle.NumRequests, "image_size": dalle.ImageSize, @@ -258,7 +263,7 @@ func (m *MetricSet) processWhisperData(events []mb.Event, data []WhisperData) { event := mb.Event{ MetricSetFields: mapstr.M{ "whisper": mapstr.M{ - "timestamp": whisper.Timestamp, + "timestamp": time.Unix(whisper.Timestamp, 0), "model_id": whisper.ModelID, "num_seconds": whisper.NumSeconds, "num_requests": whisper.NumRequests, @@ -284,7 +289,7 @@ func (m *MetricSet) processTTSData(events []mb.Event, data []TtsData) { event := mb.Event{ MetricSetFields: mapstr.M{ "tts": mapstr.M{ - "timestamp": tts.Timestamp, + "timestamp": time.Unix(tts.Timestamp, 0), "model_id": tts.ModelID, "num_characters": tts.NumCharacters, "num_requests": tts.NumRequests, diff --git a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go new file mode 100644 index 000000000000..e7754939845b --- /dev/null +++ b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go @@ -0,0 +1,175 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License; +// you may not use this file except in compliance with the Elastic License. + +//go:build integration + +package usage + +import ( + "fmt" + "net/http" + "net/http/httptest" + "testing" + + mbtest "github.com/elastic/beats/v7/metricbeat/mb/testing" + "github.com/stretchr/testify/assert" +) + +func TestFetch(t *testing.T) { + apiKey := "most_secure_token" + usagePath := "/usage" + server := initServer(usagePath, apiKey) + defer server.Close() + + f := mbtest.NewReportingMetricSetV2Error(t, getConfig(server.URL+"/usage", apiKey)) + + events, errs := mbtest.ReportingFetchV2Error(f) + if len(errs) > 0 { + t.Fatalf("Expected 0 error, has %d: %v", len(errs), errs) + } + + assert.NotEmpty(t, events) + +} + +func TestData(t *testing.T) { + apiKey := "most_secure_token" + usagePath := "/usage" + server := initServer(usagePath, apiKey) + defer server.Close() + + f := mbtest.NewReportingMetricSetV2Error(t, getConfig(server.URL+"/usage", apiKey)) + + err := mbtest.WriteEventsReporterV2Error(f, t, "") + if !assert.NoError(t, err) { + t.FailNow() + } +} + +func getConfig(url, apiKey string) map[string]interface{} { + return map[string]interface{}{ + "module": "openai", + "metricsets": []string{"usage"}, + "enabled": true, + "period": "1h", + "api_url": url, + "api_keys": []map[string]interface{}{ + {"key": apiKey}, + }, + "rate_limit": map[string]interface{}{ + "limit": 60, + "burst": 5, + }, + "collection": map[string]interface{}{ + "lookback_days": 1, + }, + } +} + +func initServer(endpoint string, api_key string) *httptest.Server { + data := []byte(`{ + "object": "list", + "data": [ + { + "organization_id": "org-dummy", + "organization_name": "Personal", + "aggregation_timestamp": 1730696460, + "n_requests": 1, + "operation": "completion-realtime", + "snapshot_id": "gpt-4o-realtime-preview-2024-10-01", + "n_context_tokens_total": 118, + "n_generated_tokens_total": 35, + "email": null, + "api_key_id": null, + "api_key_name": null, + "api_key_redacted": null, + "api_key_type": null, + "project_id": null, + "project_name": null, + "request_type": "", + "n_cached_context_tokens_total": 0 + }, + { + "organization_id": "org-dummy", + "organization_name": "Personal", + "aggregation_timestamp": 1730696460, + "n_requests": 1, + "operation": "completion", + "snapshot_id": "gpt-4o-2024-08-06", + "n_context_tokens_total": 31, + "n_generated_tokens_total": 12, + "email": null, + "api_key_id": null, + "api_key_name": null, + "api_key_redacted": null, + "api_key_type": null, + "project_id": null, + "project_name": null, + "request_type": "", + "n_cached_context_tokens_total": 0 + }, + { + "organization_id": "org-dummy", + "organization_name": "Personal", + "aggregation_timestamp": 1730697540, + "n_requests": 1, + "operation": "completion", + "snapshot_id": "ft:gpt-3.5-turbo-0125:personal:yay-renew:APjjyG8E:ckpt-step-84", + "n_context_tokens_total": 13, + "n_generated_tokens_total": 9, + "email": null, + "api_key_id": null, + "api_key_name": null, + "api_key_redacted": null, + "api_key_type": null, + "project_id": null, + "project_name": null, + "request_type": "", + "n_cached_context_tokens_total": 0 + } + ], + "ft_data": [], + "dalle_api_data": [], + "whisper_api_data": [ + { + "timestamp": 1730696460, + "model_id": "whisper-1", + "num_seconds": 2, + "num_requests": 1, + "user_id": null, + "organization_id": "org-dummy", + "api_key_id": null, + "api_key_name": null, + "api_key_redacted": null, + "api_key_type": null, + "organization_name": "Personal", + "project_id": null, + "project_name": null + } + ], + "tts_api_data": [], + "assistant_code_interpreter_data": [], + "retrieval_storage_data": [] +}`) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Validate Bearer token + authHeader := r.Header.Get("Authorization") + expectedToken := fmt.Sprintf("Bearer %s", api_key) + + if authHeader != expectedToken { + w.WriteHeader(http.StatusUnauthorized) + return + } + + // Validate the endpoint + if r.URL.Path == endpoint { + w.WriteHeader(http.StatusOK) + w.Write(data) + } else { + w.WriteHeader(http.StatusNotFound) + } + })) + return server +} From a3c46b7d302c3ab517042f432e38d42f4c642fcf Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 5 Nov 2024 20:51:59 +0530 Subject: [PATCH 05/36] update module --- metricbeat/docs/modules/openai.asciidoc | 3 ++- x-pack/metricbeat/metricbeat.reference.yml | 1 + x-pack/metricbeat/module/openai/_meta/config.yml | 2 +- x-pack/metricbeat/modules.d/openai.yml.disabled | 2 +- 4 files changed, 5 insertions(+), 3 deletions(-) diff --git a/metricbeat/docs/modules/openai.asciidoc b/metricbeat/docs/modules/openai.asciidoc index f961b1e09fbf..ad0d7f97712f 100644 --- a/metricbeat/docs/modules/openai.asciidoc +++ b/metricbeat/docs/modules/openai.asciidoc @@ -54,7 +54,8 @@ metricbeat.modules: # # Data Collection Configuration # collection: # ## Number of days to look back when collecting usage data - # lookback_days: 30---- + # lookback_days: 30 +---- [float] === Metricsets diff --git a/x-pack/metricbeat/metricbeat.reference.yml b/x-pack/metricbeat/metricbeat.reference.yml index 48fc2eefa302..263d3df085c1 100644 --- a/x-pack/metricbeat/metricbeat.reference.yml +++ b/x-pack/metricbeat/metricbeat.reference.yml @@ -1286,6 +1286,7 @@ metricbeat.modules: # collection: # ## Number of days to look back when collecting usage data # lookback_days: 30 + #----------------------------- Openmetrics Module ----------------------------- - module: openmetrics metricsets: ['collector'] diff --git a/x-pack/metricbeat/module/openai/_meta/config.yml b/x-pack/metricbeat/module/openai/_meta/config.yml index dadcd35b1bae..1290889640b9 100644 --- a/x-pack/metricbeat/module/openai/_meta/config.yml +++ b/x-pack/metricbeat/module/openai/_meta/config.yml @@ -25,4 +25,4 @@ # # Data Collection Configuration # collection: # ## Number of days to look back when collecting usage data - # lookback_days: 30 \ No newline at end of file + # lookback_days: 30 diff --git a/x-pack/metricbeat/modules.d/openai.yml.disabled b/x-pack/metricbeat/modules.d/openai.yml.disabled index 41d604c3ea9c..10a9fb485770 100644 --- a/x-pack/metricbeat/modules.d/openai.yml.disabled +++ b/x-pack/metricbeat/modules.d/openai.yml.disabled @@ -28,4 +28,4 @@ # # Data Collection Configuration # collection: # ## Number of days to look back when collecting usage data - # lookback_days: 30 \ No newline at end of file + # lookback_days: 30 From 72488f71cf76ce7f71fbfb1bd1276ae8fb391e45 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 5 Nov 2024 21:04:10 +0530 Subject: [PATCH 06/36] update module --- x-pack/metricbeat/module/openai/doc.go | 4 ++++ x-pack/metricbeat/module/openai/usage/usage.go | 2 +- .../metricbeat/module/openai/usage/usage_integration_test.go | 5 +++-- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/x-pack/metricbeat/module/openai/doc.go b/x-pack/metricbeat/module/openai/doc.go index b6d7a70b1215..5f2f07fb0bca 100644 --- a/x-pack/metricbeat/module/openai/doc.go +++ b/x-pack/metricbeat/module/openai/doc.go @@ -1,2 +1,6 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License; +// you may not use this file except in compliance with the Elastic License. + // Package openai is a Metricbeat module that contains MetricSets. package openai diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index 92319af0c557..4e9b8e88e5b7 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -155,7 +155,7 @@ func (m *MetricSet) fetchSingleDay(dateStr, apiKey string, httpClient *RLHTTPCli } func (m *MetricSet) createRequest(dateStr, apiKey string) (*http.Request, error) { - req, err := http.NewRequest(http.MethodGet, m.config.APIURL, nil) + req, err := http.NewRequestWithContext(context.TODO(), http.MethodGet, m.config.APIURL, nil) if err != nil { return nil, fmt.Errorf("error creating request: %w", err) } diff --git a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go index e7754939845b..a39eb8700a9e 100644 --- a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go +++ b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go @@ -12,8 +12,9 @@ import ( "net/http/httptest" "testing" - mbtest "github.com/elastic/beats/v7/metricbeat/mb/testing" "github.com/stretchr/testify/assert" + + mbtest "github.com/elastic/beats/v7/metricbeat/mb/testing" ) func TestFetch(t *testing.T) { @@ -166,7 +167,7 @@ func initServer(endpoint string, api_key string) *httptest.Server { // Validate the endpoint if r.URL.Path == endpoint { w.WriteHeader(http.StatusOK) - w.Write(data) + _ = w.Write(data) } else { w.WriteHeader(http.StatusNotFound) } From 2d4027dc85be76dd3e3f991a3ae5b24d492547ac Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 5 Nov 2024 22:18:48 +0530 Subject: [PATCH 07/36] fix bug --- x-pack/metricbeat/module/openai/usage/usage_integration_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go index a39eb8700a9e..8a0c9b70090c 100644 --- a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go +++ b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go @@ -167,7 +167,7 @@ func initServer(endpoint string, api_key string) *httptest.Server { // Validate the endpoint if r.URL.Path == endpoint { w.WriteHeader(http.StatusOK) - _ = w.Write(data) + _, _ = w.Write(data) } else { w.WriteHeader(http.StatusNotFound) } From e3f8fe26648b5adb3e81e254c172c4dc71d46827 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Wed, 6 Nov 2024 17:24:58 +0530 Subject: [PATCH 08/36] update module --- metricbeat/docs/modules/openai.asciidoc | 7 +- x-pack/metricbeat/metricbeat.reference.yml | 7 +- .../metricbeat/module/openai/_meta/config.yml | 6 ++ .../module/openai/usage/_meta/data.json | 2 +- .../metricbeat/module/openai/usage/config.go | 6 +- .../module/openai/usage/persistcache.go | 33 +++++++-- .../metricbeat/module/openai/usage/usage.go | 71 +++++++++++++------ .../openai/usage/usage_integration_test.go | 5 +- .../metricbeat/modules.d/openai.yml.disabled | 6 ++ 9 files changed, 109 insertions(+), 34 deletions(-) diff --git a/metricbeat/docs/modules/openai.asciidoc b/metricbeat/docs/modules/openai.asciidoc index ad0d7f97712f..6cb33cb60a3b 100644 --- a/metricbeat/docs/modules/openai.asciidoc +++ b/metricbeat/docs/modules/openai.asciidoc @@ -55,7 +55,12 @@ metricbeat.modules: # collection: # ## Number of days to look back when collecting usage data # lookback_days: 30 ----- + # ## Whether to collect usage data in realtime. Defaults to false as how + # # OpenAI usage data is collected will end up adding duplicate data to ES + # # and also making it harder to do analytics. Best approach is to avoid + # # realtime collection and collect only upto last day (in UTC). So, there's + # # at most 24h delay. + # realtime: false---- [float] === Metricsets diff --git a/x-pack/metricbeat/metricbeat.reference.yml b/x-pack/metricbeat/metricbeat.reference.yml index 263d3df085c1..a54ee2fd62e7 100644 --- a/x-pack/metricbeat/metricbeat.reference.yml +++ b/x-pack/metricbeat/metricbeat.reference.yml @@ -1286,7 +1286,12 @@ metricbeat.modules: # collection: # ## Number of days to look back when collecting usage data # lookback_days: 30 - + # ## Whether to collect usage data in realtime. Defaults to false as how + # # OpenAI usage data is collected will end up adding duplicate data to ES + # # and also making it harder to do analytics. Best approach is to avoid + # # realtime collection and collect only upto last day (in UTC). So, there's + # # at most 24h delay. + # realtime: false #----------------------------- Openmetrics Module ----------------------------- - module: openmetrics metricsets: ['collector'] diff --git a/x-pack/metricbeat/module/openai/_meta/config.yml b/x-pack/metricbeat/module/openai/_meta/config.yml index 1290889640b9..c3ab46a59932 100644 --- a/x-pack/metricbeat/module/openai/_meta/config.yml +++ b/x-pack/metricbeat/module/openai/_meta/config.yml @@ -26,3 +26,9 @@ # collection: # ## Number of days to look back when collecting usage data # lookback_days: 30 + # ## Whether to collect usage data in realtime. Defaults to false as how + # # OpenAI usage data is collected will end up adding duplicate data to ES + # # and also making it harder to do analytics. Best approach is to avoid + # # realtime collection and collect only upto last day (in UTC). So, there's + # # at most 24h delay. + # realtime: false \ No newline at end of file diff --git a/x-pack/metricbeat/module/openai/usage/_meta/data.json b/x-pack/metricbeat/module/openai/usage/_meta/data.json index 7a6355842f0e..8775b4700357 100644 --- a/x-pack/metricbeat/module/openai/usage/_meta/data.json +++ b/x-pack/metricbeat/module/openai/usage/_meta/data.json @@ -12,7 +12,7 @@ "openai": { "usage": { "data": { - "aggregation_timestamp": 1730696460, + "aggregation_timestamp": "2024-11-04T05:01:00Z", "api_key_id": null, "api_key_name": null, "api_key_redacted": null, diff --git a/x-pack/metricbeat/module/openai/usage/config.go b/x-pack/metricbeat/module/openai/usage/config.go index 5908bdd57f92..17efbf290d42 100644 --- a/x-pack/metricbeat/module/openai/usage/config.go +++ b/x-pack/metricbeat/module/openai/usage/config.go @@ -28,7 +28,8 @@ type apiKeyConfig struct { } type collectionConfig struct { - LookbackDays int `config:"lookback_days"` + LookbackDays int `config:"lookback_days"` + Realtime bool `config:"realtime"` } func defaultConfig() Config { @@ -40,7 +41,8 @@ func defaultConfig() Config { Burst: ptr(5), }, Collection: collectionConfig{ - LookbackDays: 5, // 5 days + LookbackDays: 0, // 0 days + Realtime: false, // avoid realtime collection by default }, } } diff --git a/x-pack/metricbeat/module/openai/usage/persistcache.go b/x-pack/metricbeat/module/openai/usage/persistcache.go index 9b7d725461a5..7c6d097d6ba3 100644 --- a/x-pack/metricbeat/module/openai/usage/persistcache.go +++ b/x-pack/metricbeat/module/openai/usage/persistcache.go @@ -11,7 +11,7 @@ import ( "sync" ) -// stateStore handles persistence of state markers using the filesystem +// stateStore handles persistence of key-value pairs using the filesystem type stateStore struct { Dir string // Base directory for storing state files sync.RWMutex // Protects access to the state store @@ -32,17 +32,38 @@ func (s *stateStore) getStatePath(name string) string { return path.Join(s.Dir, name) } -// Put creates a state marker file for the given key -func (s *stateStore) Put(key string) error { +// Put stores a value in a file named by the key +func (s *stateStore) Put(key string, value string) error { s.Lock() defer s.Unlock() filePath := s.getStatePath(key) + + // In case the file already exists, file is truncated. f, err := os.Create(filePath) if err != nil { return fmt.Errorf("creating state file: %w", err) } - return f.Close() + defer f.Close() + + _, err = f.WriteString(value) + if err != nil { + return fmt.Errorf("writing value to state file: %w", err) + } + return nil +} + +// Get retrieves the value stored in the file named by the key +func (s *stateStore) Get(key string) (string, error) { + s.RLock() + defer s.RUnlock() + + filePath := s.getStatePath(key) + data, err := os.ReadFile(filePath) + if err != nil { + return "", fmt.Errorf("reading state file: %w", err) + } + return string(data), nil } // Has checks if a state exists for the given key @@ -55,7 +76,7 @@ func (s *stateStore) Has(key string) bool { return err == nil } -// Remove deletes the state marker file for the given key +// Remove deletes the state file for the given key func (s *stateStore) Remove(key string) error { s.Lock() defer s.Unlock() @@ -67,7 +88,7 @@ func (s *stateStore) Remove(key string) error { return nil } -// Clear removes all state markers by deleting and recreating the state directory +// Clear removes all state files by deleting and recreating the state directory func (s *stateStore) Clear() error { s.Lock() defer s.Unlock() diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index 4e9b8e88e5b7..fb3a6a55cfe2 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -73,6 +73,13 @@ func New(base mb.BaseMetricSet) (mb.MetricSet, error) { // Fetch method implements the data gathering and data conversion to the right // format. It publishes the event which is then forwarded to the output. In case // of an error set the Error field of mb.Event or simply call report.Error(). +// +// 1. Creates a rate-limited HTTP client with configured timeout and burst settings +// 2. Sets up the time range for data collection: +// i. End date is current UTC time for realtime collection +// ii. End date is previous day for non-realtime collection +// iii. Start date is calculated based on configured lookback days +// 3. Fetches usage data for each day in the date range func (m *MetricSet) Fetch(report mb.ReporterV2) error { httpClient := newClient( context.TODO(), @@ -87,38 +94,61 @@ func (m *MetricSet) Fetch(report mb.ReporterV2) error { m.report = report endDate := time.Now().UTC() + + if !m.config.Collection.Realtime { + // If we're not collecting realtime data, then just pull until + // yesterday (in UTC). + endDate = endDate.AddDate(0, 0, -1) + } + startDate := endDate.AddDate(0, 0, -m.config.Collection.LookbackDays) return m.fetchDateRange(startDate, endDate, httpClient) } +// fetchDateRange retrieves OpenAI API usage data for each configured API key within a specified date range. +// +// For each API key: +// 1. Generates a secure SHA-256 hash of the key for state tracking +// 2. Checks the state store for the last processed date +// 3. Adjusts start date if previous state exists to avoid duplicate collection +// 4. Iterates through each day in the range, collecting usage data +// 5. Updates state store with the latest processed date func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLHTTPClient) error { for _, apiKey := range m.config.APIKeys { - // SHA-256 produces a fixed-length (64 characters) hexadecimal string - // that is safe for filenames across all major platforms. Hex encoding - // ensures that the hash is safe for use in file paths as it uses only - // alphanumeric characters. + // SHA-256 is a cryptographic hash function that generates a 256-bit (32-byte) digest, + // represented as a 64-character hexadecimal string. The hash function is deterministic, + // ensuring the same input consistently produces identical output, while being computationally + // infeasible to reverse. Its strong collision resistance makes it highly suitable for + // secure key storage. // - // Also, SHA-256 is a strong cryptographic hash function that is - // deterministic, meaning that the same input will always produce - // the same output and it is an one-way function, meaning that it is - // computationally infeasible to reverse the hash to obtain the - // original. + // The hexadecimal representation uses only alphanumeric characters [0-9a-f], making it + // ideal for cross-platform filename compatibility. The fixed-length output provides + // predictable storage requirements and consistent behavior across different systems. hasher := sha256.New() hasher.Write([]byte(apiKey.Key)) hashedKey := hex.EncodeToString(hasher.Sum(nil)) stateKey := "state_" + hashedKey - // If state exists, only fetch current day if m.stateStore.Has(stateKey) { - currentDay := endDate.Format("2006-01-02") - if err := m.fetchSingleDay(currentDay, apiKey.Key, httpClient); err != nil { - m.logger.Errorf("Error fetching data for date %s: %v", currentDay, err) + lastProcessedDate, err := m.stateStore.Get(stateKey) + if err != nil { + m.logger.Errorf("Error reading state for API key: %v", err) + continue + } + + lastDate, err := time.Parse("2006-01-02", lastProcessedDate) + if err != nil { + m.logger.Errorf("Error parsing last processed date: %v", err) + continue + } + + startDate = lastDate.AddDate(0, 0, 1) + if startDate.After(endDate) { + continue } - continue } - // First run for this API key - fetch historical data for d := startDate; !d.After(endDate); d = d.AddDate(0, 0, 1) { dateStr := d.Format("2006-01-02") if err := m.fetchSingleDay(dateStr, apiKey.Key, httpClient); err != nil { @@ -127,8 +157,7 @@ func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLH } } - // Mark this API key as processed - if err := m.stateStore.Put(stateKey); err != nil { + if err := m.stateStore.Put(stateKey, endDate.Format("2006-01-02")); err != nil { m.logger.Errorf("Error storing state for API key: %v", err) } } @@ -207,7 +236,7 @@ func (m *MetricSet) processUsageData(events []mb.Event, data []UsageData) { "data": mapstr.M{ "organization_id": usage.OrganizationID, "organization_name": usage.OrganizationName, - "aggregation_timestamp": time.Unix(usage.AggregationTimestamp, 0), + "aggregation_timestamp": time.Unix(usage.AggregationTimestamp, 0).UTC(), // epoch time to time.Time (UTC) "n_requests": usage.NRequests, "operation": usage.Operation, "snapshot_id": usage.SnapshotID, @@ -235,7 +264,7 @@ func (m *MetricSet) processDalleData(events []mb.Event, data []DalleData) { event := mb.Event{ MetricSetFields: mapstr.M{ "dalle": mapstr.M{ - "timestamp": time.Unix(dalle.Timestamp, 0), + "timestamp": time.Unix(dalle.Timestamp, 0).UTC(), // epoch time to time.Time (UTC) "num_images": dalle.NumImages, "num_requests": dalle.NumRequests, "image_size": dalle.ImageSize, @@ -263,7 +292,7 @@ func (m *MetricSet) processWhisperData(events []mb.Event, data []WhisperData) { event := mb.Event{ MetricSetFields: mapstr.M{ "whisper": mapstr.M{ - "timestamp": time.Unix(whisper.Timestamp, 0), + "timestamp": time.Unix(whisper.Timestamp, 0).UTC(), // epoch time to time.Time (UTC) "model_id": whisper.ModelID, "num_seconds": whisper.NumSeconds, "num_requests": whisper.NumRequests, @@ -289,7 +318,7 @@ func (m *MetricSet) processTTSData(events []mb.Event, data []TtsData) { event := mb.Event{ MetricSetFields: mapstr.M{ "tts": mapstr.M{ - "timestamp": time.Unix(tts.Timestamp, 0), + "timestamp": time.Unix(tts.Timestamp, 0).UTC(), // epoch time to time.Time (UTC) "model_id": tts.ModelID, "num_characters": tts.NumCharacters, "num_requests": tts.NumRequests, diff --git a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go index 8a0c9b70090c..54a779be7edf 100644 --- a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go +++ b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go @@ -11,6 +11,7 @@ import ( "net/http" "net/http/httptest" "testing" + "time" "github.com/stretchr/testify/assert" @@ -18,7 +19,7 @@ import ( ) func TestFetch(t *testing.T) { - apiKey := "most_secure_token" + apiKey := time.Now().String() // to generate a unique API key everytime usagePath := "/usage" server := initServer(usagePath, apiKey) defer server.Close() @@ -35,7 +36,7 @@ func TestFetch(t *testing.T) { } func TestData(t *testing.T) { - apiKey := "most_secure_token" + apiKey := time.Now().String() // to generate a unique API key everytime usagePath := "/usage" server := initServer(usagePath, apiKey) defer server.Close() diff --git a/x-pack/metricbeat/modules.d/openai.yml.disabled b/x-pack/metricbeat/modules.d/openai.yml.disabled index 10a9fb485770..8dba9ea0d1c9 100644 --- a/x-pack/metricbeat/modules.d/openai.yml.disabled +++ b/x-pack/metricbeat/modules.d/openai.yml.disabled @@ -29,3 +29,9 @@ # collection: # ## Number of days to look back when collecting usage data # lookback_days: 30 + # ## Whether to collect usage data in realtime. Defaults to false as how + # # OpenAI usage data is collected will end up adding duplicate data to ES + # # and also making it harder to do analytics. Best approach is to avoid + # # realtime collection and collect only upto last day (in UTC). So, there's + # # at most 24h delay. + # realtime: false \ No newline at end of file From 3259612fe21f991a437c27a04a6c997456746046 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Mon, 25 Nov 2024 15:11:21 +0530 Subject: [PATCH 09/36] Address review comments --- .../metricbeat/module/openai/usage/client.go | 16 ++-- .../metricbeat/module/openai/usage/config.go | 49 +++++----- .../metricbeat/module/openai/usage/errors.go | 11 +++ .../module/openai/usage/persistcache.go | 60 +++++++++++- .../metricbeat/module/openai/usage/usage.go | 91 ++++++++----------- .../openai/usage/usage_integration_test.go | 45 +++++---- 6 files changed, 173 insertions(+), 99 deletions(-) create mode 100644 x-pack/metricbeat/module/openai/usage/errors.go diff --git a/x-pack/metricbeat/module/openai/usage/client.go b/x-pack/metricbeat/module/openai/usage/client.go index c8f24aa8a50c..aedbd8b605ba 100644 --- a/x-pack/metricbeat/module/openai/usage/client.go +++ b/x-pack/metricbeat/module/openai/usage/client.go @@ -5,7 +5,6 @@ package usage import ( - "context" "net/http" "time" @@ -17,7 +16,6 @@ import ( // RLHTTPClient implements a rate-limited HTTP client that wraps the standard http.Client // with a rate limiter to control API request frequency. type RLHTTPClient struct { - ctx context.Context client *http.Client logger *logp.Logger Ratelimiter *rate.Limiter @@ -27,21 +25,25 @@ type RLHTTPClient struct { // It waits for rate limit token before proceeding with the request. // Returns the HTTP response and any error encountered. func (c *RLHTTPClient) Do(req *http.Request) (*http.Response, error) { - c.logger.Warn("Waiting for rate limit token") - err := c.Ratelimiter.Wait(context.TODO()) + c.logger.Debug("Waiting for rate limit token") + start := time.Now() + err := c.Ratelimiter.Wait(req.Context()) + waitDuration := time.Since(start) if err != nil { return nil, err } - c.logger.Warn("Rate limit token acquired") + c.logger.Debug("Rate limit token acquired") + if waitDuration > time.Minute { + c.logger.Infof("Rate limit wait exceeded threshold: %v", waitDuration) + } return c.client.Do(req) } // newClient creates a new rate-limited HTTP client with specified rate limiter and timeout. -func newClient(ctx context.Context, logger *logp.Logger, rl *rate.Limiter, timeout time.Duration) *RLHTTPClient { +func newClient(logger *logp.Logger, rl *rate.Limiter, timeout time.Duration) *RLHTTPClient { var client = http.DefaultClient client.Timeout = timeout return &RLHTTPClient{ - ctx: ctx, client: client, logger: logger, Ratelimiter: rl, diff --git a/x-pack/metricbeat/module/openai/usage/config.go b/x-pack/metricbeat/module/openai/usage/config.go index 17efbf290d42..ae8870e0d99e 100644 --- a/x-pack/metricbeat/module/openai/usage/config.go +++ b/x-pack/metricbeat/module/openai/usage/config.go @@ -5,6 +5,7 @@ package usage import ( + "errors" "fmt" "time" ) @@ -48,35 +49,39 @@ func defaultConfig() Config { } func (c *Config) Validate() error { - switch { - case len(c.APIKeys) == 0: - return fmt.Errorf("at least one API key must be configured") + var errs []error - case c.APIURL == "": - return fmt.Errorf("api_url cannot be empty") - - case c.RateLimit == nil: - return fmt.Errorf("rate_limit must be configured") - - case c.RateLimit.Limit == nil: - return fmt.Errorf("rate_limit.limit must be configured") - - case c.RateLimit.Burst == nil: - return fmt.Errorf("rate_limit.burst must be configured") - - case c.Timeout <= 0: - return fmt.Errorf("timeout must be greater than 0") - - case c.Collection.LookbackDays < 0: - return fmt.Errorf("lookback_days must be >= 0") + if len(c.APIKeys) == 0 { + errs = append(errs, errors.New("at least one API key must be configured")) + } + if c.APIURL == "" { + errs = append(errs, errors.New("api_url cannot be empty")) + } + if c.RateLimit == nil { + errs = append(errs, errors.New("rate_limit must be configured")) + } else { + if c.RateLimit.Limit == nil { + errs = append(errs, errors.New("rate_limit.limit must be configured")) + } + if c.RateLimit.Burst == nil { + errs = append(errs, errors.New("rate_limit.burst must be configured")) + } + } + if c.Timeout <= 0 { + errs = append(errs, errors.New("timeout must be greater than 0")) + } + if c.Collection.LookbackDays < 0 { + errs = append(errs, errors.New("lookback_days must be >= 0")) } - // API keys validation in a separate loop since it needs iteration for i, apiKey := range c.APIKeys { if apiKey.Key == "" { - return fmt.Errorf("API key at position %d cannot be empty", i) + errs = append(errs, fmt.Errorf("API key at position %d cannot be empty", i)) } } + if len(errs) > 0 { + return fmt.Errorf("validation failed: %v", errors.Join(errs...)) + } return nil } diff --git a/x-pack/metricbeat/module/openai/usage/errors.go b/x-pack/metricbeat/module/openai/usage/errors.go new file mode 100644 index 000000000000..54a54e362843 --- /dev/null +++ b/x-pack/metricbeat/module/openai/usage/errors.go @@ -0,0 +1,11 @@ +package usage + +import "errors" + +var ( + // ErrNoState indicates no previous state exists for the given API key + ErrNoState = errors.New("no previous state found") + + // ErrHTTPClientTimeout indicates request timeout + ErrHTTPClientTimeout = errors.New("http client request timeout") +) diff --git a/x-pack/metricbeat/module/openai/usage/persistcache.go b/x-pack/metricbeat/module/openai/usage/persistcache.go index 7c6d097d6ba3..def92ad32777 100644 --- a/x-pack/metricbeat/module/openai/usage/persistcache.go +++ b/x-pack/metricbeat/module/openai/usage/persistcache.go @@ -5,16 +5,40 @@ package usage import ( + "crypto/sha256" + "encoding/hex" "fmt" "os" "path" "sync" + "time" ) +// stateManager handles the storage and retrieval of state data with key hashing and caching capabilities +type stateManager struct { + store *stateStore + keyPrefix string + hashCache sync.Map +} + // stateStore handles persistence of key-value pairs using the filesystem type stateStore struct { - Dir string // Base directory for storing state files sync.RWMutex // Protects access to the state store + Dir string // Base directory for storing state files +} + +// newStateManager creates a new state manager instance with the given storage path +func newStateManager(storePath string) (*stateManager, error) { + store, err := newStateStore(storePath) + if err != nil { + return nil, fmt.Errorf("create state store: %w", err) + } + + return &stateManager{ + store: store, + keyPrefix: "state_", + hashCache: sync.Map{}, + }, nil } // newStateStore creates a new state store instance at the specified path @@ -98,3 +122,37 @@ func (s *stateStore) Clear() error { } return os.MkdirAll(s.Dir, 0o755) } + +// GetLastProcessedDate retrieves and parses the last processed date for a given API key +func (s *stateManager) GetLastProcessedDate(apiKey string) (time.Time, error) { + hashedKey := s.hashKey(apiKey) + stateKey := s.keyPrefix + hashedKey + + if !s.store.Has(stateKey) { + return time.Time{}, ErrNoState + } + + dateStr, err := s.store.Get(stateKey) + if err != nil { + return time.Time{}, fmt.Errorf("get state: %w", err) + } + + return time.Parse("2006-01-02", dateStr) +} + +// hashKey generates and caches a SHA-256 hash of the provided API key +func (s *stateManager) hashKey(apiKey string) string { + // Check cache first to avoid recomputing hashes + if hashedKey, ok := s.hashCache.Load(apiKey); ok { + return hashedKey.(string) + } + + // Generate SHA-256 hash and hex encode for safe filename usage + hasher := sha256.New() + hasher.Write([]byte(apiKey)) + hashedKey := hex.EncodeToString(hasher.Sum(nil)) + + // Cache the computed hash for future lookups + s.hashCache.Store(apiKey, hashedKey) + return hashedKey +} diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index fb3a6a55cfe2..e4c8a236c66f 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -6,10 +6,10 @@ package usage import ( "context" - "crypto/sha256" - "encoding/hex" "encoding/json" "fmt" + "io" + "net" "net/http" "path" "time" @@ -37,10 +37,11 @@ func init() { // interface methods except for Fetch. type MetricSet struct { mb.BaseMetricSet - logger *logp.Logger - config Config - report mb.ReporterV2 - stateStore *stateStore + httpClient *RLHTTPClient + logger *logp.Logger + config Config + report mb.ReporterV2 + stateManager *stateManager } // New creates a new instance of the MetricSet. New is responsible for unpacking @@ -57,16 +58,28 @@ func New(base mb.BaseMetricSet) (mb.MetricSet, error) { return nil, err } - st, err := newStateStore(paths.Resolve(paths.Data, path.Join(base.Module().Name(), base.Name()))) + sm, err := newStateManager(paths.Resolve(paths.Data, path.Join(base.Module().Name(), base.Name()))) if err != nil { - return nil, fmt.Errorf("creating state store: %w", err) + return nil, fmt.Errorf("create state manager: %w", err) } + logger := logp.NewLogger("openai.usage") + + httpClient := newClient( + logger, + rate.NewLimiter( + rate.Every(time.Duration(*config.RateLimit.Limit)*time.Second), + *config.RateLimit.Burst, + ), + config.Timeout, + ) + return &MetricSet{ BaseMetricSet: base, - logger: logp.NewLogger("openai.usage"), + httpClient: httpClient, + logger: logger, config: config, - stateStore: st, + stateManager: sm, }, nil } @@ -81,18 +94,6 @@ func New(base mb.BaseMetricSet) (mb.MetricSet, error) { // iii. Start date is calculated based on configured lookback days // 3. Fetches usage data for each day in the date range func (m *MetricSet) Fetch(report mb.ReporterV2) error { - httpClient := newClient( - context.TODO(), - m.logger, - rate.NewLimiter( - rate.Every(time.Duration(*m.config.RateLimit.Limit)*time.Second), - *m.config.RateLimit.Burst, - ), - m.config.Timeout, - ) - - m.report = report - endDate := time.Now().UTC() if !m.config.Collection.Realtime { @@ -103,7 +104,8 @@ func (m *MetricSet) Fetch(report mb.ReporterV2) error { startDate := endDate.AddDate(0, 0, -m.config.Collection.LookbackDays) - return m.fetchDateRange(startDate, endDate, httpClient) + m.report = report + return m.fetchDateRange(startDate, endDate, m.httpClient) } // fetchDateRange retrieves OpenAI API usage data for each configured API key within a specified date range. @@ -114,36 +116,13 @@ func (m *MetricSet) Fetch(report mb.ReporterV2) error { // 3. Adjusts start date if previous state exists to avoid duplicate collection // 4. Iterates through each day in the range, collecting usage data // 5. Updates state store with the latest processed date +// Update the fetchDateRange method to use the new stateManager methods func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLHTTPClient) error { for _, apiKey := range m.config.APIKeys { - // SHA-256 is a cryptographic hash function that generates a 256-bit (32-byte) digest, - // represented as a 64-character hexadecimal string. The hash function is deterministic, - // ensuring the same input consistently produces identical output, while being computationally - // infeasible to reverse. Its strong collision resistance makes it highly suitable for - // secure key storage. - // - // The hexadecimal representation uses only alphanumeric characters [0-9a-f], making it - // ideal for cross-platform filename compatibility. The fixed-length output provides - // predictable storage requirements and consistent behavior across different systems. - hasher := sha256.New() - hasher.Write([]byte(apiKey.Key)) - hashedKey := hex.EncodeToString(hasher.Sum(nil)) - stateKey := "state_" + hashedKey - - if m.stateStore.Has(stateKey) { - lastProcessedDate, err := m.stateStore.Get(stateKey) - if err != nil { - m.logger.Errorf("Error reading state for API key: %v", err) - continue - } - - lastDate, err := time.Parse("2006-01-02", lastProcessedDate) - if err != nil { - m.logger.Errorf("Error parsing last processed date: %v", err) - continue - } - - startDate = lastDate.AddDate(0, 0, 1) + lastProcessedDate, err := m.stateManager.GetLastProcessedDate(apiKey.Key) + if err == nil { + // We have previous state, adjust start date + startDate = lastProcessedDate.AddDate(0, 0, 1) if startDate.After(endDate) { continue } @@ -157,7 +136,9 @@ func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLH } } - if err := m.stateStore.Put(stateKey, endDate.Format("2006-01-02")); err != nil { + // Store using stateManager's key prefix and hashing + stateKey := m.stateManager.keyPrefix + m.stateManager.hashKey(apiKey.Key) + if err := m.stateManager.store.Put(stateKey, endDate.Format("2006-01-02")); err != nil { m.logger.Errorf("Error storing state for API key: %v", err) } } @@ -172,12 +153,16 @@ func (m *MetricSet) fetchSingleDay(dateStr, apiKey string, httpClient *RLHTTPCli resp, err := httpClient.Do(req) if err != nil { + if err, ok := err.(net.Error); ok && err.Timeout() { + return ErrHTTPClientTimeout + } return fmt.Errorf("error making request: %w", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { - return fmt.Errorf("error response from API: %s", resp.Status) + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("error response from API: status=%s, body=%s", resp.Status, string(body)) } return m.processResponse(resp, dateStr) diff --git a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go index 54a779be7edf..62d822481d92 100644 --- a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go +++ b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go @@ -13,7 +13,7 @@ import ( "testing" "time" - "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" mbtest "github.com/elastic/beats/v7/metricbeat/mb/testing" ) @@ -27,12 +27,8 @@ func TestFetch(t *testing.T) { f := mbtest.NewReportingMetricSetV2Error(t, getConfig(server.URL+"/usage", apiKey)) events, errs := mbtest.ReportingFetchV2Error(f) - if len(errs) > 0 { - t.Fatalf("Expected 0 error, has %d: %v", len(errs), errs) - } - - assert.NotEmpty(t, events) - + require.Empty(t, errs, "Expected no errors") + require.NotEmpty(t, events, "Expected events to be returned") } func TestData(t *testing.T) { @@ -44,9 +40,7 @@ func TestData(t *testing.T) { f := mbtest.NewReportingMetricSetV2Error(t, getConfig(server.URL+"/usage", apiKey)) err := mbtest.WriteEventsReporterV2Error(f, t, "") - if !assert.NoError(t, err) { - t.FailNow() - } + require.NoError(t, err, "Writing events should not return an error") } func getConfig(url, apiKey string) map[string]interface{} { @@ -132,7 +126,25 @@ func initServer(endpoint string, api_key string) *httptest.Server { } ], "ft_data": [], - "dalle_api_data": [], + "dalle_api_data": [ + { + "timestamp": 1730696460, + "num_images": 1, + "num_requests": 1, + "image_size": "1024x1024", + "operation": "generations", + "user_id": "subham.sarkar@elastic.co", + "organization_id": "org-FCp10pUDIN4slA4kNZK6UKkX", + "api_key_id": "key_10xSzP3zPsz8zB5O", + "api_key_name": "project_key", + "api_key_redacted": "sk-...zkA", + "api_key_type": "organization", + "organization_name": "Personal", + "model_id": "dall-e-3", + "project_id": "Default Project", + "project_name": "Default Project" + } + ], "whisper_api_data": [ { "timestamp": 1730696460, @@ -165,13 +177,14 @@ func initServer(endpoint string, api_key string) *httptest.Server { return } - // Validate the endpoint - if r.URL.Path == endpoint { - w.WriteHeader(http.StatusOK) - _, _ = w.Write(data) - } else { + // If it doesn't match the expected endpoint, return 404 + if r.URL.Path != endpoint { w.WriteHeader(http.StatusNotFound) + return } + + w.WriteHeader(http.StatusOK) + _, _ = w.Write(data) })) return server } From 40b1131587b4eb7bf951ba8f29d5bf19cec3c63b Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Mon, 25 Nov 2024 15:32:26 +0530 Subject: [PATCH 10/36] Address review comments --- x-pack/metricbeat/module/openai/usage/errors.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/x-pack/metricbeat/module/openai/usage/errors.go b/x-pack/metricbeat/module/openai/usage/errors.go index 54a54e362843..8aec742de34b 100644 --- a/x-pack/metricbeat/module/openai/usage/errors.go +++ b/x-pack/metricbeat/module/openai/usage/errors.go @@ -1,3 +1,7 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License; +// you may not use this file except in compliance with the Elastic License. + package usage import "errors" From 29365403de59e7590f480bbf0a3120df5110f8c5 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Mon, 25 Nov 2024 16:30:57 +0530 Subject: [PATCH 11/36] Address review comments --- x-pack/metricbeat/module/openai/usage/usage.go | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index e4c8a236c66f..c87d3f7db580 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -42,6 +42,7 @@ type MetricSet struct { config Config report mb.ReporterV2 stateManager *stateManager + headers map[string]string } // New creates a new instance of the MetricSet. New is responsible for unpacking @@ -80,6 +81,7 @@ func New(base mb.BaseMetricSet) (mb.MetricSet, error) { logger: logger, config: config, stateManager: sm, + headers: processHeaders(config.Headers), }, nil } @@ -179,7 +181,7 @@ func (m *MetricSet) createRequest(dateStr, apiKey string) (*http.Request, error) req.URL.RawQuery = q.Encode() req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", apiKey)) - for key, value := range processHeaders(m.config.Headers) { + for key, value := range m.headers { req.Header.Add(key, value) } From 5061e4630da29999d5721f92d1fea5b24533d04a Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Mon, 25 Nov 2024 16:38:54 +0530 Subject: [PATCH 12/36] Address review comments --- .../metricbeat/module/openai/usage/usage.go | 31 +++++++++---------- 1 file changed, 15 insertions(+), 16 deletions(-) diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index c87d3f7db580..0156a8b2fe82 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -85,16 +85,13 @@ func New(base mb.BaseMetricSet) (mb.MetricSet, error) { }, nil } -// Fetch method implements the data gathering and data conversion to the right -// format. It publishes the event which is then forwarded to the output. In case -// of an error set the Error field of mb.Event or simply call report.Error(). +// Fetch collects OpenAI API usage data for the configured time range. // -// 1. Creates a rate-limited HTTP client with configured timeout and burst settings -// 2. Sets up the time range for data collection: -// i. End date is current UTC time for realtime collection -// ii. End date is previous day for non-realtime collection -// iii. Start date is calculated based on configured lookback days -// 3. Fetches usage data for each day in the date range +// The collection process: +// 1. Determines the time range based on realtime/non-realtime configuration +// 2. Calculates start date using configured lookback days +// 3. Fetches usage data for each day in the range +// 4. Reports collected metrics through the mb.ReporterV2 func (m *MetricSet) Fetch(report mb.ReporterV2) error { endDate := time.Now().UTC() @@ -110,15 +107,14 @@ func (m *MetricSet) Fetch(report mb.ReporterV2) error { return m.fetchDateRange(startDate, endDate, m.httpClient) } -// fetchDateRange retrieves OpenAI API usage data for each configured API key within a specified date range. +// fetchDateRange retrieves OpenAI API usage data for each configured API key within a date range. // // For each API key: -// 1. Generates a secure SHA-256 hash of the key for state tracking -// 2. Checks the state store for the last processed date -// 3. Adjusts start date if previous state exists to avoid duplicate collection -// 4. Iterates through each day in the range, collecting usage data -// 5. Updates state store with the latest processed date -// Update the fetchDateRange method to use the new stateManager methods +// 1. Retrieves last processed date from state store +// 2. Adjusts collection range to avoid duplicates +// 3. Collects daily usage data +// 4. Updates state store with latest processed date +// 5. Handles errors per day without failing entire range func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLHTTPClient) error { for _, apiKey := range m.config.APIKeys { lastProcessedDate, err := m.stateManager.GetLastProcessedDate(apiKey.Key) @@ -147,6 +143,7 @@ func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLH return nil } +// fetchSingleDay retrieves usage data for a specific date and API key. func (m *MetricSet) fetchSingleDay(dateStr, apiKey string, httpClient *RLHTTPClient) error { req, err := m.createRequest(dateStr, apiKey) if err != nil { @@ -170,6 +167,7 @@ func (m *MetricSet) fetchSingleDay(dateStr, apiKey string, httpClient *RLHTTPCli return m.processResponse(resp, dateStr) } +// createRequest builds an HTTP request for the OpenAI usage API. func (m *MetricSet) createRequest(dateStr, apiKey string) (*http.Request, error) { req, err := http.NewRequestWithContext(context.TODO(), http.MethodGet, m.config.APIURL, nil) if err != nil { @@ -188,6 +186,7 @@ func (m *MetricSet) createRequest(dateStr, apiKey string) (*http.Request, error) return req, nil } +// processResponse handles the API response and processes the usage data. func (m *MetricSet) processResponse(resp *http.Response, dateStr string) error { var usageResponse UsageResponse if err := json.NewDecoder(resp.Body).Decode(&usageResponse); err != nil { From 34a7128d48f1f02dc5c680863f204bfa7bfb96fc Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 26 Nov 2024 02:13:30 +0530 Subject: [PATCH 13/36] Improvements --- x-pack/metricbeat/metricbeat.yml | 17 ++++++--- .../metricbeat/module/openai/usage/config.go | 13 ++++--- .../metricbeat/module/openai/usage/usage.go | 12 +++--- .../metricbeat/modules.d/openai.yml.disabled | 37 ------------------- 4 files changed, 24 insertions(+), 55 deletions(-) delete mode 100644 x-pack/metricbeat/modules.d/openai.yml.disabled diff --git a/x-pack/metricbeat/metricbeat.yml b/x-pack/metricbeat/metricbeat.yml index a148cfb3b517..9bbf792640e5 100644 --- a/x-pack/metricbeat/metricbeat.yml +++ b/x-pack/metricbeat/metricbeat.yml @@ -89,13 +89,18 @@ setup.kibana: # Configure what output to use when sending the data collected by the beat. # ---------------------------- Elasticsearch Output ---------------------------- -output.elasticsearch: - # Array of hosts to connect to. - hosts: ["localhost:9200"] - # Performance preset - one of "balanced", "throughput", "scale", - # "latency", or "custom". - preset: balanced +output.file: + path: "." + filename: openai-logs + +# output.elasticsearch: +# # Array of hosts to connect to. +# hosts: ["localhost:9200"] + +# # Performance preset - one of "balanced", "throughput", "scale", +# # "latency", or "custom". +# preset: balanced # Protocol - either `http` (default) or `https`. #protocol: "https" diff --git a/x-pack/metricbeat/module/openai/usage/config.go b/x-pack/metricbeat/module/openai/usage/config.go index ae8870e0d99e..9258f867fde4 100644 --- a/x-pack/metricbeat/module/openai/usage/config.go +++ b/x-pack/metricbeat/module/openai/usage/config.go @@ -12,24 +12,24 @@ import ( type Config struct { APIKeys []apiKeyConfig `config:"api_keys" validate:"required"` - APIURL string `config:"api_url"` + APIURL string `config:"api_url" validate:"required"` Headers []string `config:"headers"` RateLimit *rateLimitConfig `config:"rate_limit"` - Timeout time.Duration `config:"timeout"` + Timeout time.Duration `config:"timeout" validate:"required"` Collection collectionConfig `config:"collection"` } type rateLimitConfig struct { - Limit *int `config:"limit"` - Burst *int `config:"burst"` + Limit *int `config:"limit" validate:"required"` + Burst *int `config:"burst" validate:"required"` } type apiKeyConfig struct { - Key string `config:"key"` + Key string `config:"key" validate:"required"` } type collectionConfig struct { - LookbackDays int `config:"lookback_days"` + LookbackDays int `config:"lookback_days" validate:"required"` Realtime bool `config:"realtime"` } @@ -83,5 +83,6 @@ func (c *Config) Validate() error { if len(errs) > 0 { return fmt.Errorf("validation failed: %v", errors.Join(errs...)) } + return nil } diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index 0156a8b2fe82..76449124f72a 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -117,6 +117,9 @@ func (m *MetricSet) Fetch(report mb.ReporterV2) error { // 5. Handles errors per day without failing entire range func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLHTTPClient) error { for _, apiKey := range m.config.APIKeys { + // stateKey using stateManager's key prefix and hashing apiKey + stateKey := m.stateManager.keyPrefix + m.stateManager.hashKey(apiKey.Key) + lastProcessedDate, err := m.stateManager.GetLastProcessedDate(apiKey.Key) if err == nil { // We have previous state, adjust start date @@ -132,12 +135,9 @@ func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLH m.logger.Errorf("Error fetching data for date %s: %v", dateStr, err) continue } - } - - // Store using stateManager's key prefix and hashing - stateKey := m.stateManager.keyPrefix + m.stateManager.hashKey(apiKey.Key) - if err := m.stateManager.store.Put(stateKey, endDate.Format("2006-01-02")); err != nil { - m.logger.Errorf("Error storing state for API key: %v", err) + if err := m.stateManager.store.Put(stateKey, dateStr); err != nil { + m.logger.Errorf("Error storing state for API key: %v", err) + } } } return nil diff --git a/x-pack/metricbeat/modules.d/openai.yml.disabled b/x-pack/metricbeat/modules.d/openai.yml.disabled deleted file mode 100644 index 8dba9ea0d1c9..000000000000 --- a/x-pack/metricbeat/modules.d/openai.yml.disabled +++ /dev/null @@ -1,37 +0,0 @@ -# Module: openai -# Docs: https://www.elastic.co/guide/en/beats/metricbeat/main/metricbeat-module-openai.html - -- module: openai - metricsets: ["usage"] - enabled: false - period: 1h - - # # Project API Keys - Multiple API keys can be specified for different projects - # api_keys: - # - key: "api_key1" - # - key: "api_key2" - - # # API Configuration - # ## Base URL for the OpenAI usage API endpoint - # api_url: "https://api.openai.com/v1/usage" - # ## Custom headers to be included in API requests - # headers: - # - "k1: v1" - # - "k2: v2" - ## Rate Limiting Configuration - # rate_limit: - # limit: 60 # requests per second - # burst: 5 # burst size - # ## Request timeout duration - # timeout: 30s - - # # Data Collection Configuration - # collection: - # ## Number of days to look back when collecting usage data - # lookback_days: 30 - # ## Whether to collect usage data in realtime. Defaults to false as how - # # OpenAI usage data is collected will end up adding duplicate data to ES - # # and also making it harder to do analytics. Best approach is to avoid - # # realtime collection and collect only upto last day (in UTC). So, there's - # # at most 24h delay. - # realtime: false \ No newline at end of file From 7f4cc75f8a05085eb6ca95dfd3e73bf7e95d36da Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 26 Nov 2024 16:08:02 +0530 Subject: [PATCH 14/36] Fix *.yml --- x-pack/metricbeat/metricbeat.yml | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/x-pack/metricbeat/metricbeat.yml b/x-pack/metricbeat/metricbeat.yml index 9bbf792640e5..a148cfb3b517 100644 --- a/x-pack/metricbeat/metricbeat.yml +++ b/x-pack/metricbeat/metricbeat.yml @@ -89,18 +89,13 @@ setup.kibana: # Configure what output to use when sending the data collected by the beat. # ---------------------------- Elasticsearch Output ---------------------------- +output.elasticsearch: + # Array of hosts to connect to. + hosts: ["localhost:9200"] -output.file: - path: "." - filename: openai-logs - -# output.elasticsearch: -# # Array of hosts to connect to. -# hosts: ["localhost:9200"] - -# # Performance preset - one of "balanced", "throughput", "scale", -# # "latency", or "custom". -# preset: balanced + # Performance preset - one of "balanced", "throughput", "scale", + # "latency", or "custom". + preset: balanced # Protocol - either `http` (default) or `https`. #protocol: "https" From 7919c8131e113be291b4b8f84227eb8827748296 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 26 Nov 2024 16:14:53 +0530 Subject: [PATCH 15/36] Address review comments --- x-pack/metricbeat/module/openai/usage/errors.go | 9 ++------- x-pack/metricbeat/module/openai/usage/usage.go | 6 ++---- 2 files changed, 4 insertions(+), 11 deletions(-) diff --git a/x-pack/metricbeat/module/openai/usage/errors.go b/x-pack/metricbeat/module/openai/usage/errors.go index 8aec742de34b..9b8b793767c0 100644 --- a/x-pack/metricbeat/module/openai/usage/errors.go +++ b/x-pack/metricbeat/module/openai/usage/errors.go @@ -6,10 +6,5 @@ package usage import "errors" -var ( - // ErrNoState indicates no previous state exists for the given API key - ErrNoState = errors.New("no previous state found") - - // ErrHTTPClientTimeout indicates request timeout - ErrHTTPClientTimeout = errors.New("http client request timeout") -) +// ErrNoState indicates no previous state exists for the given API key +var ErrNoState = errors.New("no previous state found") diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index 76449124f72a..0d193c507fbd 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -8,7 +8,6 @@ import ( "context" "encoding/json" "fmt" - "io" "net" "net/http" "path" @@ -153,15 +152,14 @@ func (m *MetricSet) fetchSingleDay(dateStr, apiKey string, httpClient *RLHTTPCli resp, err := httpClient.Do(req) if err != nil { if err, ok := err.(net.Error); ok && err.Timeout() { - return ErrHTTPClientTimeout + return fmt.Errorf("request timed out with configured timeout: %v and error: %w", m.config.Timeout, err) } return fmt.Errorf("error making request: %w", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) - return fmt.Errorf("error response from API: status=%s, body=%s", resp.Status, string(body)) + return fmt.Errorf("error response from API: status=%s", resp.Status) } return m.processResponse(resp, dateStr) From 2a4bb56c2068ab811c787c743bdcc14871fa136e Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 26 Nov 2024 16:26:24 +0530 Subject: [PATCH 16/36] Improvements --- .../module/openai/usage/persistcache.go | 20 ++++++++++++++++--- .../metricbeat/module/openai/usage/usage.go | 4 ++-- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/x-pack/metricbeat/module/openai/usage/persistcache.go b/x-pack/metricbeat/module/openai/usage/persistcache.go index def92ad32777..81a19b2c63b6 100644 --- a/x-pack/metricbeat/module/openai/usage/persistcache.go +++ b/x-pack/metricbeat/module/openai/usage/persistcache.go @@ -7,9 +7,11 @@ package usage import ( "crypto/sha256" "encoding/hex" + "errors" "fmt" "os" "path" + "strings" "sync" "time" ) @@ -29,6 +31,10 @@ type stateStore struct { // newStateManager creates a new state manager instance with the given storage path func newStateManager(storePath string) (*stateManager, error) { + if strings.TrimSpace(storePath) == "" { + return nil, errors.New("empty path provided") + } + store, err := newStateStore(storePath) if err != nil { return nil, fmt.Errorf("create state store: %w", err) @@ -74,6 +80,11 @@ func (s *stateStore) Put(key string, value string) error { if err != nil { return fmt.Errorf("writing value to state file: %w", err) } + + if err = f.Sync(); err != nil { + return fmt.Errorf("syncing state file: %w", err) + } + return nil } @@ -125,8 +136,7 @@ func (s *stateStore) Clear() error { // GetLastProcessedDate retrieves and parses the last processed date for a given API key func (s *stateManager) GetLastProcessedDate(apiKey string) (time.Time, error) { - hashedKey := s.hashKey(apiKey) - stateKey := s.keyPrefix + hashedKey + stateKey := s.GetStateKey(apiKey) if !s.store.Has(stateKey) { return time.Time{}, ErrNoState @@ -149,10 +159,14 @@ func (s *stateManager) hashKey(apiKey string) string { // Generate SHA-256 hash and hex encode for safe filename usage hasher := sha256.New() - hasher.Write([]byte(apiKey)) + _, _ = hasher.Write([]byte(apiKey)) hashedKey := hex.EncodeToString(hasher.Sum(nil)) // Cache the computed hash for future lookups s.hashCache.Store(apiKey, hashedKey) return hashedKey } + +func (s *stateManager) GetStateKey(apiKey string) string { + return s.keyPrefix + s.hashKey(apiKey) +} diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index 0d193c507fbd..bb43641e3664 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -58,7 +58,7 @@ func New(base mb.BaseMetricSet) (mb.MetricSet, error) { return nil, err } - sm, err := newStateManager(paths.Resolve(paths.Data, path.Join(base.Module().Name(), base.Name()))) + sm, err := newStateManager(paths.Resolve(paths.Data, path.Join("state", base.Module().Name(), base.Name()))) if err != nil { return nil, fmt.Errorf("create state manager: %w", err) } @@ -117,7 +117,7 @@ func (m *MetricSet) Fetch(report mb.ReporterV2) error { func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLHTTPClient) error { for _, apiKey := range m.config.APIKeys { // stateKey using stateManager's key prefix and hashing apiKey - stateKey := m.stateManager.keyPrefix + m.stateManager.hashKey(apiKey.Key) + stateKey := m.stateManager.GetStateKey(apiKey.Key) lastProcessedDate, err := m.stateManager.GetLastProcessedDate(apiKey.Key) if err == nil { From 217f7d563d9c3948fb7c6b4cb0a03f6fbe81e0ad Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 26 Nov 2024 16:46:01 +0530 Subject: [PATCH 17/36] Make linter happy --- x-pack/metricbeat/module/openai/usage/config.go | 2 +- x-pack/metricbeat/module/openai/usage/usage.go | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/x-pack/metricbeat/module/openai/usage/config.go b/x-pack/metricbeat/module/openai/usage/config.go index 9258f867fde4..94d80ee9c6d7 100644 --- a/x-pack/metricbeat/module/openai/usage/config.go +++ b/x-pack/metricbeat/module/openai/usage/config.go @@ -81,7 +81,7 @@ func (c *Config) Validate() error { } if len(errs) > 0 { - return fmt.Errorf("validation failed: %v", errors.Join(errs...)) + return fmt.Errorf("validation failed: %w", errors.Join(errs...)) } return nil diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index bb43641e3664..20c5ba21dad9 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -7,6 +7,7 @@ package usage import ( "context" "encoding/json" + "errors" "fmt" "net" "net/http" @@ -151,7 +152,8 @@ func (m *MetricSet) fetchSingleDay(dateStr, apiKey string, httpClient *RLHTTPCli resp, err := httpClient.Do(req) if err != nil { - if err, ok := err.(net.Error); ok && err.Timeout() { + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { return fmt.Errorf("request timed out with configured timeout: %v and error: %w", m.config.Timeout, err) } return fmt.Errorf("error making request: %w", err) From b5dbfb031bbb2d9c3e7c54d1707273a37e793829 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Wed, 27 Nov 2024 13:31:06 +0530 Subject: [PATCH 18/36] Address review comments --- .../metricbeat/module/openai/usage/client.go | 23 +++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/x-pack/metricbeat/module/openai/usage/client.go b/x-pack/metricbeat/module/openai/usage/client.go index aedbd8b605ba..327189aef0ed 100644 --- a/x-pack/metricbeat/module/openai/usage/client.go +++ b/x-pack/metricbeat/module/openai/usage/client.go @@ -25,24 +25,39 @@ type RLHTTPClient struct { // It waits for rate limit token before proceeding with the request. // Returns the HTTP response and any error encountered. func (c *RLHTTPClient) Do(req *http.Request) (*http.Response, error) { - c.logger.Debug("Waiting for rate limit token") start := time.Now() + + c.logger.Debug("Waiting for rate limit token") + err := c.Ratelimiter.Wait(req.Context()) - waitDuration := time.Since(start) if err != nil { return nil, err } + c.logger.Debug("Rate limit token acquired") + + waitDuration := time.Since(start) + if waitDuration > time.Minute { c.logger.Infof("Rate limit wait exceeded threshold: %v", waitDuration) } + return c.client.Do(req) } // newClient creates a new rate-limited HTTP client with specified rate limiter and timeout. func newClient(logger *logp.Logger, rl *rate.Limiter, timeout time.Duration) *RLHTTPClient { - var client = http.DefaultClient - client.Timeout = timeout + transport := &http.Transport{ + MaxIdleConns: 100, + MaxIdleConnsPerHost: 100, + IdleConnTimeout: 90 * time.Second, + } + + client := &http.Client{ + Timeout: timeout, + Transport: transport, + } + return &RLHTTPClient{ client: client, logger: logger, From 492eca1dc322434f6b4a094bc32cf6a78da756d6 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Fri, 29 Nov 2024 00:28:36 +0530 Subject: [PATCH 19/36] gofumpt'ed --- x-pack/metricbeat/module/openai/usage/persistcache.go | 2 +- x-pack/metricbeat/module/openai/usage/usage.go | 1 - x-pack/metricbeat/module/openai/usage/usage_integration_test.go | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/x-pack/metricbeat/module/openai/usage/persistcache.go b/x-pack/metricbeat/module/openai/usage/persistcache.go index 81a19b2c63b6..bf72df9df448 100644 --- a/x-pack/metricbeat/module/openai/usage/persistcache.go +++ b/x-pack/metricbeat/module/openai/usage/persistcache.go @@ -63,7 +63,7 @@ func (s *stateStore) getStatePath(name string) string { } // Put stores a value in a file named by the key -func (s *stateStore) Put(key string, value string) error { +func (s *stateStore) Put(key, value string) error { s.Lock() defer s.Unlock() diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index 20c5ba21dad9..2b105abdf676 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -215,7 +215,6 @@ func (m *MetricSet) processResponse(resp *http.Response, dateStr string) error { } func (m *MetricSet) processUsageData(events []mb.Event, data []UsageData) { - for _, usage := range data { event := mb.Event{ MetricSetFields: mapstr.M{ diff --git a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go index 62d822481d92..6bf869dfea88 100644 --- a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go +++ b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go @@ -63,7 +63,7 @@ func getConfig(url, apiKey string) map[string]interface{} { } } -func initServer(endpoint string, api_key string) *httptest.Server { +func initServer(endpoint, api_key string) *httptest.Server { data := []byte(`{ "object": "list", "data": [ From 413010f3d3c41052f0e878a9d9ddd44f3f703fea Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Fri, 29 Nov 2024 01:08:12 +0530 Subject: [PATCH 20/36] Address review comments --- .../metricbeat/module/openai/usage/client.go | 11 ++- .../metricbeat/module/openai/usage/config.go | 6 ++ .../metricbeat/module/openai/usage/helper.go | 11 ++- .../module/openai/usage/persistcache.go | 2 +- .../metricbeat/module/openai/usage/schema.go | 79 +++++++------------ .../metricbeat/module/openai/usage/usage.go | 2 +- 6 files changed, 55 insertions(+), 56 deletions(-) diff --git a/x-pack/metricbeat/module/openai/usage/client.go b/x-pack/metricbeat/module/openai/usage/client.go index 327189aef0ed..161da993e382 100644 --- a/x-pack/metricbeat/module/openai/usage/client.go +++ b/x-pack/metricbeat/module/openai/usage/client.go @@ -5,6 +5,7 @@ package usage import ( + "fmt" "net/http" "time" @@ -13,8 +14,7 @@ import ( "github.com/elastic/elastic-agent-libs/logp" ) -// RLHTTPClient implements a rate-limited HTTP client that wraps the standard http.Client -// with a rate limiter to control API request frequency. +// RLHTTPClient wraps the standard http.Client with a rate limiter to control API request frequency. type RLHTTPClient struct { client *http.Client logger *logp.Logger @@ -42,7 +42,12 @@ func (c *RLHTTPClient) Do(req *http.Request) (*http.Response, error) { c.logger.Infof("Rate limit wait exceeded threshold: %v", waitDuration) } - return c.client.Do(req) + resp, err := c.client.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to execute request: %w", err) + } + + return resp, nil } // newClient creates a new rate-limited HTTP client with specified rate limiter and timeout. diff --git a/x-pack/metricbeat/module/openai/usage/config.go b/x-pack/metricbeat/module/openai/usage/config.go index 94d80ee9c6d7..1be571920a87 100644 --- a/x-pack/metricbeat/module/openai/usage/config.go +++ b/x-pack/metricbeat/module/openai/usage/config.go @@ -7,6 +7,7 @@ package usage import ( "errors" "fmt" + "net/url" "time" ) @@ -56,6 +57,11 @@ func (c *Config) Validate() error { } if c.APIURL == "" { errs = append(errs, errors.New("api_url cannot be empty")) + } else { + _, err := url.ParseRequestURI(c.APIURL) + if err != nil { + errs = append(errs, fmt.Errorf("invalid api_url format: %w", err)) + } } if c.RateLimit == nil { errs = append(errs, errors.New("rate_limit must be configured")) diff --git a/x-pack/metricbeat/module/openai/usage/helper.go b/x-pack/metricbeat/module/openai/usage/helper.go index f2d3adadd01e..992c025df8cc 100644 --- a/x-pack/metricbeat/module/openai/usage/helper.go +++ b/x-pack/metricbeat/module/openai/usage/helper.go @@ -6,6 +6,9 @@ package usage import "strings" +// dateFormatForStateStore is used to parse and format dates in the YYYY-MM-DD format +const dateFormatForStateStore = "2006-01-02" + func ptr[T any](value T) *T { return &value } @@ -13,11 +16,15 @@ func ptr[T any](value T) *T { func processHeaders(headers []string) map[string]string { headersMap := make(map[string]string, len(headers)) for _, header := range headers { - parts := strings.Split(header, ":") + parts := strings.SplitN(header, ":", 2) if len(parts) != 2 { continue } - headersMap[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1]) + k, v := strings.TrimSpace(parts[0]), strings.TrimSpace(parts[1]) + if k == "" || v == "" { + continue + } + headersMap[k] = v } return headersMap } diff --git a/x-pack/metricbeat/module/openai/usage/persistcache.go b/x-pack/metricbeat/module/openai/usage/persistcache.go index bf72df9df448..8fa8564292ec 100644 --- a/x-pack/metricbeat/module/openai/usage/persistcache.go +++ b/x-pack/metricbeat/module/openai/usage/persistcache.go @@ -147,7 +147,7 @@ func (s *stateManager) GetLastProcessedDate(apiKey string) (time.Time, error) { return time.Time{}, fmt.Errorf("get state: %w", err) } - return time.Parse("2006-01-02", dateStr) + return time.Parse(dateFormatForStateStore, dateStr) } // hashKey generates and caches a SHA-256 hash of the provided API key diff --git a/x-pack/metricbeat/module/openai/usage/schema.go b/x-pack/metricbeat/module/openai/usage/schema.go index 668d5b03370c..0fe130496134 100644 --- a/x-pack/metricbeat/module/openai/usage/schema.go +++ b/x-pack/metricbeat/module/openai/usage/schema.go @@ -4,6 +4,18 @@ package usage +type BaseData struct { + OrganizationID string `json:"organization_id"` + OrganizationName string `json:"organization_name"` + UserID *string `json:"user_id"` + 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"` +} + type UsageResponse struct { Object string `json:"object"` Data []UsageData `json:"data"` @@ -16,8 +28,7 @@ type UsageResponse struct { } type UsageData struct { - OrganizationID string `json:"organization_id"` - OrganizationName string `json:"organization_name"` + BaseData AggregationTimestamp int64 `json:"aggregation_timestamp"` NRequests int `json:"n_requests"` Operation string `json:"operation"` @@ -25,62 +36,32 @@ type UsageData struct { 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 DalleData struct { - Timestamp int64 `json:"timestamp"` - NumImages int `json:"num_images"` - NumRequests int `json:"num_requests"` - ImageSize string `json:"image_size"` - Operation string `json:"operation"` - UserID *string `json:"user_id"` - OrganizationID string `json:"organization_id"` - ApiKeyID *string `json:"api_key_id"` - ApiKeyName *string `json:"api_key_name"` - ApiKeyRedacted *string `json:"api_key_redacted"` - ApiKeyType *string `json:"api_key_type"` - OrganizationName string `json:"organization_name"` - ModelID string `json:"model_id"` - ProjectID *string `json:"project_id"` - ProjectName *string `json:"project_name"` + BaseData + Timestamp int64 `json:"timestamp"` + NumImages int `json:"num_images"` + NumRequests int `json:"num_requests"` + ImageSize string `json:"image_size"` + Operation string `json:"operation"` + ModelID string `json:"model_id"` } type WhisperData struct { - Timestamp int64 `json:"timestamp"` - ModelID string `json:"model_id"` - NumSeconds int `json:"num_seconds"` - NumRequests int `json:"num_requests"` - UserID *string `json:"user_id"` - OrganizationID string `json:"organization_id"` - ApiKeyID *string `json:"api_key_id"` - ApiKeyName *string `json:"api_key_name"` - ApiKeyRedacted *string `json:"api_key_redacted"` - ApiKeyType *string `json:"api_key_type"` - OrganizationName string `json:"organization_name"` - ProjectID *string `json:"project_id"` - ProjectName *string `json:"project_name"` + BaseData + Timestamp int64 `json:"timestamp"` + ModelID string `json:"model_id"` + NumSeconds int `json:"num_seconds"` + NumRequests int `json:"num_requests"` } type TtsData struct { - Timestamp int64 `json:"timestamp"` - ModelID string `json:"model_id"` - NumCharacters int `json:"num_characters"` - NumRequests int `json:"num_requests"` - UserID *string `json:"user_id"` - OrganizationID string `json:"organization_id"` - ApiKeyID *string `json:"api_key_id"` - ApiKeyName *string `json:"api_key_name"` - ApiKeyRedacted *string `json:"api_key_redacted"` - ApiKeyType *string `json:"api_key_type"` - OrganizationName string `json:"organization_name"` - ProjectID *string `json:"project_id"` - ProjectName *string `json:"project_name"` + BaseData + Timestamp int64 `json:"timestamp"` + ModelID string `json:"model_id"` + NumCharacters int `json:"num_characters"` + NumRequests int `json:"num_requests"` } diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index 2b105abdf676..018a9965d2ae 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -193,7 +193,7 @@ func (m *MetricSet) processResponse(resp *http.Response, dateStr string) error { return fmt.Errorf("error decoding response: %w", err) } - m.logger.Info("Fetched usage metrics for date:", dateStr) + m.logger.Infof("Fetched usage metrics for date: %s", dateStr) events := make([]mb.Event, 0, len(usageResponse.Data)) From 92f5062822f59f31bea1f5fa3a22337c1f1f1e29 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Fri, 29 Nov 2024 01:14:00 +0530 Subject: [PATCH 21/36] Address review comments --- x-pack/metricbeat/module/openai/usage/client.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/metricbeat/module/openai/usage/client.go b/x-pack/metricbeat/module/openai/usage/client.go index 161da993e382..3995386c30e7 100644 --- a/x-pack/metricbeat/module/openai/usage/client.go +++ b/x-pack/metricbeat/module/openai/usage/client.go @@ -31,7 +31,7 @@ func (c *RLHTTPClient) Do(req *http.Request) (*http.Response, error) { err := c.Ratelimiter.Wait(req.Context()) if err != nil { - return nil, err + return nil, fmt.Errorf("failed to acquire rate limit token: %w", err) } c.logger.Debug("Rate limit token acquired") From 9c28e4d5bf48f2bac9b7a9f001e04fef19d041d2 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Wed, 4 Dec 2024 14:08:12 +0530 Subject: [PATCH 22/36] Address review comments --- .../module/openai/usage/_meta/data.json | 19 +- .../metricbeat/module/openai/usage/config.go | 2 +- .../metricbeat/module/openai/usage/schema.go | 4 +- .../metricbeat/module/openai/usage/usage.go | 85 +++---- .../openai/usage/usage_integration_test.go | 207 +++++++++++++++++- 5 files changed, 249 insertions(+), 68 deletions(-) diff --git a/x-pack/metricbeat/module/openai/usage/_meta/data.json b/x-pack/metricbeat/module/openai/usage/_meta/data.json index 8775b4700357..e80e9552ba7b 100644 --- a/x-pack/metricbeat/module/openai/usage/_meta/data.json +++ b/x-pack/metricbeat/module/openai/usage/_meta/data.json @@ -11,25 +11,24 @@ }, "openai": { "usage": { + "api_key_id": null, + "api_key_name": null, + "api_key_redacted": null, + "api_key_type": null, "data": { - "aggregation_timestamp": "2024-11-04T05:01:00Z", - "api_key_id": null, - "api_key_name": null, - "api_key_redacted": null, - "api_key_type": null, "email": null, "n_cached_context_tokens_total": 0, "n_context_tokens_total": 118, "n_generated_tokens_total": 35, "n_requests": 1, "operation": "completion-realtime", - "organization_id": "org-dummy", - "organization_name": "Personal", - "project_id": null, - "project_name": null, "request_type": "", "snapshot_id": "gpt-4o-realtime-preview-2024-10-01" - } + }, + "organization_id": "org-dummy", + "organization_name": "Personal", + "project_id": null, + "project_name": null } }, "service": { diff --git a/x-pack/metricbeat/module/openai/usage/config.go b/x-pack/metricbeat/module/openai/usage/config.go index 1be571920a87..c82bb6271149 100644 --- a/x-pack/metricbeat/module/openai/usage/config.go +++ b/x-pack/metricbeat/module/openai/usage/config.go @@ -30,7 +30,7 @@ type apiKeyConfig struct { } type collectionConfig struct { - LookbackDays int `config:"lookback_days" validate:"required"` + LookbackDays int `config:"lookback_days"` Realtime bool `config:"realtime"` } diff --git a/x-pack/metricbeat/module/openai/usage/schema.go b/x-pack/metricbeat/module/openai/usage/schema.go index 0fe130496134..64f22393e590 100644 --- a/x-pack/metricbeat/module/openai/usage/schema.go +++ b/x-pack/metricbeat/module/openai/usage/schema.go @@ -7,7 +7,6 @@ package usage type BaseData struct { OrganizationID string `json:"organization_id"` OrganizationName string `json:"organization_name"` - UserID *string `json:"user_id"` ApiKeyID *string `json:"api_key_id"` ApiKeyName *string `json:"api_key_name"` ApiKeyRedacted *string `json:"api_key_redacted"` @@ -48,6 +47,7 @@ type DalleData struct { ImageSize string `json:"image_size"` Operation string `json:"operation"` ModelID string `json:"model_id"` + UserID string `json:"user_id"` } type WhisperData struct { @@ -56,6 +56,7 @@ type WhisperData struct { ModelID string `json:"model_id"` NumSeconds int `json:"num_seconds"` NumRequests int `json:"num_requests"` + UserID string `json:"user_id"` } type TtsData struct { @@ -64,4 +65,5 @@ type TtsData struct { ModelID string `json:"model_id"` NumCharacters int `json:"num_characters"` NumRequests int `json:"num_requests"` + UserID string `json:"user_id"` } diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index 018a9965d2ae..f3e604297c0c 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -214,31 +214,37 @@ func (m *MetricSet) processResponse(resp *http.Response, dateStr string) error { return nil } +func getBaseFields(data BaseData) mapstr.M { + return mapstr.M{ + "organization_id": data.OrganizationID, + "organization_name": data.OrganizationName, + "api_key_id": data.ApiKeyID, + "api_key_name": data.ApiKeyName, + "api_key_redacted": data.ApiKeyRedacted, + "api_key_type": data.ApiKeyType, + "project_id": data.ProjectID, + "project_name": data.ProjectName, + } +} + func (m *MetricSet) processUsageData(events []mb.Event, data []UsageData) { for _, usage := range data { event := mb.Event{ + Timestamp: time.Unix(usage.AggregationTimestamp, 0).UTC(), // epoch time to time.Time (UTC) MetricSetFields: mapstr.M{ "data": mapstr.M{ - "organization_id": usage.OrganizationID, - "organization_name": usage.OrganizationName, - "aggregation_timestamp": time.Unix(usage.AggregationTimestamp, 0).UTC(), // epoch time to time.Time (UTC) "n_requests": usage.NRequests, "operation": usage.Operation, "snapshot_id": usage.SnapshotID, "n_context_tokens_total": usage.NContextTokensTotal, "n_generated_tokens_total": usage.NGeneratedTokensTotal, "email": usage.Email, - "api_key_id": usage.ApiKeyID, - "api_key_name": usage.ApiKeyName, - "api_key_redacted": usage.ApiKeyRedacted, - "api_key_type": usage.ApiKeyType, - "project_id": usage.ProjectID, - "project_name": usage.ProjectName, "request_type": usage.RequestType, "n_cached_context_tokens_total": usage.NCachedContextTokensTotal, }, }, } + event.MetricSetFields.DeepUpdate(getBaseFields(usage.BaseData)) events = append(events, event) } m.processEvents(events) @@ -247,26 +253,19 @@ func (m *MetricSet) processUsageData(events []mb.Event, data []UsageData) { func (m *MetricSet) processDalleData(events []mb.Event, data []DalleData) { for _, dalle := range data { event := mb.Event{ + Timestamp: time.Unix(dalle.Timestamp, 0).UTC(), // epoch time to time.Time (UTC) MetricSetFields: mapstr.M{ "dalle": mapstr.M{ - "timestamp": time.Unix(dalle.Timestamp, 0).UTC(), // epoch time to time.Time (UTC) - "num_images": dalle.NumImages, - "num_requests": dalle.NumRequests, - "image_size": dalle.ImageSize, - "operation": dalle.Operation, - "user_id": dalle.UserID, - "organization_id": dalle.OrganizationID, - "api_key_id": dalle.ApiKeyID, - "api_key_name": dalle.ApiKeyName, - "api_key_redacted": dalle.ApiKeyRedacted, - "api_key_type": dalle.ApiKeyType, - "organization_name": dalle.OrganizationName, - "model_id": dalle.ModelID, - "project_id": dalle.ProjectID, - "project_name": dalle.ProjectName, + "num_images": dalle.NumImages, + "num_requests": dalle.NumRequests, + "image_size": dalle.ImageSize, + "operation": dalle.Operation, + "user_id": dalle.UserID, + "model_id": dalle.ModelID, }, }, } + event.MetricSetFields.DeepUpdate(getBaseFields(dalle.BaseData)) events = append(events, event) } m.processEvents(events) @@ -275,24 +274,17 @@ func (m *MetricSet) processDalleData(events []mb.Event, data []DalleData) { func (m *MetricSet) processWhisperData(events []mb.Event, data []WhisperData) { for _, whisper := range data { event := mb.Event{ + Timestamp: time.Unix(whisper.Timestamp, 0).UTC(), // epoch time to time.Time (UTC) MetricSetFields: mapstr.M{ "whisper": mapstr.M{ - "timestamp": time.Unix(whisper.Timestamp, 0).UTC(), // epoch time to time.Time (UTC) - "model_id": whisper.ModelID, - "num_seconds": whisper.NumSeconds, - "num_requests": whisper.NumRequests, - "user_id": whisper.UserID, - "organization_id": whisper.OrganizationID, - "api_key_id": whisper.ApiKeyID, - "api_key_name": whisper.ApiKeyName, - "api_key_redacted": whisper.ApiKeyRedacted, - "api_key_type": whisper.ApiKeyType, - "organization_name": whisper.OrganizationName, - "project_id": whisper.ProjectID, - "project_name": whisper.ProjectName, + "model_id": whisper.ModelID, + "num_seconds": whisper.NumSeconds, + "num_requests": whisper.NumRequests, + "user_id": whisper.UserID, }, }, } + event.MetricSetFields.DeepUpdate(getBaseFields(whisper.BaseData)) events = append(events, event) } m.processEvents(events) @@ -301,24 +293,17 @@ func (m *MetricSet) processWhisperData(events []mb.Event, data []WhisperData) { func (m *MetricSet) processTTSData(events []mb.Event, data []TtsData) { for _, tts := range data { event := mb.Event{ + Timestamp: time.Unix(tts.Timestamp, 0).UTC(), // epoch time to time.Time (UTC) MetricSetFields: mapstr.M{ "tts": mapstr.M{ - "timestamp": time.Unix(tts.Timestamp, 0).UTC(), // epoch time to time.Time (UTC) - "model_id": tts.ModelID, - "num_characters": tts.NumCharacters, - "num_requests": tts.NumRequests, - "user_id": tts.UserID, - "organization_id": tts.OrganizationID, - "api_key_id": tts.ApiKeyID, - "api_key_name": tts.ApiKeyName, - "api_key_redacted": tts.ApiKeyRedacted, - "api_key_type": tts.ApiKeyType, - "organization_name": tts.OrganizationName, - "project_id": tts.ProjectID, - "project_name": tts.ProjectName, + "model_id": tts.ModelID, + "num_characters": tts.NumCharacters, + "num_requests": tts.NumRequests, + "user_id": tts.UserID, }, }, } + event.MetricSetFields.DeepUpdate(getBaseFields(tts.BaseData)) events = append(events, event) } diff --git a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go index 6bf869dfea88..03d8f78240ad 100644 --- a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go +++ b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go @@ -15,20 +15,214 @@ import ( "github.com/stretchr/testify/require" + "github.com/elastic/beats/v7/metricbeat/mb" mbtest "github.com/elastic/beats/v7/metricbeat/mb/testing" + "github.com/elastic/elastic-agent-libs/mapstr" ) func TestFetch(t *testing.T) { - apiKey := time.Now().String() // to generate a unique API key everytime + apiKey := time.Now().String() // to generate a unique API key everytime to ignore the stateStore usagePath := "/usage" server := initServer(usagePath, apiKey) defer server.Close() - f := mbtest.NewReportingMetricSetV2Error(t, getConfig(server.URL+"/usage", apiKey)) + tests := []struct { + name string + expected mb.Event + }{ + { + name: "tc: #1", + expected: mb.Event{ + RootFields: nil, + ModuleFields: nil, + MetricSetFields: mapstr.M{ + "api_key_id": (*string)(nil), + "api_key_name": (*string)(nil), + "api_key_redacted": (*string)(nil), + "api_key_type": (*string)(nil), + "data": mapstr.M{ + "email": (*string)(nil), + "n_cached_context_tokens_total": 0, + "n_context_tokens_total": 118, + "n_generated_tokens_total": 35, + "n_requests": 1, + "operation": "completion-realtime", + "request_type": "", + "snapshot_id": "gpt-4o-realtime-preview-2024-10-01", + }, + "organization_id": "org-dummy", + "organization_name": "Personal", + "project_id": (*string)(nil), + "project_name": (*string)(nil), + }, + Index: "", + ID: "", + Namespace: "", + Timestamp: time.Date(2024, time.November, 4, 5, 1, 0, 0, time.UTC), + Error: error(nil), + Host: "", + Service: "", + Took: 0, + Period: 0, + DisableTimeSeries: false, + }, + }, + { + name: "tc: #2", + expected: mb.Event{ + RootFields: nil, + ModuleFields: nil, + MetricSetFields: mapstr.M{ + "api_key_id": (*string)(nil), + "api_key_name": (*string)(nil), + "api_key_redacted": (*string)(nil), + "api_key_type": (*string)(nil), + "data": mapstr.M{ + "email": (*string)(nil), + "n_cached_context_tokens_total": 0, + "n_context_tokens_total": 31, + "n_generated_tokens_total": 12, + "n_requests": 1, + "operation": "completion", + "request_type": "", + "snapshot_id": "gpt-4o-2024-08-06", + }, + "organization_id": "org-dummy", + "organization_name": "Personal", + "project_id": (*string)(nil), + "project_name": (*string)(nil), + }, + Index: "", + ID: "", + Namespace: "", + Timestamp: time.Date(2024, time.November, 4, 5, 1, 0, 0, time.UTC), + Error: error(nil), + Host: "", + Service: "", + Took: 0, + Period: 0, + DisableTimeSeries: false, + }, + }, + { + name: "tc: #3", + expected: mb.Event{ + RootFields: nil, + ModuleFields: nil, + MetricSetFields: mapstr.M{ + "api_key_id": (*string)(nil), + "api_key_name": (*string)(nil), + "api_key_redacted": (*string)(nil), + "api_key_type": (*string)(nil), + "data": mapstr.M{ + "email": (*string)(nil), + "n_cached_context_tokens_total": 0, + "n_context_tokens_total": 13, + "n_generated_tokens_total": 9, + "n_requests": 1, + "operation": "completion", + "request_type": "", + "snapshot_id": "ft:gpt-3.5-turbo-0125:personal:yay-renew:APjjyG8E:ckpt-step-84", + }, + "organization_id": "org-dummy", + "organization_name": "Personal", + "project_id": (*string)(nil), + "project_name": (*string)(nil), + }, + Index: "", + ID: "", + Namespace: "", + Timestamp: time.Date(2024, time.November, 4, 5, 19, 0, 0, time.UTC), + Error: error(nil), + Host: "", + Service: "", + Took: 0, + Period: 0, + DisableTimeSeries: false, + }, + }, + { + name: "tc: #4", + expected: mb.Event{ + RootFields: nil, + ModuleFields: nil, + MetricSetFields: mapstr.M{ + "api_key_id": ptr("key_sha_id_random"), + "api_key_name": ptr("project_key"), + "api_key_redacted": ptr("sk-...zkA"), + "api_key_type": ptr("organization"), + "dalle": mapstr.M{ + "image_size": "1024x1024", + "model_id": "dall-e-3", + "num_images": 1, + "num_requests": 1, + "operation": "generations", + "user_id": "subham.sarkar@elastic.co", + }, + "organization_id": "org-dummy", + "organization_name": "Personal", + "project_id": ptr("Default Project"), + "project_name": ptr("Default Project"), + }, + Index: "", + ID: "", + Namespace: "", + Timestamp: time.Date(2024, time.November, 4, 5, 1, 0, 0, time.UTC), + Error: error(nil), + Host: "", + Service: "", + Took: 0, + Period: 0, + DisableTimeSeries: false, + }, + }, + { + name: "tc: #5", + expected: mb.Event{ + RootFields: nil, + ModuleFields: nil, + MetricSetFields: mapstr.M{ + "api_key_id": (*string)(nil), + "api_key_name": (*string)(nil), + "api_key_redacted": (*string)(nil), + "api_key_type": (*string)(nil), + "whisper": mapstr.M{ + "model_id": "whisper-1", + "num_requests": 1, + "num_seconds": 2, + "user_id": "", + }, + "organization_id": "org-dummy", + "organization_name": "Personal", + "project_id": (*string)(nil), + "project_name": (*string)(nil), + }, + Index: "", + ID: "", + Namespace: "", + Timestamp: time.Date(2024, time.November, 4, 5, 1, 0, 0, time.UTC), + Error: error(nil), + Host: "", + Service: "", + Took: 0, + Period: 0, + DisableTimeSeries: false, + }, + }, + } + f := mbtest.NewReportingMetricSetV2Error(t, getConfig(server.URL+"/usage", apiKey)) events, errs := mbtest.ReportingFetchV2Error(f) + require.Empty(t, errs, "Expected no errors") require.NotEmpty(t, events, "Expected events to be returned") + require.Equal(t, len(tests), len(events)) + + for i, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.expected, events[i]) + }) + } } func TestData(t *testing.T) { @@ -48,7 +242,7 @@ func getConfig(url, apiKey string) map[string]interface{} { "module": "openai", "metricsets": []string{"usage"}, "enabled": true, - "period": "1h", + "period": "24h", "api_url": url, "api_keys": []map[string]interface{}{ {"key": apiKey}, @@ -58,7 +252,8 @@ func getConfig(url, apiKey string) map[string]interface{} { "burst": 5, }, "collection": map[string]interface{}{ - "lookback_days": 1, + "lookback_days": 0, + "realtime": false, }, } } @@ -134,8 +329,8 @@ func initServer(endpoint, api_key string) *httptest.Server { "image_size": "1024x1024", "operation": "generations", "user_id": "subham.sarkar@elastic.co", - "organization_id": "org-FCp10pUDIN4slA4kNZK6UKkX", - "api_key_id": "key_10xSzP3zPsz8zB5O", + "organization_id": "org-dummy", + "api_key_id": "key_sha_id_random", "api_key_name": "project_key", "api_key_redacted": "sk-...zkA", "api_key_type": "organization", From f6067b8dd848f2f7a021dd34baabf2429184c7ea Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Wed, 4 Dec 2024 14:33:14 +0530 Subject: [PATCH 23/36] Add more dummy data --- .../openai/usage/usage_integration_test.go | 103 +++++++++++++++++- 1 file changed, 100 insertions(+), 3 deletions(-) diff --git a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go index 03d8f78240ad..efa0ca96320c 100644 --- a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go +++ b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go @@ -157,7 +157,7 @@ func TestFetch(t *testing.T) { "num_images": 1, "num_requests": 1, "operation": "generations", - "user_id": "subham.sarkar@elastic.co", + "user_id": "hello-test@elastic.co", }, "organization_id": "org-dummy", "organization_name": "Personal", @@ -209,6 +209,72 @@ func TestFetch(t *testing.T) { DisableTimeSeries: false, }, }, + { + name: "tc: #6", + expected: mb.Event{ + RootFields: nil, + ModuleFields: nil, + MetricSetFields: mapstr.M{ + "api_key_id": ptr("key_fake_id"), + "api_key_name": ptr("project_key"), + "api_key_redacted": ptr("sk-...zkA"), + "api_key_type": ptr("organization"), + "tts": mapstr.M{ + "model_id": "tts-1", + "num_characters": 90, + "num_requests": 2, + "user_id": "hello-test@elastic.co", + }, + "organization_id": "org-fake", + "organization_name": "Personal", + "project_id": ptr("Default Project"), + "project_name": ptr("Default Project"), + }, + Index: "", + ID: "", + Namespace: "", + Timestamp: time.Date(2024, time.September, 4, 0, 0, 0, 0, time.UTC), + Error: error(nil), + Host: "", + Service: "", + Took: 0, + Period: 0, + DisableTimeSeries: false, + }, + }, + { + name: "tc: #7", + expected: mb.Event{ + RootFields: nil, + ModuleFields: nil, + MetricSetFields: mapstr.M{ + "api_key_id": ptr("key_fake_id"), + "api_key_name": ptr("fake_key"), + "api_key_redacted": ptr("sk-...FIA"), + "api_key_type": ptr("project"), + "tts": mapstr.M{ + "model_id": "tts-1", + "num_characters": 45, + "num_requests": 1, + "user_id": "hello-test@elastic.co", + }, + "organization_id": "org-fake", + "organization_name": "Personal", + "project_id": ptr("proj_fake_id"), + "project_name": ptr("fake_proj"), + }, + Index: "", + ID: "", + Namespace: "", + Timestamp: time.Date(2024, time.September, 5, 0, 0, 0, 0, time.UTC), + Error: error(nil), + Host: "", + Service: "", + Took: 0, + Period: 0, + DisableTimeSeries: false, + }, + }, } f := mbtest.NewReportingMetricSetV2Error(t, getConfig(server.URL+"/usage", apiKey)) @@ -328,7 +394,7 @@ func initServer(endpoint, api_key string) *httptest.Server { "num_requests": 1, "image_size": "1024x1024", "operation": "generations", - "user_id": "subham.sarkar@elastic.co", + "user_id": "hello-test@elastic.co", "organization_id": "org-dummy", "api_key_id": "key_sha_id_random", "api_key_name": "project_key", @@ -357,7 +423,38 @@ func initServer(endpoint, api_key string) *httptest.Server { "project_name": null } ], - "tts_api_data": [], + "tts_api_data": [ + { + "timestamp": 1725408000, + "model_id": "tts-1", + "num_characters": 90, + "num_requests": 2, + "user_id": "hello-test@elastic.co", + "organization_id": "org-fake", + "api_key_id": "key_fake_id", + "api_key_name": "project_key", + "api_key_redacted": "sk-...zkA", + "api_key_type": "organization", + "organization_name": "Personal", + "project_id": "Default Project", + "project_name": "Default Project" + }, + { + "timestamp": 1725494400, + "model_id": "tts-1", + "num_characters": 45, + "num_requests": 1, + "user_id": "hello-test@elastic.co", + "organization_id": "org-fake", + "api_key_id": "key_fake_id", + "api_key_name": "fake_key", + "api_key_redacted": "sk-...FIA", + "api_key_type": "project", + "organization_name": "Personal", + "project_id": "proj_fake_id", + "project_name": "fake_proj" + } + ], "assistant_code_interpreter_data": [], "retrieval_storage_data": [] }`) From ab983695bc817590c4edfd56beff4cfe5c400414 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Wed, 4 Dec 2024 14:38:44 +0530 Subject: [PATCH 24/36] Better prealloc --- x-pack/metricbeat/module/openai/usage/usage.go | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index f3e604297c0c..a4e67f1361aa 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -195,7 +195,11 @@ func (m *MetricSet) processResponse(resp *http.Response, dateStr string) error { m.logger.Infof("Fetched usage metrics for date: %s", dateStr) - events := make([]mb.Event, 0, len(usageResponse.Data)) + events := make([]mb.Event, 0, len(usageResponse.Data)+ + len(usageResponse.DalleApiData)+ + len(usageResponse.WhisperApiData)+ + len(usageResponse.TtsApiData), + ) m.processUsageData(events, usageResponse.Data) m.processDalleData(events, usageResponse.DalleApiData) From 120184533b591138ea9d5d7e053bdddad242393f Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Mon, 9 Dec 2024 02:00:23 +0530 Subject: [PATCH 25/36] Address review comments --- x-pack/metricbeat/module/openai/fields.go | 2 +- .../module/openai/usage/_meta/data.json | 8 +- .../module/openai/usage/_meta/fields.yml | 153 +++++------------- .../module/openai/usage/persistcache.go | 39 +++-- .../metricbeat/module/openai/usage/usage.go | 97 ++++++----- .../openai/usage/usage_integration_test.go | 72 ++++----- 6 files changed, 154 insertions(+), 217 deletions(-) diff --git a/x-pack/metricbeat/module/openai/fields.go b/x-pack/metricbeat/module/openai/fields.go index fa62e586ec53..de875dbdaab1 100644 --- a/x-pack/metricbeat/module/openai/fields.go +++ b/x-pack/metricbeat/module/openai/fields.go @@ -19,5 +19,5 @@ func init() { // AssetOpenai returns asset data. // This is the base64 encoded zlib format compressed contents of module/openai. func AssetOpenai() string { - return "eJzsms9uozwUxfd5iqvu8wJZfFL0/VOlzrRqM5olcvEN8QRsxr5Mmz79yAEnQIBA4qbtFC8hnN/h+uJjRKawxs0MVIqSiQkACYpxBlf5gasJgMYYmcEZPCKxCQBHE2qRklByBn9NAKC4GhLFsxgnAEuBMTez7bkpSJZgiWAHbVKcQaRVlhZHGlSrOmWtzLAId0eb5OyoG3ejEZWP2xTl/Brmd9c5AxIkLUIDTHIwxEgYEqEpXVO3WLbJWYXb7vSIKzv+R4maxYUtq+y81X7aZKgyDzpiUrwwCwoEP/id87jGzZPSTecrTm9LciA4ShJLgbof3h7yaqBR0KFZFGmMcjKJBA2xpD4Lezxn1OStwl44FVDLfE5KjFYfMtD4M0ND9anbw2Mlo2Pwr1nyiNqSnRwkjLffvkpRNxsbWnEntL2slWckS81KkYcmeyik+jSYDEIlCZ8pILVGaQJSxOKzKr2wCiB39S4AkAMgM3h4B3s70fbJJeSvZ2iHKCx1FYeFK+QXqNGWUytVqy9MmGjn9+ySbwZ1i9JuCUhFsMaNh5a0EbHGTZ+OdFAfq53Ddi90BVAjZyE1NOdA6H2h4+hHwY3rwkDoYpOibaRjzFSrHxj6WGXucqU+U+qgPqbUYTuntFjivVa20Nz9vGUPE8d14smbmH/mNzfTfw+3VwO3MK8Q3vVa1JEySwKRsAh9ZXYutl+5O8mvtF9oZW7NBUa8nN1rD+IFqwnVUsW32qVkBrWHtWObPUO3v5fffY/5dwr0XeffG75QJYpj7KGXvlidIbH7Z2V9Y+4+rYRJD1ydnLzfc7mPGL0X7zKbuAZDJbmvwC3UbCuFaDpfF98g7ccIPI08RmAH+JNE4CdKJDpYQk5OowU+05TU9CFFDFdjKPUMpXDFtH0Kta9o2AuO0TRG0xhNYzQ1Mt99NC0p8Pmt9T8hcUqZFDI670uriITs+K6iHu39NZzOTwTDHgP2BMua8e6qMWOEISYpCBXHQEhCnWokf2+dc0eAvxVHuN4Tzsr7yxd2VyqwpYJSqXrUWdubxF8sDgwpXf3nwlnlvXfCUAh/sKrqA//bYv4OAAD//+T56FM=" + return "eJzsWMtu4zgQvPsrGtnL7sHYuw8BjOwDAbKbIMlgjgIttmWOKVJDtuI4Xz+gJNp6UHZsyRMEGB5FsqpZrK6YmcIatzPQGSomJgAkSOIMrsoPVxMAgxKZxRkskNgEgKONjchIaDWD6wkAVLsh1TyXOAFYCpTczoq5KSiWYo3BDdpmOIPE6DyrvgRQmzh1rNyyBHdfQ3ButAv3I0hVjvsM1fwW5g+3JQekSEbEFpjiYImRsCRiW9vTLtGN3+BGp6lWsGAWqxXACIzWBBJfUDaW7xQyCVPijbnCIsEba/wZ17jdaNOea5zovgYDgqMisRRojlO6T6OQdoA8HctEtMbtgMO5m1nj9ti5PNGQI3mqo6cxyFlMePaZHqv9nvEgmUM9l+h5myHo5UGezOhvGNOAK3ooEY5dkScackWeqoHRbcVMolv/582KUdXXnDUiYV9WZ6IvXTrFXLcmAf5FhYbJGqXPk9bSUIjUizL4PUdLNiJNTHaW+RKlVklgslHl/3m6QONc4EEhZbx9A43YNkVX97KG76lDfO+BIGDhPZ9VLLMrHTDgyYxPFVSfE+u0sVaErxSRXqMaQ+dnhwBqp3ZFACUB5LaTF/tiksI3hPxy5ewoqoL6hWHxCnl0eX0KnpZMvVVhykQ//zv98cWi6UFqNV4odU+m8+lbYfbl1V/zu7vp3yBSlxnVNbmuCSXHPrSkDIfoGalVFdD5DXRiZqk8jYpDdG/xvLwqwfbO/Zi07GUtyouseBtslCfxhs0e7dHxozI6t2hGyOei/96RzanmKEfg+8/hhAhbHfh1JWyGBmyGGK+mpKdFJB1qwE25ZawW9BUM7cGLK9cmdE1vMdaKj9X1FZr7uRijPfRX82Pa/sK90LLmM76SM+RTYc2DlqRO0WfbscX6KV0Zr5hxrywzljH3gL+8WXhzzrlwG5gsnzqOw8Lvhm3AkjYswT+CRl1SNOaj6x+hcEq5EioZ8uTSRiRCHbgrvXBPz8B0ORGdJPEj28CyVXhY593/A6wVlpiiKNYcI6EITWaQxvsjNPcMcKM5wu2eYVD3/3xhd1KBkwpqUr1DZ+MOiS9MRpWJx5L30QP77vhkqppO/YWYPwIAAP//Z/rS7Q==" } diff --git a/x-pack/metricbeat/module/openai/usage/_meta/data.json b/x-pack/metricbeat/module/openai/usage/_meta/data.json index e80e9552ba7b..da78c73e036f 100644 --- a/x-pack/metricbeat/module/openai/usage/_meta/data.json +++ b/x-pack/metricbeat/module/openai/usage/_meta/data.json @@ -16,13 +16,13 @@ "api_key_redacted": null, "api_key_type": null, "data": { + "cached_context_tokens_total": 0, + "context_tokens_total": 118, "email": null, - "n_cached_context_tokens_total": 0, - "n_context_tokens_total": 118, - "n_generated_tokens_total": 35, - "n_requests": 1, + "generated_tokens_total": 35, "operation": "completion-realtime", "request_type": "", + "requests_total": 1, "snapshot_id": "gpt-4o-realtime-preview-2024-10-01" }, "organization_id": "org-dummy", diff --git a/x-pack/metricbeat/module/openai/usage/_meta/fields.yml b/x-pack/metricbeat/module/openai/usage/_meta/fields.yml index 203dacaced0d..c25fe30c17bb 100644 --- a/x-pack/metricbeat/module/openai/usage/_meta/fields.yml +++ b/x-pack/metricbeat/module/openai/usage/_meta/fields.yml @@ -4,21 +4,39 @@ description: > OpenAI API usage metrics and statistics fields: + # Common base fields at root level + - name: organization_id + type: keyword + description: Organization identifier + - name: organization_name + type: keyword + description: Organization name + - name: api_key_id + type: keyword + description: API key identifier + - name: api_key_name + type: keyword + description: API key name + - name: api_key_redacted + type: keyword + description: Redacted API key + - name: api_key_type + type: keyword + description: Type of API key + - name: project_id + type: keyword + description: Project identifier + - name: project_name + type: keyword + description: Project name + + # Completion/Chat usage data - name: data type: group description: > General usage data metrics fields: - - name: organization_id - type: keyword - description: Organization identifier - - name: organization_name - type: keyword - description: Organization name - - name: aggregation_timestamp - type: date - description: Timestamp of data aggregation - - name: n_requests + - name: requests_total type: long description: Number of requests made - name: operation @@ -27,52 +45,32 @@ - name: snapshot_id type: keyword description: Snapshot identifier - - name: n_context_tokens_total + - name: context_tokens_total type: long description: Total number of context tokens used - - name: n_generated_tokens_total + - name: generated_tokens_total type: long description: Total number of generated tokens - - name: n_cached_context_tokens_total + - name: cached_context_tokens_total type: long description: Total number of cached context tokens - name: email type: keyword description: User email - - name: api_key_id - type: keyword - description: API key identifier - - name: api_key_name - type: keyword - description: API key name - - name: api_key_redacted - type: keyword - description: Redacted API key - - name: api_key_type - type: keyword - description: Type of API key - - name: project_id - type: keyword - description: Project identifier - - name: project_name - type: keyword - description: Project name - name: request_type type: keyword description: Type of request + # DALL-E image generation metrics - name: dalle type: group description: > DALL-E API usage metrics fields: - - name: timestamp - type: date - description: Timestamp of request - name: num_images type: long description: Number of images generated - - name: num_requests + - name: requests_total type: long description: Number of requests - name: image_size @@ -84,124 +82,49 @@ - name: user_id type: keyword description: User identifier - - name: organization_id - type: keyword - description: Organization identifier - - name: api_key_id - type: keyword - description: API key identifier - - name: api_key_name - type: keyword - description: API key name - - name: api_key_redacted - type: keyword - description: Redacted API key - - name: api_key_type - type: keyword - description: Type of API key - - name: organization_name - type: keyword - description: Organization name - name: model_id type: keyword description: Model identifier - - name: project_id - type: keyword - description: Project identifier - - name: project_name - type: keyword - description: Project name + # Whisper speech-to-text metrics - name: whisper type: group description: > Whisper API usage metrics fields: - - name: timestamp - type: date - description: Timestamp of request - name: model_id type: keyword description: Model identifier - name: num_seconds type: long description: Number of seconds processed - - name: num_requests + - name: requests_total type: long description: Number of requests - name: user_id type: keyword description: User identifier - - name: organization_id - type: keyword - description: Organization identifier - - name: api_key_id - type: keyword - description: API key identifier - - name: api_key_name - type: keyword - description: API key name - - name: api_key_redacted - type: keyword - description: Redacted API key - - name: api_key_type - type: keyword - description: Type of API key - - name: organization_name - type: keyword - description: Organization name - - name: project_id - type: keyword - description: Project identifier - - name: project_name - type: keyword - description: Project name + # Text-to-Speech metrics - name: tts type: group description: > Text-to-Speech API usage metrics fields: - - name: timestamp - type: date - description: Timestamp of request - name: model_id type: keyword description: Model identifier - name: num_characters type: long description: Number of characters processed - - name: num_requests + - name: requests_total type: long description: Number of requests - name: user_id type: keyword description: User identifier - - name: organization_id - type: keyword - description: Organization identifier - - name: api_key_id - type: keyword - description: API key identifier - - name: api_key_name - type: keyword - description: API key name - - name: api_key_redacted - type: keyword - description: Redacted API key - - name: api_key_type - type: keyword - description: Type of API key - - name: organization_name - type: keyword - description: Organization name - - name: project_id - type: keyword - description: Project identifier - - name: project_name - type: keyword - description: Project name + # Additional data types (raw storage) - name: ft_data type: group description: > diff --git a/x-pack/metricbeat/module/openai/usage/persistcache.go b/x-pack/metricbeat/module/openai/usage/persistcache.go index 8fa8564292ec..362bdd9cab02 100644 --- a/x-pack/metricbeat/module/openai/usage/persistcache.go +++ b/x-pack/metricbeat/module/openai/usage/persistcache.go @@ -18,6 +18,7 @@ import ( // stateManager handles the storage and retrieval of state data with key hashing and caching capabilities type stateManager struct { + mu sync.RWMutex store *stateStore keyPrefix string hashCache sync.Map @@ -25,8 +26,8 @@ type stateManager struct { // stateStore handles persistence of key-value pairs using the filesystem type stateStore struct { - sync.RWMutex // Protects access to the state store - Dir string // Base directory for storing state files + Dir string // Base directory for storing state files + mu sync.RWMutex } // newStateManager creates a new state manager instance with the given storage path @@ -41,6 +42,7 @@ func newStateManager(storePath string) (*stateManager, error) { } return &stateManager{ + mu: sync.RWMutex{}, store: store, keyPrefix: "state_", hashCache: sync.Map{}, @@ -64,8 +66,8 @@ func (s *stateStore) getStatePath(name string) string { // Put stores a value in a file named by the key func (s *stateStore) Put(key, value string) error { - s.Lock() - defer s.Unlock() + s.mu.Lock() + defer s.mu.Unlock() filePath := s.getStatePath(key) @@ -90,8 +92,8 @@ func (s *stateStore) Put(key, value string) error { // Get retrieves the value stored in the file named by the key func (s *stateStore) Get(key string) (string, error) { - s.RLock() - defer s.RUnlock() + s.mu.RLock() + defer s.mu.RUnlock() filePath := s.getStatePath(key) data, err := os.ReadFile(filePath) @@ -103,8 +105,8 @@ func (s *stateStore) Get(key string) (string, error) { // Has checks if a state exists for the given key func (s *stateStore) Has(key string) bool { - s.RLock() - defer s.RUnlock() + s.mu.RLock() + defer s.mu.RUnlock() filePath := s.getStatePath(key) _, err := os.Stat(filePath) @@ -113,8 +115,8 @@ func (s *stateStore) Has(key string) bool { // Remove deletes the state file for the given key func (s *stateStore) Remove(key string) error { - s.Lock() - defer s.Unlock() + s.mu.Lock() + defer s.mu.Unlock() filePath := s.getStatePath(key) if err := os.Remove(filePath); err != nil && !os.IsNotExist(err) { @@ -125,8 +127,8 @@ func (s *stateStore) Remove(key string) error { // Clear removes all state files by deleting and recreating the state directory func (s *stateStore) Clear() error { - s.Lock() - defer s.Unlock() + s.mu.Lock() + defer s.mu.Unlock() if err := os.RemoveAll(s.Dir); err != nil { return fmt.Errorf("clearing state directory: %w", err) @@ -136,6 +138,9 @@ func (s *stateStore) Clear() error { // GetLastProcessedDate retrieves and parses the last processed date for a given API key func (s *stateManager) GetLastProcessedDate(apiKey string) (time.Time, error) { + s.mu.RLock() + defer s.mu.RUnlock() + stateKey := s.GetStateKey(apiKey) if !s.store.Has(stateKey) { @@ -150,6 +155,15 @@ func (s *stateManager) GetLastProcessedDate(apiKey string) (time.Time, error) { return time.Parse(dateFormatForStateStore, dateStr) } +// SaveState saves the last processed date for a given API key +func (sm *stateManager) SaveState(apiKey, dateStr string) error { + sm.mu.Lock() + defer sm.mu.Unlock() + + stateKey := sm.GetStateKey(apiKey) + return sm.store.Put(stateKey, dateStr) +} + // hashKey generates and caches a SHA-256 hash of the provided API key func (s *stateManager) hashKey(apiKey string) string { // Check cache first to avoid recomputing hashes @@ -167,6 +181,7 @@ func (s *stateManager) hashKey(apiKey string) string { return hashedKey } +// GetStateKey generates a unique state key for a given API key func (s *stateManager) GetStateKey(apiKey string) string { return s.keyPrefix + s.hashKey(apiKey) } diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index a4e67f1361aa..eea6b860bb20 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -117,9 +117,6 @@ func (m *MetricSet) Fetch(report mb.ReporterV2) error { // 5. Handles errors per day without failing entire range func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLHTTPClient) error { for _, apiKey := range m.config.APIKeys { - // stateKey using stateManager's key prefix and hashing apiKey - stateKey := m.stateManager.GetStateKey(apiKey.Key) - lastProcessedDate, err := m.stateManager.GetLastProcessedDate(apiKey.Key) if err == nil { // We have previous state, adjust start date @@ -135,7 +132,8 @@ func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLH m.logger.Errorf("Error fetching data for date %s: %v", dateStr, err) continue } - if err := m.stateManager.store.Put(stateKey, dateStr); err != nil { + + if err := m.stateManager.SaveState(apiKey.Key, dateStr); err != nil { m.logger.Errorf("Error storing state for API key: %v", err) } } @@ -193,27 +191,21 @@ func (m *MetricSet) processResponse(resp *http.Response, dateStr string) error { return fmt.Errorf("error decoding response: %w", err) } - m.logger.Infof("Fetched usage metrics for date: %s", dateStr) - - events := make([]mb.Event, 0, len(usageResponse.Data)+ - len(usageResponse.DalleApiData)+ - len(usageResponse.WhisperApiData)+ - len(usageResponse.TtsApiData), - ) + m.logger.Infof("Fetching usage metrics for date: %s", dateStr) - m.processUsageData(events, usageResponse.Data) - m.processDalleData(events, usageResponse.DalleApiData) - m.processWhisperData(events, usageResponse.WhisperApiData) - m.processTTSData(events, usageResponse.TtsApiData) + m.processUsageData(usageResponse.Data) + m.processDalleData(usageResponse.DalleApiData) + m.processWhisperData(usageResponse.WhisperApiData) + m.processTTSData(usageResponse.TtsApiData) // Process additional data. // // NOTE(shmsr): During testing, could not get the usage data for the following // and found no documentation, example responses, etc. That's why let's store them // as it is so that we can use processors later on to process them as needed. - m.processFTData(events, usageResponse.FtData) - m.processAssistantCodeInterpreterData(events, usageResponse.AssistantCodeInterpreterData) - m.processRetrievalStorageData(events, usageResponse.RetrievalStorageData) + m.processFTData(usageResponse.FtData) + m.processAssistantCodeInterpreterData(usageResponse.AssistantCodeInterpreterData) + m.processRetrievalStorageData(usageResponse.RetrievalStorageData) return nil } @@ -231,20 +223,21 @@ func getBaseFields(data BaseData) mapstr.M { } } -func (m *MetricSet) processUsageData(events []mb.Event, data []UsageData) { +func (m *MetricSet) processUsageData(data []UsageData) { + events := make([]mb.Event, 0, len(data)) for _, usage := range data { event := mb.Event{ Timestamp: time.Unix(usage.AggregationTimestamp, 0).UTC(), // epoch time to time.Time (UTC) MetricSetFields: mapstr.M{ "data": mapstr.M{ - "n_requests": usage.NRequests, - "operation": usage.Operation, - "snapshot_id": usage.SnapshotID, - "n_context_tokens_total": usage.NContextTokensTotal, - "n_generated_tokens_total": usage.NGeneratedTokensTotal, - "email": usage.Email, - "request_type": usage.RequestType, - "n_cached_context_tokens_total": usage.NCachedContextTokensTotal, + "requests_total": usage.NRequests, + "operation": usage.Operation, + "snapshot_id": usage.SnapshotID, + "context_tokens_total": usage.NContextTokensTotal, + "generated_tokens_total": usage.NGeneratedTokensTotal, + "email": usage.Email, + "request_type": usage.RequestType, + "cached_context_tokens_total": usage.NCachedContextTokensTotal, }, }, } @@ -254,18 +247,19 @@ func (m *MetricSet) processUsageData(events []mb.Event, data []UsageData) { m.processEvents(events) } -func (m *MetricSet) processDalleData(events []mb.Event, data []DalleData) { +func (m *MetricSet) processDalleData(data []DalleData) { + events := make([]mb.Event, 0, len(data)) for _, dalle := range data { event := mb.Event{ Timestamp: time.Unix(dalle.Timestamp, 0).UTC(), // epoch time to time.Time (UTC) MetricSetFields: mapstr.M{ "dalle": mapstr.M{ - "num_images": dalle.NumImages, - "num_requests": dalle.NumRequests, - "image_size": dalle.ImageSize, - "operation": dalle.Operation, - "user_id": dalle.UserID, - "model_id": dalle.ModelID, + "num_images": dalle.NumImages, + "requests_total": dalle.NumRequests, + "image_size": dalle.ImageSize, + "operation": dalle.Operation, + "user_id": dalle.UserID, + "model_id": dalle.ModelID, }, }, } @@ -275,16 +269,17 @@ func (m *MetricSet) processDalleData(events []mb.Event, data []DalleData) { m.processEvents(events) } -func (m *MetricSet) processWhisperData(events []mb.Event, data []WhisperData) { +func (m *MetricSet) processWhisperData(data []WhisperData) { + events := make([]mb.Event, 0, len(data)) for _, whisper := range data { event := mb.Event{ Timestamp: time.Unix(whisper.Timestamp, 0).UTC(), // epoch time to time.Time (UTC) MetricSetFields: mapstr.M{ "whisper": mapstr.M{ - "model_id": whisper.ModelID, - "num_seconds": whisper.NumSeconds, - "num_requests": whisper.NumRequests, - "user_id": whisper.UserID, + "model_id": whisper.ModelID, + "num_seconds": whisper.NumSeconds, + "requests_total": whisper.NumRequests, + "user_id": whisper.UserID, }, }, } @@ -294,7 +289,8 @@ func (m *MetricSet) processWhisperData(events []mb.Event, data []WhisperData) { m.processEvents(events) } -func (m *MetricSet) processTTSData(events []mb.Event, data []TtsData) { +func (m *MetricSet) processTTSData(data []TtsData) { + events := make([]mb.Event, 0, len(data)) for _, tts := range data { event := mb.Event{ Timestamp: time.Unix(tts.Timestamp, 0).UTC(), // epoch time to time.Time (UTC) @@ -302,7 +298,7 @@ func (m *MetricSet) processTTSData(events []mb.Event, data []TtsData) { "tts": mapstr.M{ "model_id": tts.ModelID, "num_characters": tts.NumCharacters, - "num_requests": tts.NumRequests, + "requests_total": tts.NumRequests, "user_id": tts.UserID, }, }, @@ -314,7 +310,8 @@ func (m *MetricSet) processTTSData(events []mb.Event, data []TtsData) { m.processEvents(events) } -func (m *MetricSet) processFTData(events []mb.Event, data []interface{}) { +func (m *MetricSet) processFTData(data []interface{}) { + events := make([]mb.Event, 0, len(data)) for _, ft := range data { event := mb.Event{ MetricSetFields: mapstr.M{ @@ -328,7 +325,8 @@ func (m *MetricSet) processFTData(events []mb.Event, data []interface{}) { m.processEvents(events) } -func (m *MetricSet) processAssistantCodeInterpreterData(events []mb.Event, data []interface{}) { +func (m *MetricSet) processAssistantCodeInterpreterData(data []interface{}) { + events := make([]mb.Event, 0, len(data)) for _, aci := range data { event := mb.Event{ MetricSetFields: mapstr.M{ @@ -342,7 +340,8 @@ func (m *MetricSet) processAssistantCodeInterpreterData(events []mb.Event, data m.processEvents(events) } -func (m *MetricSet) processRetrievalStorageData(events []mb.Event, data []interface{}) { +func (m *MetricSet) processRetrievalStorageData(data []interface{}) { + events := make([]mb.Event, 0, len(data)) for _, rs := range data { event := mb.Event{ MetricSetFields: mapstr.M{ @@ -357,10 +356,10 @@ func (m *MetricSet) processRetrievalStorageData(events []mb.Event, data []interf } func (m *MetricSet) processEvents(events []mb.Event) { - if len(events) > 0 { - for i := range events { - m.report.Event(events[i]) - } + if len(events) == 0 { + return + } + for i := range events { + m.report.Event(events[i]) } - clear(events) } diff --git a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go index efa0ca96320c..8cb842419955 100644 --- a/x-pack/metricbeat/module/openai/usage/usage_integration_test.go +++ b/x-pack/metricbeat/module/openai/usage/usage_integration_test.go @@ -41,14 +41,14 @@ func TestFetch(t *testing.T) { "api_key_redacted": (*string)(nil), "api_key_type": (*string)(nil), "data": mapstr.M{ - "email": (*string)(nil), - "n_cached_context_tokens_total": 0, - "n_context_tokens_total": 118, - "n_generated_tokens_total": 35, - "n_requests": 1, - "operation": "completion-realtime", - "request_type": "", - "snapshot_id": "gpt-4o-realtime-preview-2024-10-01", + "email": (*string)(nil), + "cached_context_tokens_total": 0, + "context_tokens_total": 118, + "generated_tokens_total": 35, + "requests_total": 1, + "operation": "completion-realtime", + "request_type": "", + "snapshot_id": "gpt-4o-realtime-preview-2024-10-01", }, "organization_id": "org-dummy", "organization_name": "Personal", @@ -78,14 +78,14 @@ func TestFetch(t *testing.T) { "api_key_redacted": (*string)(nil), "api_key_type": (*string)(nil), "data": mapstr.M{ - "email": (*string)(nil), - "n_cached_context_tokens_total": 0, - "n_context_tokens_total": 31, - "n_generated_tokens_total": 12, - "n_requests": 1, - "operation": "completion", - "request_type": "", - "snapshot_id": "gpt-4o-2024-08-06", + "email": (*string)(nil), + "cached_context_tokens_total": 0, + "context_tokens_total": 31, + "generated_tokens_total": 12, + "requests_total": 1, + "operation": "completion", + "request_type": "", + "snapshot_id": "gpt-4o-2024-08-06", }, "organization_id": "org-dummy", "organization_name": "Personal", @@ -115,14 +115,14 @@ func TestFetch(t *testing.T) { "api_key_redacted": (*string)(nil), "api_key_type": (*string)(nil), "data": mapstr.M{ - "email": (*string)(nil), - "n_cached_context_tokens_total": 0, - "n_context_tokens_total": 13, - "n_generated_tokens_total": 9, - "n_requests": 1, - "operation": "completion", - "request_type": "", - "snapshot_id": "ft:gpt-3.5-turbo-0125:personal:yay-renew:APjjyG8E:ckpt-step-84", + "email": (*string)(nil), + "cached_context_tokens_total": 0, + "context_tokens_total": 13, + "generated_tokens_total": 9, + "requests_total": 1, + "operation": "completion", + "request_type": "", + "snapshot_id": "ft:gpt-3.5-turbo-0125:personal:yay-renew:APjjyG8E:ckpt-step-84", }, "organization_id": "org-dummy", "organization_name": "Personal", @@ -152,12 +152,12 @@ func TestFetch(t *testing.T) { "api_key_redacted": ptr("sk-...zkA"), "api_key_type": ptr("organization"), "dalle": mapstr.M{ - "image_size": "1024x1024", - "model_id": "dall-e-3", - "num_images": 1, - "num_requests": 1, - "operation": "generations", - "user_id": "hello-test@elastic.co", + "image_size": "1024x1024", + "model_id": "dall-e-3", + "num_images": 1, + "requests_total": 1, + "operation": "generations", + "user_id": "hello-test@elastic.co", }, "organization_id": "org-dummy", "organization_name": "Personal", @@ -187,10 +187,10 @@ func TestFetch(t *testing.T) { "api_key_redacted": (*string)(nil), "api_key_type": (*string)(nil), "whisper": mapstr.M{ - "model_id": "whisper-1", - "num_requests": 1, - "num_seconds": 2, - "user_id": "", + "model_id": "whisper-1", + "requests_total": 1, + "num_seconds": 2, + "user_id": "", }, "organization_id": "org-dummy", "organization_name": "Personal", @@ -222,7 +222,7 @@ func TestFetch(t *testing.T) { "tts": mapstr.M{ "model_id": "tts-1", "num_characters": 90, - "num_requests": 2, + "requests_total": 2, "user_id": "hello-test@elastic.co", }, "organization_id": "org-fake", @@ -255,7 +255,7 @@ func TestFetch(t *testing.T) { "tts": mapstr.M{ "model_id": "tts-1", "num_characters": 45, - "num_requests": 1, + "requests_total": 1, "user_id": "hello-test@elastic.co", }, "organization_id": "org-fake", From 90c2680e0141f0c6b93eb0114bb91a1fc9444af7 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Mon, 9 Dec 2024 11:46:44 +0530 Subject: [PATCH 26/36] make update --- metricbeat/docs/fields.asciidoc | 338 ++++---------------------------- 1 file changed, 43 insertions(+), 295 deletions(-) diff --git a/metricbeat/docs/fields.asciidoc b/metricbeat/docs/fields.asciidoc index 1705b79fe707..cfd7fa17af84 100644 --- a/metricbeat/docs/fields.asciidoc +++ b/metricbeat/docs/fields.asciidoc @@ -56675,14 +56675,7 @@ OpenAI API usage metrics and statistics -[float] -=== data - -General usage data metrics - - - -*`openai.usage.data.organization_id`*:: +*`openai.usage.organization_id`*:: + -- Organization identifier @@ -56691,7 +56684,7 @@ type: keyword -- -*`openai.usage.data.organization_name`*:: +*`openai.usage.organization_name`*:: + -- Organization name @@ -56700,127 +56693,125 @@ type: keyword -- -*`openai.usage.data.aggregation_timestamp`*:: +*`openai.usage.api_key_id`*:: + -- -Timestamp of data aggregation +API key identifier -type: date +type: keyword -- -*`openai.usage.data.n_requests`*:: +*`openai.usage.api_key_name`*:: + -- -Number of requests made +API key name -type: long +type: keyword -- -*`openai.usage.data.operation`*:: +*`openai.usage.api_key_redacted`*:: + -- -Operation type +Redacted API key type: keyword -- -*`openai.usage.data.snapshot_id`*:: +*`openai.usage.api_key_type`*:: + -- -Snapshot identifier +Type of API key type: keyword -- -*`openai.usage.data.n_context_tokens_total`*:: +*`openai.usage.project_id`*:: + -- -Total number of context tokens used +Project identifier -type: long +type: keyword -- -*`openai.usage.data.n_generated_tokens_total`*:: +*`openai.usage.project_name`*:: + -- -Total number of generated tokens +Project name -type: long +type: keyword -- -*`openai.usage.data.n_cached_context_tokens_total`*:: -+ --- -Total number of cached context tokens +[float] +=== data -type: long +General usage data metrics --- -*`openai.usage.data.email`*:: + +*`openai.usage.data.requests_total`*:: + -- -User email +Number of requests made -type: keyword +type: long -- -*`openai.usage.data.api_key_id`*:: +*`openai.usage.data.operation`*:: + -- -API key identifier +Operation type type: keyword -- -*`openai.usage.data.api_key_name`*:: +*`openai.usage.data.snapshot_id`*:: + -- -API key name +Snapshot identifier type: keyword -- -*`openai.usage.data.api_key_redacted`*:: +*`openai.usage.data.context_tokens_total`*:: + -- -Redacted API key +Total number of context tokens used -type: keyword +type: long -- -*`openai.usage.data.api_key_type`*:: +*`openai.usage.data.generated_tokens_total`*:: + -- -Type of API key +Total number of generated tokens -type: keyword +type: long -- -*`openai.usage.data.project_id`*:: +*`openai.usage.data.cached_context_tokens_total`*:: + -- -Project identifier +Total number of cached context tokens -type: keyword +type: long -- -*`openai.usage.data.project_name`*:: +*`openai.usage.data.email`*:: + -- -Project name +User email type: keyword @@ -56842,15 +56833,6 @@ DALL-E API usage metrics -*`openai.usage.dalle.timestamp`*:: -+ --- -Timestamp of request - -type: date - --- - *`openai.usage.dalle.num_images`*:: + -- @@ -56860,7 +56842,7 @@ type: long -- -*`openai.usage.dalle.num_requests`*:: +*`openai.usage.dalle.requests_total`*:: + -- Number of requests @@ -56896,60 +56878,6 @@ type: keyword -- -*`openai.usage.dalle.organization_id`*:: -+ --- -Organization identifier - -type: keyword - --- - -*`openai.usage.dalle.api_key_id`*:: -+ --- -API key identifier - -type: keyword - --- - -*`openai.usage.dalle.api_key_name`*:: -+ --- -API key name - -type: keyword - --- - -*`openai.usage.dalle.api_key_redacted`*:: -+ --- -Redacted API key - -type: keyword - --- - -*`openai.usage.dalle.api_key_type`*:: -+ --- -Type of API key - -type: keyword - --- - -*`openai.usage.dalle.organization_name`*:: -+ --- -Organization name - -type: keyword - --- - *`openai.usage.dalle.model_id`*:: + -- @@ -56959,24 +56887,6 @@ type: keyword -- -*`openai.usage.dalle.project_id`*:: -+ --- -Project identifier - -type: keyword - --- - -*`openai.usage.dalle.project_name`*:: -+ --- -Project name - -type: keyword - --- - [float] === whisper @@ -56984,15 +56894,6 @@ Whisper API usage metrics -*`openai.usage.whisper.timestamp`*:: -+ --- -Timestamp of request - -type: date - --- - *`openai.usage.whisper.model_id`*:: + -- @@ -57011,7 +56912,7 @@ type: long -- -*`openai.usage.whisper.num_requests`*:: +*`openai.usage.whisper.requests_total`*:: + -- Number of requests @@ -57029,78 +56930,6 @@ type: keyword -- -*`openai.usage.whisper.organization_id`*:: -+ --- -Organization identifier - -type: keyword - --- - -*`openai.usage.whisper.api_key_id`*:: -+ --- -API key identifier - -type: keyword - --- - -*`openai.usage.whisper.api_key_name`*:: -+ --- -API key name - -type: keyword - --- - -*`openai.usage.whisper.api_key_redacted`*:: -+ --- -Redacted API key - -type: keyword - --- - -*`openai.usage.whisper.api_key_type`*:: -+ --- -Type of API key - -type: keyword - --- - -*`openai.usage.whisper.organization_name`*:: -+ --- -Organization name - -type: keyword - --- - -*`openai.usage.whisper.project_id`*:: -+ --- -Project identifier - -type: keyword - --- - -*`openai.usage.whisper.project_name`*:: -+ --- -Project name - -type: keyword - --- - [float] === tts @@ -57108,15 +56937,6 @@ Text-to-Speech API usage metrics -*`openai.usage.tts.timestamp`*:: -+ --- -Timestamp of request - -type: date - --- - *`openai.usage.tts.model_id`*:: + -- @@ -57135,7 +56955,7 @@ type: long -- -*`openai.usage.tts.num_requests`*:: +*`openai.usage.tts.requests_total`*:: + -- Number of requests @@ -57153,78 +56973,6 @@ type: keyword -- -*`openai.usage.tts.organization_id`*:: -+ --- -Organization identifier - -type: keyword - --- - -*`openai.usage.tts.api_key_id`*:: -+ --- -API key identifier - -type: keyword - --- - -*`openai.usage.tts.api_key_name`*:: -+ --- -API key name - -type: keyword - --- - -*`openai.usage.tts.api_key_redacted`*:: -+ --- -Redacted API key - -type: keyword - --- - -*`openai.usage.tts.api_key_type`*:: -+ --- -Type of API key - -type: keyword - --- - -*`openai.usage.tts.organization_name`*:: -+ --- -Organization name - -type: keyword - --- - -*`openai.usage.tts.project_id`*:: -+ --- -Project identifier - -type: keyword - --- - -*`openai.usage.tts.project_name`*:: -+ --- -Project name - -type: keyword - --- - [float] === ft_data From e1ea8d2c68f23d10af06024564467119529c1a5f Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Mon, 9 Dec 2024 12:38:02 +0530 Subject: [PATCH 27/36] Include OpenAI with Agentbeat --- x-pack/agentbeat/agentbeat.spec.yml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/x-pack/agentbeat/agentbeat.spec.yml b/x-pack/agentbeat/agentbeat.spec.yml index a2af5ce6bf8f..5b6c33e4ceb1 100644 --- a/x-pack/agentbeat/agentbeat.spec.yml +++ b/x-pack/agentbeat/agentbeat.spec.yml @@ -507,6 +507,11 @@ inputs: platforms: *platforms outputs: *outputs command: *metricbeat_command + - name: openai/metrics + description: "OpenAI metrics" + platforms: *platforms + outputs: *outputs + command: *metricbeat_command - name: panw/metrics description: "Palo Alto Networks metrics" platforms: *platforms From e907d984eea7fdc33021ffbd07519fdd3f327226 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 10 Dec 2024 03:00:58 +0530 Subject: [PATCH 28/36] More changes --- .../metricbeat/module/openai/_meta/config.yml | 4 +- .../metricbeat/module/openai/usage/config.go | 4 +- .../metricbeat/module/openai/usage/usage.go | 59 ++++++++++++------- .../metricbeat/modules.d/openai.yml.disabled | 37 ++++++++++++ 4 files changed, 79 insertions(+), 25 deletions(-) create mode 100644 x-pack/metricbeat/modules.d/openai.yml.disabled diff --git a/x-pack/metricbeat/module/openai/_meta/config.yml b/x-pack/metricbeat/module/openai/_meta/config.yml index c3ab46a59932..52636fe191db 100644 --- a/x-pack/metricbeat/module/openai/_meta/config.yml +++ b/x-pack/metricbeat/module/openai/_meta/config.yml @@ -17,8 +17,8 @@ # - "k2: v2" ## Rate Limiting Configuration # rate_limit: - # limit: 60 # requests per second - # burst: 5 # burst size + # limit: 12 # requests per second + # burst: 1 # burst size # ## Request timeout duration # timeout: 30s diff --git a/x-pack/metricbeat/module/openai/usage/config.go b/x-pack/metricbeat/module/openai/usage/config.go index c82bb6271149..f930a0137476 100644 --- a/x-pack/metricbeat/module/openai/usage/config.go +++ b/x-pack/metricbeat/module/openai/usage/config.go @@ -39,8 +39,8 @@ func defaultConfig() Config { APIURL: "https://api.openai.com/v1/usage", Timeout: 30 * time.Second, RateLimit: &rateLimitConfig{ - Limit: ptr(60), - Burst: ptr(5), + Limit: ptr(12), + Burst: ptr(1), }, Collection: collectionConfig{ LookbackDays: 0, // 0 days diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index eea6b860bb20..d8347c2754aa 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -14,6 +14,7 @@ import ( "path" "time" + "golang.org/x/sync/errgroup" "golang.org/x/time/rate" "github.com/elastic/beats/v7/libbeat/common/cfgwarn" @@ -116,33 +117,49 @@ func (m *MetricSet) Fetch(report mb.ReporterV2) error { // 4. Updates state store with latest processed date // 5. Handles errors per day without failing entire range func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLHTTPClient) error { - for _, apiKey := range m.config.APIKeys { - lastProcessedDate, err := m.stateManager.GetLastProcessedDate(apiKey.Key) - if err == nil { - // We have previous state, adjust start date - startDate = lastProcessedDate.AddDate(0, 0, 1) - if startDate.After(endDate) { - continue + g, ctx := errgroup.WithContext(context.TODO()) + + for i := range m.config.APIKeys { + apiKey := m.config.APIKeys[i] + apiKeyIdx := i + 1 + g.Go(func() error { + lastProcessedDate, err := m.stateManager.GetLastProcessedDate(apiKey.Key) + if err == nil { + currentStartDate := lastProcessedDate.AddDate(0, 0, 1) + if currentStartDate.After(endDate) { + return nil + } + startDate = currentStartDate } - } - for d := startDate; !d.After(endDate); d = d.AddDate(0, 0, 1) { - dateStr := d.Format("2006-01-02") - if err := m.fetchSingleDay(dateStr, apiKey.Key, httpClient); err != nil { - m.logger.Errorf("Error fetching data for date %s: %v", dateStr, err) - continue + for d := startDate; !d.After(endDate); d = d.AddDate(0, 0, 1) { + select { + case <-ctx.Done(): + return ctx.Err() + default: + dateStr := d.Format("2006-01-02") + if err := m.fetchSingleDay(apiKeyIdx, dateStr, apiKey.Key, httpClient); err != nil { + m.logger.Errorf("Error fetching data (api key #%d) for date %s: %v", apiKeyIdx, dateStr, err) + continue + } + if err := m.stateManager.SaveState(apiKey.Key, dateStr); err != nil { + m.logger.Errorf("Error storing state for API key: %v at index %d", err, apiKeyIdx) + } + } } + return nil + }) + } - if err := m.stateManager.SaveState(apiKey.Key, dateStr); err != nil { - m.logger.Errorf("Error storing state for API key: %v", err) - } - } + if err := g.Wait(); err != nil { + m.logger.Errorf("Error fetching data: %v", err) } + return nil } // fetchSingleDay retrieves usage data for a specific date and API key. -func (m *MetricSet) fetchSingleDay(dateStr, apiKey string, httpClient *RLHTTPClient) error { +func (m *MetricSet) fetchSingleDay(apiKeyIdx int, dateStr, apiKey string, httpClient *RLHTTPClient) error { req, err := m.createRequest(dateStr, apiKey) if err != nil { return fmt.Errorf("error creating request: %w", err) @@ -162,7 +179,7 @@ func (m *MetricSet) fetchSingleDay(dateStr, apiKey string, httpClient *RLHTTPCli return fmt.Errorf("error response from API: status=%s", resp.Status) } - return m.processResponse(resp, dateStr) + return m.processResponse(apiKeyIdx, resp, dateStr) } // createRequest builds an HTTP request for the OpenAI usage API. @@ -185,13 +202,13 @@ func (m *MetricSet) createRequest(dateStr, apiKey string) (*http.Request, error) } // processResponse handles the API response and processes the usage data. -func (m *MetricSet) processResponse(resp *http.Response, dateStr string) error { +func (m *MetricSet) processResponse(apiKeyIdx int, resp *http.Response, dateStr string) error { var usageResponse UsageResponse if err := json.NewDecoder(resp.Body).Decode(&usageResponse); err != nil { return fmt.Errorf("error decoding response: %w", err) } - m.logger.Infof("Fetching usage metrics for date: %s", dateStr) + m.logger.Infof("Fetching usage metrics (api key #%d) for date: %s", apiKeyIdx, dateStr) m.processUsageData(usageResponse.Data) m.processDalleData(usageResponse.DalleApiData) diff --git a/x-pack/metricbeat/modules.d/openai.yml.disabled b/x-pack/metricbeat/modules.d/openai.yml.disabled new file mode 100644 index 000000000000..0ca3855f7b71 --- /dev/null +++ b/x-pack/metricbeat/modules.d/openai.yml.disabled @@ -0,0 +1,37 @@ +# Module: openai +# Docs: https://www.elastic.co/guide/en/beats/metricbeat/main/metricbeat-module-openai.html + +- module: openai + metricsets: ["usage"] + enabled: false + period: 1h + + # # Project API Keys - Multiple API keys can be specified for different projects + # api_keys: + # - key: "api_key1" + # - key: "api_key2" + + # # API Configuration + # ## Base URL for the OpenAI usage API endpoint + # api_url: "https://api.openai.com/v1/usage" + # ## Custom headers to be included in API requests + # headers: + # - "k1: v1" + # - "k2: v2" + ## Rate Limiting Configuration + # rate_limit: + # limit: 12 # requests per second + # burst: 1 # burst size + # ## Request timeout duration + # timeout: 30s + + # # Data Collection Configuration + # collection: + # ## Number of days to look back when collecting usage data + # lookback_days: 30 + # ## Whether to collect usage data in realtime. Defaults to false as how + # # OpenAI usage data is collected will end up adding duplicate data to ES + # # and also making it harder to do analytics. Best approach is to avoid + # # realtime collection and collect only upto last day (in UTC). So, there's + # # at most 24h delay. + # realtime: false \ No newline at end of file From 942454de8f767bd90e8d3180c01515e60abc3538 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 10 Dec 2024 03:06:28 +0530 Subject: [PATCH 29/36] make check --- x-pack/metricbeat/metricbeat.reference.yml | 6 +++--- x-pack/metricbeat/module/openai/_meta/config.yml | 6 +++--- x-pack/metricbeat/modules.d/openai.yml.disabled | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/x-pack/metricbeat/metricbeat.reference.yml b/x-pack/metricbeat/metricbeat.reference.yml index 4cf44666d39b..8bf85bd4d7b5 100644 --- a/x-pack/metricbeat/metricbeat.reference.yml +++ b/x-pack/metricbeat/metricbeat.reference.yml @@ -1296,9 +1296,9 @@ metricbeat.modules: # - "k2: v2" ## Rate Limiting Configuration # rate_limit: - # limit: 60 # requests per second - # burst: 5 # burst size - # ## Request timeout duration + # limit: 12 # seconds between requests + # burst: 1 # max concurrent requests + # ## Request Timeout Duration # timeout: 30s # # Data Collection Configuration diff --git a/x-pack/metricbeat/module/openai/_meta/config.yml b/x-pack/metricbeat/module/openai/_meta/config.yml index 52636fe191db..b4048313a1c9 100644 --- a/x-pack/metricbeat/module/openai/_meta/config.yml +++ b/x-pack/metricbeat/module/openai/_meta/config.yml @@ -17,9 +17,9 @@ # - "k2: v2" ## Rate Limiting Configuration # rate_limit: - # limit: 12 # requests per second - # burst: 1 # burst size - # ## Request timeout duration + # limit: 12 # seconds between requests + # burst: 1 # max concurrent requests + # ## Request Timeout Duration # timeout: 30s # # Data Collection Configuration diff --git a/x-pack/metricbeat/modules.d/openai.yml.disabled b/x-pack/metricbeat/modules.d/openai.yml.disabled index 0ca3855f7b71..b70f4127264b 100644 --- a/x-pack/metricbeat/modules.d/openai.yml.disabled +++ b/x-pack/metricbeat/modules.d/openai.yml.disabled @@ -20,9 +20,9 @@ # - "k2: v2" ## Rate Limiting Configuration # rate_limit: - # limit: 12 # requests per second - # burst: 1 # burst size - # ## Request timeout duration + # limit: 12 # seconds between requests + # burst: 1 # max concurrent requests + # ## Request Timeout Duration # timeout: 30s # # Data Collection Configuration From be9b5f2527e5ed5e2babb9427251b3b332d6a6a6 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 10 Dec 2024 11:35:21 +0530 Subject: [PATCH 30/36] make update --- metricbeat/docs/modules/openai.asciidoc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/metricbeat/docs/modules/openai.asciidoc b/metricbeat/docs/modules/openai.asciidoc index 6cb33cb60a3b..c211765ef31d 100644 --- a/metricbeat/docs/modules/openai.asciidoc +++ b/metricbeat/docs/modules/openai.asciidoc @@ -46,9 +46,9 @@ metricbeat.modules: # - "k2: v2" ## Rate Limiting Configuration # rate_limit: - # limit: 60 # requests per second - # burst: 5 # burst size - # ## Request timeout duration + # limit: 12 # seconds between requests + # burst: 1 # max concurrent requests + # ## Request Timeout Duration # timeout: 30s # # Data Collection Configuration From aff36b559417b719738f52834c928a2f00a3fca7 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Tue, 10 Dec 2024 11:41:46 +0530 Subject: [PATCH 31/36] nitpick --- x-pack/metricbeat/module/openai/usage/usage.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index d8347c2754aa..1766cfa73674 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -137,7 +137,7 @@ func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLH case <-ctx.Done(): return ctx.Err() default: - dateStr := d.Format("2006-01-02") + dateStr := d.Format(dateFormatForStateStore) if err := m.fetchSingleDay(apiKeyIdx, dateStr, apiKey.Key, httpClient); err != nil { m.logger.Errorf("Error fetching data (api key #%d) for date %s: %v", apiKeyIdx, dateStr, err) continue From 549f26e13fcf2488ae999bb970f2458b9742cfbb Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Wed, 11 Dec 2024 00:44:53 +0530 Subject: [PATCH 32/36] Logging and other basic improvements --- x-pack/metricbeat/module/openai/usage/usage.go | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/x-pack/metricbeat/module/openai/usage/usage.go b/x-pack/metricbeat/module/openai/usage/usage.go index 1766cfa73674..86dbe76ce815 100644 --- a/x-pack/metricbeat/module/openai/usage/usage.go +++ b/x-pack/metricbeat/module/openai/usage/usage.go @@ -94,7 +94,7 @@ func New(base mb.BaseMetricSet) (mb.MetricSet, error) { // 3. Fetches usage data for each day in the range // 4. Reports collected metrics through the mb.ReporterV2 func (m *MetricSet) Fetch(report mb.ReporterV2) error { - endDate := time.Now().UTC() + endDate := time.Now().UTC().Truncate(time.Hour * 24) // truncate to day as we only collect daily data if !m.config.Collection.Realtime { // If we're not collecting realtime data, then just pull until @@ -127,11 +127,14 @@ func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLH if err == nil { currentStartDate := lastProcessedDate.AddDate(0, 0, 1) if currentStartDate.After(endDate) { + m.logger.Infof("Skipping API key #%d as current start date (%s) is after end date (%s)", apiKeyIdx, currentStartDate, endDate) return nil } startDate = currentStartDate } + m.logger.Debugf("Fetching data for API key #%d from %s to %s", apiKeyIdx, startDate, endDate) + for d := startDate; !d.After(endDate); d = d.AddDate(0, 0, 1) { select { case <-ctx.Done(): @@ -139,6 +142,8 @@ func (m *MetricSet) fetchDateRange(startDate, endDate time.Time, httpClient *RLH default: dateStr := d.Format(dateFormatForStateStore) if err := m.fetchSingleDay(apiKeyIdx, dateStr, apiKey.Key, httpClient); err != nil { + // If there's an error, log it and continue to the next day. + // In this case, we are not saving the state. m.logger.Errorf("Error fetching data (api key #%d) for date %s: %v", apiKeyIdx, dateStr, err) continue } From 33daed7d39824d8125dffec93fef809495cc9962 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Thu, 12 Dec 2024 12:23:07 +0530 Subject: [PATCH 33/36] Update metricbeat/docs/modules/openai.asciidoc Co-authored-by: Brandon Morelli --- metricbeat/docs/modules/openai.asciidoc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/metricbeat/docs/modules/openai.asciidoc b/metricbeat/docs/modules/openai.asciidoc index c211765ef31d..ce5fd3e0ccc5 100644 --- a/metricbeat/docs/modules/openai.asciidoc +++ b/metricbeat/docs/modules/openai.asciidoc @@ -60,7 +60,8 @@ metricbeat.modules: # # and also making it harder to do analytics. Best approach is to avoid # # realtime collection and collect only upto last day (in UTC). So, there's # # at most 24h delay. - # realtime: false---- + # realtime: false +---- [float] === Metricsets From 2b836c28e7f98e03c0cfeaa0771089014b091de1 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Thu, 12 Dec 2024 12:56:53 +0530 Subject: [PATCH 34/36] Fix docs --- x-pack/metricbeat/metricbeat.reference.yml | 1 + x-pack/metricbeat/module/openai/_meta/config.yml | 2 +- x-pack/metricbeat/modules.d/openai.yml.disabled | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/x-pack/metricbeat/metricbeat.reference.yml b/x-pack/metricbeat/metricbeat.reference.yml index 28f4dd4546bd..a3dad62e8a6e 100644 --- a/x-pack/metricbeat/metricbeat.reference.yml +++ b/x-pack/metricbeat/metricbeat.reference.yml @@ -1314,6 +1314,7 @@ metricbeat.modules: # # realtime collection and collect only upto last day (in UTC). So, there's # # at most 24h delay. # realtime: false + #----------------------------- Openmetrics Module ----------------------------- - module: openmetrics metricsets: ['collector'] diff --git a/x-pack/metricbeat/module/openai/_meta/config.yml b/x-pack/metricbeat/module/openai/_meta/config.yml index b4048313a1c9..a34fd7b183de 100644 --- a/x-pack/metricbeat/module/openai/_meta/config.yml +++ b/x-pack/metricbeat/module/openai/_meta/config.yml @@ -31,4 +31,4 @@ # # and also making it harder to do analytics. Best approach is to avoid # # realtime collection and collect only upto last day (in UTC). So, there's # # at most 24h delay. - # realtime: false \ No newline at end of file + # realtime: false diff --git a/x-pack/metricbeat/modules.d/openai.yml.disabled b/x-pack/metricbeat/modules.d/openai.yml.disabled index b70f4127264b..6a881ae86419 100644 --- a/x-pack/metricbeat/modules.d/openai.yml.disabled +++ b/x-pack/metricbeat/modules.d/openai.yml.disabled @@ -34,4 +34,4 @@ # # and also making it harder to do analytics. Best approach is to avoid # # realtime collection and collect only upto last day (in UTC). So, there's # # at most 24h delay. - # realtime: false \ No newline at end of file + # realtime: false From 32a9593ec2b97bde2f36b77aa89616faa8542c20 Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Thu, 12 Dec 2024 14:26:28 +0530 Subject: [PATCH 35/36] Update CODEOWNER --- .github/CODEOWNERS | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index a06730335729..dca850c88650 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -222,6 +222,7 @@ CHANGELOG* /x-pack/metricbeat/module/iis @elastic/obs-infraobs-integrations /x-pack/metricbeat/module/istio/ @elastic/obs-cloudnative-monitoring /x-pack/metricbeat/module/mssql @elastic/obs-infraobs-integrations +/x-pack/metricbeat/module/openai @elastic/obs-infraobs-integrations /x-pack/metricbeat/module/oracle @elastic/obs-infraobs-integrations /x-pack/metricbeat/module/panw @elastic/obs-infraobs-integrations /x-pack/metricbeat/module/prometheus/ @elastic/obs-cloudnative-monitoring From ae85fa804f94e672051d44ea508bb37a6f0014da Mon Sep 17 00:00:00 2001 From: subham sarkar Date: Fri, 13 Dec 2024 13:04:24 +0530 Subject: [PATCH 36/36] Make linter happy --- x-pack/metricbeat/module/openai/usage/persistcache.go | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/x-pack/metricbeat/module/openai/usage/persistcache.go b/x-pack/metricbeat/module/openai/usage/persistcache.go index 362bdd9cab02..51a8d6fb0f89 100644 --- a/x-pack/metricbeat/module/openai/usage/persistcache.go +++ b/x-pack/metricbeat/module/openai/usage/persistcache.go @@ -156,12 +156,12 @@ func (s *stateManager) GetLastProcessedDate(apiKey string) (time.Time, error) { } // SaveState saves the last processed date for a given API key -func (sm *stateManager) SaveState(apiKey, dateStr string) error { - sm.mu.Lock() - defer sm.mu.Unlock() +func (s *stateManager) SaveState(apiKey, dateStr string) error { + s.mu.Lock() + defer s.mu.Unlock() - stateKey := sm.GetStateKey(apiKey) - return sm.store.Put(stateKey, dateStr) + stateKey := s.GetStateKey(apiKey) + return s.store.Put(stateKey, dateStr) } // hashKey generates and caches a SHA-256 hash of the provided API key