From 68b00718615c1ae9cc5c73bbc65575870c5cfda6 Mon Sep 17 00:00:00 2001 From: Nathan Weinberg <31703736+nathan-weinberg@users.noreply.github.com> Date: Mon, 4 Aug 2025 16:12:02 -0400 Subject: [PATCH] chore: standardize session not found error (#3031) # What does this PR do? 1. Creates a new `SessionNotFoundError` class 2. Implements the new class where appropriate Relates to #2379 Signed-off-by: Nathan Weinberg --- llama_stack/apis/common/errors.py | 8 ++++++++ .../agents/meta_reference/agent_instance.py | 3 ++- .../inline/agents/meta_reference/agents.py | 5 ----- .../agents/meta_reference/persistence.py | 19 ++++++++++--------- 4 files changed, 20 insertions(+), 15 deletions(-) diff --git a/llama_stack/apis/common/errors.py b/llama_stack/apis/common/errors.py index f742513c3..bef048191 100644 --- a/llama_stack/apis/common/errors.py +++ b/llama_stack/apis/common/errors.py @@ -50,3 +50,11 @@ class ToolGroupNotFoundError(ValueError): f"Tool group '{toolgroup_name}' not found. Use client.toolgroups.list() to list available tool groups." ) super().__init__(message) + + +class SessionNotFoundError(ValueError): + """raised when Llama Stack cannot find a referenced session or access is denied""" + + def __init__(self, session_name: str) -> None: + message = f"Session '{session_name}' not found or access denied." + super().__init__(message) diff --git a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py index fb22e10cc..5f7c90879 100644 --- a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py +++ b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py @@ -44,6 +44,7 @@ from llama_stack.apis.common.content_types import ( ToolCallDelta, ToolCallParseStatus, ) +from llama_stack.apis.common.errors import SessionNotFoundError from llama_stack.apis.inference import ( ChatCompletionResponseEventType, CompletionMessage, @@ -214,7 +215,7 @@ class ChatAgent(ShieldRunnerMixin): is_resume = isinstance(request, AgentTurnResumeRequest) session_info = await self.storage.get_session_info(request.session_id) if session_info is None: - raise ValueError(f"Session {request.session_id} not found") + raise SessionNotFoundError(request.session_id) turns = await self.storage.get_session_turns(request.session_id) if is_resume and len(turns) == 0: diff --git a/llama_stack/providers/inline/agents/meta_reference/agents.py b/llama_stack/providers/inline/agents/meta_reference/agents.py index b326a95b4..15695ec48 100644 --- a/llama_stack/providers/inline/agents/meta_reference/agents.py +++ b/llama_stack/providers/inline/agents/meta_reference/agents.py @@ -230,8 +230,6 @@ class MetaReferenceAgentsImpl(Agents): agent = await self._get_agent_impl(agent_id) session_info = await agent.storage.get_session_info(session_id) - if session_info is None: - raise ValueError(f"Session {session_id} not found") turns = await agent.storage.get_session_turns(session_id) if turn_ids: turns = [turn for turn in turns if turn.turn_id in turn_ids] @@ -244,9 +242,6 @@ class MetaReferenceAgentsImpl(Agents): async def delete_agents_session(self, agent_id: str, session_id: str) -> None: agent = await self._get_agent_impl(agent_id) - session_info = await agent.storage.get_session_info(session_id) - if session_info is None: - raise ValueError(f"Session {session_id} not found") # Delete turns first, then the session await agent.storage.delete_session_turns(session_id) diff --git a/llama_stack/providers/inline/agents/meta_reference/persistence.py b/llama_stack/providers/inline/agents/meta_reference/persistence.py index 8d4713cd2..7a8d99b78 100644 --- a/llama_stack/providers/inline/agents/meta_reference/persistence.py +++ b/llama_stack/providers/inline/agents/meta_reference/persistence.py @@ -10,6 +10,7 @@ import uuid from datetime import UTC, datetime from llama_stack.apis.agents import AgentConfig, Session, ToolExecutionStep, Turn +from llama_stack.apis.common.errors import SessionNotFoundError from llama_stack.core.access_control.access_control import AccessDeniedError, is_action_allowed from llama_stack.core.access_control.datatypes import AccessRule from llama_stack.core.datatypes import User @@ -61,12 +62,12 @@ class AgentPersistence: ) return session_id - async def get_session_info(self, session_id: str) -> AgentSessionInfo | None: + async def get_session_info(self, session_id: str) -> AgentSessionInfo: value = await self.kvstore.get( key=f"session:{self.agent_id}:{session_id}", ) if not value: - return None + raise SessionNotFoundError(session_id) session_info = AgentSessionInfo(**json.loads(value)) @@ -95,7 +96,7 @@ class AgentPersistence: async def add_vector_db_to_session(self, session_id: str, vector_db_id: str): session_info = await self.get_session_if_accessible(session_id) if session_info is None: - raise ValueError(f"Session {session_id} not found or access denied") + raise SessionNotFoundError(session_id) session_info.vector_db_id = vector_db_id await self.kvstore.set( @@ -105,7 +106,7 @@ class AgentPersistence: async def add_turn_to_session(self, session_id: str, turn: Turn): if not await self.get_session_if_accessible(session_id): - raise ValueError(f"Session {session_id} not found or access denied") + raise SessionNotFoundError(session_id) await self.kvstore.set( key=f"session:{self.agent_id}:{session_id}:{turn.turn_id}", @@ -114,7 +115,7 @@ class AgentPersistence: async def get_session_turns(self, session_id: str) -> list[Turn]: if not await self.get_session_if_accessible(session_id): - raise ValueError(f"Session {session_id} not found or access denied") + raise SessionNotFoundError(session_id) values = await self.kvstore.values_in_range( start_key=f"session:{self.agent_id}:{session_id}:", @@ -137,7 +138,7 @@ class AgentPersistence: async def get_session_turn(self, session_id: str, turn_id: str) -> Turn | None: if not await self.get_session_if_accessible(session_id): - raise ValueError(f"Session {session_id} not found or access denied") + raise SessionNotFoundError(session_id) value = await self.kvstore.get( key=f"session:{self.agent_id}:{session_id}:{turn_id}", @@ -148,7 +149,7 @@ class AgentPersistence: async def set_in_progress_tool_call_step(self, session_id: str, turn_id: str, step: ToolExecutionStep): if not await self.get_session_if_accessible(session_id): - raise ValueError(f"Session {session_id} not found or access denied") + raise SessionNotFoundError(session_id) await self.kvstore.set( key=f"in_progress_tool_call_step:{self.agent_id}:{session_id}:{turn_id}", @@ -166,7 +167,7 @@ class AgentPersistence: async def set_num_infer_iters_in_turn(self, session_id: str, turn_id: str, num_infer_iters: int): if not await self.get_session_if_accessible(session_id): - raise ValueError(f"Session {session_id} not found or access denied") + raise SessionNotFoundError(session_id) await self.kvstore.set( key=f"num_infer_iters_in_turn:{self.agent_id}:{session_id}:{turn_id}", @@ -218,6 +219,6 @@ class AgentPersistence: """ session_info = await self.get_session_info(session_id) if session_info is None: - raise ValueError(f"Session {session_id} not found") + raise SessionNotFoundError(session_id) await self.kvstore.delete(key=f"session:{self.agent_id}:{session_id}")