Avoid code duplication at retrieve_response for responses API (v2)

Author: luis5tb

src/app/endpoints/query_v2.py

@@ -17,7 +17,7 @@
 from authentication import get_auth_dependency
 from authentication.interface import AuthTuple
 from authorization.middleware import authorize
-from configuration import configuration
+from configuration import AppConfig, configuration
 import metrics
 from models.config import Action
 from models.requests import QueryRequest
@@ -190,31 +190,9 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
         validate_attachments_metadata(query_request.attachments)
 
     # Prepare tools for responses API
-    toolgroups: list[dict[str, Any]] | None = None
-    if not query_request.no_tools:
-        toolgroups = []
-        # Get vector stores for RAG tools
-        vector_store_ids = [
-            vector_store.id for vector_store in (await client.vector_stores.list()).data
-        ]
-
-        # Add RAG tools if vector stores are available
-        rag_tools = get_rag_tools(vector_store_ids)
-        if rag_tools:
-            toolgroups.extend(rag_tools)
-
-        # Add MCP server tools
-        mcp_tools = get_mcp_tools(configuration.mcp_servers, token)
-        if mcp_tools:
-            toolgroups.extend(mcp_tools)
-            logger.debug(
-                "Configured %d MCP tools: %s",
-                len(mcp_tools),
-                [tool.get("server_label", "unknown") for tool in mcp_tools],
-            )
-        # Convert empty list to None for consistency with existing behavior
-        if not toolgroups:
-            toolgroups = None
+    toolgroups = await prepare_tools_for_responses_api(
+        client, query_request, token, configuration
+    )
 
     # Prepare input for Responses API
     # Convert attachments to text and concatenate with query
@@ -470,3 +448,55 @@ def get_mcp_tools(mcp_servers: list, token: str | None = None) -> list[dict[str,
 
         tools.append(tool_def)
     return tools
+
+
+async def prepare_tools_for_responses_api(
+    client: AsyncLlamaStackClient,
+    query_request: QueryRequest,
+    token: str,
+    config: AppConfig,
+) -> list[dict[str, Any]] | None:
+    """
+    Prepare tools for Responses API including RAG and MCP tools.
+
+    This function retrieves vector stores and combines them with MCP
+    server tools to create a unified toolgroups list for the Responses API.
+
+    Args:
+        client: The Llama Stack client instance
+        query_request: The user's query request
+        token: Authentication token for MCP tools
+        config: Configuration object containing MCP server settings
+
+    Returns:
+        list[dict[str, Any]] | None: List of tool configurations for the
+        Responses API, or None if no_tools is True or no tools are available
+    """
+    if query_request.no_tools:
+        return None
+
+    toolgroups = []
+    # Get vector stores for RAG tools
+    vector_store_ids = [
+        vector_store.id for vector_store in (await client.vector_stores.list()).data
+    ]
+
+    # Add RAG tools if vector stores are available
+    rag_tools = get_rag_tools(vector_store_ids)
+    if rag_tools:
+        toolgroups.extend(rag_tools)
+
+    # Add MCP server tools
+    mcp_tools = get_mcp_tools(config.mcp_servers, token)
+    if mcp_tools:
+        toolgroups.extend(mcp_tools)
+        logger.debug(
+            "Configured %d MCP tools: %s",
+            len(mcp_tools),
+            [tool.get("server_label", "unknown") for tool in mcp_tools],
+        )
+    # Convert empty list to None for consistency with existing behavior
+    if not toolgroups:
+        return None
+
+    return toolgroups

src/app/endpoints/streaming_query_v2.py

@@ -21,9 +21,8 @@
 )
 from app.endpoints.query_v2 import (
     extract_token_usage_from_responses_api,
-    get_mcp_tools,
-    get_rag_tools,
     get_topic_summary,
+    prepare_tools_for_responses_api,
 )
 from app.endpoints.streaming_query import (
     format_stream_data,
@@ -448,31 +447,9 @@ async def retrieve_response(
         validate_attachments_metadata(query_request.attachments)
 
     # Prepare tools for responses API
-    toolgroups: list[dict[str, Any]] | None = None
-    if not query_request.no_tools:
-        toolgroups = []
-        # Get vector stores for RAG tools
-        vector_store_ids = [
-            vector_store.id for vector_store in (await client.vector_stores.list()).data
-        ]
-
-        # Add RAG tools if vector stores are available
-        rag_tools = get_rag_tools(vector_store_ids)
-        if rag_tools:
-            toolgroups.extend(rag_tools)
-
-        # Add MCP server tools
-        mcp_tools = get_mcp_tools(configuration.mcp_servers, token)
-        if mcp_tools:
-            toolgroups.extend(mcp_tools)
-            logger.debug(
-                "Configured %d MCP tools: %s",
-                len(mcp_tools),
-                [tool.get("server_label", "unknown") for tool in mcp_tools],
-            )
-        # Convert empty list to None for consistency with existing behavior
-        if not toolgroups:
-            toolgroups = None
+    toolgroups = await prepare_tools_for_responses_api(
+        client, query_request, token, configuration
+    )
 
     response = await client.responses.create(
         input=query_request.query,