[FIX] Use actual model name for Azure OpenAI cost tracking#1805
Open
pk-zipstack wants to merge 3 commits intomainfrom
Open
[FIX] Use actual model name for Azure OpenAI cost tracking#1805pk-zipstack wants to merge 3 commits intomainfrom
pk-zipstack wants to merge 3 commits intomainfrom
Conversation
Azure OpenAI uses deployment_name for LiteLLM API routing, but cost tracking needs the real model name (e.g., gpt-4o) to match pricing table entries. Preserve the user-provided model as cost_model in validate() output and use it for usage recording and embedding metadata. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review infoConfiguration used: Organization UI Review profile: CHILL Plan: Pro Cache: Disabled due to Reviews > Disable Cache setting Knowledge base: Disabled due to 📒 Files selected for processing (1)
Summary by CodeRabbit
WalkthroughCapture user-provided model names as a separate cost_model during Azure parameter validation, store cost_model on Embedding/LLM instances (removing it from kwargs), and propagate cost_model into usage/logging while deployment_name continues to be used for routing. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant Adapter as AzureAdapter
participant Embedding
participant LLM
participant Usage as UsageLogger
User->>AzureAdapter: provide model (may be deployment_name)
AzureAdapter->>AzureAdapter: normalize deployment_name\ncapture original model -> cost_model
AzureAdapter->>Embedding: pass kwargs (with cost_model)
Embedding->>Embedding: pop store cost_model as _cost_model
User->>LLM: request completion
LLM->>LLM: use deployment_name for routing\nuse _cost_model (or fallback) for logging
LLM->>Usage: emit usage with model = cost_model
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 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 |
for more information, see https://pre-commit.ci
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@unstract/sdk1/src/unstract/sdk1/embedding.py`:
- Around line 195-199: Move the pop of "cost_model" into Embedding.__init__: in
the Embedding class constructor, pop "cost_model" from self.kwargs (or kwargs)
and store it as self._cost_model so it is removed before any calls to
litellm.embedding(); remove the current pop in the code that sets
self.model_name (the snippet using
self._embedding_instance.kwargs.pop("cost_model", None)), and update
EmbeddingCompat to read adapter metadata from Embedding._cost_model instead of
relying on post-init popping; ensure get_embeddings(), get_aembedding(), and
get_aembeddings() no longer pass cost_model through to litellm.embedding(),
while get_embedding() behavior remains unchanged.
ℹ️ Review info
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to Reviews > Disable Cache setting
Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting
📒 Files selected for processing (3)
unstract/sdk1/src/unstract/sdk1/adapters/base1.pyunstract/sdk1/src/unstract/sdk1/embedding.pyunstract/sdk1/src/unstract/sdk1/llm.py
pk-zipstack
requested review from
Deepak-Kesavan,
chandrasekharan-zipstack and
gaya3-zipstack
Move cost_model removal from EmbeddingCompat into Embedding.__init__ so it is popped before the test connection call and never passed to litellm.embedding() in any method (get_embeddings, get_aembedding, get_aembeddings). EmbeddingCompat now reads the stored _cost_model attribute instead of post-init popping from kwargs. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
|
Contributor
Test ResultsSummary
Runner Tests - Full Report
SDK1 Tests - Full Report
|
chandrasekharan-zipstack
approved these changes
Feb 25, 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
2 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
modelfield (actual model name likegpt-4o) ascost_modelin Azure OpenAI adapter validation outputcost_modelfor usage recording and embedding metadata instead of the deployment nameWhy
deployment_name(e.g.,my-gpt4-prod) as themodelparameter for LiteLLM routing (azure/my-gpt4-prod), which is correct for API callsmodel_nameagainst a pricing table keyed by actual model names (e.g.,azure/gpt-4o)model_nameisazure/my-gpt4-prod, no pricing entry matches and costs are not trackedHow
base1.py—AzureOpenAILLMParameters.validate(): Capture originalmodelvalue beforevalidate_model()overwrites it withdeployment_name. Addcost_model(withazure/prefix) to the validated output dict when a model name was providedbase1.py—AzureOpenAIEmbeddingParameters.validate(): Same pattern for embedding adaptersllm.py—__init__(): Popcost_modelfromself.kwargsintoself._cost_modelso it doesn't get passed to litellmllm.py—complete(),stream_complete(),acomplete(): Popcost_modelfrom re-validatedcompletion_kwargsbefore passing to litellm. Useself._cost_model(falling back toself.kwargs["model"]) in_record_usage()callsembedding.py—EmbeddingCompat.__init__(): Popcost_modelfrom kwargs and use it forself.model_name, falling back to the routing model nameCan 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)
cost_modelin theirvalidate()output, soself._cost_modelisNoneandself.kwargs["model"]is used as before — zero impact. For Azure adapters, when themodelfield is empty,cost_modelis not set and behavior falls back to current (deployment name used for cost). Thecost_modelkey is always popped before passing kwargs to litellm, so no unknown parameter errors.Database Migrations
Env Config
Relevant Docs
Related Issues or PRs
Dependencies Versions
Notes on Testing
model=gpt-4oanddeployment_name=my-deploygpt-4o(matching the pricing table) instead ofmy-deployScreenshots
N/A
Checklist
I have read and understood the Contribution Guidelines.
🤖 Generated with Claude Code