Merge pull request lightspeed-core#908 from JslYoon/JslYoon-RAG_toggle

Variable RAG DB setup for query, query v2, and streaming query

Author: tisnik

src/app/endpoints/query.py

@@ -751,9 +751,14 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
             ),
         }
 
-        vector_db_ids = [
-            vector_store.id for vector_store in (await client.vector_stores.list()).data
-        ]
+        # Use specified vector stores or fetch all available ones
+        if query_request.vector_store_ids:
+            vector_db_ids = query_request.vector_store_ids
+        else:
+            vector_db_ids = [
+                vector_store.id
+                for vector_store in (await client.vector_stores.list()).data
+            ]
         toolgroups = (get_rag_toolgroups(vector_db_ids) or []) + [
             mcp_server.name for mcp_server in configuration.mcp_servers
         ]

src/app/endpoints/query_v2.py

@@ -741,10 +741,13 @@ async def prepare_tools_for_responses_api(
         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
-    ]
+    # Get vector stores for RAG tools - use specified ones or fetch all
+    if query_request.vector_store_ids:
+        vector_store_ids = query_request.vector_store_ids
+    else:
+        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)

src/app/endpoints/streaming_query.py

@@ -1088,9 +1088,14 @@ async def retrieve_response(
             ),
         }
 
-        vector_db_ids = [
-            vector_store.id for vector_store in (await client.vector_stores.list()).data
-        ]
+        # Use specified vector stores or fetch all available ones
+        if query_request.vector_store_ids:
+            vector_db_ids = query_request.vector_store_ids
+        else:
+            vector_db_ids = [
+                vector_store.id
+                for vector_store in (await client.vector_stores.list()).data
+            ]
         toolgroups = (get_rag_toolgroups(vector_db_ids) or []) + [
             mcp_server.name for mcp_server in configuration.mcp_servers
         ]

src/models/requests.py

@@ -83,6 +83,7 @@ class QueryRequest(BaseModel):
         no_tools: Whether to bypass all tools and MCP servers (default: False).
         generate_topic_summary: Whether to generate topic summary for new conversations.
         media_type: The optional media type for response format (application/json or text/plain).
+        vector_store_ids: The optional list of specific vector store IDs to query for RAG.
 
     Example:
         ```python
@@ -159,6 +160,13 @@ class QueryRequest(BaseModel):
         examples=[MEDIA_TYPE_JSON, MEDIA_TYPE_TEXT],
     )
 
+    vector_store_ids: Optional[list[str]] = Field(
+        None,
+        description="Optional list of specific vector store IDs to query for RAG. "
+        "If not provided, all available vector stores will be queried.",
+        examples=["ocp_docs", "knowledge_base", "vector_db_1"],
+    )
+
     # provides examples for /docs endpoint
     model_config = {
         "extra": "forbid",
@@ -172,6 +180,7 @@ class QueryRequest(BaseModel):
                     "system_prompt": "You are a helpful assistant",
                     "no_tools": False,
                     "generate_topic_summary": True,
+                    "vector_store_ids": ["ocp_docs", "knowledge_base"],
                     "attachments": [
                         {
                             "attachment_type": "log",