Skip to content

feat(litellm): return token usage and cost with the response - #6678

Merged
shahargl merged 1 commit into
keephq:mainfrom
artemrootman:feat/litellm-return-usage
Aug 6, 2026
Merged

feat(litellm): return token usage and cost with the response#6678
shahargl merged 1 commit into
keephq:mainfrom
artemrootman:feat/litellm-return-usage

Conversation

@artemrootman

Copy link
Copy Markdown
Contributor

Problem

See #6677. LitellmProvider._query returns only {"response": ...} and discards the usage block, so a workflow cannot tell how many tokens an analysis consumed or what it cost. The only way to get that today is to skip the provider and call the API from a python step.

Fix

Return the usage fields next to the response:

  • prompt_tokens, completion_tokens, total_tokens, cost — read through .get() so a backend that reports no usage yields None rather than raising;
  • model as reported by the API, which can differ from the requested one when a proxy routes the call.

response keeps its meaning, so existing workflows are unaffected.

Tests

Two added to tests/providers/litellm_provider/test_litellm_response_parsing.py: usage and model surfaced from a normal response, and a response without a usage block leaving the fields None while the text still comes through. 9 tests in the file, all passing; black clean.

Fixes #6677

@dosubot dosubot Bot added size:XS This PR changes 0-9 lines, ignoring generated files. Enhancement New feature or request Provider Providers related issues labels Aug 6, 2026
@artemrootman

Copy link
Copy Markdown
Contributor Author

The unit-tests failure is unrelated to this change: the run reports 1078 passed, 84 skipped, 1 error, and the error is at setup of tests/test_auth.py::test_deleted_api_key_authentication[MULTI_TENANT], with HTTPSConnection(host='https', port=443) — a real network call. This PR only touches the return value of LitellmProvider._query; its two new tests are among the 1078 that passed.

The cause looks reproducible rather than random, if it helps:

  1. tests/fixtures/client.py sets AUTH0_DOMAIN to https://auth0domain.com — with the scheme included — for the MULTI_TENANT string parameter.
  2. _mock_oidc_discovery() is only applied when is_auth0 is true, which is set from request.param.get("AUTH_TYPE") == "AUTH0" and therefore only for the dict-style parameters. The MULTI_TENANT string path gets nullcontext().
  3. ee/identitymanager/identity_managers/auth0/auth0_authverifier.py then runs at import time and builds the JWKS URL as f"https://{auth_domain}/.well-known/jwks.json", which with the scheme already present becomes https://https://auth0domain.com/.... urllib3 parses the host out of that as https, hence the connection attempt in the traceback.

So whenever that module is imported unmocked in this parameterisation, the setup depends on an outbound HTTPS request, which is why it surfaces intermittently across workers. Applying the discovery mock for MULTI_TENANT as well, or storing AUTH0_DOMAIN without the scheme, would remove the network dependency.

Happy to send that as a separate PR if it is useful. Could you re-run the job in the meantime? I cannot trigger it from a fork.

@shahargl shahargl left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

lgtm

@dosubot dosubot Bot added the lgtm This PR has been approved by a maintainer label Aug 6, 2026
@shahargl
shahargl merged commit 8cca2c3 into keephq:main Aug 6, 2026
14 of 16 checks passed
@artemrootman
artemrootman deleted the feat/litellm-return-usage branch August 6, 2026 14:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Enhancement New feature or request lgtm This PR has been approved by a maintainer Provider Providers related issues size:XS This PR changes 0-9 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

LiteLLM provider discards token usage, so a workflow cannot report what a run cost

2 participants