[REFACTOR] Replace CostCalculationHelper with litellm.cost_per_token#1906
Open
chandrasekharan-zipstack wants to merge 2 commits intomainfrom
Open
[REFACTOR] Replace CostCalculationHelper with litellm.cost_per_token#1906chandrasekharan-zipstack wants to merge 2 commits intomainfrom
chandrasekharan-zipstack wants to merge 2 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
Summary by CodeRabbit
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 moved from platform-service to SDK; litellm.cost_per_token called before provider-prefix stripping; embedding completion tokens correctly zeroed; exception fallback to 0.0 is safe. |
| platform-service/src/unstract/platform_service/controller/platform.py | CostCalculationHelper removed; cost_in_dollars now read directly from payload with a safe 0.0 default for backward compatibility with older SDK clients. |
| platform-service/src/unstract/platform_service/helper/cost_calculation.py | Deleted — external HTTP pricing fetch, file-storage TTL cache, and manual model-price lookup replaced by litellm's bundled pricing database. |
| platform-service/src/unstract/platform_service/env.py | MODEL_PRICES_URL, MODEL_PRICES_TTL_IN_DAYS, and MODEL_PRICES_FILE_PATH env vars removed; no remaining callers. |
| platform-service/src/unstract/platform_service/utils.py | Orphaned format_float_positional helper removed; no remaining callers after cost_calculation.py deletion. |
Sequence Diagram
sequenceDiagram
participant SDK as sdk1 Audit
participant LiteLLM as litellm.cost_per_token
participant PS as platform-service /usage
SDK->>LiteLLM: cost_per_token(model="azure/gpt-4o", prompt_tokens, completion_tokens)
LiteLLM-->>SDK: (prompt_cost, completion_cost)
Note over SDK: cost_in_dollars = prompt_cost + completion_cost<br/>display_model_name = "gpt-4o" (prefix stripped)
SDK->>PS: POST /usage { model_name, cost_in_dollars, ... }
PS-->>SDK: 200 OK
Prompt To Fix All With AI
This is a comment left during a code review.
Path: unstract/sdk1/src/unstract/sdk1/audit.py
Line: 104-107
Comment:
**Exception details not captured in cost-lookup failure log**
The bare `except Exception` silently swallows the error type and message. When `cost_per_token` fails for a new or custom model name, the only log entry is the model name — the root cause (e.g. `NotFoundError`, `BudgetExceededError`, malformed model string) is lost. Adding `exc_info=True` makes these failures much easier to diagnose without changing the fallback behaviour.
```suggestion
except Exception:
logger.debug(
"Cost lookup failed for model %s, defaulting to 0",
model_name,
exc_info=True,
)
```
How can I resolve this? If you propose a fix, please make it concise.Reviews (2): Last reviewed commit: "[REFACTOR] Zero out completion_tokens fo..." | 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
Test ResultsSummary
Runner Tests - Full Report
SDK1 Tests - Full Report
|
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
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
3 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.