refactor conversation to agent and conversation mapping #317
Conversation
WalkthroughThis change removes the global conversation-to-agent ID mapping and explicit TTLCache-based agent caching from the codebase. Instead, agent retrieval and session management now rely directly on client API calls. Agent IDs are assigned from conversation IDs, and session IDs are explicitly handled. All related test code is updated to eliminate references to the removed caches and mappings. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant API_Endpoint
participant Agent
Client->>API_Endpoint: Request (with conversation_id)
API_Endpoint->>Agent: get_agent(conversation_id) via client API
Note right of API_Endpoint: agent_id = conversation_id\nsession_id retrieved from agent sessions
Agent-->>API_Endpoint: Agent instance, conversation_id, session_id
API_Endpoint->>Agent: agent.create_turn(session_id, ...)
API_Endpoint-->>Client: Response
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Poem
Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🔇 Additional comments (11)
✨ 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 (
|
maorfr
force-pushed
the
agent-conversation-mapping
branch
from
5152a7d to
0c16c16
Compare
manstis
left a comment
manstis
left a comment
There was a problem hiding this comment.
I welcome this improvement; but wonder if we can go a step further.
See my comments.
eranco74
left a comment
eranco74
left a comment
There was a problem hiding this comment.
Looks great!
I do think we should keep the agent cache, but I guess we can start without and see if this has an impact on performance
eranco74
left a comment
eranco74
left a comment
There was a problem hiding this comment.
You have an issue with the session!
The conversation id you return from get_agent is later used as the session_id.
Instead you should list the agent sessions and use the one you find, if you didn't find an existing session -> create a new one.
maorfr
force-pushed
the
agent-conversation-mapping
branch
from
dec00e9 to
7f61218
Compare
|
while this code is ugly - it works. |
![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
🧹 Nitpick comments (1)
src/app/endpoints/streaming_query.py (1)
559-559: Logging shows conversation_id instead of session_idThe debug log still shows "Session ID: %s" but logs
conversation_idinstead ofsession_id. This could be confusing for debugging.Apply this diff to fix the logging:
- logger.debug("Session ID: %s", conversation_id) + logger.debug("Session ID: %s", session_id)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
src/app/endpoints/conversations.py(3 hunks)src/app/endpoints/query.py(5 hunks)src/app/endpoints/streaming_query.py(5 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- src/app/endpoints/conversations.py
- src/app/endpoints/query.py
🧰 Additional context used
🧬 Code Graph Analysis (1)
src/app/endpoints/streaming_query.py (2)
src/utils/suid.py (1)
get_suid(6-12)src/app/endpoints/query.py (1)
get_agent(73-109)
🔇 Additional comments (5)
src/app/endpoints/streaming_query.py (5)
4-4: Import added for error suppressionThe
suppressimport is correctly added to handle theValueErrorwhen retrieving non-existent agents.
58-58: Function signature updated for explicit session handlingThe return type change from
tuple[AsyncAgent, str]totuple[AsyncAgent, str, str]correctly reflects the new explicit session ID handling, maintaining consistency with the synchronous version inquery.py.
60-64: Agent retrieval logic correctly implementedThe async agent retrieval with error suppression is properly implemented. The
suppress(ValueError)appropriately handles cases where the agent doesn't exist, allowing the flow to continue with new agent creation.
530-537: Function call correctly updated for three-value returnThe unpacking of the
get_agentreturn value is correctly updated to handle the new session_id parameter.
568-568: Session ID correctly used in agent.create_turnThe explicit use of
session_idinstead ofconversation_idas the session parameter is correct and aligns with the PR objectives of proper session management.
umago
left a comment
umago
left a comment
There was a problem hiding this comment.
Thanks for the patch, a few other comments inline. Maybe I'm missing something but, the logic here seems quite confusing to me
src/app/endpoints/query.py
Outdated
| existing_agent_id = None | ||
| if conversation_id: | ||
| with suppress(ValueError): | ||
| existing_agent_id = client.agents.retrieve(agent_id=conversation_id).agent_id |
There was a problem hiding this comment.
I'm trying to get my head around the logic here. If conversation_id is provided and an Agent already exists that matches that conversation_id (agent_id == conversation_id). Why not return the existing agent at this point ?
Why do we need to create a new agent, replace it's agent_id with conversation_id, delete the old agent and then return ?
There was a problem hiding this comment.
client.agents.retrieve returns the agent configuration, not an agent instance.
We need to call this method to verify that the agent exists in the llama-stack, but we still need to create the agent instance later.
| if existing_agent_id and conversation_id: | ||
| orphan_agent_id = agent.agent_id | ||
| agent.agent_id = conversation_id | ||
| client.agents.delete(agent_id=orphan_agent_id) |
There was a problem hiding this comment.
Shouldn't this delete happens before we change the ID of the newly created agent to conversation_id ? Otherwise, this delete could also affect the new agent as it's ID will now match here ?
There was a problem hiding this comment.
Note agent_id=orphan_agent_id, this deletes the newly created agent from the llama-stack server DB
There was a problem hiding this comment.
@umago IDK if this could help explain: #317 (comment)
There was a problem hiding this comment.
Thanks! I was on PTO yesterday, reading thru it now. It does makes sense. Thank you
maorfr
force-pushed
the
agent-conversation-mapping
branch
from
7f61218 to
8ae4803
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
src/app/endpoints/conversations.py(3 hunks)src/app/endpoints/query.py(5 hunks)src/app/endpoints/streaming_query.py(5 hunks)tests/unit/app/endpoints/test_conversations.py(0 hunks)tests/unit/app/endpoints/test_query.py(0 hunks)tests/unit/app/endpoints/test_streaming_query.py(0 hunks)
💤 Files with no reviewable changes (3)
- tests/unit/app/endpoints/test_query.py
- tests/unit/app/endpoints/test_conversations.py
- tests/unit/app/endpoints/test_streaming_query.py
🚧 Files skipped from review as they are similar to previous changes (2)
- src/app/endpoints/conversations.py
- src/app/endpoints/streaming_query.py
🧰 Additional context used
🧬 Code Graph Analysis (1)
src/app/endpoints/query.py (2)
src/utils/suid.py (1)
get_suid(6-12)src/app/endpoints/streaming_query.py (1)
get_agent(51-89)
⏰ 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 (5)
src/app/endpoints/query.py (5)
3-3: LGTM: Import necessary for new error handling.The
suppressimport is correctly added to support the new agent retrieval logic that gracefully handles cases where an agent doesn't exist.
81-81: LGTM: Function signature updated correctly.The return type change to include
session_idaligns with the PR objectives and is consistent with similar changes in the streaming endpoint.
83-86: LGTM: Clean agent retrieval implementation.The logic correctly checks for existing agents using the client API instead of cache lookup. The use of
suppress(ValueError)is appropriate for handling non-existent agents gracefully.
284-284: LGTM: Function call updated correctly.The unpacking correctly handles the new return signature of
get_agentthat now includessession_id.
328-328: LGTM: Correct session management.Using
session_idinstead ofconversation_idin thecreate_turncall is the correct approach for explicit session management as outlined in the PR objectives.
maorfr
force-pushed
the
agent-conversation-mapping
branch
from
a15c75c to
f055c77
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (4)
tests/unit/app/endpoints/test_query.py (3)
1084-1112: Consider renaming the test function for clarity.The function name
test_get_agent_cache_hitis misleading since the cache was removed in this refactor. Consider renaming totest_get_agent_existing_conversationor similar to better reflect that it tests retrieving an existing agent by conversation ID.The test logic correctly verifies the new three-tuple return and session handling.
1114-1157: Consider renaming the test function for clarity.The function name
test_get_agent_cache_miss_with_conversation_idis misleading since the cache was removed. Consider renaming totest_get_agent_nonexistent_conversationto better reflect that it tests the fallback behavior when an agent cannot be retrieved by conversation ID.The test logic correctly verifies that a new agent is created when the existing one is not found.
1084-1708: Consider adding test coverage for session retrieval edge cases.While the existing test updates are comprehensive and correctly handle the new
get_agentsignature, consider adding test coverage for:
- What happens when
client.agents.session.listreturns empty data or multiple sessions- Error handling when session retrieval fails
- Verification that
client.agents.deleteis called when reusing an existing agent (as mentioned in the relevant code snippet)The current changes are well-executed and maintain consistency across all test functions.
tests/unit/app/endpoints/test_streaming_query.py (1)
1214-1227: Consider updating the test name for clarity.The test name
test_get_agent_cache_hitsuggests cache functionality, but according to the PR objectives, the explicit caching mechanism was removed in favor of direct agent retrieval. The test logic correctly mocksAsyncAgentand unpacks the three return values, but the name might be misleading.Consider renaming to something like
test_get_agent_with_existing_conversation_idto better reflect the current functionality.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
src/app/endpoints/query.py(5 hunks)src/app/endpoints/streaming_query.py(5 hunks)tests/unit/app/endpoints/test_conversations.py(4 hunks)tests/unit/app/endpoints/test_query.py(40 hunks)tests/unit/app/endpoints/test_streaming_query.py(38 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
- tests/unit/app/endpoints/test_conversations.py
- src/app/endpoints/streaming_query.py
- src/app/endpoints/query.py
🧰 Additional context used
🧬 Code Graph Analysis (1)
tests/unit/app/endpoints/test_query.py (1)
src/app/endpoints/query.py (1)
get_agent(71-110)
⏰ 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 (13)
tests/unit/app/endpoints/test_query.py (5)
380-382: LGTM!Mock correctly updated to return the new three-tuple
(agent, conversation_id, session_id)matching the refactoredget_agentfunction signature.
395-395: LGTM!Assertions correctly updated to verify the
conversation_idreturn value from the refactoredget_agentfunction. All assertions consistently check for"fake_conversation_id"which matches the mock setup.Also applies to: 430-430, 478-478, 529-529, 582-582, 637-637, 690-690, 744-744, 810-810, 884-884, 962-962
417-419: LGTM!All mock patches consistently updated to return the new three-tuple format. The consistent use of the same mock values across all tests ensures maintainability and reduces confusion.
Also applies to: 465-467, 516-518, 569-571, 624-626, 677-679, 731-733, 797-799, 857-859, 949-951, 1499-1501, 1548-1550
1195-1207: LGTM!Test functions correctly updated to handle the new three-tuple return from
get_agent. The assertions properly verify that:
conversation_idequalsagent.agent_id(aligning with the PR objective of using agent ID as conversation ID)session_idmatches the expected mock valueAlso applies to: 1246-1258, 1301-1313
1611-1624: LGTM!Test functions correctly verify the
no_toolsparameter functionality while properly handling the new three-tuple return fromget_agent. The assertions appropriately validate both the agent creation and the return values.Also applies to: 1670-1683
tests/unit/app/endpoints/test_streaming_query.py (8)
317-317: LGTM!The mock return value correctly reflects the updated
get_agentfunction signature that now returns(agent, conversation_id, session_id)instead of the previous(agent, conversation_id).
334-334: LGTM!The test correctly uses the session_id returned by
get_agentin thecreate_turncall, reflecting the updated session management approach.
354-354: LGTM!All these changes consistently follow the updated pattern:
- Mock return values for
get_agentnow include the third elementsession_idcreate_turncalls properly use thesession_idparameterThe test updates correctly reflect the refactored session management.
Also applies to: 371-371, 404-404, 419-419, 455-455, 470-470, 508-508, 535-535, 563-563, 578-578, 616-616, 631-631
1248-1250: LGTM!The mock correctly simulates the scenario where agent retrieval fails, triggering the fallback to create a new agent. This tests the error handling path appropriately.
1275-1275: LGTM!The test correctly implements the unified ID approach where:
- The function returns three values as expected
- The conversation_id is unified with the agent_id
- The session_id is properly tracked and asserted
This aligns with the PR objective of simplifying the conversation-agent-session relationship.
Also applies to: 1286-1287
1332-1332: LGTM!These changes consistently implement:
- Correct unpacking of the three-value return from
get_agent- Proper assertions validating the unified conversation_id/agent_id approach
- Appropriate session_id handling
All changes align with the refactoring objectives.
Also applies to: 1343-1344, 1389-1389, 1400-1401, 1448-1448, 1459-1460
1561-1561: Verify the return value assignment.The variable name
result_session_idsuggests the function should return a session_id, butstreaming_query_endpoint_handlertypically returns aStreamingResponseobject. This might be an incorrect variable name or an issue with the test logic.Please verify that the
streaming_query_endpoint_handlerfunction signature and return type are correctly reflected in this test.
1674-1674: LGTM!These final changes maintain consistency with the refactoring approach:
- Proper 3-tuple returns from mocked
get_agentcalls- Correct conversation_id assertions reflecting the unified ID approach
- Appropriate session_id handling throughout
All changes support the simplified conversation-agent-session relationship as intended.
Also applies to: 1686-1686, 1724-1724, 1736-1736, 1790-1790, 1802-1803, 1854-1854, 1866-1867
maorfr
force-pushed
the
agent-conversation-mapping
branch
from
eb35db5 to
baab097
Compare
|
@eranco74 also PTAL |
| assert exc_info.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE | ||
| assert "Unable to connect to Llama Stack" in exc_info.value.detail["response"] | ||
|
|
||
| def test_llama_stack_not_found_error(self, mocker, setup_configuration): |
|
|
||
|
|
||
| def test_get_agent_cache_hit(prepare_agent_mocks): | ||
| def test_get_agent_cache_hit(prepare_agent_mocks, mocker): |
There was a problem hiding this comment.
we no longer have cache...
yet we want to test this case (rehydration) just need to rename the test
| conversation_id=conversation_id, | ||
| ) | ||
|
|
||
| # Assert cached agent is returned |
| ): | ||
| """Test get_agent function when conversation_id is provided but agent not in cache.""" | ||
| mock_client, mock_agent = prepare_agent_mocks | ||
| mock_client.agents.retrieve.side_effect = ValueError( |
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (3)
tests/unit/app/endpoints/test_query.py (3)
1084-1084: Rename test to reflect agent rehydration instead of cache hitThe test name still references "cache hit" but the cache has been removed from the implementation. Consider renaming to
test_get_agent_rehydrationor similar.
1106-1107: Update comment to remove cache referenceThe comment "Assert cached agent is returned" is misleading since the cache has been removed. Update to reflect the actual behavior.
1112-1115: Rename test to remove cache miss referenceThe test name references "cache miss" but caching has been removed. Consider renaming to
test_get_agent_with_existing_conversation_idor similar.
🧹 Nitpick comments (1)
tests/unit/app/endpoints/test_conversations.py (1)
264-281: Inconsistent mock setup - unused mock object.The test creates
mock_session_objwith amodel_dumpmethod (lines 259-260) but then usesmock_session_datadirectly in the list return value. Since the implementation now usessession.list, the mock should return the raw data format that the actual API would return.Consider removing the unused mock object:
- # Mock session data with model_dump method - mock_session_obj = mocker.Mock() - mock_session_obj.model_dump.return_value = mock_session_data - # Mock LlamaStackClientHolder mock_client = mocker.Mock() mock_client.agents.session.list.return_value = mocker.Mock( data=[mock_session_data] )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tests/unit/app/endpoints/test_conversations.py(6 hunks)tests/unit/app/endpoints/test_query.py(40 hunks)tests/unit/app/endpoints/test_streaming_query.py(38 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/unit/app/endpoints/test_streaming_query.py
🔇 Additional comments (3)
tests/unit/app/endpoints/test_conversations.py (3)
118-135: LGTM! Test updates align with function signature changes.The changes correctly reflect the updated
simplify_session_datafunction that now accepts dictionary data directly instead of objects with amodel_dump()method. The test logic and assertions remain sound.
214-214: LGTM! Mock updated to reflect new client API usage.The change from
session.retrievetosession.listcorrectly aligns with the refactored implementation that switched to using the list method for session retrieval.
400-400: LGTM! Correctly reflects the new ID mapping approach.The change from
VALID_AGENT_IDtoVALID_CONVERSATION_IDcorrectly implements the refactored architecture where conversation IDs are used directly as agent IDs, eliminating the separate mapping.
|
/lgtm |
Signed-off-by: Eran Cohen <[email protected]>
eranco74
force-pushed
the
agent-conversation-mapping
branch
from
df8f5d1 to
0d5306e
Compare
|
assisted-chat bump: rh-ecosystem-edge/assisted-chat#88 |
5 participants
part of https://issues.redhat.com/browse/MGMT-21329
Since a conversation is mapped to a single agent, and that agent uses the convesation_id as the session_id we can simplify the flow to use a single ID.
conversation id == agent id, each agent will have a single session_id, the session id is internal detail that we can retrive using:
Proposed flow:
User: Initial query (no conversation id)
LSC: query without conversatin id
User: query with conversation id
LSC: query with conversatin id
we keep using an agent cache with
@lru_cacheprevious art: https://gitlab.cee.redhat.com/ask-red-hat/backend/ifd-backend-api/-/blob/develop/ifd/core/agents/llama_stack.py?ref_type=heads#L73-80
Summary by CodeRabbit
Refactor
Tests