Skip to content

Commit 1e8c1c9

Browse files
Include the token usage for every conversation and workspace (#788)
* Include the token usage for every conversation and workspace Related: #418 This PR does introduces the changes necessary to track the used tokens per request and then process them to return them in the API. Specific changes: - Make sure we process all the stream and record at the very end - Include the flag `"stream_options": {"include_usage": True},` so the providers respond with the tokens - Added the necessary processing for the API - Modified the initial API models to display correctly the tokens and its price * Moved token recording to DB * Changed token usage code to get info from file and added GHA to get file periodically * formatting changes * Move model cost to dedicated folder * Fix problems with copilot streaming
1 parent 8d0432d commit 1e8c1c9

File tree

22 files changed

+8650
-92
lines changed

22 files changed

+8650
-92
lines changed
+59
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
# This workflow updates the file in /model_cost_data/model_prices_and_context_window.json
2+
name: Update model prices and context window JSON file
3+
4+
on:
5+
workflow_call:
6+
schedule:
7+
- cron: '0 2 * * 0' # Run every Sunday at 2:00 AM
8+
9+
jobs:
10+
update_model_prices:
11+
runs-on: ubuntu-latest
12+
permissions:
13+
contents: write
14+
pull-requests: write
15+
16+
steps:
17+
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4
18+
19+
- name: Get the latest file
20+
run: |
21+
curl -Ss 'https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json' > model_cost_data/model_prices_and_context_window.json
22+
23+
- name: Check if file changed
24+
id: check-model-prices
25+
run: |
26+
if ! git diff --quiet model_cost_data/model_prices_and_context_window.json ; then
27+
echo "changed=true" >> "$GITHUB_OUTPUT"
28+
else
29+
echo "changed=false" >> "$GITHUB_OUTPUT"
30+
fi
31+
32+
- name: Set git config
33+
run: |
34+
git config --local user.email "github-actions[bot]@users.noreply.github.com"
35+
git config --local user.name "github-actions[bot]"
36+
37+
- name: Get current date
38+
id: date
39+
run: |
40+
echo "CURRENT_DATETIME=$(date +%Y-%m-%d)" >> $GITHUB_ENV
41+
42+
- name: Generate PR if needed
43+
if: steps.check-model-prices.outputs.changed == 'true'
44+
run: |
45+
git checkout -b update-model-prices-$GITHUB_SHA
46+
47+
git add model_cost_data/model_prices_and_context_window.json
48+
git commit -m "Update model_prices_and_context_window.json to version generated on ${{ env.CURRENT_DATETIME }}"
49+
50+
echo "Pushing branch so we can create a PR..."
51+
git push --set-upstream origin update-model-prices-$GITHUB_SHA
52+
53+
gh pr create --title "Update model_prices_and_context_window.json" \
54+
--body "This PR updates the model_prices_and_context_window.json definition to the version generated on ${{ env.CURRENT_DATETIME }}" \
55+
--repo "$GITHUB_REPOSITORY" \
56+
--base main \
57+
--head update-model-prices-$GITHUB_SHA
58+
env:
59+
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
"""add token usage columns
2+
3+
Revision ID: 0c3539f66339
4+
Revises: 0f9b8edc8e46
5+
Create Date: 2025-01-28 09:15:54.767311+00:00
6+
7+
"""
8+
9+
from typing import Sequence, Union
10+
11+
from alembic import op
12+
13+
# revision identifiers, used by Alembic.
14+
revision: str = "0c3539f66339"
15+
down_revision: Union[str, None] = "0f9b8edc8e46"
16+
branch_labels: Union[str, Sequence[str], None] = None
17+
depends_on: Union[str, Sequence[str], None] = None
18+
19+
20+
def upgrade() -> None:
21+
# Begin transaction
22+
op.execute("BEGIN TRANSACTION;")
23+
24+
# We add the columns to the outputs table
25+
# Add the columns with default values to avoid issues with the existing data
26+
# The prices of the tokens may change in the future,
27+
# so we need to store the cost of the tokens at the time of the request
28+
op.execute("ALTER TABLE outputs ADD COLUMN input_tokens INT DEFAULT NULL;")
29+
op.execute("ALTER TABLE outputs ADD COLUMN output_tokens INT DEFAULT NULL;")
30+
op.execute("ALTER TABLE outputs ADD COLUMN input_cost FLOAT DEFAULT NULL;")
31+
op.execute("ALTER TABLE outputs ADD COLUMN output_cost FLOAT DEFAULT NULL;")
32+
33+
# Finish transaction
34+
op.execute("COMMIT;")
35+
36+
37+
def downgrade() -> None:
38+
# Begin transaction
39+
op.execute("BEGIN TRANSACTION;")
40+
41+
op.execute("ALTER TABLE outputs DROP COLUMN input_tokens;")
42+
op.execute("ALTER TABLE outputs DROP COLUMN output_tokens;")
43+
op.execute("ALTER TABLE outputs DROP COLUMN input_cost;")
44+
op.execute("ALTER TABLE outputs DROP COLUMN output_cost;")
45+
46+
# Finish transaction
47+
op.execute("COMMIT;")

0 commit comments

Comments
 (0)