x-pack/metricbeat/module/openai: Add new module #41516

shmsr · 2024-11-04T20:26:34Z

Proposed commit message

Implement a new module for OpenAI usage collection. This module operates on https://api.openai.com/v1/usage (by default; also configurable for Proxy URLs, etc.) and collects the limited set of usage metrics emitted from the undocumented endpoint.

Example how the usage endpoints emits metrics:

Given timestamps t0, t1, t2, ... tn in ascending order:

At t0 (first collection):

   usage_metrics_1: *

At t1 (after new API usage):

   usage_metrics_1: *
   usage_metrics_2: *

At t2 (continuous collection):

   usage_metrics_1: *
   usage_metrics_2: *
   usage_metrics_3: *

and so on.

Example response:

{
  "object": "list",
  "data": [
    {
      "organization_id": "org-xxx",
      "organization_name": "Personal",
      "aggregation_timestamp": 1725389580,
      "n_requests": 1,
      "operation": "completion",
      "snapshot_id": "gpt-4o-mini-2024-07-18",
      "n_context_tokens_total": 62,
      "n_generated_tokens_total": 21,
      "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": ""
    },
    {
      "organization_id": "org-xxx",
      "organization_name": "Personal",
      "aggregation_timestamp": 1725389640,
      "n_requests": 1,
      "operation": "completion",
      "snapshot_id": "gpt-4o-mini-2024-07-18",
      "n_context_tokens_total": 97,
      "n_generated_tokens_total": 17,
      "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": ""
    }
  ],
  "tpm_data": [
    {
      "organization_id": "org-xxx",
      "organization_name": "Personal",
      "day_timestamp": 1725321600,
      "snapshot_id": "gpt-4o-mini-2024-07-18",
      "operation": "completion",
      "p90_context_tpm": 97,
      "p90_generated_tpm": 21,
      "p90_provisioned_context_tpm": 0,
      "p90_provisioned_generated_tpm": 0,
      "max_context_tpm": 97,
      "max_generated_tpm": 21,
      "max_provisioned_context_tpm": 0,
      "max_provisioned_generated_tpm": 0
    }
  ],
  "ft_data": [],
  "dalle_api_data": [],
  "whisper_api_data": [],
  "tts_api_data": [],
  "assistant_code_interpreter_data": [],
  "retrieval_storage_data": []
}

As soon as the API is used, usage is generated after a few times. So, if collecting using the module real-time and that too multiple times of the day, it would collect duplicates and it is not good for storage as well as analytics of the usage data.

It's better to collect time.Now() (in UTC) - 24h so that we get full usage collection of the past day (in UTC) and it avoids duplication. So that's why I have introduced a config realtime and set it to false as the collection is 24h delayed; we are now getting daily data. realtime: true will work as any other normal collection where metrics are fetched in set intervals. Our recommendation is to keep realtime: false.

As this is a metricbeat module, we do not have existing package that gives us support to store the cursor. So, in order to avoid pulling already pulled data, timestamps are being stored per API key. Logic for the same is commented in the code on how it is stored. We are using a new custom code to store the state in order to store the cursor and begin from the next available date.

Checklist

My code follows the style guidelines of this project
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
I have made corresponding change to the default configuration files
I have added tests that prove my fix is effective or that my feature works
I have added an entry in CHANGELOG.next.asciidoc or CHANGELOG-developer.next.asciidoc.

Author's Checklist

Check the state store
Validate with usage dashboard of OpenAI

How to test this PR locally

Run metricbeat and use your OpenAI's API key to collect usage metrics.

mergify · 2024-11-04T20:27:13Z

This pull request does not have a backport label.
If this is a bug or security fix, could you label this PR @shmsr? 🙏.
For such, you'll need to label your PR with:

The upcoming major version of the Elastic Stack
The upcoming minor version of the Elastic Stack (if you're not pushing a breaking change)

To fixup this pull request, you need to add the backport labels for the needed
branches, such as:

backport-8./d is the label to automatically backport to the 8./d branch. /d is the digit

mergify · 2024-11-04T20:27:13Z

backport-8.x has been added to help with the transition to the new branch 8.x.
If you don't need it please use backport-skip label and remove the backport-8.x label.

shmsr · 2024-11-05T17:47:25Z

Getting hit by this error: #41174 (comment) and hence the CI is failing. Rest all okay.

shmsr · 2024-11-06T09:38:51Z

To continue with my testing and to avoid: Limit of total fields [10000] has been exceeded error until it is fixed for 8.15.x and older, I am using: setup.template.fields where I specify a new field that only contains the ecs fields and openai fields from fields.yml and nothing else.

See this: https://www.elastic.co/guide/en/beats/metricbeat/current/configuration-template.html

So, this has currently unblocked me but yes we definitely need a fix for this.

elasticmachine · 2024-11-12T07:32:31Z

Pinging @elastic/elastic-agent-data-plane (Team:Elastic-Agent-Data-Plane)

shmsr · 2024-11-12T09:11:52Z

I've explain the complicated collection mechanism in the PR description itself. Rest is self-explanatory from the code. Please let me know, if anything needs further clarification.

x-pack/metricbeat/module/openai/usage/_meta/fields.yml

x-pack/metricbeat/module/openai/usage/client.go

x-pack/metricbeat/module/openai/usage/config.go

x-pack/metricbeat/module/openai/usage/usage.go

x-pack/metricbeat/module/openai/usage/usage_integration_test.go

AndersonQ · 2024-11-13T14:40:15Z

x-pack/metricbeat/module/openai/usage/usage_integration_test.go

+    }
+  ],
+  "ft_data": [],
+  "dalle_api_data": [],


It'd be good to have data for each data set

Yeah, I tried generating ft (fine-tuning) data but it doesn't seem to work. As OpenAI provides this API as undocumented, I couldn't find a single source with any samples. Not even sure they even populate for the response of this particular endpoint. For dalle_api_data, I'll add.

x-pack/metricbeat/module/openai/usage/usage_integration_test.go

AndersonQ

You forgot to revert the changes you made to the default config file

x-pack/metricbeat/metricbeat.yml

x-pack/metricbeat/module/openai/usage/usage.go

AndersonQ · 2024-11-26T07:41:18Z

x-pack/metricbeat/module/openai/usage/usage.go

 		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))


be careful here to do not leak sensitive information. I usually try to log the body in its own log, without including the body in the error message

I was trying this morning only. I did found one case where the logging the body is not ideal. In case of invalid API keys, the body shows the API key that it's invalid to connect to OpenAI. So, in case there's a typo when typing the API key, most part of the API key is exposed. So, I'll remove this in order to avoid any sensitive info leakage.

AndersonQ

Sorry I spotted a wee issue on how you create the http client and measure how long the request take, as soon as it's fixed, it's good to go.

x-pack/metricbeat/module/openai/usage/client.go

AndersonQ · 2024-11-27T06:38:32Z

x-pack/metricbeat/module/openai/usage/client.go

+	err := c.Ratelimiter.Wait(req.Context())
+	waitDuration := time.Since(start)
+	if err != nil {
+		return nil, err
+	}


[NIT|Go convention]

Suggested change

err := c.Ratelimiter.Wait(req.Context())

waitDuration := time.Since(start)

if err != nil {

return nil, err

}

err := c.Ratelimiter.Wait(req.Context())

if err != nil {

return nil, err

}

waitDuration := time.Since(start)

AndersonQ · 2024-11-27T06:40:37Z

x-pack/metricbeat/module/openai/usage/client.go

+	var client = http.DefaultClient
+	client.Timeout = timeout


[Blocker]
You should not change the default cleint, you should create a new one.

var client = http.DefaultClient will make client to reference the http.DefaultClient, so any change in client will take effect for everyone using http.DefaultClient

Suggested change

var client = http.DefaultClient

client.Timeout = timeout

client := http.Client{Timeout: timeout}

This is a very good point. Thanks.

ishleenk17 · 2024-11-27T07:04:48Z

@shmsr : Started reviewing the module. Can you please link the document where this approach and the metrics we are collcting was decided ?

shmsr · 2024-11-27T07:49:02Z

@shmsr : Started reviewing the module. Can you please link the document where this approach and the metrics we are collcting was decided ?

As it is a private repo, I've linked reversely linked the issue as per recommendation. Also, the research there will show the research on which API to use. And the decision on how to process the metrics (more implementation specific details) is defined here in the description itself: #41516 (comment); @muthu-mps and I decided it on a call.

shmsr · 2024-11-27T09:47:07Z

Thanks @AndersonQ. Very helpful review comments; appreciate it! ✨

muthu-mps · 2024-11-27T10:04:26Z

x-pack/metricbeat/module/openai/usage/_meta/data.json

+                "api_key_redacted": null,
+                "api_key_type": null,
+                "email": null,
+                "n_cached_context_tokens_total": 0,


@shmsr , @ishleenk17 - Are we okay with the field names? IMO, Updating the field names something similar to
n_context_tokens_total -> tokens.total or tokens_total in the context of usability. But still it is also good to keep as it is to relate the ES fields with the openai API metrics. WDYT?

Yeah, we should keep fields more readable and inline with other LLM Integrations

I was thinking to keep it like this here and change the names in ingest pipelines in integrations. What do you think?

x-pack/metricbeat/module/openai/usage/client.go

devamanv · 2024-11-27T09:54:04Z

x-pack/metricbeat/module/openai/usage/config.go

+	"time"
+)
+
+type Config struct {


Although not necessary but the name Config sounds a little vague, and doesn't clearly explain the purpose. Consider renaming it for clarity.

Do you have any suggestion? The Config is like that in many modules, so kept it like that only

Search for type Config struct (case-insensitive search) in metricbeat/module/*

x-pack/metricbeat/module/openai/usage/helper.go

devamanv · 2024-11-27T10:31:07Z

x-pack/metricbeat/module/openai/usage/helper.go

+		if len(parts) != 2 {
+			continue
+		}
+		headersMap[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1])


Nit: Extract them separately for clarity.

x-pack/metricbeat/module/openai/usage/persistcache.go

x-pack/metricbeat/module/openai/usage/schema.go

x-pack/metricbeat/module/openai/usage/usage.go

ishleenk17 · 2024-11-27T12:31:02Z

x-pack/metricbeat/module/openai/usage/usage.go

+					"project_name":                  usage.ProjectName,
+					"request_type":                  usage.RequestType,
+					"n_cached_context_tokens_total": usage.NCachedContextTokensTotal,
+				},


Does this API provide a division of Input and Output tokens?
If yes, we can include that also.

I didn't get this. Can you please elaborate?

x-pack/metricbeat/module/openai/usage/client.go

x-pack/metricbeat/module/openai/usage/config.go

ishleenk17 · 2024-11-27T12:54:13Z

x-pack/metricbeat/module/openai/usage/config.go

+		}
+	}
+	if c.Timeout <= 0 {
+		errs = append(errs, errors.New("timeout must be greater than 0"))


Can this happen ? Since we are having a default value for it always ?

Ideally, no. Yes, because of default. Just kept it there as it's a just a one time check and to avoid issues if someday someone changes the default i.e., 0 (or lower).

Oh, yes also in the case if the user has configured the timeout as 0 or lower.

ishleenk17 · 2024-11-27T12:56:36Z

x-pack/metricbeat/module/openai/usage/config.go

+		APIURL:  "https://api.openai.com/v1/usage",
+		Timeout: 30 * time.Second,
+		RateLimit: &rateLimitConfig{
+			Limit: ptr(60),


What defines these limits ?
If you could point me to documentation due to which we have chosen these defaults ?

The timeout is a sane default. And the rate-limit is decided as per testing. If you see the meta ticket, you'd find an issue there where we have researched about rate-limits of this API. Not linking it here as it's private. Can share over Slack.

ishleenk17 · 2024-11-27T13:03:47Z

x-pack/metricbeat/module/openai/usage/persistcache.go

+}
+
+// Clear removes all state files by deleting and recreating the state directory
+func (s *stateStore) Clear() error {


If our intent here is to remove the files, since we are again recreating the directory. Should we delete just the files instead?
It might avoid potential permission issues or race conditions. WDYT ?

Eg:

files, err := os.ReadDir(s.Dir) if err != nil { return fmt.Errorf("reading state directory: %w", err) } for _, file := range files { if err := os.Remove(path.Join(s.Dir, file.Name())); err != nil { return fmt.Errorf("removing state file: %w", err) } }

Actually, when I was writing this I was creating it to be good package for reuse later as well and this method was inteded to completely reset the state. Plan earlier was to add subdirectories too inside the state folder because what if in future we add a new dataset. For now, we just have "usage", but in case we add more.

But as it just resides in this package, we are not even using this method. It is unused. We can even remove this method.

So, that's why this. It's a complete reset of state.

ishleenk17 · 2024-11-27T13:05:14Z

x-pack/metricbeat/module/openai/usage/persistcache.go

+		return nil, errors.New("empty path provided")
+	}
+
+	store, err := newStateStore(storePath)


Do we need to check if the storepath indeed exists ?

we are creating the storePath. Part of it already there and the rest we create with MkdirAll.

See this: newStateManager(paths.Resolve(paths.Data, path.Join("state", base.Module().Name(), base.Name())))

paths.Data is supplied by Agent using --paths.data flag. So, even if the directory doesn't exist, we create them.

x-pack/metricbeat/module/openai/usage/persistcache.go

x-pack/metricbeat/module/openai/usage/schema.go

x-pack/metricbeat/module/openai: Add new module

0e1175c

botelastic bot added the needs_team Indicates that the issue/PR needs a Team:* label label Nov 4, 2024

mergify bot assigned shmsr Nov 4, 2024

mergify bot added the backport-8.x Automated backport to the 8.x branch with mergify label Nov 4, 2024

shmsr added Module: openai Team:Obs-InfraObs Label for the Observability Infrastructure Monitoring team labels Nov 4, 2024

botelastic bot removed the needs_team Indicates that the issue/PR needs a Team:* label label Nov 4, 2024

shmsr added 2 commits November 5, 2024 01:59

Merge branch 'main' into openai-metrics

bf45bf6

update module

8725cbb

shmsr requested a review from a team November 5, 2024 07:51

shmsr added 6 commits November 5, 2024 17:35

update module

03995f6

update module

406c1a4

update module

a3c46b7

update module

72488f7

Merge branch 'main' into openai-metrics

b5c750a

fix bug

2d4027d

update module

e3f8fe2

shmsr marked this pull request as ready for review November 12, 2024 07:28

shmsr requested a review from a team as a code owner November 12, 2024 07:28

shmsr requested review from AndersonQ and belimawr November 12, 2024 07:28

pierrehilbert added the Team:Elastic-Agent-Data-Plane Label for the Agent Data Plane team label Nov 12, 2024

Merge branch 'main' into openai-metrics

82cd979

AndersonQ reviewed Nov 13, 2024

View reviewed changes

shmsr added 6 commits November 25, 2024 15:11

Address review comments

3259612

Address review comments

40b1131

Address review comments

2936540

Merge branch 'main' into openai-metrics

850fc43

Address review comments

5061e46

Improvements

34a7128

AndersonQ requested changes Nov 26, 2024

View reviewed changes

shmsr added 3 commits November 26, 2024 16:08

Fix *.yml

7f4cc75

Address review comments

7919c81

Improvements

2a4bb56

shmsr requested review from devamanv and AndersonQ November 26, 2024 10:57

shmsr added 2 commits November 26, 2024 16:28

Merge branch 'main' into openai-metrics

4899f3a

Make linter happy

217f7d5

AndersonQ requested changes Nov 27, 2024

View reviewed changes

Address review comments

b5dbfb0

AndersonQ approved these changes Nov 27, 2024

View reviewed changes

muthu-mps reviewed Nov 27, 2024

View reviewed changes

devamanv reviewed Nov 27, 2024

View reviewed changes

ishleenk17 reviewed Nov 27, 2024

View reviewed changes

x-pack/metricbeat/module/openai/usage/schema.go Show resolved Hide resolved

shmsr added 5 commits November 29, 2024 00:28

gofumpt'ed

492eca1

Address review comments

413010f

Address review comments

92f5062

Merge branch 'main' into openai-metrics

fdb81c9

Merge branch 'main' into openai-metrics

78bbd11

	var client = http.DefaultClient
	client.Timeout = timeout
	client := http.Client{Timeout: timeout}

x-pack/metricbeat/module/openai: Add new module #41516

Are you sure you want to change the base?

x-pack/metricbeat/module/openai: Add new module #41516

Conversation

shmsr commented Nov 4, 2024 • edited Loading

Proposed commit message

Checklist

Author's Checklist

How to test this PR locally

mergify bot commented Nov 4, 2024

mergify bot commented Nov 4, 2024

shmsr commented Nov 5, 2024 • edited Loading

shmsr commented Nov 6, 2024

elasticmachine commented Nov 12, 2024

shmsr commented Nov 12, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

AndersonQ left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

AndersonQ left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ishleenk17 commented Nov 27, 2024

shmsr commented Nov 27, 2024

shmsr commented Nov 27, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shmsr Nov 28, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shmsr Nov 28, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shmsr commented Nov 4, 2024 •

edited

Loading

shmsr commented Nov 5, 2024 •

edited

Loading

shmsr Nov 28, 2024 •

edited

Loading

shmsr Nov 28, 2024 •

edited

Loading