Enable RAG tool when a vector DB is available (#166)

* Enable RAG tool when a vector DB is available

* Replace typing.List with list

Author: TamiTakamiya

src/app/endpoints/query.py

@@ -10,6 +10,10 @@
 
 from llama_stack_client import LlamaStackClient  # type: ignore
 from llama_stack_client.types import UserMessage  # type: ignore
+from llama_stack_client.types.agents.turn_create_params import (
+    ToolgroupAgentToolGroupWithArgs,
+    Toolgroup,
+)
 from llama_stack_client.types.model_list_response import ModelListResponse
 
 from fastapi import APIRouter, HTTPException, status, Depends
@@ -182,11 +186,13 @@ def retrieve_response(
     )
     session_id = agent.create_session("chat_session")
     logger.debug("Session ID: %s", session_id)
+    vector_db_ids = [vector_db.identifier for vector_db in client.vector_dbs.list()]
     response = agent.create_turn(
         messages=[UserMessage(role="user", content=query_request.query)],
         session_id=session_id,
         documents=query_request.get_documents(),
         stream=False,
+        toolgroups=get_rag_toolgroups(vector_db_ids),
     )
     return str(response.output_message.content)  # type: ignore[union-attr]
 
@@ -282,3 +288,21 @@ def store_transcript(  # pylint: disable=too-many-arguments,too-many-positional-
         json.dump(data_to_store, transcript_file)
 
     logger.info("Transcript successfully stored at: %s", transcript_file_path)
+
+
+def get_rag_toolgroups(
+    vector_db_ids: list[str],
+) -> list[Toolgroup] | None:
+    """Return a list of RAG Tool groups if the given vector DB list is not empty."""
+    return (
+        [
+            ToolgroupAgentToolGroupWithArgs(
+                name="builtin::rag/knowledge_search",
+                args={
+                    "vector_db_ids": vector_db_ids,
+                },
+            )
+        ]
+        if vector_db_ids
+        else None
+    )

src/app/endpoints/streaming_query.py

@@ -20,6 +20,7 @@
 
 
 from app.endpoints.query import (
+    get_rag_toolgroups,
     is_transcripts_enabled,
     retrieve_conversation_id,
     store_transcript,
@@ -202,11 +203,15 @@ async def retrieve_response(
     )
     session_id = await agent.create_session("chat_session")
     logger.debug("Session ID: %s", session_id)
+    vector_db_ids = [
+        vector_db.identifier for vector_db in await client.vector_dbs.list()
+    ]
     response = await agent.create_turn(
         messages=[UserMessage(role="user", content=query_request.query)],
         session_id=session_id,
         documents=query_request.get_documents(),
         stream=True,
+        toolgroups=get_rag_toolgroups(vector_db_ids),
     )
 
     return response

tests/unit/app/endpoints/test_query.py

@@ -12,6 +12,7 @@
     is_transcripts_enabled,
     construct_transcripts_path,
     store_transcript,
+    get_rag_toolgroups,
 )
 from models.requests import QueryRequest, Attachment
 from models.config import ModelContextProtocolServer
@@ -277,12 +278,45 @@ def test_validate_attachments_metadata_invalid_content_type():
     )
 
 
+def test_retrieve_response_vector_db_available(mocker):
+    """Test the retrieve_response function."""
+    mock_agent = mocker.Mock()
+    mock_agent.create_turn.return_value.output_message.content = "LLM answer"
+    mock_client = mocker.Mock()
+    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 configuration with empty MCP servers
+    mock_config = mocker.Mock()
+    mock_config.mcp_servers = []
+    mocker.patch("app.endpoints.query.configuration", mock_config)
+    mocker.patch("app.endpoints.query.Agent", return_value=mock_agent)
+
+    query_request = QueryRequest(query="What is OpenStack?")
+    model_id = "fake_model_id"
+    access_token = "test_token"
+
+    response = retrieve_response(mock_client, model_id, query_request, access_token)
+
+    assert response == "LLM answer"
+    mock_agent.create_turn.assert_called_once_with(
+        messages=[UserMessage(content="What is OpenStack?", role="user", context=None)],
+        session_id=mocker.ANY,
+        documents=[],
+        stream=False,
+        toolgroups=get_rag_toolgroups(["VectorDB-1"]),
+    )
+
+
 def test_retrieve_response_no_available_shields(mocker):
     """Test the retrieve_response function."""
     mock_agent = mocker.Mock()
     mock_agent.create_turn.return_value.output_message.content = "LLM answer"
     mock_client = mocker.Mock()
     mock_client.shields.list.return_value = []
+    mock_client.vector_dbs.list.return_value = []
 
     # Mock configuration with empty MCP servers
     mock_config = mocker.Mock()
@@ -302,6 +336,7 @@ def test_retrieve_response_no_available_shields(mocker):
         session_id=mocker.ANY,
         documents=[],
         stream=False,
+        toolgroups=None,
     )
 
 
@@ -319,6 +354,7 @@ def identifier(self):
     mock_agent.create_turn.return_value.output_message.content = "LLM answer"
     mock_client = mocker.Mock()
     mock_client.shields.list.return_value = [MockShield("shield1")]
+    mock_client.vector_dbs.list.return_value = []
 
     # Mock configuration with empty MCP servers
     mock_config = mocker.Mock()
@@ -338,6 +374,7 @@ def identifier(self):
         session_id=mocker.ANY,
         documents=[],
         stream=False,
+        toolgroups=None,
     )
 
 
@@ -358,6 +395,7 @@ def identifier(self):
         MockShield("shield1"),
         MockShield("shield2"),
     ]
+    mock_client.vector_dbs.list.return_value = []
 
     # Mock configuration with empty MCP servers
     mock_config = mocker.Mock()
@@ -377,6 +415,7 @@ def identifier(self):
         session_id=mocker.ANY,
         documents=[],
         stream=False,
+        toolgroups=None,
     )
 
 
@@ -386,6 +425,7 @@ def test_retrieve_response_with_one_attachment(mocker):
     mock_agent.create_turn.return_value.output_message.content = "LLM answer"
     mock_client = mocker.Mock()
     mock_client.shields.list.return_value = []
+    mock_client.vector_dbs.list.return_value = []
 
     # Mock configuration with empty MCP servers
     mock_config = mocker.Mock()
@@ -418,6 +458,7 @@ def test_retrieve_response_with_one_attachment(mocker):
                 "mime_type": "text/plain",
             },
         ],
+        toolgroups=None,
     )
 
 
@@ -427,6 +468,7 @@ def test_retrieve_response_with_two_attachments(mocker):
     mock_agent.create_turn.return_value.output_message.content = "LLM answer"
     mock_client = mocker.Mock()
     mock_client.shields.list.return_value = []
+    mock_client.vector_dbs.list.return_value = []
 
     # Mock configuration with empty MCP servers
     mock_config = mocker.Mock()
@@ -468,6 +510,7 @@ def test_retrieve_response_with_two_attachments(mocker):
                 "mime_type": "application/yaml",
             },
         ],
+        toolgroups=None,
     )
 
 
@@ -477,6 +520,7 @@ def test_retrieve_response_with_mcp_servers(mocker):
     mock_agent.create_turn.return_value.output_message.content = "LLM answer"
     mock_client = mocker.Mock()
     mock_client.shields.list.return_value = []
+    mock_client.vector_dbs.list.return_value = []
 
     # Mock configuration with MCP servers
     mcp_servers = [
@@ -536,6 +580,7 @@ def test_retrieve_response_with_mcp_servers_empty_token(mocker):
     mock_agent.create_turn.return_value.output_message.content = "LLM answer"
     mock_client = mocker.Mock()
     mock_client.shields.list.return_value = []
+    mock_client.vector_dbs.list.return_value = []
 
     # Mock configuration with MCP servers
     mcp_servers = [
@@ -646,3 +691,16 @@ def test_store_transcript(mocker):
         },
         mocker.ANY,
     )
+
+
+def test_get_rag_toolgroups(mocker):
+    """Test get_rag_toolgroups function."""
+    vector_db_ids = []
+    result = get_rag_toolgroups(vector_db_ids)
+    assert result is None
+
+    vector_db_ids = ["Vector-DB-1", "Vector-DB-2"]
+    result = get_rag_toolgroups(vector_db_ids)
+    assert len(result) == 1
+    assert result[0]["name"] == "builtin::rag/knowledge_search"
+    assert result[0]["args"]["vector_db_ids"] == vector_db_ids

tests/unit/app/endpoints/test_streaming_query.py

@@ -1,5 +1,6 @@
 import pytest
 
+from app.endpoints.query import get_rag_toolgroups
 from app.endpoints.streaming_query import (
     streaming_query_endpoint_handler,
     retrieve_response,
@@ -114,12 +115,41 @@ async def test_streaming_query_endpoint_handler_store_transcript(mocker):
     await _test_streaming_query_endpoint_handler(mocker, store_transcript=True)
 
 
+async def test_retrieve_response_vector_db_available(mocker):
+    """Test the retrieve_response function."""
+    mock_agent = mocker.AsyncMock()
+    mock_agent.create_turn.return_value.output_message.content = "LLM answer"
+    mock_client = mocker.AsyncMock()
+    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]
+
+    mocker.patch("app.endpoints.streaming_query.AsyncAgent", return_value=mock_agent)
+
+    query_request = QueryRequest(query="What is OpenStack?")
+    model_id = "fake_model_id"
+
+    response = await retrieve_response(mock_client, model_id, query_request)
+
+    # For streaming, the response should be the streaming object
+    assert response is not None
+    mock_agent.create_turn.assert_called_once_with(
+        messages=[UserMessage(content="What is OpenStack?", role="user", context=None)],
+        session_id=mocker.ANY,
+        documents=[],
+        stream=True,  # Should be True for streaming endpoint
+        toolgroups=get_rag_toolgroups(["VectorDB-1"]),
+    )
+
+
 async def test_retrieve_response_no_available_shields(mocker):
     """Test the retrieve_response function."""
     mock_agent = mocker.AsyncMock()
     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 = []
 
     mocker.patch("app.endpoints.streaming_query.AsyncAgent", return_value=mock_agent)
 
@@ -135,6 +165,7 @@ async def test_retrieve_response_no_available_shields(mocker):
         session_id=mocker.ANY,
         documents=[],
         stream=True,  # Should be True for streaming endpoint
+        toolgroups=None,
     )
 
 
@@ -166,6 +197,7 @@ def identifier(self):
         session_id=mocker.ANY,
         documents=[],
         stream=True,  # Should be True for streaming endpoint
+        toolgroups=None,
     )
 
 
@@ -186,6 +218,7 @@ def identifier(self):
         MockShield("shield1"),
         MockShield("shield2"),
     ]
+    mock_client.vector_dbs.list.return_value = []
 
     mocker.patch("app.endpoints.streaming_query.AsyncAgent", return_value=mock_agent)
 
@@ -200,6 +233,7 @@ def identifier(self):
         session_id=mocker.ANY,
         documents=[],
         stream=True,  # Should be True for streaming endpoint
+        toolgroups=None,
     )
 
 
@@ -209,6 +243,7 @@ async def test_retrieve_response_with_one_attachment(mocker):
     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 = []
 
     attachments = [
         Attachment(
@@ -235,6 +270,7 @@ async def test_retrieve_response_with_one_attachment(mocker):
                 "mime_type": "text/plain",
             },
         ],
+        toolgroups=None,
     )
 
 
@@ -244,6 +280,7 @@ async def test_retrieve_response_with_two_attachments(mocker):
     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 = []
 
     attachments = [
         Attachment(
@@ -279,6 +316,7 @@ async def test_retrieve_response_with_two_attachments(mocker):
                 "mime_type": "application/yaml",
             },
         ],
+        toolgroups=None,
     )