Move query and streaming_query to vector_stores

Author: luis5tb

src/app/endpoints/query.py

@@ -732,10 +732,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
@@ -831,30 +831,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_store_ids": vector_store_ids,
                 },
             )
         ]
-        if vector_db_ids
+        if vector_store_ids
         else None
     )

src/app/endpoints/streaming_query.py

@@ -1031,10 +1031,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

tests/unit/app/endpoints/test_query.py

@@ -1414,16 +1414,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_store_ids"] == vector_store_ids
 
 
 @pytest.mark.asyncio

tests/unit/app/endpoints/test_streaming_query.py

@@ -453,9 +453,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()
@@ -498,7 +500,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()
@@ -554,7 +558,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()
@@ -611,7 +617,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()
@@ -670,7 +678,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()
@@ -723,7 +733,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()
@@ -777,7 +789,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()
@@ -1188,7 +1202,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 = [
@@ -1267,7 +1283,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 = [
@@ -1331,7 +1349,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 = [
@@ -1587,9 +1607,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 = [
@@ -1637,9 +1659,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 = [