Skip to content

Conversation

@lbrdnk
Copy link
Contributor

@lbrdnk lbrdnk commented Sep 16, 2025

Closes https://linear.app/metabase/issue/BOT-410/implement-token-usage-tracking-for-semantic-search

This PR adds tracking of tokens consumed by embeddings api calls initiated from semantic search. Number of tokens of every request is stored in appdb. The tables is trimmed every day, storing rolling 2 months of data.

This PR

  • add migration with the new appdb table,
  • adds new model, :model/SemanticSearchTokenTracking,
  • modifies get-embedding... implementations to take opts,
  • passes opts from near toplevel (upsert-index!, query-index) with appropriate request :type, that is stored with token count in the new table,
  • adds new daily job to trim that table,
  • adds tests.

@metabase-bot metabase-bot bot added the .Team/Metabot Metabot team label Sep 16, 2025
@lbrdnk lbrdnk force-pushed the sem-search-token-tracking branch from 3b8e11b to f39c72d Compare September 17, 2025 14:14
@lbrdnk lbrdnk force-pushed the sem-search-token-tracking branch from f39c72d to ba98ce2 Compare September 19, 2025 14:01
@lbrdnk lbrdnk added the backport Automatically create PR on current release branch on merge label Sep 19, 2025
@lbrdnk lbrdnk changed the title [WIP] Semantic search token usage tracking Semantic search token usage tracking Sep 19, 2025
@lbrdnk lbrdnk requested a review from a team September 19, 2025 15:37
@github-actions
Copy link
Contributor

github-actions bot commented Sep 19, 2025

e2e tests failed on b8c2364b945cc6f2e1af9b401adb46b76f2baddd-3

e2e test run

File Test Name
embedding-reproductions.cy.spec.js (flaky) issue 40660 > static dashboard content shouldn't overflow its container (#40660)

constraints:
nullable: false
- column:
name: total_tokens
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we not want to store in/out separately? :)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@piranha probably no need as it's an embedding model, so tracking total tokens is enough

Copy link
Contributor

@piranha piranha left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So the logic seems fine, but I kind of want confirmation that total is fine and we don't want separate in/out figures.

Comment on lines +282 to +286
(is (= 1 (t2/count :model/SemanticSearchTokenTracking)))
(let [{:keys [request_type total_tokens]}
(t2/select-one :model/SemanticSearchTokenTracking)]
(is (= :index request_type))
(is (= 13 total_tokens))))
Copy link
Contributor

@piranha piranha Sep 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is by no means a request to change, but I personally find it easier to read =? stuff, like this:

(is (=? [{:request_type :index :total_tokens 13}]
        (t2/select :model/SemanticSearchTokenTracking)))

@lbrdnk lbrdnk merged commit 63cc5c3 into master Sep 23, 2025
506 of 516 checks passed
@lbrdnk lbrdnk deleted the sem-search-token-tracking branch September 23, 2025 14:50
github-automation-metabase pushed a commit that referenced this pull request Sep 23, 2025
* Pass type of embedding request

* Initial migration

* Add remaining opts args

* Add SemanticSearchTokenTracking module

* Connect token tracking to ai-service impl

* Add test for token tracking writes.

* Remove prompt_tokens including migration

* Add usage trimmer job

* Trimmer test

* Activate trimmer job

* test

* Record tokens for openai

* Update test

* Use total-tokens

* Comment

* linter

* Add index

* Exclude from copy
metamben pushed a commit that referenced this pull request Sep 23, 2025
* Pass type of embedding request

* Initial migration

* Add remaining opts args

* Add SemanticSearchTokenTracking module

* Connect token tracking to ai-service impl

* Add test for token tracking writes.

* Remove prompt_tokens including migration

* Add usage trimmer job

* Trimmer test

* Activate trimmer job

* test

* Record tokens for openai

* Update test

* Use total-tokens

* Comment

* linter

* Add index

* Exclude from copy
github-automation-metabase added a commit that referenced this pull request Sep 23, 2025
* Pass type of embedding request

* Initial migration

* Add remaining opts args

* Add SemanticSearchTokenTracking module

* Connect token tracking to ai-service impl

* Add test for token tracking writes.

* Remove prompt_tokens including migration

* Add usage trimmer job

* Trimmer test

* Activate trimmer job

* test

* Record tokens for openai

* Update test

* Use total-tokens

* Comment

* linter

* Add index

* Exclude from copy

Co-authored-by: lbrdnk <lbrdnk@users.noreply.github.com>
@github-actions github-actions bot added this to the 0.56.7 milestone Sep 23, 2025
galdre pushed a commit that referenced this pull request Sep 24, 2025
* Pass type of embedding request

* Initial migration

* Add remaining opts args

* Add SemanticSearchTokenTracking module

* Connect token tracking to ai-service impl

* Add test for token tracking writes.

* Remove prompt_tokens including migration

* Add usage trimmer job

* Trimmer test

* Activate trimmer job

* test

* Record tokens for openai

* Update test

* Use total-tokens

* Comment

* linter

* Add index

* Exclude from copy
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

backport Automatically create PR on current release branch on merge .Team/Metabot Metabot team

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants