Move query and streaming_query to vector_stores

Author: luis5tb

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
     )

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

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 = [