UN-3415 [REFACTOR] Replace CostCalculationHelper with litellm.cost_per_token#1906
Merged
muhammad-ali-e merged 3 commits intomainfrom Apr 17, 2026
Merged
UN-3415 [REFACTOR] Replace CostCalculationHelper with litellm.cost_per_token#1906muhammad-ali-e merged 3 commits intomainfrom
muhammad-ali-e merged 3 commits intomainfrom
Conversation
Move cost calculation from platform-service to sdk1's Audit class, using litellm's built-in cost_per_token() instead of a custom helper that fetched pricing data from an external URL. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
WalkthroughServer-side cost calculation and related utilities were removed; cost is now computed in the SDK using LiteLLM and sent in the usage payload. The platform controller accepts Changes
Sequence DiagramsequenceDiagram
participant SDK as SDK (Client)
participant LiteLLM as LiteLLM
participant PlatformService as Platform Service
rect rgba(100, 150, 200, 0.5)
Note over SDK,PlatformService: Old Flow (Server-side cost calc)
SDK->>PlatformService: POST /usage (model, tokens)
PlatformService->>PlatformService: Load pricing (cache/URL)
PlatformService->>PlatformService: Calculate cost
PlatformService-->>SDK: Response (includes computed cost)
end
rect rgba(150, 200, 100, 0.5)
Note over SDK,PlatformService: New Flow (Client-side cost calc)
SDK->>LiteLLM: cost_per_token(model, prompt_tokens, completion_tokens)
LiteLLM-->>SDK: cost_in_dollars
SDK->>PlatformService: POST /usage (model, tokens, cost_in_dollars)
PlatformService->>PlatformService: Accept cost from payload
PlatformService-->>SDK: Response
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Contributor
|
| Filename | Overview |
|---|---|
| unstract/sdk1/src/unstract/sdk1/audit.py | Cost calculation added using litellm.cost_per_token(); embedding events correctly zero completion_tokens; full model name used for cost lookup before stripping provider prefix for DB storage. |
| platform-service/src/unstract/platform_service/controller/platform.py | Simplified /usage endpoint: reads pre-computed cost_in_dollars from payload (default 0.0), removing CostCalculationHelper and input-token branching logic. |
| platform-service/src/unstract/platform_service/helper/cost_calculation.py | Deleted — external HTTP fetch, file-storage TTL caching, and manual pricing logic fully replaced by litellm's bundled pricing database. |
| platform-service/src/unstract/platform_service/env.py | Removed three required env vars (MODEL_PRICES_URL, MODEL_PRICES_TTL_IN_DAYS, MODEL_PRICES_FILE_PATH) that are no longer needed. |
| platform-service/src/unstract/platform_service/utils.py | Removed the now-orphaned format_float_positional helper function that was only used by the deleted CostCalculationHelper. |
Sequence Diagram
sequenceDiagram
participant Tool as SDK1 Tool
participant Audit as Audit.push_usage_data()
participant LiteLLM as litellm cost_per_token
participant Platform as platform-service /usage
Tool->>Audit: push_usage_data(model_name, token_counter, event_type)
Audit->>Audit: compute prompt_tokens and completion_tokens
Note over Audit: For embedding events, completion_tokens set to 0
alt model_name is set and known
Audit->>LiteLLM: cost_per_token(model, prompt_tokens, completion_tokens)
LiteLLM-->>Audit: (prompt_cost, completion_cost)
Audit->>Audit: cost_in_dollars = sum of costs
else unknown or no model
Audit->>Audit: cost_in_dollars = 0.0 (exception fallback)
end
Audit->>Audit: strip provider prefix from model_name for DB storage
Audit->>Platform: POST /usage with cost_in_dollars in payload
Platform->>Platform: read cost_in_dollars from payload (default 0.0)
Platform->>Platform: INSERT into token_usage table
Platform-->>Audit: 200 OK
Reviews (3): Last reviewed commit: "Merge branch 'main' into refactor/litell..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
platform-service/src/unstract/platform_service/controller/platform.py (1)
219-219: Consider validating thecost_in_dollarsvalue from the payload.The endpoint now trusts the client-provided
cost_in_dollarsvalue without validation. While the endpoint is protected by authentication, consider adding basic type validation to ensure data integrity:
- The value could be a non-numeric type (string, None beyond default, dict, etc.)
- The value could be negative, which doesn't make sense for costs
💡 Optional validation
- cost_in_dollars = payload.get("cost_in_dollars", 0.0) + cost_in_dollars = payload.get("cost_in_dollars", 0.0) + if not isinstance(cost_in_dollars, (int, float)) or cost_in_dollars < 0: + cost_in_dollars = 0.0🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@platform-service/src/unstract/platform_service/controller/platform.py` at line 219, Validate the payload's cost_in_dollars after the line where cost_in_dollars = payload.get("cost_in_dollars", 0.0): ensure it's a numeric value and non-negative by attempting to coerce to float (or checking isinstance int/float) and rejecting invalid input; if coercion fails or the value is < 0, return a 400-style validation error (or set a safe default and log) from the controller function that contains this variable so downstream logic only sees a validated non-negative float.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@platform-service/src/unstract/platform_service/controller/platform.py`:
- Line 219: Validate the payload's cost_in_dollars after the line where
cost_in_dollars = payload.get("cost_in_dollars", 0.0): ensure it's a numeric
value and non-negative by attempting to coerce to float (or checking isinstance
int/float) and rejecting invalid input; if coercion fails or the value is < 0,
return a 400-style validation error (or set a safe default and log) from the
controller function that contains this variable so downstream logic only sees a
validated non-negative float.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 2b5e73f7-f615-4b4d-937b-8c09f0b2db10
📒 Files selected for processing (5)
platform-service/src/unstract/platform_service/controller/platform.pyplatform-service/src/unstract/platform_service/env.pyplatform-service/src/unstract/platform_service/helper/cost_calculation.pyplatform-service/src/unstract/platform_service/utils.pyunstract/sdk1/src/unstract/sdk1/audit.py
💤 Files with no reviewable changes (3)
- platform-service/src/unstract/platform_service/utils.py
- platform-service/src/unstract/platform_service/env.py
- platform-service/src/unstract/platform_service/helper/cost_calculation.py
Explicitly set completion_tokens to 0 for embedding events before calling cost_per_token, making the assumption that embeddings have no completion tokens explicit rather than relying on the counter always being zero. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
unstract/sdk1/src/unstract/sdk1/audit.py (1)
104-107: Consider logging the exception details for easier debugging.The exception is caught but the actual error message is not logged, making it harder to diagnose failures. For unknown models this is fine, but other failures (e.g., litellm API changes, unexpected types) would be harder to debug.
Proposed improvement
except Exception: logger.debug( - "Cost lookup failed for model %s, defaulting to 0", model_name + "Cost lookup failed for model %s, defaulting to 0", model_name, + exc_info=True, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@unstract/sdk1/src/unstract/sdk1/audit.py` around lines 104 - 107, The except block in the cost lookup (in unstract.sdk1.audit.py, around the function performing model cost lookup) currently swallows all exceptions and only logs the model_name; modify the exception handler to include the actual exception details—either by capturing the exception as e and adding it to the log message (e.g., include str(e)) or by passing exc_info=True to logger.debug—so failures from litellm API changes or unexpected types are recorded alongside the existing "Cost lookup failed for model %s, defaulting to 0" message.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@unstract/sdk1/src/unstract/sdk1/audit.py`:
- Around line 104-107: The except block in the cost lookup (in
unstract.sdk1.audit.py, around the function performing model cost lookup)
currently swallows all exceptions and only logs the model_name; modify the
exception handler to include the actual exception details—either by capturing
the exception as e and adding it to the log message (e.g., include str(e)) or by
passing exc_info=True to logger.debug—so failures from litellm API changes or
unexpected types are recorded alongside the existing "Cost lookup failed for
model %s, defaulting to 0" message.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: a41f4ae9-3710-44c0-8d58-4b068323c771
📒 Files selected for processing (1)
unstract/sdk1/src/unstract/sdk1/audit.py
Deepak-Kesavan
approved these changes
Apr 9, 2026
chandrasekharan-zipstack
changed the title
Contributor
Test ResultsSummary
Runner Tests - Full Report
SDK1 Tests - Full Report
|
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
CostCalculationHelperin platform-service with litellm's built-incost_per_token(), moving cost calculation to sdk1'sAuditclass (caller-side).Why
CostCalculationHelperfetched pricing data from an external URL, cached it in file storage with a TTL, and did manual price lookups — all of which litellm already handles natively with its bundled pricing database (2645+ models).How
audit.py(sdk1): Compute cost vialitellm.cost_per_token()using the full model name (e.g.azure/gpt-4o) before stripping the provider prefix for DB storage. Send pre-computedcost_in_dollarsin the payload to platform-service.platform.py(platform-service): Readcost_in_dollarsdirectly from the payload instead of computing it. RemovedCostCalculationHelperimport,providervariable, and input token branching logic.cost_calculation.py: No longer needed — was only used by the/usageendpoint.env.py: RemovedMODEL_PRICES_URL,MODEL_PRICES_TTL_IN_DAYS,MODEL_PRICES_FILE_PATHenv vars.utils.py: Removed orphanedformat_float_positionalfunction.Can this PR break any existing features. If yes, please list possible items. If no, please explain why. (PS: Admins do not merge the PR without this section filled)
gpt-4o, notazure/gpt-4o). API deployment responses and dashboard queries are unaffected.cost=0.0(same behavior as before via theexcept Exceptionfallback)./usageendpoint is backward-compatible: ifcost_in_dollarsis not in the payload (e.g. from an older SDK), it defaults to0.0.Database Migrations
Env Config
MODEL_PRICES_URL,MODEL_PRICES_TTL_IN_DAYS,MODEL_PRICES_FILE_PATH(no longer needed)Relevant Docs
Related Issues or PRs
Dependencies Versions
litellmis already a transitive dependency of platform-service viaunstract-sdk1.Notes on Testing
/usageendpoint accepts requests and stores cost correctly.litellm.cost_per_token()against previousCostCalculationHelperoutput forazure/gpt-4o— values match.gpt-4o,gpt-4o-miniazure/gpt-4oclaude-sonnet-4-20250514text-embedding-3-smallcost=0.0Screenshots
N/A
Checklist
I have read and understood the Contribution Guidelines.