From c21be66427a7432c802f7a6e256367dfd771f7cc Mon Sep 17 00:00:00 2001 From: JR Boos Date: Wed, 6 Aug 2025 14:09:42 -0400 Subject: [PATCH] Enhance get_conversation_endpoint_handler to retrieve session data correctly - Updated the endpoint handler to first list agent sessions and then retrieve the specific session data using the session ID. - Modified unit tests to mock the session retrieval process, ensuring the model_dump method is called correctly. This change improves the accuracy of conversation data retrieval. --- src/app/endpoints/conversations.py | 8 +++++++- tests/unit/app/endpoints/test_conversations.py | 6 ++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/src/app/endpoints/conversations.py b/src/app/endpoints/conversations.py index 59df19e17..f7133c959 100644 --- a/src/app/endpoints/conversations.py +++ b/src/app/endpoints/conversations.py @@ -134,7 +134,13 @@ async def get_conversation_endpoint_handler( try: client = AsyncLlamaStackClientHolder().get_client() - session_data = (await client.agents.session.list(agent_id=agent_id)).data[0] + agent_sessions = (await client.agents.session.list(agent_id=agent_id)).data + session_id = str(agent_sessions[0].get("session_id")) + + session_response = await client.agents.session.retrieve( + agent_id=agent_id, session_id=session_id + ) + session_data = session_response.model_dump() logger.info("Successfully retrieved conversation %s", conversation_id) diff --git a/tests/unit/app/endpoints/test_conversations.py b/tests/unit/app/endpoints/test_conversations.py index 8b8954475..1958d131a 100644 --- a/tests/unit/app/endpoints/test_conversations.py +++ b/tests/unit/app/endpoints/test_conversations.py @@ -305,6 +305,12 @@ async def test_successful_conversation_retrieval( mock_client.agents.session.list.return_value = mocker.Mock( data=[mock_session_data] ) + + # Mock session.retrieve to return an object with model_dump() method + mock_session_retrieve_result = mocker.Mock() + mock_session_retrieve_result.model_dump.return_value = mock_session_data + mock_client.agents.session.retrieve.return_value = mock_session_retrieve_result + mock_client_holder = mocker.patch( "app.endpoints.conversations.AsyncLlamaStackClientHolder" )