Store the model ID without the provider ID in DB #379
Conversation
In 3f7ed75 we accidentally stored the model ID as "<provider>/<model>", which is not what we want. We want to store the model ID as is, so we can re-use it later when the user omits it, without string manipulation.
WalkthroughThe model selection utility now returns (llama_stack_model_id, model_id, provider_id). Both query and streaming_query handlers unpack the triple and pass llama_stack_model_id to retrieve_response, while continuing to use model_id/provider_id for metrics and persistence. Tests were updated to reflect the new return signature and assertions. Changes
Sequence Diagram(s)sequenceDiagram
actor Client
participant API as Query/Streaming Handlers
participant Selector as select_model_and_provider_id
participant Retriever as retrieve_response
Client->>API: request(model_id?, provider_id?)
API->>Selector: select_model_and_provider_id(models, model_id, provider_id)
Selector-->>API: (llama_stack_model_id, model_id, provider_id)
API->>Retriever: retrieve_response(model_id=llama_stack_model_id, ...)
Retriever-->>API: response
API-->>Client: response (metrics use model_id/provider_id)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~15–20 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (1)
src/app/endpoints/query.py (1)
263-272: Bug: fallback returns provider-prefixed identifier as model_id when selecting first LLMWhen no model/provider hints are provided and you pick the first available LLM, you set:
model_id = model.identifierreturn model_id, model_id, provider_idIf
model.identifieris "provider/model", this persists "provider/model" as the plainmodel_id, reintroducing the bug this PR aims to fix. Extract the plain model name and return the proper triple.Apply this diff:
try: model = next( m for m in models if m.model_type == "llm" # pyright: ignore[reportAttributeAccessIssue] ) - model_id = model.identifier - provider_id = model.provider_id - logger.info("Selected model: %s", model) - return model_id, model_id, provider_id + # Fully-qualified identifier as exposed by LlamaStack + llama_stack_model_id = model.identifier + provider_id = model.provider_id + # Persist the plain model id (without provider prefix) to DB/metrics + plain_model_id = llama_stack_model_id.split("/", 1)[-1] + logger.info("Selected model: %s", model) + return llama_stack_model_id, plain_model_id, provider_idThis ensures:
- llama_stack_model_id: provider/model
- model_id (plain): model
- provider_id: provider
🧹 Nitpick comments (7)
src/app/endpoints/streaming_query.py (1)
424-429: Triple unpack aligns with the new contractDestructuring
(llama_stack_model_id, model_id, provider_id)here is correct and keepsmodel_id/provider_idavailable for persistence/metrics.Add a brief inline comment to clarify the semantics of each tuple element to avoid future misuse.
src/app/endpoints/query.py (2)
242-243: Signature now returns a triple (llama_stack_model_id, model_id, provider_id)Type hint update looks good. Consider clarifying the docstring to explicitly define each element’s semantics (full vs. plain identifiers).
300-301: Return ordering and values are correct in the validated pathReturning
(llama_stack_model_id, model_id, provider_id)after validation matches the new contract and preserves the plainmodel_idfor persistence.Optionally, defensively guard against
model_idcontaining '/' here by normalizing it before persistence to avoid future regressions if callers pass a provider-prefixed model.tests/unit/app/endpoints/test_streaming_query.py (3)
268-269: Strengthen the mock to better validate tuple semanticsCurrently the mock returns identical values for the first two positions, making it hard to catch regressions if the wrong element is passed to
retrieve_response. Use distinct values and (in a test) assert the first element is used for the agent call.Consider:
- return_value=("fake_model_id", "fake_model_id", "fake_provider_id"), + # First element is the fully-qualified llama id; second is the plain model id + return_value=("provX/modelX", "modelX", "provX"),Additionally, in
test_auth_tuple_unpacking_in_streaming_query_endpoint_handler, add an assertion thatretrieve_responsewas called with the fully-qualified id:# After the handler call: # args: (client, model_id, query_request, token, ...) assert mock_retrieve_response.call_args[0][1] == "test_model" # or "provX/modelX" if you adjust the mock here too
1281-1283: Updated mock triple looks consistentReturn shape matches the new contract. Consider the suggestion above to use distinct values for first/second to catch mis-wiring.
1327-1329: Updated mock triple looks consistentSame note: distinct first/second values would improve regression detection.
tests/unit/app/endpoints/test_query.py (1)
135-136: Mock triple matches the new return signatureGood. To make tests more robust, prefer using a distinct first/second element (e.g., "provX/modelX" vs "modelX") when you want to assert which one is used where.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
src/app/endpoints/query.py(4 hunks)src/app/endpoints/streaming_query.py(1 hunks)tests/unit/app/endpoints/test_query.py(7 hunks)tests/unit/app/endpoints/test_streaming_query.py(4 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: e2e_tests
🔇 Additional comments (8)
src/app/endpoints/streaming_query.py (1)
430-436: Correct identifier used for LlamaStack callPassing
llama_stack_model_idintoretrieve_responseensures the agent receives the fully-qualified identifier while DB/metrics continue to use the plain model and provider IDs.src/app/endpoints/query.py (1)
186-195: Using llama_stack_model_id for retrieval is correct; keep model/provider for metricsThe triple destructuring and forwarding
llama_stack_model_idintoretrieve_responsepreserves the intent: API calls use the fully-qualified id, while DB/metrics use the plain model and provider.tests/unit/app/endpoints/test_streaming_query.py (1)
1374-1376: Updated mock triple looks consistentNo issues.
tests/unit/app/endpoints/test_query.py (5)
217-223: Asserting tuple decomposition from request inputs is correctThe assertions reflect the desired behavior: llama id includes provider/model, plain model is stored separately.
253-260: Asserting tuple decomposition from configuration defaults is correctLooks good and aligned with the contract.
1140-1142: Mock triple matches the new return signatureConsistent with updated code paths.
1179-1181: Mock triple matches the new return signatureConsistent with updated code paths.
1218-1220: Mock triple matches the new return signatureConsistent with updated code paths.
| llama_stack_model_id, model_id, provider_id = select_model_and_provider_id( | ||
| model_list, query_request.model, query_request.provider | ||
| ) | ||
|
|
||
| # Assert return the first available LLM model when no model/provider is | ||
| # specified in the request or in the configuration | ||
| assert llama_stack_model_id == "first_model" | ||
| assert model_id == "first_model" | ||
| assert provider_id == "provider1" |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Add coverage for fallback with provider-prefixed identifiers
This test uses identifiers without provider prefixes, so it won’t catch the regression where the fallback returns a provider-prefixed model as the plain model. Add a companion test where the first available LLM has identifier "provider1/model1" and assert that the plain model is "model1".
Proposed additional test:
def test_select_model_and_provider_id_first_from_list_with_prefixed_identifier(mocker):
"""Ensure fallback splits provider/model correctly for persistence."""
model_list = [
mocker.Mock(identifier="provider1/model1", model_type="llm", provider_id="provider1"),
mocker.Mock(identifier="provider2/model2", model_type="llm", provider_id="provider2"),
]
query_request = QueryRequest(query="What is OpenStack?")
llama_stack_model_id, model_id, provider_id = select_model_and_provider_id(
model_list, query_request.model, query_request.provider
)
assert llama_stack_model_id == "provider1/model1"
assert model_id == "model1" # plain model id
assert provider_id == "provider1"🏁 Script executed:
#!/bin/bash
# Locate the implementation of select_model_and_provider_id
rg "def select_model_and_provider_id" -n -A 10Length of output: 920
🏁 Script executed:
#!/bin/bash
# Display the full implementation of select_model_and_provider_id for review
sed -n '240,340p' src/app/endpoints/query.pyLength of output: 3986
Ensure fallback handles provider-prefixed model identifiers
The current fallback branch in select_model_and_provider_id unconditionally assigns both returned IDs to model.identifier, so if that identifier already contains a provider prefix (e.g. "provider1/model1"), the plain model_id also comes back as "provider1/model1". We should:
- Add a unit test to catch this case.
- Update the fallback implementation to split out the plain model ID.
Locations to update:
tests/unit/app/endpoints/test_query.pysrc/app/endpoints/query.py(fallback branch around line 250)
Proposed test addition:
def test_select_model_and_provider_id_first_with_prefixed_identifier(mocker):
"""Fallback should split provider/model correctly when identifier is prefixed."""
model_list = [
mocker.Mock(identifier="provider1/model1", model_type="llm", provider_id="provider1"),
mocker.Mock(identifier="provider2/model2", model_type="llm", provider_id="provider2"),
]
query_request = QueryRequest(query="foo")
llama_stack_model_id, model_id, provider_id = select_model_and_provider_id(
model_list, query_request.model, query_request.provider
)
assert llama_stack_model_id == "provider1/model1"
assert model_id == "model1" # plain model id
assert provider_id == "provider1"Proposed implementation diff in src/app/endpoints/query.py:
model = next(m for m in models if m.model_type == "llm")
- model_id = model.identifier
- provider_id = model.provider_id
- return model_id, model_id, provider_id
+ full_id = model.identifier
+ provider_id = model.provider_id
+ # split off plain model ID if a prefix is present
+ plain_id = full_id.split("/", 1)[1] if "/" in full_id else full_id
+ return full_id, plain_id, provider_id🤖 Prompt for AI Agents
In tests/unit/app/endpoints/test_query.py around lines 279-287, add a unit test
that supplies a model_list whose first model.identifier is provider-prefixed
(e.g. "provider1/model1") and asserts that select_model_and_provider_id returns
llama_stack_model_id as the full "provider1/model1", model_id as the plain
"model1", and provider_id as "provider1"; and in src/app/endpoints/query.py
around line ~250 update the fallback branch so that when using model.identifier
you detect a provider prefix (split on the first '/' if present), return the
full identifier as llama_stack_model_id, the right-hand segment as model_id, and
the provider part as provider_id (fall back to existing model.provider_id when
no prefix), keeping other behavior unchanged.
2 participants
Description
In 3f7ed75 we accidentally stored the model ID as "<provider>/<model>", which is not what we want. We want to store the model ID as is, so we can re-use it later when the user omits it, without string manipulation.
Type of change
Related Tickets & Documents
3f7ed75
Checklist before requesting a review
Testing
Modified unit tests. Made sure omitting the model ID and provider ID when giving the conversation ID actually works now.
Summary by CodeRabbit
New Features
Refactor
Tests