feat: Enable multi-collection query generation by updating context handling to pass a list of selected collections to the backend.

Author: ChingEnLin

backend/models/schemas.py

@@ -30,9 +30,7 @@ class QueryPrompt(BaseModel):
     user_input: str
     account_id: str  # Added for cross-collection schema fetching
     db_context: DBContext
-    collection_context: CollectionContext | None = (
-        None  # Optional context for specific collection
-    )
+    collection_context: list[CollectionContext] = []  # List of contexts for selected collections
     intermediate_context: object | None = (
         None  # Optional intermediate context for complex queries
     )

backend/routes/query.py

@@ -31,10 +31,20 @@ def nl2query(prompt: QueryPrompt = Body(...), authorization: str = Header(...)):
     try:
         user_token = authorization.replace("Bearer ", "")
         access_token = exchange_token_obo(user_token)
-        # Fetch schema summary
-        # Note: We need the connection string to fetch the schema
-        # Ideally we might cache this, but for now we fetch it live
-        schema_summary = get_database_schema_summary(prompt.account_id, prompt.db_context.name, access_token)
+        # Use provided contexts if available to avoid re-fetch and ensure consistency
+        if prompt.collection_context:
+            summary = []
+            for ctx in prompt.collection_context:
+                 doc_str = str(ctx.sampleDocument) if ctx.sampleDocument else "No documents found"
+                 summary.append(f"Collection: {ctx.name}\nSample Document: {doc_str}")
+            schema_summary = "\n\n".join(summary)
+        # Fallback: fetch schema summary from DB
+        else:
+             schema_summary = get_database_schema_summary(
+                prompt.account_id, 
+                prompt.db_context.name, 
+                access_token
+             )
     except Exception as e:
         print(f"Error fetching schema context: {e}")
         schema_summary = "Could not fetch schema summary."
@@ -44,8 +54,8 @@ def nl2query(prompt: QueryPrompt = Body(...), authorization: str = Header(...)):
         prompt.user_input,
         collections,
         prompt.db_context.name,
-        prompt.collection_context,
-        prompt.intermediate_context,
+        collection_context=None,
+        intermediate_context=prompt.intermediate_context,
         all_collections_schema=schema_summary
     )
 

frontend/pages/QueryGeneratorPage.tsx

@@ -619,7 +619,19 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
         throw new Error("No account ID available for query generation.");
       }
 
-      const result = await generateMongoQuery(prompt, accountId, connectedDbInfo ?? undefined, collectionCtx, intermediateContext?.data);
+      // Map selected collection names to their full info objects
+      const selectedCollectionInfos = selectedCollections
+        .map(name => collectionDetailsMap[name])
+        .filter((info): info is CollectionInfo => !!info);
+
+      const result = await generateMongoQuery(
+        prompt,
+        accountId,
+        connectedDbInfo ?? undefined,
+        collectionCtx,
+        intermediateContext?.data,
+        selectedCollectionInfos // Pass full info objects
+      );
       setQueryResult(result);
       setIntermediateContext(null); // Clear context after use
 
@@ -1058,7 +1070,7 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
     if (isPromptUnchanged) return 'Query Generated';
     if (selectedCollections.length > 0) {
       if (selectedCollections.length === 1) return `Generate Query for ${selectedCollections[0]} collection`;
-      return `Generate Query for ${selectedCollections.length} collections`;
+      return `Generate Query across ${selectedCollections.length} collections`;
     }
     return 'Generate Query';
   }, [isLoading, selectedCollections, isPromptUnchanged]);

frontend/services/geminiService.ts

@@ -18,8 +18,9 @@ export const generateMongoQuery = async (
     userInput: string,
     accountId: string,
     dbInfo?: DbInfo,
-    collectionContext?: CollectionInfo,
+    collectionContext?: CollectionInfo, // Kept for backward compatibility/single select
     intermediateContext?: any,
+    selectedCollections: CollectionInfo[] = [],
 ): Promise<QueryResultData> => {
     // --- DEVELOPMENT MOCK ---
     if (!USE_MSAL_AUTH) {
@@ -66,7 +67,11 @@ export const generateMongoQuery = async (
             user_input: userInput,
             account_id: accountId,
             db_context: dbInfo, // Send DB context to the backend for more accurate queries
-            collection_context: collectionContext, // Optional: send collection context if available
+            // Send mapping of selected collections as context
+            collection_context: selectedCollections.length > 0 ? selectedCollections.map(col => ({
+                name: col.name || "",
+                sampleDocument: col.sampleDocument
+            })) : (collectionContext ? [collectionContext] : []), // Fallback to single context as array
             intermediate_context: intermediateContext, // Optional: send data from a previous query
         }),
     });