diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index 62cdb878..c21648a4 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -733,10 +733,10 @@ async def retrieve_response( # pylint: disable=too-many-locals,too-many-branche ), } - vector_db_ids = [ - vector_db.identifier for vector_db in await client.vector_dbs.list() + vector_store_ids = [ + vector_store.id for vector_store in (await client.vector_stores.list()).data ] - toolgroups = (get_rag_toolgroups(vector_db_ids) or []) + [ + toolgroups = (get_rag_toolgroups(vector_store_ids) or []) + [ mcp_server.name for mcp_server in configuration.mcp_servers ] # Convert empty list to None for consistency with existing behavior @@ -830,30 +830,30 @@ def validate_attachments_metadata(attachments: list[Attachment]) -> None: def get_rag_toolgroups( - vector_db_ids: list[str], + vector_store_ids: list[str], ) -> list[Toolgroup] | None: """ - Return a list of RAG Tool groups if the given vector DB list is not empty. + Return a list of RAG Tool groups if the given vector store list is not empty. Generate a list containing a RAG knowledge search toolgroup if - vector database IDs are provided. + vector store IDs are provided. Parameters: - vector_db_ids (list[str]): List of vector database identifiers to include in the toolgroup. + vector_store_ids (list[str]): List of vector store identifiers to include in the toolgroup. Returns: list[Toolgroup] | None: A list with a single RAG toolgroup if - vector_db_ids is non-empty; otherwise, None. + vector_store_ids is non-empty; otherwise, None. """ return ( [ ToolgroupAgentToolGroupWithArgs( name="builtin::rag/knowledge_search", args={ - "vector_db_ids": vector_db_ids, + "vector_db_ids": vector_store_ids, }, ) ] - if vector_db_ids + if vector_store_ids else None ) diff --git a/src/app/endpoints/streaming_query.py b/src/app/endpoints/streaming_query.py index 1b440a33..50e39cd5 100644 --- a/src/app/endpoints/streaming_query.py +++ b/src/app/endpoints/streaming_query.py @@ -1061,10 +1061,10 @@ async def retrieve_response( ), } - vector_db_ids = [ - vector_db.identifier for vector_db in await client.vector_dbs.list() + vector_store_ids = [ + vector_store.id for vector_store in (await client.vector_stores.list()).data ] - toolgroups = (get_rag_toolgroups(vector_db_ids) or []) + [ + toolgroups = (get_rag_toolgroups(vector_store_ids) or []) + [ mcp_server.name for mcp_server in configuration.mcp_servers ] # Convert empty list to None for consistency with existing behavior diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index 54b46a3c..4b86c245 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -536,9 +536,11 @@ async def test_retrieve_response_no_returned_message( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message = None mock_client.shields.list.return_value = [] - mock_vector_db = mocker.Mock() - mock_vector_db.identifier = "VectorDB-1" - mock_client.vector_dbs.list.return_value = [mock_vector_db] + mock_vector_store = mocker.Mock() + mock_vector_store.id = "VectorDB-1" + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [mock_vector_store] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -574,9 +576,11 @@ async def test_retrieve_response_message_without_content( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = None mock_client.shields.list.return_value = [] - mock_vector_db = mocker.Mock() - mock_vector_db.identifier = "VectorDB-1" - mock_client.vector_dbs.list.return_value = [mock_vector_db] + mock_vector_store = mocker.Mock() + mock_vector_store.id = "VectorDB-1" + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [mock_vector_store] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -613,9 +617,11 @@ async def test_retrieve_response_vector_db_available( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_vector_db = mocker.Mock() - mock_vector_db.identifier = "VectorDB-1" - mock_client.vector_dbs.list.return_value = [mock_vector_db] + mock_vector_store = mocker.Mock() + mock_vector_store.id = "VectorDB-1" + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [mock_vector_store] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -660,7 +666,9 @@ async def test_retrieve_response_no_available_shields( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_client.vector_dbs.list.return_value = [] + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -716,7 +724,9 @@ def __repr__(self) -> str: mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [MockShield("shield1")] - mock_client.vector_dbs.list.return_value = [] + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -775,7 +785,9 @@ def __repr__(self) -> str: MockShield("shield1"), MockShield("shield2"), ] - mock_client.vector_dbs.list.return_value = [] + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -836,7 +848,9 @@ def __repr__(self) -> str: MockShield("output_shield3"), MockShield("inout_shield4"), ] - mock_client.vector_dbs.list.return_value = [] + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -891,7 +905,9 @@ async def test_retrieve_response_with_one_attachment( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_client.vector_dbs.list.return_value = [] + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -947,7 +963,9 @@ async def test_retrieve_response_with_two_attachments( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_client.vector_dbs.list.return_value = [] + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -1129,7 +1147,9 @@ async def test_retrieve_response_with_mcp_servers( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_client.vector_dbs.list.return_value = [] + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with MCP servers mcp_servers = [ @@ -1210,7 +1230,9 @@ async def test_retrieve_response_with_mcp_servers_empty_token( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_client.vector_dbs.list.return_value = [] + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with MCP servers mcp_servers = [ @@ -1269,7 +1291,9 @@ async def test_retrieve_response_with_mcp_servers_and_mcp_headers( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_client.vector_dbs.list.return_value = [] + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with MCP servers mcp_servers = [ @@ -1380,9 +1404,11 @@ async def test_retrieve_response_shield_violation( text="LLM answer", type="text" ) mock_client.shields.list.return_value = [] - mock_vector_db = mocker.Mock() - mock_vector_db.identifier = "VectorDB-1" - mock_client.vector_dbs.list.return_value = [mock_vector_db] + mock_vector_store = mocker.Mock() + mock_vector_store.id = "VectorDB-1" + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [mock_vector_store] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -1419,16 +1445,16 @@ async def test_retrieve_response_shield_violation( def test_get_rag_toolgroups() -> None: """Test get_rag_toolgroups function.""" - vector_db_ids: list[str] = [] - result = get_rag_toolgroups(vector_db_ids) + vector_store_ids: list[str] = [] + result = get_rag_toolgroups(vector_store_ids) assert result is None - vector_db_ids = ["Vector-DB-1", "Vector-DB-2"] - result = get_rag_toolgroups(vector_db_ids) + vector_store_ids = ["Vector-DB-1", "Vector-DB-2"] + result = get_rag_toolgroups(vector_store_ids) assert result is not None assert len(result) == 1 assert result[0]["name"] == "builtin::rag/knowledge_search" - assert result[0]["args"]["vector_db_ids"] == vector_db_ids + assert result[0]["args"]["vector_db_ids"] == vector_store_ids @pytest.mark.asyncio @@ -1649,9 +1675,11 @@ async def test_retrieve_response_no_tools_bypasses_mcp_and_rag( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_vector_db = mocker.Mock() - mock_vector_db.identifier = "VectorDB-1" - mock_client.vector_dbs.list.return_value = [mock_vector_db] + mock_vector_store = mocker.Mock() + mock_vector_store.id = "VectorDB-1" + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [mock_vector_store] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with MCP servers mcp_servers = [ @@ -1704,9 +1732,11 @@ async def test_retrieve_response_no_tools_false_preserves_functionality( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_vector_db = mocker.Mock() - mock_vector_db.identifier = "VectorDB-1" - mock_client.vector_dbs.list.return_value = [mock_vector_db] + mock_vector_store = mocker.Mock() + mock_vector_store.id = "VectorDB-1" + mock_vector_stores_response = mocker.Mock() + mock_vector_stores_response.data = [mock_vector_store] + mock_client.vector_stores.list.return_value = mock_vector_stores_response # Mock configuration with MCP servers mcp_servers = [ diff --git a/tests/unit/app/endpoints/test_streaming_query.py b/tests/unit/app/endpoints/test_streaming_query.py index e93f9566..e00afcb4 100644 --- a/tests/unit/app/endpoints/test_streaming_query.py +++ b/tests/unit/app/endpoints/test_streaming_query.py @@ -376,9 +376,11 @@ async def test_retrieve_response_vector_db_available( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_vector_db = mocker.Mock() - mock_vector_db.identifier = "VectorDB-1" - mock_client.vector_dbs.list.return_value = [mock_vector_db] + mock_vector_store = mocker.Mock() + mock_vector_store.id = "VectorDB-1" + mock_list_response = mocker.Mock() + mock_list_response.data = [mock_vector_store] + mock_client.vector_stores.list.return_value = mock_list_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -421,7 +423,9 @@ async def test_retrieve_response_no_available_shields( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_client.vector_dbs.list.return_value = [] + mock_list_response = mocker.Mock() + mock_list_response.data = [] + mock_client.vector_stores.list.return_value = mock_list_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -477,7 +481,9 @@ def __repr__(self) -> str: mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [MockShield("shield1")] - mock_client.vector_dbs.list.return_value = [] + mock_list_response = mocker.Mock() + mock_list_response.data = [] + mock_client.vector_stores.list.return_value = mock_list_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -534,7 +540,9 @@ def __repr__(self) -> str: MockShield("shield1"), MockShield("shield2"), ] - mock_client.vector_dbs.list.return_value = [] + mock_list_response = mocker.Mock() + mock_list_response.data = [] + mock_client.vector_stores.list.return_value = mock_list_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -593,7 +601,9 @@ def __repr__(self) -> str: MockShield("output_shield3"), MockShield("inout_shield4"), ] - mock_client.vector_dbs.list.return_value = [] + mock_list_response = mocker.Mock() + mock_list_response.data = [] + mock_client.vector_stores.list.return_value = mock_list_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -646,7 +656,9 @@ async def test_retrieve_response_with_one_attachment( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_client.vector_dbs.list.return_value = [] + mock_list_response = mocker.Mock() + mock_list_response.data = [] + mock_client.vector_stores.list.return_value = mock_list_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -700,7 +712,9 @@ async def test_retrieve_response_with_two_attachments( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_client.vector_dbs.list.return_value = [] + mock_list_response = mocker.Mock() + mock_list_response.data = [] + mock_client.vector_stores.list.return_value = mock_list_response # Mock configuration with empty MCP servers mock_config = mocker.Mock() @@ -1111,7 +1125,9 @@ async def test_retrieve_response_with_mcp_servers( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_client.vector_dbs.list.return_value = [] + mock_list_response = mocker.Mock() + mock_list_response.data = [] + mock_client.vector_stores.list.return_value = mock_list_response # Mock configuration with MCP servers mcp_servers = [ @@ -1190,7 +1206,9 @@ async def test_retrieve_response_with_mcp_servers_empty_token( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_client.vector_dbs.list.return_value = [] + mock_list_response = mocker.Mock() + mock_list_response.data = [] + mock_client.vector_stores.list.return_value = mock_list_response # Mock configuration with MCP servers mcp_servers = [ @@ -1254,7 +1272,9 @@ async def test_retrieve_response_with_mcp_servers_and_mcp_headers( mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client = mocker.AsyncMock() mock_client.shields.list.return_value = [] - mock_client.vector_dbs.list.return_value = [] + mock_list_response = mocker.Mock() + mock_list_response.data = [] + mock_client.vector_stores.list.return_value = mock_list_response # Mock configuration with MCP servers mcp_servers = [ @@ -1510,9 +1530,11 @@ async def test_retrieve_response_no_tools_bypasses_mcp_and_rag( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_vector_db = mocker.Mock() - mock_vector_db.identifier = "VectorDB-1" - mock_client.vector_dbs.list.return_value = [mock_vector_db] + mock_vector_store = mocker.Mock() + mock_vector_store.id = "VectorDB-1" + mock_list_response = mocker.Mock() + mock_list_response.data = [mock_vector_store] + mock_client.vector_stores.list.return_value = mock_list_response # Mock configuration with MCP servers mcp_servers = [ @@ -1560,9 +1582,11 @@ async def test_retrieve_response_no_tools_false_preserves_functionality( mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] - mock_vector_db = mocker.Mock() - mock_vector_db.identifier = "VectorDB-1" - mock_client.vector_dbs.list.return_value = [mock_vector_db] + mock_vector_store = mocker.Mock() + mock_vector_store.id = "VectorDB-1" + mock_list_response = mocker.Mock() + mock_list_response.data = [mock_vector_store] + mock_client.vector_stores.list.return_value = mock_list_response # Mock configuration with MCP servers mcp_servers = [