-
Notifications
You must be signed in to change notification settings - Fork 85
LCORE-1880: Refactor of responses models dumping #1645
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
asimurka
wants to merge
1
commit into
lightspeed-core:main
Choose a base branch
from
asimurka:refactor_responses_endpoint
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
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
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
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,39 @@ | ||
| """Utility functions for models.""" | ||
|
|
||
| from typing import Any | ||
|
|
||
| from llama_stack_api.openai_responses import OpenAIResponseInputTool as InputTool | ||
|
|
||
|
|
||
| def add_mcp_authorizations( | ||
| dumped_tools: list[dict[str, Any]], | ||
| tools: list[InputTool], | ||
| ) -> list[dict[str, Any]]: | ||
| """Merge MCP authorization into serialized tool dicts keyed by server_label. | ||
|
|
||
| Args: | ||
| dumped_tools: Serialized tools. | ||
| tools: Live tool models. MCP entries with authorization are mapped by | ||
| server_label. | ||
|
|
||
| Returns: | ||
| A new list of dicts. For MCP rows, authorization is set only when a | ||
| matching non-None token exists. | ||
| """ | ||
| authorizations = { | ||
| tool.server_label: tool.authorization | ||
| for tool in tools | ||
| if tool.type == "mcp" and tool.authorization is not None | ||
| } # server_labels are unique by design | ||
| result: list[dict[str, Any]] = [] | ||
| for dumped in dumped_tools: | ||
| row = dict(dumped) | ||
| if ( | ||
| row.get("type") == "mcp" | ||
| and (label := row.get("server_label")) is not None | ||
| and (token := authorizations.get(label)) is not None | ||
| ): | ||
| row["authorization"] = token | ||
|
|
||
| result.append(row) | ||
| return result |
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,111 @@ | ||
| """Unit tests for models.utils (mirrors src/models/utils.py).""" | ||
|
|
||
| from llama_stack_api.openai_responses import ( | ||
| OpenAIResponseInputToolFileSearch as InputToolFileSearch, | ||
| ) | ||
| from llama_stack_api.openai_responses import ( | ||
| OpenAIResponseInputToolMCP as InputToolMCP, | ||
| ) | ||
|
|
||
| from models.utils import add_mcp_authorizations | ||
|
|
||
|
|
||
| class TestAddMcpAuthorizations: | ||
| """Tests for add_mcp_authorizations with realistic MCP tool rows. | ||
|
|
||
| Assumes server_label is present on MCP dicts and unique across configured | ||
| servers; see InputToolMCP in llama-stack-api. | ||
| """ | ||
|
|
||
| def test_merges_authorization_by_server_label(self) -> None: | ||
| """MCP model_dump omits authorization; the helper restores it by server_label.""" | ||
| live = InputToolMCP( | ||
| server_label="alpha", | ||
| server_url="http://alpha", | ||
| require_approval="never", | ||
| authorization="secret-token", | ||
| ) | ||
| dumped = [live.model_dump()] | ||
| assert "authorization" not in dumped[0] | ||
|
|
||
| out = add_mcp_authorizations(dumped, [live]) | ||
| assert len(out) == 1 | ||
| assert out[0]["authorization"] == "secret-token" | ||
| assert out[0]["server_label"] == "alpha" | ||
|
|
||
| def test_two_mcp_servers_distinct_tokens(self) -> None: | ||
| """Each server_label receives its own authorization.""" | ||
| a = InputToolMCP( | ||
| server_label="srv-a", | ||
| server_url="http://a", | ||
| require_approval="never", | ||
| authorization="token-a", | ||
| ) | ||
| b = InputToolMCP( | ||
| server_label="srv-b", | ||
| server_url="http://b", | ||
| require_approval="never", | ||
| authorization="token-b", | ||
| ) | ||
| dumped = [a.model_dump(), b.model_dump()] | ||
| assert "authorization" not in dumped[0] | ||
| assert "authorization" not in dumped[1] | ||
|
|
||
| out = add_mcp_authorizations(dumped, [a, b]) | ||
| assert out[0]["authorization"] == "token-a" | ||
| assert out[1]["authorization"] == "token-b" | ||
|
|
||
| def test_file_search_row_unchanged_no_authorization_merge(self) -> None: | ||
| """Non-MCP rows are copied; MCP row still gets auth from live list.""" | ||
| mcp = InputToolMCP( | ||
| server_label="m", | ||
| server_url="http://m", | ||
| require_approval="never", | ||
| authorization="mcp-secret", | ||
| ) | ||
| fs = InputToolFileSearch(type="file_search", vector_store_ids=["vs-1"]) | ||
| dumped = [fs.model_dump(), mcp.model_dump()] | ||
| assert "authorization" not in dumped[1] | ||
|
|
||
| out = add_mcp_authorizations(dumped, [fs, mcp]) | ||
| assert out[0]["type"] == "file_search" | ||
| assert "authorization" not in out[0] | ||
| assert out[1]["authorization"] == "mcp-secret" | ||
|
|
||
| def test_subset_dumped_rows_still_match_live_by_label(self) -> None: | ||
| """When only some MCP tools appear in dumped_tools, labels still align.""" | ||
| first = InputToolMCP( | ||
| server_label="one", | ||
| server_url="http://one", | ||
| require_approval="never", | ||
| authorization="tok-one", | ||
| ) | ||
| second = InputToolMCP( | ||
| server_label="two", | ||
| server_url="http://two", | ||
| require_approval="never", | ||
| authorization="tok-two", | ||
| ) | ||
| dumped = [second.model_dump()] | ||
| assert "authorization" not in dumped[0] | ||
|
|
||
| out = add_mcp_authorizations(dumped, [first, second]) | ||
| assert len(out) == 1 | ||
| assert out[0]["authorization"] == "tok-two" | ||
|
|
||
| def test_does_not_mutate_input_list_or_dicts(self) -> None: | ||
| """Output is new containers; inputs stay as provided.""" | ||
| live = InputToolMCP( | ||
| server_label="s", | ||
| server_url="http://s", | ||
| require_approval="never", | ||
| authorization="t", | ||
| ) | ||
| dumped = [live.model_dump()] | ||
| row = dumped[0] | ||
| assert "authorization" not in row | ||
|
|
||
| out = add_mcp_authorizations(dumped, [live]) | ||
| assert out is not dumped | ||
| assert out[0] is not row | ||
| assert "authorization" not in row |
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
Oops, something went wrong.
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.
Uh oh!
There was an error while loading. Please reload this page.