refactor conversation to agent and conversation mapping

Author: maorfr

src/app/endpoints/conversations.py

@@ -18,8 +18,6 @@
 router = APIRouter(tags=["conversations"])
 auth_dependency = get_auth_dependency()
 
-conversation_id_to_agent_id: dict[str, str] = {}
-
 conversation_responses: dict[int | str, dict[str, Any]] = {
     200: {
         "conversation_id": "123e4567-e89b-12d3-a456-426614174000",
@@ -126,17 +124,7 @@ def get_conversation_endpoint_handler(
             },
         )
 
-    agent_id = conversation_id_to_agent_id.get(conversation_id)
-    if not agent_id:
-        logger.error("Agent ID not found for conversation %s", conversation_id)
-        raise HTTPException(
-            status_code=status.HTTP_404_NOT_FOUND,
-            detail={
-                "response": "conversation ID not found",
-                "cause": f"conversation ID {conversation_id} not found!",
-            },
-        )
-
+    agent_id = conversation_id
     logger.info("Retrieving conversation %s", conversation_id)
 
     try:
@@ -206,16 +194,7 @@ def delete_conversation_endpoint_handler(
                 "cause": f"Conversation ID {conversation_id} is not a valid UUID",
             },
         )
-    agent_id = conversation_id_to_agent_id.get(conversation_id)
-    if not agent_id:
-        logger.error("Agent ID not found for conversation %s", conversation_id)
-        raise HTTPException(
-            status_code=status.HTTP_404_NOT_FOUND,
-            detail={
-                "response": "conversation ID not found",
-                "cause": f"conversation ID {conversation_id} not found!",
-            },
-        )
+    agent_id = conversation_id
     logger.info("Deleting conversation %s", conversation_id)
 
     try:

src/app/endpoints/query.py

@@ -1,6 +1,7 @@
 """Handler for REST API call to provide answer to query."""
 
 from datetime import datetime, UTC
+from functools import lru_cache
 import json
 import logging
 import os
@@ -23,7 +24,6 @@
 
 from client import LlamaStackClientHolder
 from configuration import configuration
-from app.endpoints.conversations import conversation_id_to_agent_id
 import metrics
 from models.responses import QueryResponse, UnauthorizedResponse, ForbiddenResponse
 from models.requests import QueryRequest, Attachment
@@ -39,9 +39,6 @@
 router = APIRouter(tags=["query"])
 auth_dependency = get_auth_dependency()
 
-# Global agent registry to persist agents across requests
-_agent_cache: TTLCache[str, Agent] = TTLCache(maxsize=1000, ttl=3600)
-
 query_response: dict[int | str, dict[str, Any]] = {
     200: {
         "conversation_id": "123e4567-e89b-12d3-a456-426614174000",
@@ -73,6 +70,7 @@ def is_transcripts_enabled() -> bool:
     return configuration.user_data_collection_configuration.transcripts_enabled
 
 
+@lru_cache
 def get_agent(  # pylint: disable=too-many-arguments,too-many-positional-arguments
     client: LlamaStackClient,
     model_id: str,
@@ -82,15 +80,6 @@ def get_agent(  # pylint: disable=too-many-arguments,too-many-positional-argumen
     conversation_id: str | None,
 ) -> tuple[Agent, str]:
     """Get existing agent or create a new one with session persistence."""
-    if conversation_id is not None:
-        agent = _agent_cache.get(conversation_id)
-        if agent:
-            logger.debug(
-                "Reusing existing agent with conversation_id: %s", conversation_id
-            )
-            return agent, conversation_id
-        logger.debug("No existing agent found for conversation_id: %s", conversation_id)
-
     logger.debug("Creating new agent")
     # TODO(lucasagomes): move to ReActAgent
     agent = Agent(
@@ -102,10 +91,11 @@ def get_agent(  # pylint: disable=too-many-arguments,too-many-positional-argumen
         tool_parser=GraniteToolParser.get_parser(model_id),
         enable_session_persistence=True,
     )
-    conversation_id = agent.create_session(get_suid())
-    logger.debug("Created new agent and conversation_id: %s", conversation_id)
-    _agent_cache[conversation_id] = agent
-    conversation_id_to_agent_id[conversation_id] = agent.agent_id
+    if conversation_id:
+        agent.agent_id = conversation_id
+    else:
+        agent.create_session(get_suid())
+        conversation_id = agent.agent_id
 
     return agent, conversation_id
 

src/app/endpoints/streaming_query.py

@@ -1,5 +1,6 @@
 """Handler for REST API call to provide answer to streaming query."""
 
+from functools import lru_cache
 import json
 import logging
 import re
@@ -31,7 +32,6 @@
 from utils.suid import get_suid
 from utils.types import GraniteToolParser
 
-from app.endpoints.conversations import conversation_id_to_agent_id
 from app.endpoints.query import (
     get_rag_toolgroups,
     is_input_shield,
@@ -46,11 +46,9 @@
 router = APIRouter(tags=["streaming_query"])
 auth_dependency = get_auth_dependency()
 
-# Global agent registry to persist agents across requests
-_agent_cache: TTLCache[str, AsyncAgent] = TTLCache(maxsize=1000, ttl=3600)
-
 
 # # pylint: disable=R0913,R0917
+@lru_cache
 async def get_agent(
     client: AsyncLlamaStackClient,
     model_id: str,
@@ -60,15 +58,6 @@ async def get_agent(
     conversation_id: str | None,
 ) -> tuple[AsyncAgent, str]:
     """Get existing agent or create a new one with session persistence."""
-    if conversation_id is not None:
-        agent = _agent_cache.get(conversation_id)
-        if agent:
-            logger.debug(
-                "Reusing existing agent with conversation_id: %s", conversation_id
-            )
-            return agent, conversation_id
-        logger.debug("No existing agent found for conversation_id: %s", conversation_id)
-
     logger.debug("Creating new agent")
     agent = AsyncAgent(
         client,  # type: ignore[arg-type]
@@ -79,10 +68,13 @@ async def get_agent(
         tool_parser=GraniteToolParser.get_parser(model_id),
         enable_session_persistence=True,
     )
-    conversation_id = await agent.create_session(get_suid())
-    logger.debug("Created new agent and conversation_id: %s", conversation_id)
-    _agent_cache[conversation_id] = agent
-    conversation_id_to_agent_id[conversation_id] = agent.agent_id
+
+    if conversation_id:
+        agent._agent_id = conversation_id
+    else:
+        conversation_id = agent.agent_id
+        await agent.create_session(get_suid())
+
     return agent, conversation_id
 
 

tests/unit/app/endpoints/test_conversations.py

@@ -7,7 +7,6 @@
 from app.endpoints.conversations import (
     get_conversation_endpoint_handler,
     delete_conversation_endpoint_handler,
-    conversation_id_to_agent_id,
     simplify_session_data,
 )
 from models.responses import ConversationResponse, ConversationDeleteResponse
@@ -48,16 +47,6 @@ def setup_configuration_fixture():
     return cfg
 
 
-@pytest.fixture(autouse=True)
-def setup_conversation_mapping():
-    """Set up and clean up the conversation ID to agent ID mapping."""
-    # Clear the mapping before each test
-    conversation_id_to_agent_id.clear()
-    yield
-    # Clean up after each test
-    conversation_id_to_agent_id.clear()
-
-
 @pytest.fixture(name="mock_session_data")
 def mock_session_data_fixture():
     """Create mock session data for testing."""
@@ -243,9 +232,6 @@ def test_llama_stack_connection_error(self, mocker, setup_configuration):
         mocker.patch("app.endpoints.conversations.configuration", setup_configuration)
         mocker.patch("app.endpoints.conversations.check_suid", return_value=True)
 
-        # Set up conversation mapping
-        conversation_id_to_agent_id[VALID_CONVERSATION_ID] = VALID_AGENT_ID
-
         # Mock LlamaStackClientHolder to raise APIConnectionError
         mock_client = mocker.Mock()
         mock_client.agents.session.retrieve.side_effect = APIConnectionError(
@@ -268,9 +254,6 @@ def test_llama_stack_not_found_error(self, mocker, setup_configuration):
         mocker.patch("app.endpoints.conversations.configuration", setup_configuration)
         mocker.patch("app.endpoints.conversations.check_suid", return_value=True)
 
-        # Set up conversation mapping
-        conversation_id_to_agent_id[VALID_CONVERSATION_ID] = VALID_AGENT_ID
-
         # Mock LlamaStackClientHolder to raise NotFoundError
         mock_client = mocker.Mock()
         mock_client.agents.session.retrieve.side_effect = NotFoundError(
@@ -294,9 +277,6 @@ def test_session_retrieve_exception(self, mocker, setup_configuration):
         mocker.patch("app.endpoints.conversations.configuration", setup_configuration)
         mocker.patch("app.endpoints.conversations.check_suid", return_value=True)
 
-        # Set up conversation mapping
-        conversation_id_to_agent_id[VALID_CONVERSATION_ID] = VALID_AGENT_ID
-
         # Mock LlamaStackClientHolder to raise a general exception
         mock_client = mocker.Mock()
         mock_client.agents.session.retrieve.side_effect = Exception(
@@ -323,9 +303,6 @@ def test_successful_conversation_retrieval(
         mocker.patch("app.endpoints.conversations.configuration", setup_configuration)
         mocker.patch("app.endpoints.conversations.check_suid", return_value=True)
 
-        # Set up conversation mapping
-        conversation_id_to_agent_id[VALID_CONVERSATION_ID] = VALID_AGENT_ID
-
         # Mock session data with model_dump method
         mock_session_obj = mocker.Mock()
         mock_session_obj.model_dump.return_value = mock_session_data
@@ -394,9 +371,6 @@ def test_llama_stack_connection_error(self, mocker, setup_configuration):
         mocker.patch("app.endpoints.conversations.configuration", setup_configuration)
         mocker.patch("app.endpoints.conversations.check_suid", return_value=True)
 
-        # Set up conversation mapping
-        conversation_id_to_agent_id[VALID_CONVERSATION_ID] = VALID_AGENT_ID
-
         # Mock LlamaStackClientHolder to raise APIConnectionError
         mock_client = mocker.Mock()
         mock_client.agents.session.delete.side_effect = APIConnectionError(request=None)
@@ -416,9 +390,6 @@ def test_llama_stack_not_found_error(self, mocker, setup_configuration):
         mocker.patch("app.endpoints.conversations.configuration", setup_configuration)
         mocker.patch("app.endpoints.conversations.check_suid", return_value=True)
 
-        # Set up conversation mapping
-        conversation_id_to_agent_id[VALID_CONVERSATION_ID] = VALID_AGENT_ID
-
         # Mock LlamaStackClientHolder to raise NotFoundError
         mock_client = mocker.Mock()
         mock_client.agents.session.delete.side_effect = NotFoundError(
@@ -442,9 +413,6 @@ def test_session_deletion_exception(self, mocker, setup_configuration):
         mocker.patch("app.endpoints.conversations.configuration", setup_configuration)
         mocker.patch("app.endpoints.conversations.check_suid", return_value=True)
 
-        # Set up conversation mapping
-        conversation_id_to_agent_id[VALID_CONVERSATION_ID] = VALID_AGENT_ID
-
         # Mock LlamaStackClientHolder to raise a general exception
         mock_client = mocker.Mock()
         mock_client.agents.session.delete.side_effect = Exception(
@@ -470,9 +438,6 @@ def test_successful_conversation_deletion(self, mocker, setup_configuration):
         mocker.patch("app.endpoints.conversations.configuration", setup_configuration)
         mocker.patch("app.endpoints.conversations.check_suid", return_value=True)
 
-        # Set up conversation mapping
-        conversation_id_to_agent_id[VALID_CONVERSATION_ID] = VALID_AGENT_ID
-
         # Mock LlamaStackClientHolder
         mock_client = mocker.Mock()
         mock_client.agents.session.delete.return_value = None  # Successful deletion

tests/unit/app/endpoints/test_query.py

@@ -20,7 +20,6 @@
     store_transcript,
     get_rag_toolgroups,
     get_agent,
-    _agent_cache,
 )
 
 from models.requests import QueryRequest, Attachment
@@ -65,8 +64,6 @@ def prepare_agent_mocks_fixture(mocker):
     mock_agent = mocker.Mock()
     mock_agent.create_turn.return_value.steps = []
     yield mock_client, mock_agent
-    # cleanup agent cache after tests
-    _agent_cache.clear()
 
 
 def test_query_endpoint_handler_configuration_not_loaded(mocker):
@@ -1065,7 +1062,6 @@ def test_get_agent_cache_hit(prepare_agent_mocks):
 
     # Set up cache with existing agent
     conversation_id = "test_conversation_id"
-    _agent_cache[conversation_id] = mock_agent
 
     result_agent, result_conversation_id = get_agent(
         client=mock_client,
@@ -1132,9 +1128,6 @@ def test_get_agent_cache_miss_with_conversation_id(
         enable_session_persistence=True,
     )
 
-    # Verify agent was stored in cache
-    assert _agent_cache["new_session_id"] == mock_agent
-
 
 def test_get_agent_no_conversation_id(setup_configuration, prepare_agent_mocks, mocker):
     """Test get_agent function when conversation_id is None."""
@@ -1185,9 +1178,6 @@ def test_get_agent_no_conversation_id(setup_configuration, prepare_agent_mocks,
         enable_session_persistence=True,
     )
 
-    # Verify agent was stored in cache
-    assert _agent_cache["new_session_id"] == mock_agent
-
 
 def test_get_agent_empty_shields(setup_configuration, prepare_agent_mocks, mocker):
     """Test get_agent function with empty shields list."""

tests/unit/app/endpoints/test_streaming_query.py

@@ -41,7 +41,6 @@
     retrieve_response,
     stream_build_event,
     get_agent,
-    _agent_cache,
 )
 from models.requests import QueryRequest, Attachment
 from models.config import ModelContextProtocolServer
@@ -109,8 +108,6 @@ def prepare_agent_mocks_fixture(mocker):
     mock_client = mocker.AsyncMock()
     mock_agent = mocker.AsyncMock()
     yield mock_client, mock_agent
-    # cleanup agent cache after tests
-    _agent_cache.clear()
 
 
 @pytest.mark.asyncio
@@ -1213,7 +1210,6 @@ async def test_get_agent_cache_hit(prepare_agent_mocks):
 
     # Set up cache with existing agent
     conversation_id = "test_conversation_id"
-    _agent_cache[conversation_id] = mock_agent
 
     result_agent, result_conversation_id = await get_agent(
         client=mock_client,
@@ -1284,9 +1280,6 @@ async def test_get_agent_cache_miss_with_conversation_id(
         enable_session_persistence=True,
     )
 
-    # Verify agent was stored in cache
-    assert _agent_cache["new_session_id"] == mock_agent
-
 
 @pytest.mark.asyncio
 async def test_get_agent_no_conversation_id(
@@ -1343,9 +1336,6 @@ async def test_get_agent_no_conversation_id(
         enable_session_persistence=True,
     )
 
-    # Verify agent was stored in cache
-    assert _agent_cache["new_session_id"] == mock_agent
-
 
 @pytest.mark.asyncio
 async def test_get_agent_empty_shields(