phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-15 14:08:00 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

change to _get_session_from_values instead of _get_sessions_from_keys

Browse source 
Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
This commit is contained in:
Francisco Javier Arceo 2025-08-12 15:10:09 -04:00
parent a4180894ec
commit 9bd642f311
2 changed files with 20 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

25
llama_stack/providers/inline/agents/meta_reference/persistence.py
View file

@ -183,37 +183,36 @@ class AgentPersistence:
) )
return int(value) if value else None 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: Args:
keys: List of keys from the KV store values: List of serialized JSON values from the KV store
Returns: Returns:
List of parsed Session objects 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 = [] sessions = []
for key in session_keys: for value in values:
try: try:
value = await self.kvstore.get(key) data = json.loads(value)
if value: if "turn_id" in data:
session_info = Session(**json.loads(value)) continue
sessions.append(session_info)
session_info = Session(**data)
sessions.append(session_info)
except Exception as e: except Exception as e:
log.error(f"Error parsing session info: {e}") log.error(f"Error parsing session info: {e}")
continue continue
return sessions return sessions
async def list_sessions(self) -> list[Session]: 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}:", start_key=f"session:{self.agent_id}:",
end_key=f"session:{self.agent_id}:\xff\xff\xff\xff", 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 return sessions
async def delete_session_turns(self, session_id: str) -> None: async def delete_session_turns(self, session_id: str) -> None:

20
tests/unit/providers/agent/test_agent_meta_reference.py
View file

@ -111,7 +111,8 @@ class TestAgentPersistenceListSessions:
if custom_data: if custom_data:
mock_data.update(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): async def mock_get(key):
return mock_data.get(key) return mock_data.get(key)
@ -237,12 +238,12 @@ class TestAgentPersistenceListSessions:
assert mock_log.error.call_count == error_scenario["expected_error_count"] assert mock_log.error.call_count == error_scenario["expected_error_count"]
async def test_list_sessions_empty(self, agent_persistence, mock_kvstore): 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() result = await agent_persistence.list_sessions()
assert result == [] 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" 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") assert hasattr(result[0], "started_at")
async def test_list_sessions_kvstore_exception(self, agent_persistence, mock_kvstore): 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"): with pytest.raises(Exception, match="KVStore error"):
await agent_persistence.list_sessions() await agent_persistence.list_sessions()
@ -313,11 +314,6 @@ class TestAgentPersistenceListSessions:
"completed_at": "2025-08-05T14:31:51.305384Z", "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 = { mock_data = {
"session:test-agent-123:1f08fd1c-5a9d-459d-a00b-36d4dfa49b7d": json.dumps(session_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( "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): async def mock_get(key):
return mock_data.get(key) return mock_data.get(key)
@ -342,10 +338,10 @@ class TestAgentPersistenceListSessions:
mock_log.error.assert_not_called() mock_log.error.assert_not_called()
async def test_list_sessions_key_range_construction(self, agent_persistence, mock_kvstore): 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() 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" start_key="session:test-agent-123:", end_key="session:test-agent-123:\xff\xff\xff\xff"
) )