From 9bd642f31156c1feed4de876648e1d5b7c9aba72 Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Tue, 12 Aug 2025 15:10:09 -0400 Subject: [PATCH] change to _get_session_from_values instead of _get_sessions_from_keys Signed-off-by: Francisco Javier Arceo --- .../agents/meta_reference/persistence.py | 25 +++++++++---------- .../agent/test_agent_meta_reference.py | 20 ++++++--------- 2 files changed, 20 insertions(+), 25 deletions(-) diff --git a/llama_stack/providers/inline/agents/meta_reference/persistence.py b/llama_stack/providers/inline/agents/meta_reference/persistence.py index 47d27cb37..7053ee43c 100644 --- a/llama_stack/providers/inline/agents/meta_reference/persistence.py +++ b/llama_stack/providers/inline/agents/meta_reference/persistence.py @@ -183,37 +183,36 @@ class AgentPersistence: ) return int(value) if value else None - async def _get_sessions_from_keys(self, keys: list[str]) -> list[Session]: + async def _get_sessions_from_values(self, values: list[str]) -> list[Session]: """ - Helper method to filter session keys and retrieve session objects. + Helper method to process values from the KV store and extract session objects. Args: - keys: List of keys from the KV store + values: List of serialized JSON values from the KV store Returns: List of parsed Session objects """ - # filter session metadata to relevant keys (3 parts: agent_id, session_id, and turn_id) - session_keys = [key for key in keys if len(key.split(":")) == 3] - sessions = [] - for key in session_keys: + for value in values: try: - value = await self.kvstore.get(key) - if value: - session_info = Session(**json.loads(value)) - sessions.append(session_info) + data = json.loads(value) + if "turn_id" in data: + continue + + session_info = Session(**data) + sessions.append(session_info) except Exception as e: log.error(f"Error parsing session info: {e}") continue return sessions async def list_sessions(self) -> list[Session]: - keys = await self.kvstore.keys_in_range( + values = await self.kvstore.values_in_range( start_key=f"session:{self.agent_id}:", end_key=f"session:{self.agent_id}:\xff\xff\xff\xff", ) - sessions = await self._get_sessions_from_keys(keys) + sessions = await self._get_sessions_from_values(values) return sessions async def delete_session_turns(self, session_id: str) -> None: diff --git a/tests/unit/providers/agent/test_agent_meta_reference.py b/tests/unit/providers/agent/test_agent_meta_reference.py index 8415a7b56..3fc60024a 100644 --- a/tests/unit/providers/agent/test_agent_meta_reference.py +++ b/tests/unit/providers/agent/test_agent_meta_reference.py @@ -111,7 +111,8 @@ class TestAgentPersistenceListSessions: if custom_data: mock_data.update(custom_data) - mock_kvstore.keys_in_range.return_value = all_keys + values_list = list(mock_data.values()) + mock_kvstore.values_in_range.return_value = values_list async def mock_get(key): return mock_data.get(key) @@ -237,12 +238,12 @@ class TestAgentPersistenceListSessions: assert mock_log.error.call_count == error_scenario["expected_error_count"] async def test_list_sessions_empty(self, agent_persistence, mock_kvstore): - mock_kvstore.keys_in_range.return_value = [] + mock_kvstore.values_in_range.return_value = [] result = await agent_persistence.list_sessions() assert result == [] - mock_kvstore.keys_in_range.assert_called_once_with( + mock_kvstore.values_in_range.assert_called_once_with( start_key="session:test-agent-123:", end_key="session:test-agent-123:\xff\xff\xff\xff" ) @@ -267,7 +268,7 @@ class TestAgentPersistenceListSessions: assert hasattr(result[0], "started_at") async def test_list_sessions_kvstore_exception(self, agent_persistence, mock_kvstore): - mock_kvstore.keys_in_range.side_effect = Exception("KVStore error") + mock_kvstore.values_in_range.side_effect = Exception("KVStore error") with pytest.raises(Exception, match="KVStore error"): await agent_persistence.list_sessions() @@ -313,11 +314,6 @@ class TestAgentPersistenceListSessions: "completed_at": "2025-08-05T14:31:51.305384Z", } - keys = [ - "session:test-agent-123:1f08fd1c-5a9d-459d-a00b-36d4dfa49b7d", - "session:test-agent-123:1f08fd1c-5a9d-459d-a00b-36d4dfa49b7d:eb7e818f-41fb-49a0-bdd6-464974a2d2ad", - ] - mock_data = { "session:test-agent-123:1f08fd1c-5a9d-459d-a00b-36d4dfa49b7d": json.dumps(session_data), "session:test-agent-123:1f08fd1c-5a9d-459d-a00b-36d4dfa49b7d:eb7e818f-41fb-49a0-bdd6-464974a2d2ad": json.dumps( @@ -325,7 +321,7 @@ class TestAgentPersistenceListSessions: ), } - mock_kvstore.keys_in_range.return_value = keys + mock_kvstore.values_in_range.return_value = list(mock_data.values()) async def mock_get(key): return mock_data.get(key) @@ -342,10 +338,10 @@ class TestAgentPersistenceListSessions: mock_log.error.assert_not_called() async def test_list_sessions_key_range_construction(self, agent_persistence, mock_kvstore): - mock_kvstore.keys_in_range.return_value = [] + mock_kvstore.values_in_range.return_value = [] await agent_persistence.list_sessions() - mock_kvstore.keys_in_range.assert_called_once_with( + mock_kvstore.values_in_range.assert_called_once_with( start_key="session:test-agent-123:", end_key="session:test-agent-123:\xff\xff\xff\xff" )