fix: Updated tools playground to allow vdb selection (meta-llama#1960)

MichaelClifford · web-flow · commit f12011794bc9 · 2025-04-17T09:29:40.000+02:00
# What does this PR do?

This PR lets users select an existing vdb to use with their agent on the
tools page of the playground. The drop down menu that lets users select
a vdb only appears when the rag tool is selected. Without this change,
there is no way for a user to specify which vdb they want their rag tool
to use on the tools page. I have intentionally left the RAG options
sparse here since the full RAG options are exposed on the RAG page.

## Test Plan

Without these changes the RAG tool will throw the following error:
`name: knowledge_search) does not have any content `

With these changes the RAG tool works as expected.

Signed-off-by: Michael Clifford &lt;mcliffor@redhat.com&gt;
diff --git a/llama_stack/distribution/ui/page/playground/tools.py b/llama_stack/distribution/ui/page/playground/tools.py
@@ -37,6 +37,17 @@ def reset_agent():
             label="Available ToolGroups", options=builtin_tools_list, selection_mode="multi", on_change=reset_agent
         )
 
+        if "builtin::rag" in toolgroup_selection:
+            vector_dbs = llama_stack_api.client.vector_dbs.list() or []
+            if not vector_dbs:
+                st.info("No vector databases available for selection.")
+            vector_dbs = [vector_db.identifier for vector_db in vector_dbs]
+            selected_vector_dbs = st.multiselect(
+                label="Select Document Collections to use in RAG queries",
+                options=vector_dbs,
+                on_change=reset_agent,
+            )
+
         st.subheader("MCP Servers")
         mcp_selection = st.pills(
             label="Available MCP Servers", options=mcp_tools_list, selection_mode="multi", on_change=reset_agent
@@ -67,6 +78,16 @@ def reset_agent():
             on_change=reset_agent,
         )
 
+    for i, tool_name in enumerate(toolgroup_selection):
+        if tool_name == "builtin::rag":
+            tool_dict = dict(
+                name="builtin::rag",
+                args={
+                    "vector_db_ids": list(selected_vector_dbs),
+                },
+            )
+            toolgroup_selection[i] = tool_dict
+
     @st.cache_resource
     def create_agent():
         return Agent(
