diff --git a/backend/models/schemas.py b/backend/models/schemas.py
index 5024b2a..73d3ef1 100644
--- a/backend/models/schemas.py
+++ b/backend/models/schemas.py
@@ -28,10 +28,11 @@ class CollectionContext(BaseModel):
 
 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
     )
@@ -54,3 +55,22 @@ class DebugQueryRequest(BaseModel):
 
 class DebugSuggestionResponse(BaseModel):
     suggestion: str
+
+
+class SchemaRelationshipsRequest(BaseModel):
+    account_id: str
+    database_name: str
+    collection_names: list[str]
+
+
+class Relationship(BaseModel):
+    source_collection: str
+    source_field: str
+    target_collection: str
+    target_field: str
+    description: str
+    confidence: float  # 0.0 to 1.0
+
+
+class SchemaRelationshipsResponse(BaseModel):
+    relationships: list[Relationship]
diff --git a/backend/routes/query.py b/backend/routes/query.py
index 529f488..af3df51 100644
--- a/backend/routes/query.py
+++ b/backend/routes/query.py
@@ -8,12 +8,19 @@
     ExecuteInput,
     DebugQueryRequest,
     DebugSuggestionResponse,
+    SchemaRelationshipsRequest,
+    SchemaRelationshipsResponse,
 )
 from services.gemini_service import (
     generate_query_from_prompt,
     generate_suggestion_from_query_error,
+    generate_schema_relationships,
+)
+from services.mongo_service import (
+    execute_mongo_query,
+    transform_mongo_result,
+    get_database_schema_summary,
 )
-from services.mongo_service import execute_mongo_query, transform_mongo_result
 from models.analyze import AnalyzeRequest, AnalyzeResponse
 from services.analyze_service import analyze_query_result
 
@@ -21,17 +28,70 @@
 
 
 @router.post("/nl2query")
-def nl2query(prompt: QueryPrompt = Body(...)):
+def nl2query(prompt: QueryPrompt = Body(...), authorization: str = Header(...)):
+    if not authorization.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Invalid token format")
+
+    try:
+        user_token = authorization.replace("Bearer ", "")
+        access_token = exchange_token_obo(user_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."
+
     collections = [col.name for col in prompt.db_context.collections]
     return generate_query_from_prompt(
         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,
     )
 
 
+@router.post("/infer-relationships", response_model=SchemaRelationshipsResponse)
+def infer_relationships(
+    request: SchemaRelationshipsRequest = Body(...), authorization: str = Header(...)
+):
+    if not authorization.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Invalid token format")
+
+    try:
+        user_token = authorization.replace("Bearer ", "")
+        access_token = exchange_token_obo(user_token)
+
+        # Fetch schema summary ONLY for correct collections
+        schema_summary = get_database_schema_summary(
+            request.account_id,
+            request.database_name,
+            access_token,
+            collection_filter=request.collection_names,
+        )
+
+        return generate_schema_relationships(schema_summary)
+
+    except Exception as e:
+        print(f"Error inferring relationships: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 @router.post("/execute")
 def execute(query: ExecuteInput = Body(...), authorization: str = Header(...)):
     if not authorization.startswith("Bearer "):
diff --git a/backend/services/gemini_service.py b/backend/services/gemini_service.py
index 91e4726..0040c6d 100644
--- a/backend/services/gemini_service.py
+++ b/backend/services/gemini_service.py
@@ -1,6 +1,11 @@
 from google import genai
 from google.genai import types
-from models.schemas import GeneratedCode, CollectionContext, DebugSuggestionResponse
+from models.schemas import (
+    GeneratedCode,
+    CollectionContext,
+    DebugSuggestionResponse,
+    SchemaRelationshipsResponse,
+)
 from pydantic import BaseModel, Field
 from typing import Optional, List
 
@@ -31,6 +36,7 @@ class AuditSummaryResponse(BaseModel):
 Database: {database}
 Available collections: {collections}
 Sample collection document (optional): {collection_context}
+Schema summary for ALL collections (for JOINs/lookups): {all_collections_schema}
 Intermediate context (optional): {intermediate_context}
 Return:
 only one line of pure pymongo query code (e.g., db["collection"].find(...))
@@ -112,6 +118,7 @@ def generate_query_from_prompt(
     database: str,
     collection_context: CollectionContext = None,
     intermediate_context: dict = None,
+    all_collections_schema: str = "",
 ) -> GeneratedCode:
     # Prune intermediate_context to remove image/large data
     safe_intermediate_context = (
@@ -124,6 +131,7 @@ def generate_query_from_prompt(
         collection_context=(
             collection_context.sampleDocument if collection_context else ""
         ),
+        all_collections_schema=all_collections_schema,
         intermediate_context=safe_intermediate_context,
     )
     client = genai.Client()
@@ -244,3 +252,58 @@ def summarize_audit_results(
             summary="Could not generate summary due to parsing error.",
             visualization=VisualizationConfig(available=False),
         )
+
+
+PROMPT_TEMPLATE_RELATIONSHIPS = """
+You are a database architect. Analyze the provided MongoDB document samples to identify likely foreign key relationships and JOIN conditions between collections.
+
+Schema/Samples:
+{schema_summary}
+
+Tasks:
+1. Identify likely relationships (e.g., `userId` in `orders` -> `_id` in `users`).
+2. Provide a confidence score (0.0 - 1.0) and a brief description for each.
+3. Return a JSON object with a "relationships" key containing a list of these findings.
+
+Output Format (Json):
+{{
+  "relationships": [
+    {{
+      "source_collection": "orders",
+      "source_field": "userId",
+      "target_collection": "users",
+      "target_field": "_id",
+      "description": "Orders belong to Users",
+      "confidence": 0.95
+    }}
+  ]
+}}
+"""
+
+
+def generate_schema_relationships(schema_summary: str) -> SchemaRelationshipsResponse:
+    from models.schemas import SchemaRelationshipsResponse
+
+    full_prompt = PROMPT_TEMPLATE_RELATIONSHIPS.format(schema_summary=schema_summary)
+    client = genai.Client()
+    response = client.models.generate_content(
+        model="gemini-2.5-flash",
+        contents=full_prompt,
+        config=types.GenerateContentConfig(
+            response_mime_type="application/json",
+            response_schema=SchemaRelationshipsResponse,
+            thinking_config=types.ThinkingConfig(thinking_budget=0),
+        ),
+    )
+
+    if hasattr(response, "parsed") and response.parsed:
+        return response.parsed
+
+    import json
+
+    try:
+        data = json.loads(response.text)
+        return SchemaRelationshipsResponse(**data)
+    except Exception as e:
+        print(f"Error parsing Gemini relationship response: {e}")
+        return SchemaRelationshipsResponse(relationships=[])
diff --git a/backend/services/mongo_service.py b/backend/services/mongo_service.py
index a0a5b56..5a60bfa 100644
--- a/backend/services/mongo_service.py
+++ b/backend/services/mongo_service.py
@@ -56,3 +56,38 @@ def transform_mongo_result(result):
     elif isinstance(result, DeleteResult):
         return {"deleted_count": result.deleted_count}
     return result
+
+
+def get_database_schema_summary(
+    account_id: str,
+    database: str,
+    access_token: str,
+    collection_filter: list[str] = None,
+) -> str:
+    from services.azure_cosmos_resources import get_connection_string
+
+    try:
+        connection_string = get_connection_string(account_id, access_token)
+        client = pymongo.MongoClient(connection_string)
+        db = client[database]
+        summary = []
+
+        # Determine which collections to scan
+        if collection_filter:
+            target_collections = collection_filter
+        else:
+            target_collections = db.list_collection_names()
+
+        for collection_name in target_collections:
+            # Skip system collections if scanning all (if explicit filter, try to fetch)
+            if not collection_filter and collection_name.startswith("system."):
+                continue
+
+            doc = db[collection_name].find_one()
+            doc_str = str(doc) if doc else "No documents found"
+            summary.append(f"Collection: {collection_name}\nSample Document: {doc_str}")
+
+        return "\n\n".join(summary)
+    except Exception as e:
+        print(f"Error fetching schema summary: {e}")
+        return "Could not fetch schema summary."
diff --git a/backend/tests/test_query_routes.py b/backend/tests/test_query_routes.py
index a13e8d5..32880b0 100644
--- a/backend/tests/test_query_routes.py
+++ b/backend/tests/test_query_routes.py
@@ -15,7 +15,10 @@
 def test_nl2query(client):
     """Test natural language to query conversion."""
     # Mock dependencies
-    with patch("routes.query.generate_query_from_prompt") as mock_generate:
+    with (
+        patch("routes.query.generate_query_from_prompt") as mock_generate,
+        patch("routes.query.exchange_token_obo") as mock_exchange,
+    ):
         mock_generate.return_value = {"generated_code": "db.users.find({})"}
 
         # Create test data
@@ -25,11 +28,15 @@ def test_nl2query(client):
 
         prompt = QueryPrompt(
             user_input="Find all users",
+            account_id="test-account",
             db_context=db_context,
-            collection_context=collection_context,
+            collection_context=[collection_context],
         )
 
-        response = client.post("/query/nl2query", json=prompt.model_dump())
+        headers = {"authorization": "Bearer valid-token"}
+        response = client.post(
+            "/query/nl2query", json=prompt.model_dump(), headers=headers
+        )
 
         assert response.status_code == 200
         data = response.json()
diff --git a/backend/tests/test_schemas.py b/backend/tests/test_schemas.py
index a4a7f0f..cd5c0ee 100644
--- a/backend/tests/test_schemas.py
+++ b/backend/tests/test_schemas.py
@@ -91,19 +91,22 @@ def test_query_prompt():
 
     prompt = QueryPrompt(
         user_input="Find all users",
+        account_id="test-account",
         db_context=db_context,
-        collection_context=collection_context,
+        collection_context=[collection_context],
         intermediate_context={"key": "value"},
     )
 
     assert prompt.user_input == "Find all users"
     assert prompt.db_context.name == "test-db"
-    assert prompt.collection_context.name == "users"
+    assert prompt.collection_context[0].name == "users"
     assert prompt.intermediate_context == {"key": "value"}
 
     # Test with minimal required fields
-    minimal_prompt = QueryPrompt(user_input="Find all users", db_context=db_context)
-    assert minimal_prompt.collection_context is None
+    minimal_prompt = QueryPrompt(
+        user_input="Find all users", account_id="test-account", db_context=db_context
+    )
+    assert minimal_prompt.collection_context == []
     assert minimal_prompt.intermediate_context is None
 
 
diff --git a/frontend/components/SchemaRelationshipGraph.tsx b/frontend/components/SchemaRelationshipGraph.tsx
new file mode 100644
index 0000000..5fe1ff3
--- /dev/null
+++ b/frontend/components/SchemaRelationshipGraph.tsx
@@ -0,0 +1,242 @@
+
+import React, { useMemo, useState } from 'react';
+import { Relationship, SchemaRelationshipsResponse } from '../types';
+
+interface SchemaRelationshipGraphProps {
+    relationships: SchemaRelationshipsResponse;
+    selectedCollections: string[];
+}
+
+interface Node {
+    id: string;
+    x: number;
+    y: number;
+}
+
+interface Edge {
+    source: Node;
+    target: Node;
+    data: Relationship;
+}
+
+const SchemaRelationshipGraph: React.FC<SchemaRelationshipGraphProps> = ({ relationships, selectedCollections }) => {
+    const [hoveredNode, setHoveredNode] = useState<string | null>(null);
+    const [hoveredEdge, setHoveredEdge] = useState<number | null>(null);
+
+    // Filter collections to only those selected (or involved in relationships between selected)
+    // Actually, we should show all "selected" collections as nodes, even if isolated.
+    const nodes: Node[] = useMemo(() => {
+        const uniqueCols = Array.from(new Set(selectedCollections));
+        const count = uniqueCols.length;
+        const radius = 120; // Radius of the circle layout
+        const centerX = 200;
+        const centerY = 150; // Reduced height
+
+        if (count === 2) {
+            return [
+                { id: uniqueCols[0], x: 80, y: centerY },
+                { id: uniqueCols[1], x: 320, y: centerY }
+            ];
+        }
+
+        return uniqueCols.map((col, i) => {
+            const angle = (i / count) * 2 * Math.PI - Math.PI / 2; // Start from top
+            return {
+                id: col,
+                x: centerX + radius * Math.cos(angle),
+                y: centerY + radius * Math.sin(angle),
+            };
+        });
+    }, [selectedCollections]);
+
+    const edges: Edge[] = useMemo(() => {
+        return relationships.relationships
+            .filter(rel => selectedCollections.includes(rel.source_collection) && selectedCollections.includes(rel.target_collection))
+            .map(rel => {
+                const sourceNode = nodes.find(n => n.id === rel.source_collection);
+                const targetNode = nodes.find(n => n.id === rel.target_collection);
+                if (!sourceNode || !targetNode) return null;
+                return { source: sourceNode, target: targetNode, data: rel };
+            })
+            .filter((e): e is Edge => e !== null);
+    }, [relationships, nodes, selectedCollections]);
+
+    // Helper to calculate bezier curve control point
+    const getPath = (source: Node, target: Node, index: number) => {
+        const dx = target.x - source.x;
+        const dy = target.y - source.y;
+        const dist = Math.sqrt(dx * dx + dy * dy);
+
+        // Curvature logic: Straight line if direct, curved if multiple or loops
+        let cx = (source.x + target.x) / 2;
+        let cy = (source.y + target.y) / 2;
+
+        // Add offset for multiple edges or just for style
+        // If it's a bidirectional pair, we need curvature.
+        // Simple logic: curve upwards/downwards based on direction or index
+
+        // A simple consistent curve offset
+        // const offset = 40; // This variable was not used, removed.
+        // Calculate perpendicular vector
+        const px = -dy / dist;
+        const py = dx / dist;
+
+        // Small randomish offset based on data string hash or index to allow separate lines for multiple rels
+        const curveAmount = 30 + (index * 20);
+
+        cx += px * curveAmount;
+        cy += py * curveAmount;
+
+        // Calculate intersection point on the edge of the target circle (radius 30 + arrow size approx 10)
+        // We want the arrow to point to the edge, efficiently.
+        // Actually, Q curves are hard to chop exactly. 
+        // Easier approach: Use a marker-end that has refX properly set, OR calculate the point on the circle.
+
+        // Let's recalculate target/source points to be on the circumference.
+        const radius = 30 + 5; // Node radius + buffer
+        const angleSource = Math.atan2(cy - source.y, cx - source.x);
+        const angleTarget = Math.atan2(cy - target.y, cx - target.x);
+
+        const sx = source.x + radius * Math.cos(angleSource);
+        const sy = source.y + radius * Math.sin(angleSource);
+
+        const tx = target.x + radius * Math.cos(angleTarget);
+        const ty = target.y + radius * Math.sin(angleTarget);
+
+        return `M ${sx} ${sy} Q ${cx} ${cy} ${tx} ${ty}`;
+    };
+
+    return (
+        <div className="w-full flex flex-col items-center bg-slate-50 dark:bg-slate-900/50 rounded-xl border border-slate-200 dark:border-slate-700 p-6 overflow-hidden relative">
+            <div className="absolute top-2 right-2 flex gap-2">
+                <div className="flex items-center gap-1 text-[10px] text-slate-400 dark:text-slate-500 bg-white dark:bg-slate-800 px-2 py-1 rounded-full border border-slate-200 dark:border-slate-700 shadow-sm">
+                    <span className="w-2 h-2 rounded-full bg-blue-500 animate-pulse"></span>
+                    AI Inferred Connection
+                </div>
+            </div>
+
+            <svg width="400" height="300" viewBox="0 0 400 300" className="w-full max-w-[500px] h-auto pointer-events-auto">
+                <defs>
+                    <marker id="arrowhead" markerWidth="10" markerHeight="7" refX="22" refY="3.5" orient="auto">
+                        <polygon points="0 0, 10 3.5, 0 7" fill="#64748b" className="dark:fill-slate-400" />
+                    </marker>
+                    <marker id="arrowhead-hover" markerWidth="10" markerHeight="7" refX="22" refY="3.5" orient="auto">
+                        <polygon points="0 0, 10 3.5, 0 7" fill="#3b82f6" />
+                    </marker>
+                    <filter id="glow" x="-20%" y="-20%" width="140%" height="140%">
+                        <feGaussianBlur stdDeviation="2.5" result="coloredBlur" />
+                        <feMerge>
+                            <feMergeNode in="coloredBlur" />
+                            <feMergeNode in="SourceGraphic" />
+                        </feMerge>
+                    </filter>
+                </defs>
+
+                {/* Edges */}
+                {edges.map((edge, i) => {
+                    const isHovered = hoveredEdge === i;
+                    const pathData = getPath(edge.source, edge.target, i);
+
+                    return (
+                        <g
+                            key={i}
+                            onMouseEnter={() => setHoveredEdge(i)}
+                            onMouseLeave={() => setHoveredEdge(null)}
+                            className="transition-opacity duration-300"
+                            style={{ opacity: (hoveredEdge !== null && !isHovered) ? 0.3 : 1 }}
+                        >
+                            {/* Invisible wideline for easier hovering */}
+                            <path d={pathData} stroke="transparent" strokeWidth="20" fill="none" className="cursor-pointer" />
+
+                            {/* Visible line - No arrowheads */}
+                            <path
+                                d={pathData}
+                                stroke={isHovered ? "#3b82f6" : "#cbd5e1"}
+                                strokeWidth={isHovered ? 3 : 2}
+                                fill="none"
+                                className={`transition-colors duration-300 ${isHovered ? '' : 'dark:stroke-slate-600'}`}
+                            />
+
+                            {/* Animated particle flow only on hover */}
+                            {isHovered && (
+                                <circle r="3" fill="#3b82f6">
+                                    <animateMotion dur="1.5s" repeatCount="indefinite" path={pathData} />
+                                </circle>
+                            )}
+                        </g>
+                    );
+                })}
+
+                {/* Nodes */}
+                {nodes.map(node => {
+                    const isHovered = hoveredNode === node.id;
+                    const isRelatedToHoveredEdge = hoveredEdge !== null && (edges[hoveredEdge!].source.id === node.id || edges[hoveredEdge!].target.id === node.id);
+
+                    return (
+                        <g
+                            key={node.id}
+                            transform={`translate(${node.x},${node.y})`}
+                            onMouseEnter={() => setHoveredNode(node.id)}
+                            onMouseLeave={() => setHoveredNode(null)}
+                            className="cursor-pointer transition-all duration-300"
+                            style={{
+                                opacity: (hoveredEdge !== null && !isRelatedToHoveredEdge) ? 0.4 : 1
+                            }}
+                        >
+                            <circle
+                                r="30"
+                                fill={isHovered || isRelatedToHoveredEdge ? "#eff6ff" : "white"}
+                                stroke={isHovered || isRelatedToHoveredEdge ? "#3b82f6" : "#cbd5e1"}
+                                strokeWidth={isHovered ? 3 : 2}
+                                className={`transition-colors duration-300 dark:bg-slate-800 dark:fill-slate-800 ${isHovered ? '' : 'dark:stroke-slate-600'}`}
+                                filter={isHovered ? "url(#glow)" : ""}
+                            />
+                            {/* Full Name */}
+                            <text
+                                y="5"
+                                textAnchor="middle"
+                                alignmentBaseline="middle"
+                                className={`text-[10px] font-bold pointer-events-none custom-text-shadow ${isHovered ? 'fill-blue-600' : 'fill-slate-600 dark:fill-slate-300'}`}
+                                style={{ textShadow: '0 1px 2px rgba(255,255,255,0.8)' }}
+                            >
+                                {node.id}
+                            </text>
+                        </g>
+                    );
+                })}
+            </svg>
+
+            {/* Detail Card Overlay - Appears when hovering an edge */}
+            <div className={`absolute bottom-4 left-0 right-0 mx-auto w-[90%] bg-white/90 dark:bg-slate-800/90 backdrop-blur-sm border border-slate-200 dark:border-slate-600 shadow-lg rounded-lg p-3 transition-transform duration-300 ${hoveredEdge !== null ? 'translate-y-0 opacity-100' : 'translate-y-10 opacity-0 pointer-events-none'}`}>
+                {hoveredEdge !== null && (
+                    <div className="flex flex-col items-center text-center">
+                        <div className="flex items-center gap-2 mb-1">
+                            <span className="font-mono bg-slate-100 dark:bg-slate-700 px-1.5 py-0.5 rounded text-xs text-slate-700 dark:text-slate-200 font-bold">
+                                {edges[hoveredEdge].data.source_collection}.{edges[hoveredEdge].data.source_field}
+                            </span>
+                            <span className="text-slate-400">→</span>
+                            <span className="font-mono bg-slate-100 dark:bg-slate-700 px-1.5 py-0.5 rounded text-xs text-slate-700 dark:text-slate-200 font-bold">
+                                {edges[hoveredEdge].data.target_collection}.{edges[hoveredEdge].data.target_field}
+                            </span>
+                        </div>
+                        <p className="text-xs text-slate-600 dark:text-slate-400">
+                            {edges[hoveredEdge].data.description}
+                        </p>
+                    </div>
+                )}
+            </div>
+
+            {/* Empty State Overlay */}
+            {edges.length === 0 && (
+                <div className="absolute inset-0 flex items-center justify-center pointer-events-none">
+                    <div className="bg-white/80 dark:bg-slate-900/80 backdrop-blur p-4 rounded-lg text-center">
+                        <p className="text-slate-500 text-sm">No connections found</p>
+                    </div>
+                </div>
+            )}
+
+        </div>
+    );
+};
+
+export default SchemaRelationshipGraph;
diff --git a/frontend/pages/QueryGeneratorPage.tsx b/frontend/pages/QueryGeneratorPage.tsx
index 9bb4823..bf45a37 100644
--- a/frontend/pages/QueryGeneratorPage.tsx
+++ b/frontend/pages/QueryGeneratorPage.tsx
@@ -1,11 +1,10 @@
-
 import React, { useState, useCallback, useEffect, useMemo, useRef } from 'react';
 import { createPortal } from 'react-dom';
-import { generateMongoQuery, debugMongoQuery, analyzeQueryResult } from '../services/geminiService';
+import { generateMongoQuery, debugMongoQuery, analyzeQueryResult, inferSchemaRelationships } from '../services/geminiService';
 import { getAzureCosmosAccounts, getDatabasesForAccount, runMongoQuery, getCollectionInfo, clearSystemCache } from '../services/dbService';
 import { getSavedQueries, saveQuery, updateSavedQuery, deleteSavedQuery } from '../services/userDataService';
 import { generateIpynbContent, downloadFile } from '../services/notebookService';
-import { QueryResultData, DbInfo, CollectionInfo, CosmosDBAccount, SelectedResource, DebuggingResult, AnalysisResult, NotebookStep, SavedQuery } from '../types';
+import { QueryResultData, DbInfo, CollectionInfo, CosmosDBAccount, SelectedResource, DebuggingResult, AnalysisResult, NotebookStep, SavedQuery, SchemaRelationshipsResponse } from '../types';
 import { mockECommerceDbInfo, mockCollectionInfoMap, mockFindUsersQuery, mockUserFindResult, mockSavedQueries } from '../services/mockData';
 import { getAuthErrorMessage, isAuthenticationExpiredError } from '../utils/authErrorHandler';
 import QueryDisplay from '../components/QueryDisplay';
@@ -14,6 +13,7 @@ import Loader from '../components/Loader';
 import Tutorial from '../components/Tutorial';
 import JsonDisplay from '../components/JsonDisplay';
 import CollectionActionPanel from '../components/CollectionActionPanel';
+import SchemaRelationshipGraph from '../components/SchemaRelationshipGraph';
 import SavedQueriesPanel from '../components/SavedQueriesPanel';
 import SaveQueryDialog from '../components/SaveQueryDialog';
 import ShareQueryDialog from '../components/ShareQueryDialog';
@@ -73,15 +73,15 @@ const UserMenu: React.FC<{
   }, [menuRef]);
 
   const getCacheButtonContent = () => {
-      if (isClearingCache) return <><SpinnerIcon className="w-4 h-4" /> Clearing...</>;
-      if (cacheClearStatus === 'success') return <><CheckIcon className="w-4 h-4" /> Cleared!</>;
-      if (cacheClearStatus === 'error') return <>Error</>;
-      return <><CachedIcon className="w-4 h-4" /> Clear Cache</>;
+    if (isClearingCache) return <><SpinnerIcon className="w-4 h-4" /> Clearing...</>;
+    if (cacheClearStatus === 'success') return <><CheckIcon className="w-4 h-4" /> Cleared!</>;
+    if (cacheClearStatus === 'error') return <>Error</>;
+    return <><CachedIcon className="w-4 h-4" /> Clear Cache</>;
   };
-  
+
   return (
     <div className="relative" ref={menuRef}>
-      <button 
+      <button
         id="user-menu-button"
         onClick={() => setIsOpen(!isOpen)}
         className="flex items-center gap-1.5 h-9 px-3 border border-slate-300 dark:border-slate-600 rounded-md text-slate-600 dark:text-slate-300 bg-white dark:bg-slate-800 hover:bg-slate-100 dark:hover:bg-slate-700 transition-colors"
@@ -108,7 +108,7 @@ const UserMenu: React.FC<{
             <button onClick={() => { onStartTutorial(); setIsOpen(false); }} className="w-full flex items-center gap-3 px-3 py-2 text-sm text-slate-700 dark:text-slate-200 rounded-md hover:bg-slate-100 dark:hover:bg-slate-700">
               <HelpIcon className="w-5 h-5" /> Start Tutorial
             </button>
-            
+
             <div className="h-px bg-slate-200 dark:bg-slate-700 my-2"></div>
 
             <button onClick={onClearCache} disabled={isClearingCache || cacheClearStatus !== 'idle'} className="w-full flex items-center gap-3 px-3 py-2 text-sm text-slate-700 dark:text-slate-200 rounded-md hover:bg-slate-100 dark:hover:bg-slate-700 disabled:opacity-50">
@@ -116,7 +116,7 @@ const UserMenu: React.FC<{
             </button>
 
             <div className="h-px bg-slate-200 dark:bg-slate-700 my-2"></div>
-            
+
             <button onClick={() => { onLogout(); setIsOpen(false); }} className="w-full flex items-center gap-3 px-3 py-2 text-sm text-red-600 dark:text-red-400 rounded-md hover:bg-red-50 dark:hover:bg-red-900/40">
               <SignOutIcon className="w-5 h-5" /> Sign Out
             </button>
@@ -138,54 +138,43 @@ interface HeaderUIProps {
   onStartTutorial: () => void;
   onShowSavedQueries: () => void;
   onShowShortcuts: () => void;
-  onNavigateToExplorer: () => void;
-  isExplorerNavEnabled: boolean;
   isUserMenuForcedOpen?: boolean;
 }
 
-const HeaderUI: React.FC<HeaderUIProps> = ({ name, onLogout, onClearCache, isClearingCache, cacheClearStatus, onStartTutorial, onShowSavedQueries, onShowShortcuts, onNavigateToExplorer, isExplorerNavEnabled, isUserMenuForcedOpen }) => {
+const HeaderUI: React.FC<HeaderUIProps> = ({ name, onLogout, onClearCache, isClearingCache, cacheClearStatus, onStartTutorial, onShowSavedQueries, onShowShortcuts, isUserMenuForcedOpen }) => {
   const { theme, toggleTheme } = useTheme();
 
   return (
     <header id="tutorial-header-actions" className="flex flex-col gap-4 sm:flex-row sm:items-center sm:justify-between mb-8">
-        <div className="flex items-center space-x-4">
-            <MongoIcon className="w-12 h-12 text-blue-500" />
-            <div>
-                <h1 className="text-3xl sm:text-4xl font-bold text-slate-900 dark:text-slate-100">QueryPal</h1>
-                <p className="text-slate-500 dark:text-slate-400 text-sm sm:text-base">Your AI-powered database assistant.</p>
-            </div>
-        </div>
-        <div className="flex items-center gap-2 flex-wrap">
-          <button
-            onClick={onNavigateToExplorer}
-            disabled={!isExplorerNavEnabled}
-            className="h-9 flex items-center justify-center gap-2 px-3 border text-sm font-medium rounded-md transition-all duration-300 bg-white dark:bg-slate-800 border-slate-300 dark:border-slate-600 text-slate-600 dark:text-slate-300 hover:bg-slate-100 dark:hover:bg-slate-700 disabled:opacity-50 disabled:cursor-not-allowed"
-            title="Explore and manage data"
-          >
-            <DataGridIcon className="w-5 h-5" />
-            <span className="hidden md:inline">Data Explorer</span>
-          </button>
-          <button
-            onClick={toggleTheme}
-            className="h-9 w-9 flex items-center justify-center border border-slate-300 dark:border-slate-600 rounded-md text-slate-600 dark:text-slate-300 bg-white dark:bg-slate-800 hover:bg-slate-100 dark:hover:bg-slate-700 transition-colors"
-            aria-label="Toggle theme"
-            title="Toggle light/dark mode"
-          >
-            {theme === 'light' ? <MoonIcon className="w-5 h-5" /> : <SunIcon className="w-5 h-5" />}
-          </button>
-          
-          <UserMenu
-            name={name}
-            onLogout={onLogout}
-            onClearCache={onClearCache}
-            isClearingCache={isClearingCache}
-            cacheClearStatus={cacheClearStatus}
-            onStartTutorial={onStartTutorial}
-            onShowSavedQueries={onShowSavedQueries}
-            onShowShortcuts={onShowShortcuts}
-            isForcedOpen={isUserMenuForcedOpen}
-          />
+      <div className="flex items-center space-x-4">
+        <MongoIcon className="w-12 h-12 text-blue-500" />
+        <div>
+          <h1 className="text-3xl sm:text-4xl font-bold text-slate-900 dark:text-slate-100">QueryPal</h1>
+          <p className="text-slate-500 dark:text-slate-400 text-sm sm:text-base">Your AI-powered database assistant.</p>
         </div>
+      </div>
+      <div className="flex items-center gap-2 flex-wrap">
+        <button
+          onClick={toggleTheme}
+          className="h-9 w-9 flex items-center justify-center border border-slate-300 dark:border-slate-600 rounded-md text-slate-600 dark:text-slate-300 bg-white dark:bg-slate-800 hover:bg-slate-100 dark:hover:bg-slate-700 transition-colors"
+          aria-label="Toggle theme"
+          title="Toggle light/dark mode"
+        >
+          {theme === 'light' ? <MoonIcon className="w-5 h-5" /> : <SunIcon className="w-5 h-5" />}
+        </button>
+
+        <UserMenu
+          name={name}
+          onLogout={onLogout}
+          onClearCache={onClearCache}
+          isClearingCache={isClearingCache}
+          cacheClearStatus={cacheClearStatus}
+          onStartTutorial={onStartTutorial}
+          onShowSavedQueries={onShowSavedQueries}
+          onShowShortcuts={onShowShortcuts}
+          isForcedOpen={isUserMenuForcedOpen}
+        />
+      </div>
     </header>
   );
 };
@@ -211,11 +200,11 @@ const NotebookStepCard: React.FC<NotebookStepCardProps> = ({ step, index, onRemo
           <div className="flex items-center gap-2 opacity-0 group-hover:opacity-100 transition-opacity">
             {step.isEditing ? (
               <button onClick={handleSave} className="flex items-center gap-1 text-xs px-2 py-1 rounded bg-blue-600 text-white hover:bg-blue-700" title="Save changes">
-                <CheckIcon className="w-3 h-3"/> Save
+                <CheckIcon className="w-3 h-3" /> Save
               </button>
             ) : (
               <button onClick={() => onSetEditing(step.id, true)} className="flex items-center gap-1 text-xs px-2 py-1 rounded bg-slate-600 text-slate-200 hover:bg-slate-500" title="Edit note">
-                <EditIcon className="w-3 h-3"/> Edit
+                <EditIcon className="w-3 h-3" /> Edit
               </button>
             )}
             <button onClick={() => onRemove(step.id)} className="p-1 rounded-full text-slate-500 hover:bg-red-900/50 hover:text-red-400" aria-label="Remove note" title="Remove step">
@@ -273,66 +262,66 @@ const NotebookStepCard: React.FC<NotebookStepCardProps> = ({ step, index, onRemo
 };
 
 interface NotebookPanelProps {
-    steps: NotebookStep[];
-    onClose: () => void;
-    onExport: () => void;
-    onClear: () => void;
-    onRemoveStep: (id: string) => void;
-    onAddNote: () => void;
-    onUpdateStep: (id: string, content: string) => void;
-    onSetEditing: (id: string, isEditing: boolean) => void;
+  steps: NotebookStep[];
+  onClose: () => void;
+  onExport: () => void;
+  onClear: () => void;
+  onRemoveStep: (id: string) => void;
+  onAddNote: () => void;
+  onUpdateStep: (id: string, content: string) => void;
+  onSetEditing: (id: string, isEditing: boolean) => void;
 }
 
 const NotebookPanel: React.FC<NotebookPanelProps> = ({ steps, onClose, onExport, onClear, onRemoveStep, onAddNote, onUpdateStep, onSetEditing }) => (
-    <>
-      <div
-        onClick={onClose}
-        className="fixed inset-0 bg-black bg-opacity-60 z-40 animate-fade-in-fast"
-        aria-hidden="true"
-      ></div>
-      <aside id="tutorial-notebook-panel" className="fixed top-0 right-0 h-full w-full md:w-[450px] bg-slate-900 shadow-2xl z-50 flex flex-col animate-slide-in-drawer">
-        <header className="flex items-center justify-between p-4 border-b border-slate-700 flex-shrink-0">
-          <h3 className="text-lg font-semibold text-white flex items-center gap-3">
-            <NotebookIcon className="w-5 h-5 text-blue-400" />
-            Query Notebook
-          </h3>
-          <button
-            onClick={onClose}
-            className="p-1.5 rounded-full text-slate-400 hover:bg-slate-700 hover:text-white transition-colors"
-            aria-label="Close notebook panel"
-            title="Close notebook"
-          >
-            <XIcon className="w-5 h-5" />
-          </button>
-        </header>
-        <div className="flex-shrink-0 p-4 border-b border-slate-700 flex items-center justify-between gap-2">
-            <div className="flex items-center gap-2">
-                <button onClick={onAddNote} className="flex items-center gap-2 px-3 py-1.5 border border-slate-600 text-sm font-medium rounded-md text-slate-300 bg-slate-800 hover:bg-slate-700 transition-colors" title="Add a new markdown note"><PlusCircleIcon className="w-4 h-4"/>Add Note</button>
-                <button onClick={onClear} disabled={steps.length === 0} className="flex items-center gap-2 px-3 py-1.5 border border-slate-600 text-sm font-medium rounded-md text-slate-300 bg-slate-800 hover:bg-slate-700 disabled:opacity-50 disabled:cursor-not-allowed transition-colors" title="Clear all steps"><TrashIcon className="w-4 h-4"/>Clear All</button>
-            </div>
-            <button onClick={onExport} disabled={steps.length === 0} className="flex items-center gap-2 px-4 py-1.5 border border-transparent text-sm font-medium rounded-md text-white bg-blue-600 hover:bg-blue-700 disabled:opacity-50 disabled:cursor-not-allowed transition-colors" title="Export as .ipynb notebook"><DownloadIcon className="w-4 h-4"/>Export .ipynb</button>
-        </div>
-        <div className="flex-grow overflow-auto p-4 space-y-4">
-          {steps.length > 0 ? (
-            steps.map((step, index) => (
-              <NotebookStepCard 
-                key={step.id} 
-                step={step} 
-                index={index} 
-                onRemove={onRemoveStep}
-                onUpdate={onUpdateStep}
-                onSetEditing={onSetEditing}
-              />
-            ))
-          ) : (
-            <div className="text-center text-slate-500 h-full flex flex-col items-center justify-center">
-                <p className="font-semibold">No steps recorded yet.</p>
-                <p className="text-sm">Run a query or add a note to begin.</p>
-            </div>
-          )}
+  <>
+    <div
+      onClick={onClose}
+      className="fixed inset-0 bg-black bg-opacity-60 z-40 animate-fade-in-fast"
+      aria-hidden="true"
+    ></div>
+    <aside id="tutorial-notebook-panel" className="fixed top-0 right-0 h-full w-full md:w-[450px] bg-slate-900 shadow-2xl z-50 flex flex-col animate-slide-in-drawer">
+      <header className="flex items-center justify-between p-4 border-b border-slate-700 flex-shrink-0">
+        <h3 className="text-lg font-semibold text-white flex items-center gap-3">
+          <NotebookIcon className="w-5 h-5 text-blue-400" />
+          Query Notebook
+        </h3>
+        <button
+          onClick={onClose}
+          className="p-1.5 rounded-full text-slate-400 hover:bg-slate-700 hover:text-white transition-colors"
+          aria-label="Close notebook panel"
+          title="Close notebook"
+        >
+          <XIcon className="w-5 h-5" />
+        </button>
+      </header>
+      <div className="flex-shrink-0 p-4 border-b border-slate-700 flex items-center justify-between gap-2">
+        <div className="flex items-center gap-2">
+          <button onClick={onAddNote} className="flex items-center gap-2 px-3 py-1.5 border border-slate-600 text-sm font-medium rounded-md text-slate-300 bg-slate-800 hover:bg-slate-700 transition-colors" title="Add a new markdown note"><PlusCircleIcon className="w-4 h-4" />Add Note</button>
+          <button onClick={onClear} disabled={steps.length === 0} className="flex items-center gap-2 px-3 py-1.5 border border-slate-600 text-sm font-medium rounded-md text-slate-300 bg-slate-800 hover:bg-slate-700 disabled:opacity-50 disabled:cursor-not-allowed transition-colors" title="Clear all steps"><TrashIcon className="w-4 h-4" />Clear All</button>
         </div>
-      </aside>
-    </>
+        <button onClick={onExport} disabled={steps.length === 0} className="flex items-center gap-2 px-4 py-1.5 border border-transparent text-sm font-medium rounded-md text-white bg-blue-600 hover:bg-blue-700 disabled:opacity-50 disabled:cursor-not-allowed transition-colors" title="Export as .ipynb notebook"><DownloadIcon className="w-4 h-4" />Export .ipynb</button>
+      </div>
+      <div className="flex-grow overflow-auto p-4 space-y-4">
+        {steps.length > 0 ? (
+          steps.map((step, index) => (
+            <NotebookStepCard
+              key={step.id}
+              step={step}
+              index={index}
+              onRemove={onRemoveStep}
+              onUpdate={onUpdateStep}
+              onSetEditing={onSetEditing}
+            />
+          ))
+        ) : (
+          <div className="text-center text-slate-500 h-full flex flex-col items-center justify-center">
+            <p className="font-semibold">No steps recorded yet.</p>
+            <p className="text-sm">Run a query or add a note to begin.</p>
+          </div>
+        )}
+      </div>
+    </aside>
+  </>
 );
 
 export interface QueryGeneratorPageProps {
@@ -347,7 +336,7 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
   const [userInput, setUserInput] = useState<string>('');
   const [isLoading, setIsLoading] = useState<boolean>(false);
   const [error, setError] = useState<string | null>(null);
-  
+
   // State for DB resources & connection
   const [azureAccounts, setAzureAccounts] = useState<CosmosDBAccount[]>([]);
   const [isLoadingAccounts, setIsLoadingAccounts] = useState<boolean>(true);
@@ -370,7 +359,7 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
   const [isExecuting, setIsExecuting] = useState<boolean>(false);
   const [executionResult, setExecutionResult] = useState<any | null>(null);
   const [executionError, setExecutionError] = useState<string | null>(null);
-  
+
   // State for intermediate context (multi-step queries)
   const [intermediateContext, setIntermediateContext] = useState<{ data: any; source: string; } | null>(null);
   const [currentQueryContextSource, setCurrentQueryContextSource] = useState<string | null>(null);
@@ -388,11 +377,18 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
   const [analysisError, setAnalysisError] = useState<string | null>(null);
 
   // State for collection details
-  const [selectedCollection, setSelectedCollection] = useState<string | null>(null);
-  const [collectionInfo, setCollectionInfo] = useState<CollectionInfo | null>(null);
-  const [isFetchingCollectionInfo, setIsFetchingCollectionInfo] = useState<boolean>(false);
+  const [selectedCollections, setSelectedCollections] = useState<string[]>([]);
+
+  const [collectionDetailsMap, setCollectionDetailsMap] = useState<Record<string, CollectionInfo>>({});
+  const [loadingCollections, setLoadingCollections] = useState<Record<string, boolean>>({});
+  const [expandedCollectionSchemas, setExpandedCollectionSchemas] = useState<Record<string, boolean>>({}); // Track open/close state for stacking
   const [collectionInfoError, setCollectionInfoError] = useState<string | null>(null);
-  
+
+  // State for relationship inference
+  const [relationships, setRelationships] = useState<SchemaRelationshipsResponse | null>(null);
+  const [isAnalyzingRelationships, setIsAnalyzingRelationships] = useState<boolean>(false);
+  const [relationshipError, setRelationshipError] = useState<string | null>(null);
+
   // State for cache clearing
   const [isClearingCache, setIsClearingCache] = useState<boolean>(false);
   const [cacheClearStatus, setCacheClearStatus] = useState<'idle' | 'success' | 'error'>('idle');
@@ -411,7 +407,7 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
   const [savedQueries, setSavedQueries] = useState<SavedQuery[]>([]);
   const [isLoadingSavedQueries, setIsLoadingSavedQueries] = useState<boolean>(false);
   const [isSavedQueriesPanelOpen, setIsSavedQueriesPanelOpen] = useState<boolean>(false);
-  const [saveDialogState, setSaveDialogState] = useState<{ isOpen: boolean; data?: Partial<SavedQuery> & { prompt: string; code: string }}>({ isOpen: false });
+  const [saveDialogState, setSaveDialogState] = useState<{ isOpen: boolean; data?: Partial<SavedQuery> & { prompt: string; code: string } }>({ isOpen: false });
   const [shareDialogState, setShareDialogState] = useState<{ isOpen: boolean; query?: SavedQuery }>({ isOpen: false });
   const [isSavingQuery, setIsSavingQuery] = useState(false);
 
@@ -423,45 +419,45 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
     if (!connectedResource) return '';
     return azureAccounts.find(acc => acc.id === connectedResource.accountId)?.name ?? 'Unknown Account';
   }, [connectedResource, azureAccounts]);
-  
+
   const fetchAccounts = useCallback(async () => {
-      setIsLoadingAccounts(true);
-      setDbError(null);
-      try {
-        const accounts = await getAzureCosmosAccounts();
-        setAzureAccounts(accounts);
-      } catch (e) {
-        if (e instanceof Error) {
-            // Check for authentication-related errors
-            if (isAuthenticationExpiredError(e)) {
-                setDbError(getAuthErrorMessage(e));
-            } else if (e.message.includes('AuthorizationFailed') || e.message.includes('403')) {
-                setDbError("Permission Denied: You may not have the required permissions to list Azure resources. Please contact your administrator.");
-            } else {
-                setDbError("Could not load Azure accounts from server. Ensure the backend is running and you have permissions.");
-            }
+    setIsLoadingAccounts(true);
+    setDbError(null);
+    try {
+      const accounts = await getAzureCosmosAccounts();
+      setAzureAccounts(accounts);
+    } catch (e) {
+      if (e instanceof Error) {
+        // Check for authentication-related errors
+        if (isAuthenticationExpiredError(e)) {
+          setDbError(getAuthErrorMessage(e));
+        } else if (e.message.includes('AuthorizationFailed') || e.message.includes('403')) {
+          setDbError("Permission Denied: You may not have the required permissions to list Azure resources. Please contact your administrator.");
         } else {
-            setDbError("An unknown error occurred while fetching Azure accounts.");
+          setDbError("Could not load Azure accounts from server. Ensure the backend is running and you have permissions.");
         }
-      } finally {
-        setIsLoadingAccounts(false);
+      } else {
+        setDbError("An unknown error occurred while fetching Azure accounts.");
       }
-    }, []);
-  
+    } finally {
+      setIsLoadingAccounts(false);
+    }
+  }, []);
+
   const fetchSavedQueries = useCallback(async () => {
     setIsLoadingSavedQueries(true);
     try {
-        const queries = await getSavedQueries();
-        setSavedQueries(queries);
-    } catch(e) {
-        // Log error details for debugging
-        if (e instanceof Error && isAuthenticationExpiredError(e)) {
-            console.error("Failed to fetch saved queries due to authentication error:", getAuthErrorMessage(e));
-        } else {
-            console.error("Failed to fetch saved queries:", e);
-        }
+      const queries = await getSavedQueries();
+      setSavedQueries(queries);
+    } catch (e) {
+      // Log error details for debugging
+      if (e instanceof Error && isAuthenticationExpiredError(e)) {
+        console.error("Failed to fetch saved queries due to authentication error:", getAuthErrorMessage(e));
+      } else {
+        console.error("Failed to fetch saved queries:", e);
+      }
     } finally {
-        setIsLoadingSavedQueries(false);
+      setIsLoadingSavedQueries(false);
     }
   }, []);
 
@@ -479,29 +475,29 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
   useEffect(() => {
     // If the tutorial is active and on a step that targets an item inside the user menu, force it open.
     if (isTutorialActive && tutorialStepIndex === 6) {
-        setIsUserMenuOpenForTutorial(true);
+      setIsUserMenuOpenForTutorial(true);
     } else {
-        setIsUserMenuOpenForTutorial(false);
+      setIsUserMenuOpenForTutorial(false);
     }
   }, [isTutorialActive, tutorialStepIndex]);
-  
+
   const handleClearCache = useCallback(async () => {
-      setIsClearingCache(true);
-      setCacheClearStatus('idle');
-      setDbError(null); // Clear old DB errors
-      try {
-          await clearSystemCache();
-          setCacheClearStatus('success');
-          // Refresh accounts list after clearing cache
-          await fetchAccounts();
-      } catch (e) {
-          if (e instanceof Error) setDbError(e.message);
-          else setDbError("An unknown error occurred while clearing the cache.");
-      } finally {
-          setIsClearingCache(false);
-          // Reset the button state after 3 seconds
-          setTimeout(() => setCacheClearStatus('idle'), 3000);
-      }
+    setIsClearingCache(true);
+    setCacheClearStatus('idle');
+    setDbError(null); // Clear old DB errors
+    try {
+      await clearSystemCache();
+      setCacheClearStatus('success');
+      // Refresh accounts list after clearing cache
+      await fetchAccounts();
+    } catch (e) {
+      if (e instanceof Error) setDbError(e.message);
+      else setDbError("An unknown error occurred while clearing the cache.");
+    } finally {
+      setIsClearingCache(false);
+      // Reset the button state after 3 seconds
+      setTimeout(() => setCacheClearStatus('idle'), 3000);
+    }
   }, [fetchAccounts]);
 
   const clearQueryState = useCallback(() => {
@@ -512,30 +508,32 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
     setExecutionError(null);
     setDebuggingResult(null);
     setDebugError(null);
+    setAnalysisResult(null);
+    setAnalysisError(null);
     setCodeHistory([]);
     setHistoryIndex(-1);
     setIntermediateContext(null);
     setQuerySourceCollection(null);
-    setAnalysisResult(null);
-    setAnalysisError(null);
     setCurrentQueryContextSource(null);
+    setRelationships(null); // Clear relationships
+    setRelationshipError(null); // Clear relationship error
   }, []);
-  
+
   const handleDisconnect = useCallback(() => {
     setConnectedDbInfo(null);
     setConnectedResource(null);
     clearQueryState();
     setUserInput('');
-    setSelectedCollection(null);
-    setCollectionInfo(null);
+    setSelectedCollections([]);
+    setCollectionDetailsMap({});
   }, [clearQueryState]);
-  
+
   const handleSelectAccount = useCallback(async (accountId: string) => {
     if (selectedAccountId === accountId) {
-        // Deselect if clicking the same account again
-        setSelectedAccountId(null);
-        setAccountDatabases([]);
-        return;
+      // Deselect if clicking the same account again
+      setSelectedAccountId(null);
+      setAccountDatabases([]);
+      return;
     }
 
     setSelectedAccountId(accountId);
@@ -545,34 +543,34 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
 
     const account = azureAccounts.find(acc => acc.id === accountId);
     if (!account) {
-        setDbError("An error occurred: Could not find the selected account details.");
-        setIsLoadingDatabases(false);
-        return;
+      setDbError("An error occurred: Could not find the selected account details.");
+      setIsLoadingDatabases(false);
+      return;
     }
-    
+
     // If an old database was connected from another account, disconnect it
-    if(connectedResource?.accountId !== account.id) {
-        handleDisconnect();
+    if (connectedResource?.accountId !== account.id) {
+      handleDisconnect();
     }
-    
+
     try {
-        const dbs = await getDatabasesForAccount(account.id);
-        setAccountDatabases(dbs);
-    } catch(e) {
-        if (e instanceof Error) {
-            // Check for authentication-related errors first
-            if (isAuthenticationExpiredError(e)) {
-                setDbError(getAuthErrorMessage(e));
-            } else if (e.message.includes('AuthorizationFailed') || e.message.includes('403')) {
-                setDbError("Permission Denied: You may not have the required Azure role (e.g., 'Cosmos DB Operator') to access databases for this account. Please check your permissions.");
-            } else {
-                setDbError(e.message);
-            }
+      const dbs = await getDatabasesForAccount(account.id);
+      setAccountDatabases(dbs);
+    } catch (e) {
+      if (e instanceof Error) {
+        // Check for authentication-related errors first
+        if (isAuthenticationExpiredError(e)) {
+          setDbError(getAuthErrorMessage(e));
+        } else if (e.message.includes('AuthorizationFailed') || e.message.includes('403')) {
+          setDbError("Permission Denied: You may not have the required Azure role (e.g., 'Cosmos DB Operator') to access databases for this account. Please check your permissions.");
         } else {
-            setDbError("Could not load databases for this account.");
+          setDbError(e.message);
         }
+      } else {
+        setDbError("Could not load databases for this account.");
+      }
     } finally {
-        setIsLoadingDatabases(false);
+      setIsLoadingDatabases(false);
     }
   }, [selectedAccountId, connectedResource, azureAccounts, handleDisconnect]);
 
@@ -581,8 +579,8 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
     if (!account) return;
 
     setConnectedResource({
-        accountId: account.id,
-        databaseName: dbInfo.name,
+      accountId: account.id,
+      databaseName: dbInfo.name,
     });
     setConnectedDbInfo(dbInfo);
     clearQueryState();
@@ -605,7 +603,24 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
 
 
     try {
-      const result = await generateMongoQuery(prompt, connectedDbInfo ?? undefined, collectionCtx, intermediateContext?.data);
+      const accountId = connectedResource?.accountId || selectedAccountId;
+      if (!accountId) {
+        throw new Error("No account ID available for query generation.");
+      }
+
+      // 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
 
@@ -623,20 +638,22 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
     } finally {
       setIsLoading(false);
     }
-  }, [connectedDbInfo, codeHistory, historyIndex, intermediateContext]);
-  
+  }, [connectedDbInfo, codeHistory, historyIndex, intermediateContext, connectedResource, selectedAccountId]);
+
   const handleGenerateQueryClick = useCallback(() => {
     if (intermediateContext) {
-        setCurrentQueryContextSource(intermediateContext.source);
+      setCurrentQueryContextSource(intermediateContext.source);
     } else {
-        setCurrentQueryContextSource(null);
+      setCurrentQueryContextSource(null);
     }
-    setQuerySourceCollection(selectedCollection);
+    // For now, if multiple are selected, we might want to prioritize one for the "source" tag or change how it works.
+    // Let's us the first selected one, or a joined string.
+    setQuerySourceCollection(selectedCollections.length > 0 ? selectedCollections.join(', ') : null);
     setLastSuccessfulPrompt(userInput);
     // If a collection is selected, pass its info as context.
-    // Otherwise, this will be undefined, and the query will be against the whole DB.
-    handleGenerateQuery(userInput, collectionInfo ?? undefined);
-  }, [userInput, intermediateContext, selectedCollection, collectionInfo, handleGenerateQuery]);
+    const primaryContext = selectedCollections.length > 0 ? collectionDetailsMap[selectedCollections[0]] : undefined;
+    handleGenerateQuery(userInput, primaryContext);
+  }, [userInput, intermediateContext, selectedCollections, collectionDetailsMap, handleGenerateQuery]);
 
   const handleRunQuery = useCallback(async () => {
     if (!editableCode.trim() || !connectedDbInfo || !connectedResource) {
@@ -652,7 +669,7 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
     setAnalysisError(null);
 
     try {
-      const result = await runMongoQuery(selectedAccountId, editableCode, connectedResource);
+      const result = await runMongoQuery(connectedResource.accountId, editableCode, connectedResource);
       setExecutionResult(result);
 
       // --- Add step to notebook ---
@@ -663,25 +680,25 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
         prompt: lastSuccessfulPrompt || 'Query executed without a new prompt.',
         query: editableCode,
         resultSample: resultSample,
-        contextSource: currentQueryContextSource,
+        contextSource: currentQueryContextSource ?? undefined,
       };
       setNotebookSteps(prev => [...prev, newStep]);
       setCurrentQueryContextSource(null); // Reset after use
 
     } catch (e) {
-        if (e instanceof Error) {
-            if (e.message.includes('AuthorizationFailed') || e.message.includes('403')) {
-                setExecutionError("Permission Denied: You do not have permission to execute queries against this database.");
-            } else {
-                setExecutionError(e.message);
-            }
+      if (e instanceof Error) {
+        if (e.message.includes('AuthorizationFailed') || e.message.includes('403')) {
+          setExecutionError("Permission Denied: You do not have permission to execute queries against this database.");
         } else {
-            setExecutionError("An unknown error occurred while running the query.");
+          setExecutionError(e.message);
         }
+      } else {
+        setExecutionError("An unknown error occurred while running the query.");
+      }
     } finally {
       setIsExecuting(false);
     }
-  }, [editableCode, connectedDbInfo, connectedResource, selectedAccountId, lastSuccessfulPrompt, currentQueryContextSource]);
+  }, [editableCode, connectedDbInfo, connectedResource, lastSuccessfulPrompt, currentQueryContextSource]);
 
   const handleDebugQuery = useCallback(async () => {
     if (!editableCode || !executionError) return;
@@ -691,13 +708,13 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
     setDebuggingResult(null);
 
     try {
-        const result = await debugMongoQuery(editableCode, executionError);
-        setDebuggingResult(result);
+      const result = await debugMongoQuery(editableCode, executionError);
+      setDebuggingResult(result);
     } catch (e) {
-        if (e instanceof Error) setDebugError(e.message);
-        else setDebugError("An unexpected error occurred while debugging.");
+      if (e instanceof Error) setDebugError(e.message);
+      else setDebugError("An unexpected error occurred while debugging.");
     } finally {
-        setIsDebugging(false);
+      setIsDebugging(false);
     }
   }, [editableCode, executionError]);
 
@@ -709,44 +726,117 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
     setAnalysisResult(null);
 
     try {
-        const result = await analyzeQueryResult(dataToAnalyze);
-        setAnalysisResult(result);
+      const result = await analyzeQueryResult(dataToAnalyze);
+      setAnalysisResult(result);
     } catch (e) {
-        if (e instanceof Error) setAnalysisError(e.message);
-        else setAnalysisError("An unexpected error occurred during AI analysis.");
+      if (e instanceof Error) setAnalysisError(e.message);
+      else setAnalysisError("An unexpected error occurred during AI analysis.");
     } finally {
-        setIsAnalyzing(false);
+      setIsAnalyzing(false);
     }
   }, []);
 
-  const handleCollectionClick = useCallback(async (collectionName: string) => {
+  const handleCollectionClick = useCallback(async (collectionName: string, event?: React.MouseEvent) => {
     if (!connectedResource) return;
-    if (selectedCollection === collectionName) {
-        setSelectedCollection(null);
-        setCollectionInfo(null);
-        return;
+
+    // Clear previous relationship analysis when selection changes
+    setRelationships(null);
+    setRelationshipError(null);
+
+    const isModifierPressed = event ? (event.ctrlKey || event.metaKey) : false;
+
+    let newSelection: string[] = [];
+
+    if (isModifierPressed) {
+      // Multi-selection logic
+      if (selectedCollections.includes(collectionName)) {
+        newSelection = selectedCollections.filter(c => c !== collectionName);
+      } else {
+        newSelection = [...selectedCollections, collectionName];
+      }
+    } else {
+      // Single selection functionality (toggle off if same clicked, otherwise set new)
+      if (selectedCollections.length === 1 && selectedCollections[0] === collectionName) {
+        newSelection = [];
+      } else {
+        newSelection = [collectionName];
+      }
     }
-    setSelectedCollection(collectionName);
-    setIsFetchingCollectionInfo(true);
-    setCollectionInfoError(null);
-    setCollectionInfo(null);
-    try {
-        const info = await getCollectionInfo(collectionName, connectedResource);
-        setCollectionInfo(info);
-    } catch (e) {
-        if (e instanceof Error) {
-            if (e.message.includes('AuthorizationFailed') || e.message.includes('403')) {
-                setCollectionInfoError("Permission Denied: You do not have permission to view details for this collection.");
-            } else {
-                setCollectionInfoError(e.message);
-            }
-        } else {
-            setCollectionInfoError("Failed to fetch collection details.");
+
+    setSelectedCollections(newSelection);
+
+    if (newSelection.length === 0) {
+      // Clear errors if any
+      return;
+    }
+
+    // Fetch details for any newly selected collection that we don't have yet
+    newSelection.forEach(async (col) => {
+      if (!collectionDetailsMap[col] && !loadingCollections[col]) {
+        setLoadingCollections(prev => ({ ...prev, [col]: true }));
+        try {
+          // Use the existing getCollectionInfo. 
+          // Note: getCollectionInfo signature might be (accountId, dbName, collectionName) or (collectionName, resource).
+          // I need to check the import or usage elsewhere. 
+          // Line 4 import says: import { ... getCollectionInfo ... } from '../services/dbService';
+          // Line 764 usage (in previous view) was: getCollectionInfo(newSelection[0], connectedResource);
+          // But typically it is (accountId, databaseName, collectionName).
+          // Let's assume (collectionName, connectedResource) based on previous usage or check dbService.
+          // Actually in Step 382 I wrote: getCollectionInfo(connectedResource!.accountId, connectedResource!.databaseName, col);
+          // Let's stick to what works. I will use the pattern that was there or verify signature.
+          // In Step 374, import is from dbService.
+          // In Step 382, I tried to use explicit args.
+          // Let's assume signature: getCollectionInfo(collectionName, connectedResource) based on line 764 of previous original code?
+          // Wait, original code was: `const info = await getCollectionInfo(newSelection[0], connectedResource);` in Step 382 diff (LEFT side).
+          // So I should use that signature.
+          const info = await getCollectionInfo(col, connectedResource);
+          setCollectionDetailsMap(prev => ({ ...prev, [col]: info }));
+        } catch (e: any) {
+          console.error(`Failed to load details for ${col}`, e);
+        } finally {
+          setLoadingCollections(prev => ({ ...prev, [col]: false }));
         }
+      }
+    });
+  }, [connectedResource, selectedCollections, collectionDetailsMap, loadingCollections]);
+
+  const handleAnalyzeRelationships = useCallback(async () => {
+    if (!connectedResource || selectedCollections.length < 2) return;
+
+    setIsAnalyzingRelationships(true);
+    setRelationshipError(null);
+    // Don't clear immediately to prevent flashing if we are just refreshing? 
+    // Actually for new selection we want to clear.
+
+    try {
+      const result = await inferSchemaRelationships(
+        connectedResource.accountId,
+        connectedResource.databaseName,
+        selectedCollections
+      );
+      setRelationships(result);
+    } catch (e: any) {
+      setRelationshipError(e.message || "Failed to analyze relationships.");
     } finally {
-        setIsFetchingCollectionInfo(false);
+      setIsAnalyzingRelationships(false);
     }
-  }, [selectedCollection, connectedResource]);
+  }, [connectedResource, selectedCollections]);
+
+  // Debounced effect to trigger analysis when selections change
+  useEffect(() => {
+    // Clean up previous state if selection is insufficient
+    if (selectedCollections.length < 2) {
+      setRelationships(null);
+      setRelationshipError(null);
+      return;
+    }
+
+    const timer = setTimeout(() => {
+      handleAnalyzeRelationships();
+    }, 800); // 800ms debounce
+
+    return () => clearTimeout(timer);
+  }, [selectedCollections, handleAnalyzeRelationships]);
 
   const handleNavigateHistory = useCallback((direction: 'prev' | 'next') => {
     const newIndex = direction === 'prev' ? historyIndex - 1 : historyIndex + 1;
@@ -755,23 +845,23 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
       setEditableCode(codeHistory[newIndex]);
     }
   }, [historyIndex, codeHistory]);
-  
+
   const handleSetIntermediateContext = useCallback((data: any, source: string) => {
-      setIntermediateContext({ data, source });
+    setIntermediateContext({ data, source });
   }, []);
 
   // --- Notebook Handlers ---
   const handleExportNotebook = useCallback(() => {
     if (notebookSteps.length === 0) return;
     // Turn off editing mode for all notes before exporting
-    const stepsToExport = notebookSteps.map(step => ({...step, isEditing: false }));
+    const stepsToExport = notebookSteps.map(step => ({ ...step, isEditing: false }));
     const dbName = connectedDbInfo?.name;
     const content = generateIpynbContent(stepsToExport, dbName);
     downloadFile(content, 'querypal-notebook.ipynb', 'application/json');
   }, [notebookSteps, connectedDbInfo]);
 
   const handleClearNotebook = () => {
-      setNotebookSteps([]);
+    setNotebookSteps([]);
   };
 
   const handleRemoveNotebookStep = useCallback((id: string) => {
@@ -786,17 +876,17 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
       isEditing: true, // Start in editing mode
     };
     // Turn off editing for all other notes
-    setNotebookSteps(prev => [...prev.map(s => ({...s, isEditing: false})), newNote]);
+    setNotebookSteps(prev => [...prev.map(s => ({ ...s, isEditing: false })), newNote]);
   }, []);
 
   const handleUpdateNotebookStep = (id: string, content: string) => {
     setNotebookSteps(prev => prev.map(s => s.id === id ? { ...s, prompt: content } : s));
   };
-  
+
   const handleSetEditingStep = (id: string, isEditing: boolean) => {
     setNotebookSteps(prev => prev.map(s => s.id === id ? { ...s, isEditing } : { ...s, isEditing: false }));
   };
-  
+
   // --- Saved Queries Handlers ---
   const handleOpenSaveDialog = useCallback(() => {
     if (!editableCode) return; // Don't open if there's no code to save
@@ -809,25 +899,25 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
   const handleEditSavedQuery = (query: SavedQuery) => {
     setSaveDialogState({ isOpen: true, data: query });
   };
-  
+
   const handleSaveOrUpdateQuery = useCallback(async (data: Pick<SavedQuery, 'name' | 'prompt' | 'code'> | SavedQuery) => {
     setIsSavingQuery(true);
     try {
-        if ('id' in data) {
-            // Update
-            const updated = await updateSavedQuery(data as SavedQuery);
-            setSavedQueries(prev => prev.map(q => q.id === updated.id ? updated : q));
-        } else {
-            // Create
-            const newQuery = await saveQuery(data as Pick<SavedQuery, 'name' | 'prompt' | 'code'>);
-            setSavedQueries(prev => [...prev, newQuery]);
-        }
-        setSaveDialogState({ isOpen: false });
-    } catch(e) {
-        console.error("Failed to save or update query:", e);
-        // Maybe set an error in the dialog in the future
+      if ('id' in data) {
+        // Update
+        const updated = await updateSavedQuery(data as SavedQuery);
+        setSavedQueries(prev => prev.map(q => q.id === updated.id ? updated : q));
+      } else {
+        // Create
+        const newQuery = await saveQuery(data as Pick<SavedQuery, 'name' | 'prompt' | 'code'>);
+        setSavedQueries(prev => [...prev, newQuery]);
+      }
+      setSaveDialogState({ isOpen: false });
+    } catch (e) {
+      console.error("Failed to save or update query:", e);
+      // Maybe set an error in the dialog in the future
     } finally {
-        setIsSavingQuery(false);
+      setIsSavingQuery(false);
     }
   }, []);
 
@@ -836,10 +926,10 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
     const originalQueries = savedQueries;
     setSavedQueries(prev => prev.filter(q => q.id !== queryId));
     try {
-        await deleteSavedQuery(queryId);
-    } catch(e) {
-        console.error("Failed to delete query:", e);
-        setSavedQueries(originalQueries); // Revert on failure
+      await deleteSavedQuery(queryId);
+    } catch (e) {
+      console.error("Failed to delete query:", e);
+      setSavedQueries(originalQueries); // Revert on failure
     }
   }, [savedQueries]);
 
@@ -848,11 +938,11 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
     setUserInput(query.prompt);
     setLastSuccessfulPrompt(query.prompt);
     setEditableCode(query.code);
-    
+
     // Set code history for the loaded query
     setCodeHistory([query.code]);
     setHistoryIndex(0);
-    
+
     setIsSavedQueriesPanelOpen(false); // Close panel after loading
   };
 
@@ -865,28 +955,55 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
     // Optimistic update for a snappy UI
     setSavedQueries(prev => prev.map(q => q.id === queryToUpdate.id ? queryToUpdate : q));
     setShareDialogState({ isOpen: false }); // Close dialog immediately
-    
+
     try {
-        // Now call the backend
-        await updateSavedQuery(queryToUpdate);
+      // Now call the backend
+      await updateSavedQuery(queryToUpdate);
     } catch (e) {
-        console.error("Failed to update sharing settings:", e);
-        // On failure, revert the change by re-fetching from the source of truth
-        fetchSavedQueries(); 
+      console.error("Failed to update sharing settings:", e);
+      // On failure, revert the change by re-fetching from the source of truth
+      fetchSavedQueries();
     }
   };
 
-  const handleNavigateToExplorer = () => {
-    if (connectedResource && connectedDbInfo) {
-      onNavigateToExplorer({
-        resource: connectedResource,
-        dbInfo: connectedDbInfo,
-        accountName: connectedAccountName,
-        availableDbs: accountDatabases,
-        availableAccounts: azureAccounts,
-      });
+  const handleLaunchExplorer = useCallback((targetDb: DbInfo, targetAccount: CosmosDBAccount, explicitDbs?: DbInfo[]) => {
+    // Determine the list of available DBs.
+    // If explicitDbs is provided (e.g. from quick explore fetch), use it.
+    // Else if we are launching for the currently selected account (in the list), use the state accountDatabases.
+    // Otherwise fallback to just the target DB (which limits switching capabilities).
+    const dbsForExplorer = explicitDbs ||
+      ((selectedAccountId === targetAccount.id && accountDatabases.length > 0)
+        ? accountDatabases
+        : [targetDb]);
+
+    onNavigateToExplorer({
+      resource: { accountId: targetAccount.id, databaseName: targetDb.name },
+      dbInfo: targetDb,
+      accountName: targetAccount.name,
+      availableDbs: dbsForExplorer,
+      availableAccounts: azureAccounts,
+    });
+  }, [onNavigateToExplorer, selectedAccountId, accountDatabases, azureAccounts]);
+
+  const handleQuickExploreAccount = useCallback(async (account: CosmosDBAccount) => {
+    // If we already have the DBs for this account in state (because it's selected)
+    if (selectedAccountId === account.id && accountDatabases.length > 0) {
+      handleLaunchExplorer(accountDatabases[0], account, accountDatabases);
+      return;
     }
-  };
+
+    try {
+      const dbs = await getDatabasesForAccount(account.id);
+      if (dbs.length > 0) {
+        handleLaunchExplorer(dbs[0], account, dbs);
+      } else {
+        setError(`No databases found in '${account.name}'. Cannot launch Data Explorer.`);
+      }
+    } catch (e) {
+      console.error("Quick explorer failed", e);
+      setError(`Failed to load databases for '${account.name}'.`);
+    }
+  }, [selectedAccountId, accountDatabases, handleLaunchExplorer]);
 
   // --- Tutorial Demo Mode Logic ---
   const isDemoModeForCollectionStep = isTutorialActive && tutorialStepIndex === 2;
@@ -898,7 +1015,7 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
   const isDemoModeForContextActiveStep = isTutorialActive && tutorialStepIndex === 12;
   const isDemoModeForNotebookButtonStep = isTutorialActive && tutorialStepIndex === 13;
   const isDemoModeForNotebookPanelStep = isTutorialActive && tutorialStepIndex === 14;
-  
+
   const demoActive = isDemoModeForCollectionStep || isDemoModeForResultsStep || isDemoModeForContextActiveStep || isDemoModeForDebugStep || isDemoModeForNotebookButtonStep || isDemoModeForNotebookPanelStep || isDemoModeForRunStep || isDemoModeForSaveStep || isDemoModeForSavedQueriesPanelStep;
 
   const isConnectedForRender = (connectedDbInfo && connectedResource) || demoActive;
@@ -907,62 +1024,72 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
   // --- Global Keyboard Shortcuts ---
   useEffect(() => {
     const handleKeyDown = (event: KeyboardEvent) => {
-        const isMac = navigator.platform.toUpperCase().indexOf('MAC') >= 0;
-        const isModifier = isMac ? event.metaKey : event.ctrlKey;
-
-        // Escape key to close modals/panels (topmost)
-        if (event.key === 'Escape') {
-            if (saveDialogState.isOpen) setSaveDialogState({ isOpen: false });
-            else if (shareDialogState.isOpen) setShareDialogState({ isOpen: false });
-            else if (isShortcutCheatsheetOpen) setIsShortcutCheatsheetOpen(false);
-            else if (isNotebookPanelOpen) setIsNotebookPanelOpen(false);
-            else if (isSavedQueriesPanelOpen) setIsSavedQueriesPanelOpen(false);
-            else if (isContextViewerOpen) setIsContextViewerOpen(false);
-            else if (isTutorialActive) setIsTutorialActive(false); // Also close tutorial
-            return;
+      const isMac = navigator.platform.toUpperCase().indexOf('MAC') >= 0;
+      const isModifier = isMac ? event.metaKey : event.ctrlKey;
+
+      // Escape key to close modals/panels (topmost)
+      if (event.key === 'Escape') {
+        if (saveDialogState.isOpen) setSaveDialogState({ isOpen: false });
+        else if (shareDialogState.isOpen) setShareDialogState({ isOpen: false });
+        else if (isShortcutCheatsheetOpen) setIsShortcutCheatsheetOpen(false);
+        else if (isNotebookPanelOpen) setIsNotebookPanelOpen(false);
+        else if (isSavedQueriesPanelOpen) setIsSavedQueriesPanelOpen(false);
+        else if (isContextViewerOpen) setIsContextViewerOpen(false);
+        else if (isTutorialActive) setIsTutorialActive(false); // Also close tutorial
+        return;
+      }
+
+      // Prevent shortcuts from firing while typing in dialogs/modals
+      const isTypingInDialog = (event.target as HTMLElement).closest('[role="dialog"]');
+      if (isTypingInDialog) return;
+
+      if (isModifier) {
+        // Cmd/Ctrl + S to Save
+        if (event.key === 's') {
+          event.preventDefault();
+          if (editableCode && !isQuerySectionDisabled) {
+            handleOpenSaveDialog();
+          }
         }
-        
-        // Prevent shortcuts from firing while typing in dialogs/modals
-        const isTypingInDialog = (event.target as HTMLElement).closest('[role="dialog"]');
-        if (isTypingInDialog) return;
-
-        if (isModifier) {
-            // Cmd/Ctrl + S to Save
-            if (event.key === 's') {
-                event.preventDefault();
-                if (editableCode && !isQuerySectionDisabled) {
-                    handleOpenSaveDialog();
-                }
-            }
-
-            // Cmd/Ctrl + / to open shortcuts
-            if (event.key === '/') {
-                event.preventDefault();
-                setIsShortcutCheatsheetOpen(true);
-            }
+
+        // Cmd/Ctrl + / to open shortcuts
+        if (event.key === '/') {
+          event.preventDefault();
+          setIsShortcutCheatsheetOpen(true);
         }
+      }
     };
 
     window.addEventListener('keydown', handleKeyDown);
     return () => {
-        window.removeEventListener('keydown', handleKeyDown);
+      window.removeEventListener('keydown', handleKeyDown);
     };
   }, [editableCode, isQuerySectionDisabled, handleOpenSaveDialog, saveDialogState.isOpen, shareDialogState.isOpen, isShortcutCheatsheetOpen, isNotebookPanelOpen, isSavedQueriesPanelOpen, isContextViewerOpen, isTutorialActive]);
 
   const dbInfoForRender = demoActive ? mockECommerceDbInfo : connectedDbInfo;
   const accountNameForRender = demoActive ? 'prod-ecommerce-db' : connectedAccountName;
 
-  const selectedCollectionForRender = isDemoModeForCollectionStep ? 'users' : selectedCollection;
-  const collectionInfoForRender = isDemoModeForCollectionStep ? mockCollectionInfoMap.get('users')! : collectionInfo;
-  const showCollectionPanel = isDemoModeForCollectionStep || isFetchingCollectionInfo || (collectionInfo && selectedCollection === collectionInfo.name) || collectionInfoError;
-  
+  const selectedCollectionsForRender = isDemoModeForCollectionStep ? ['users'] : selectedCollections;
+
+  // For demo/tutorial purposes, ensure appropriate data is in the map if simulated
+  const collectionDetailsMapForRender = useMemo(() => {
+    if (isDemoModeForCollectionStep) {
+      return { 'users': mockCollectionInfoMap.get('users')! };
+    }
+    return collectionDetailsMap;
+  }, [isDemoModeForCollectionStep, collectionDetailsMap]);
+
+  const isPromptUnchanged = userInput.trim() === lastSuccessfulPrompt.trim() && !!editableCode;
+
   const generateButtonText = useMemo(() => {
     if (isLoading) return 'Generating...';
-    if (selectedCollection) {
-        return `Generate Query for ${selectedCollection} collection`;
+    if (isPromptUnchanged) return 'Query Generated';
+    if (selectedCollections.length > 0) {
+      if (selectedCollections.length === 1) return `Generate Query for ${selectedCollections[0]} collection`;
+      return `Generate Query across ${selectedCollections.length} collections`;
     }
     return 'Generate Query';
-  }, [isLoading, selectedCollection]);
+  }, [isLoading, selectedCollections, isPromptUnchanged]);
 
   // --- Demo Mode Context Banner ---
   const showContextBanner = intermediateContext || isDemoModeForContextActiveStep;
@@ -997,12 +1124,12 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
     </>,
     document.body
   ) : null;
-  
+
   // --- Demo Mode Notebook Panel ---
   const showNotebookPanel = isNotebookPanelOpen || isDemoModeForNotebookPanelStep;
   const demoNotebookSteps: NotebookStep[] = [
-      { id: 'demo-1', type: 'query', prompt: 'Find all users from Canada', query: "db['users'].find({'country': 'Canada'})", resultSample: mockUserFindResult.slice(0, 1) },
-      { id: 'demo-2', type: 'query', prompt: 'From those users, find the ones named Alice', query: "db['users'].find({'name': 'Alice'})", resultSample: mockUserFindResult.slice(0, 1), contextSource: "'users' collection" },
+    { id: 'demo-1', type: 'query', prompt: 'Find all users from Canada', query: "db['users'].find({'country': 'Canada'})", resultSample: mockUserFindResult.slice(0, 1) },
+    { id: 'demo-2', type: 'query', prompt: 'From those users, find the ones named Alice', query: "db['users'].find({'name': 'Alice'})", resultSample: mockUserFindResult.slice(0, 1), contextSource: "'users' collection" },
   ];
 
   const notebookPanelDrawer = showNotebookPanel ? createPortal(
@@ -1022,64 +1149,62 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
   const showSavedQueriesPanel = isSavedQueriesPanelOpen || isDemoModeForSavedQueriesPanelStep;
   const savedQueriesPanelDrawer = showSavedQueriesPanel ? createPortal(
     <SavedQueriesPanel
-        onClose={() => setIsSavedQueriesPanelOpen(false)}
-        queries={isDemoModeForSavedQueriesPanelStep ? mockSavedQueries : savedQueries}
-        onLoad={handleLoadSavedQuery}
-        onEdit={handleEditSavedQuery}
-        onDelete={handleDeleteSavedQuery}
-        onShare={handleOpenShareDialog}
-        isLoading={isDemoModeForSavedQueriesPanelStep ? false : isLoadingSavedQueries}
-        currentUserEmail={email || 'dev.user@example.com'}
+      onClose={() => setIsSavedQueriesPanelOpen(false)}
+      queries={isDemoModeForSavedQueriesPanelStep ? mockSavedQueries : savedQueries}
+      onLoad={handleLoadSavedQuery}
+      onEdit={handleEditSavedQuery}
+      onDelete={handleDeleteSavedQuery}
+      onShare={handleOpenShareDialog}
+      isLoading={isDemoModeForSavedQueriesPanelStep ? false : isLoadingSavedQueries}
+      currentUserEmail={email || 'dev.user@example.com'}
     />,
     document.body
   ) : null;
 
   const saveQueryDialog = saveDialogState.isOpen ? createPortal(
     <SaveQueryDialog
-        isOpen={saveDialogState.isOpen}
-        onClose={() => setSaveDialogState({ isOpen: false })}
-        onSave={handleSaveOrUpdateQuery}
-        isSaving={isSavingQuery}
-        initialData={saveDialogState.data!}
+      isOpen={saveDialogState.isOpen}
+      onClose={() => setSaveDialogState({ isOpen: false })}
+      onSave={handleSaveOrUpdateQuery}
+      isSaving={isSavingQuery}
+      initialData={saveDialogState.data!}
     />,
     document.body
   ) : null;
 
   const shareQueryDialog = shareDialogState.isOpen && shareDialogState.query ? createPortal(
     <ShareQueryDialog
-        isOpen={shareDialogState.isOpen}
-        onClose={() => setShareDialogState({ isOpen: false })}
-        onSave={handleUpdateSharing}
-        query={shareDialogState.query}
+      isOpen={shareDialogState.isOpen}
+      onClose={() => setShareDialogState({ isOpen: false })}
+      onSave={handleUpdateSharing}
+      query={shareDialogState.query}
     />,
     document.body
   ) : null;
 
   const shortcutCheatsheet = createPortal(
-      <ShortcutCheatsheet 
-          isOpen={isShortcutCheatsheetOpen}
-          onClose={() => setIsShortcutCheatsheetOpen(false)}
-      />,
-      document.body
+    <ShortcutCheatsheet
+      isOpen={isShortcutCheatsheetOpen}
+      onClose={() => setIsShortcutCheatsheetOpen(false)}
+    />,
+    document.body
   );
 
 
   return (
     <div className="min-h-screen bg-slate-50 dark:bg-slate-900 text-slate-800 dark:text-slate-200 font-sans p-4 sm:p-6 lg:p-8">
       <div className="max-w-4xl mx-auto">
-        
-        <HeaderUI 
-            name={name}
-            onLogout={onLogout}
-            onClearCache={handleClearCache}
-            isClearingCache={isClearingCache}
-            cacheClearStatus={cacheClearStatus}
-            onStartTutorial={() => setIsTutorialActive(true)}
-            onShowSavedQueries={() => setIsSavedQueriesPanelOpen(true)}
-            onShowShortcuts={() => setIsShortcutCheatsheetOpen(true)}
-            onNavigateToExplorer={handleNavigateToExplorer}
-            isExplorerNavEnabled={!!connectedResource}
-            isUserMenuForcedOpen={isUserMenuOpenForTutorial}
+
+        <HeaderUI
+          name={name}
+          onLogout={onLogout}
+          onClearCache={handleClearCache}
+          isClearingCache={isClearingCache}
+          cacheClearStatus={cacheClearStatus}
+          onStartTutorial={() => setIsTutorialActive(true)}
+          onShowSavedQueries={() => setIsSavedQueriesPanelOpen(true)}
+          onShowShortcuts={() => setIsShortcutCheatsheetOpen(true)}
+          isUserMenuForcedOpen={isUserMenuOpenForTutorial}
         />
 
         <main className="space-y-8">
@@ -1094,13 +1219,23 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
                       Connected to: {accountNameForRender} / <span className="font-bold">{dbInfoForRender.name}</span>
                     </p>
                   </div>
-                  <button
-                    onClick={handleDisconnect}
-                    className="px-3 py-1.5 border border-red-300 text-sm font-medium rounded-md text-red-600 hover:bg-red-50 hover:text-red-700 dark:border-red-500/50 dark:text-red-400 dark:hover:bg-red-900/40 transition-colors"
-                    title="Disconnect from database"
-                  >
-                    Disconnect
-                  </button>
+                  <div className="flex items-center gap-2">
+                    <button
+                      onClick={() => connectedDbInfo && connectedResource && handleLaunchExplorer(connectedDbInfo, azureAccounts.find(a => a.id === connectedResource.accountId) || azureAccounts[0])}
+                      className="px-3 py-1.5 border border-slate-300 dark:border-slate-600 text-sm font-medium rounded-md text-slate-700 dark:text-slate-300 hover:bg-slate-100 dark:hover:bg-slate-700 transition-colors flex items-center gap-2"
+                      title="Open in Data Explorer"
+                    >
+                      <DataGridIcon className="w-4 h-4" />
+                      Explorer
+                    </button>
+                    <button
+                      onClick={handleDisconnect}
+                      className="px-3 py-1.5 border border-red-300 text-sm font-medium rounded-md text-red-600 hover:bg-red-50 hover:text-red-700 dark:border-red-500/50 dark:text-red-400 dark:hover:bg-red-900/40 transition-colors"
+                      title="Disconnect from database"
+                    >
+                      Disconnect
+                    </button>
+                  </div>
                 </div>
                 <div className="mt-4 grid grid-cols-1 md:grid-cols-3 gap-4 text-sm">
                   <div className="bg-slate-100 dark:bg-slate-700/50 p-3 rounded-lg">
@@ -1112,95 +1247,186 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
                     <p className="text-slate-900 dark:text-slate-100 font-semibold text-lg">{dbInfoForRender.size ?? 'N/A'}</p>
                   </div>
                   <div className="bg-slate-100 dark:bg-slate-700/50 p-3 rounded-lg col-span-1 md:col-span-3">
-                     <p className="text-slate-600 dark:text-slate-300 mb-2 flex items-center gap-2 font-medium"><ServerIcon className="w-4 h-4" /> Collections</p>
-                     <div className="flex flex-wrap gap-2">
-                       {dbInfoForRender.collections.map(col => (
-                         <button 
-                            key={col.name} 
-                            onClick={() => handleCollectionClick(col.name)}
-                            className={`text-xs font-mono px-3 py-1 rounded-full transition-all duration-200 ${selectedCollectionForRender === col.name ? 'bg-blue-500 text-white font-bold ring-2 ring-blue-300 dark:bg-blue-600 dark:ring-blue-500' : 'bg-slate-200 text-slate-700 hover:bg-slate-300 dark:bg-slate-700 dark:text-slate-200 dark:hover:bg-slate-600'}`}
-                            title={`View details for the ${col.name} collection`}
+                    <p className="text-slate-600 dark:text-slate-300 mb-2 flex items-center gap-2 font-medium">
+                      <ServerIcon className="w-4 h-4" /> Collections
+                      <span className="text-xs font-normal text-slate-400 dark:text-slate-500">(Hold Ctrl/Cmd to select multiple)</span>
+                    </p>
+                    <div className="flex flex-wrap gap-2">
+                      {dbInfoForRender.collections.map(col => (
+                        <button
+                          key={col.name}
+                          onClick={(e) => handleCollectionClick(col.name, e)}
+                          className={`text-xs font-mono px-3 py-1 rounded-full transition-all duration-200 ${selectedCollectionsForRender.includes(col.name) ? 'bg-blue-500 text-white font-bold ring-2 ring-blue-300 dark:bg-blue-600 dark:ring-blue-500' : 'bg-slate-200 text-slate-700 hover:bg-slate-300 dark:bg-slate-700 dark:text-slate-200 dark:hover:bg-slate-600'}`}
+                          title={`View details for the ${col.name} collection`}
                         >
-                            {col.name}
-                         </button>
-                       ))}
-                     </div>
+                          {col.name}
+                        </button>
+                      ))}
+                    </div>
                   </div>
                 </div>
-                 {/* Collection Action Panel */}
-                 {showCollectionPanel && (
-                    <div id="tutorial-collection-panel" className="mt-4">
-                        {isFetchingCollectionInfo && !isDemoModeForCollectionStep && <div className="text-center p-4 text-slate-500 dark:text-slate-400">Fetching collection details...</div>}
-                        {collectionInfoError && !isDemoModeForCollectionStep && (
-                            <div className="text-red-600 bg-red-50 border border-red-200 text-sm mt-4 p-3 rounded-md dark:bg-red-900/30 dark:border-red-500/50 dark:text-red-300">
-                                {collectionInfoError}
-                            </div>
-                        )}
-                        {collectionInfoForRender && selectedCollectionForRender === collectionInfoForRender.name && (
-                            <CollectionActionPanel
-                                info={collectionInfoForRender}
-                                onClose={() => { setSelectedCollection(null); setCollectionInfo(null); }}
-                            />
-                        )}
+
+                {/* Multi-Collection Analysis Panel */}
+                {selectedCollectionsForRender.length > 1 && (
+                  <div className="mt-6 bg-slate-50 dark:bg-slate-700/30 rounded-lg p-4 border border-slate-200 dark:border-slate-700 animate-fade-in">
+                    <div className="flex items-center justify-between mb-4">
+                      <h3 className="text-md font-bold text-slate-800 dark:text-slate-200 flex items-center gap-2">
+                        <span className="text-blue-500">🔗</span> Schema Connections
+                      </h3>
+                      <div className="bg-blue-50 dark:bg-blue-900/20 text-blue-600 dark:text-blue-300 text-xs px-2 py-1 rounded border border-blue-100 dark:border-blue-800/30 flex items-center gap-1">
+                        <span>✨ AI Inferred</span>
+                      </div>
                     </div>
-                 )}
+
+                    <p className="text-xs text-slate-500 dark:text-slate-400 mb-4 italic">
+                      Note: These connections are inferred by AI based on the schema and may not reflect strict foreign key constraints.
+                    </p>
+
+                    {(isAnalyzingRelationships || !relationships) && !relationshipError && (
+                      <div className="flex items-center justify-center p-6 text-slate-500 dark:text-slate-400">
+                        <div className="animate-spin rounded-full h-5 w-5 border-b-2 border-blue-500 mr-2"></div>
+                        Searching for connections...
+                      </div>
+                    )}
+
+                    {relationshipError && (
+                      <div className="text-red-600 bg-red-50 border border-red-200 text-sm p-3 rounded-md dark:bg-red-900/30 dark:border-red-500/50 dark:text-red-300">
+                        {relationshipError}
+                      </div>
+                    )}
+
+                    {relationships && !isAnalyzingRelationships && (
+                      <div className="animate-fade-in w-full flex justify-center">
+                        <SchemaRelationshipGraph
+                          relationships={relationships}
+                          selectedCollections={selectedCollections}
+                        />
+                      </div>
+                    )}
+                  </div>
+                )}
+
+                {/* Collection Action Panel (Single or Stacked) */}
+                {selectedCollectionsForRender.length > 0 && (
+                  <div id="tutorial-collection-panel" className="mt-4 space-y-2">
+                    {/* Error Display */}
+                    {collectionInfoError && (
+                      <div className="text-red-600 bg-red-50 border border-red-200 text-sm p-3 rounded-md dark:bg-red-900/30 dark:border-red-500/50 dark:text-red-300">
+                        {collectionInfoError}
+                      </div>
+                    )}
+
+                    {selectedCollectionsForRender.map(colName => {
+                      const info = collectionDetailsMapForRender[colName];
+                      const isLoading = loadingCollections[colName];
+                      const isExpanded = expandedCollectionSchemas[colName] || (selectedCollectionsForRender.length === 1); // Default expanded if single
+
+                      return (
+                        <div key={colName} className="bg-slate-50 dark:bg-slate-800/30 rounded-lg border border-slate-200 dark:border-slate-700 overflow-hidden">
+                          <button
+                            onClick={() => setExpandedCollectionSchemas(prev => ({ ...prev, [colName]: !prev[colName] }))}
+                            className="w-full flex items-center justify-between p-3 text-left hover:bg-slate-100 dark:hover:bg-slate-700/50 transition-colors"
+                          >
+                            <div className="flex items-center gap-2">
+                              <ChevronDownIcon className={`w-5 h-5 text-slate-400 transition-transform ${isExpanded ? 'rotate-180' : ''}`} />
+                              <span className="font-bold text-slate-700 dark:text-slate-200">{colName}</span>
+                              {isLoading && <SpinnerIcon className="w-4 h-4 animate-spin text-blue-500" />}
+                            </div>
+                            <div className="flex items-center gap-2 text-xs text-slate-500">
+                              {info && <span>{info.documentCount.toLocaleString()} docs</span>}
+                            </div>
+                          </button>
+
+                          {isExpanded && (
+                            <div className="p-3 border-t border-slate-200 dark:border-slate-700 bg-white dark:bg-slate-800">
+                              {isLoading ? (
+                                <div className="py-4 text-center text-slate-500 text-sm">Loading details...</div>
+                              ) : info ? (
+                                <CollectionActionPanel
+                                  info={info}
+                                  onClose={() => {
+                                    // Deselect this collection
+                                    setSelectedCollections(prev => prev.filter(c => c !== colName));
+                                  }}
+                                />
+                              ) : (
+                                <div className="py-4 text-center text-red-500 text-sm">Failed to load details.</div>
+                              )}
+                            </div>
+                          )}
+                        </div>
+                      );
+                    })}
+                  </div>
+                )}
               </div>
             ) : (
               <div>
-                <h2 className="text-xl font-bold text-slate-900 dark:text-slate-100 flex items-center gap-2"><DatabaseIcon className="w-6 h-6 text-blue-500"/> Select a Database to Connect</h2>
-                 {dbError && !isLoadingAccounts && !isLoadingDatabases && (
-                     <p className="text-red-600 bg-red-50 border border-red-200 text-sm mt-4 p-3 rounded-md dark:bg-red-900/30 dark:border-red-500/50 dark:text-red-300">{dbError}</p>
-                 )}
+                <h2 className="text-xl font-bold text-slate-900 dark:text-slate-100 flex items-center gap-2"><DatabaseIcon className="w-6 h-6 text-blue-500" /> Select a Database to Connect</h2>
+                {dbError && !isLoadingAccounts && !isLoadingDatabases && (
+                  <p className="text-red-600 bg-red-50 border border-red-200 text-sm mt-4 p-3 rounded-md dark:bg-red-900/30 dark:border-red-500/50 dark:text-red-300">{dbError}</p>
+                )}
                 {isLoadingAccounts ? (
-                    <div className="text-center p-8 text-slate-500 dark:text-slate-400">Loading your Azure accounts...</div>
+                  <div className="text-center p-8 text-slate-500 dark:text-slate-400">Loading your Azure accounts...</div>
                 ) : !dbError && azureAccounts.length === 0 ? (
-                    <div className="text-center p-8 text-slate-500 dark:text-slate-400 border-2 border-dashed border-slate-300 dark:border-slate-600 rounded-lg mt-4">
-                        No accessible Cosmos DB accounts found.
-                        <br/>
-                        <span className="text-xs">Ensure your account has Reader permissions on the resources.</span>
-                    </div>
+                  <div className="text-center p-8 text-slate-500 dark:text-slate-400 border-2 border-dashed border-slate-300 dark:border-slate-600 rounded-lg mt-4">
+                    No accessible Cosmos DB accounts found.
+                    <br />
+                    <span className="text-xs">Ensure your account has Reader permissions on the resources.</span>
+                  </div>
                 ) : (
-                    <div className="mt-4 space-y-4">
-                        {azureAccounts.map(account => (
-                            <div key={account.id} className="bg-slate-50 dark:bg-slate-800/50 p-4 rounded-lg border dark:border-slate-700">
-                               <button 
-                                    onClick={() => handleSelectAccount(account.id)}
-                                    disabled={isLoadingDatabases}
-                                    className="w-full text-left font-bold text-slate-800 dark:text-slate-100 flex items-center gap-2 disabled:cursor-not-allowed disabled:opacity-60"
-                                    title={`Load databases for ${account.name}`}
-                                >
-                                    <CloudIcon className="w-5 h-5 text-slate-500" />
-                                    {account.name}
-                               </button>
-                               {selectedAccountId === account.id && (
-                                   <div className="mt-3 pl-7 animate-fade-in">
-                                        {isLoadingDatabases ? (
-                                             <div className="text-sm text-slate-500 dark:text-slate-400 py-2">Loading databases...</div>
-                                        ): dbError ? (
-                                             <div className="text-red-600 bg-red-50 border border-red-200 text-sm mt-2 p-3 rounded-md dark:bg-red-900/30 dark:border-red-500/50 dark:text-red-300">
-                                                {dbError}
-                                             </div>
-                                        ) : accountDatabases.length > 0 ? (
-                                            <div className="flex flex-wrap gap-2">
-                                                {accountDatabases.map(db => (
-                                                    <button
-                                                        key={db.name}
-                                                        onClick={() => handleConnectDatabase(db)}
-                                                        className="px-4 py-2 border border-transparent text-sm font-medium rounded-md text-white bg-blue-600 hover:bg-blue-700 focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-offset-slate-50 dark:focus:ring-offset-slate-800 focus:ring-blue-500 transition-colors"
-                                                        title={`Connect to the ${db.name} database`}
-                                                    >
-                                                        {db.name}
-                                                    </button>
-                                                ))}
-                                            </div>
-                                        ) : (
-                                            <p className="text-xs text-slate-500 dark:text-slate-400">No databases found in this account.</p>
-                                        )}
-                                   </div>
-                               )}
-                            </div>
-                        ))}
-                    </div>
+                  <div className="mt-4 space-y-4">
+                    {azureAccounts.map(account => (
+                      <div key={account.id} className="bg-slate-50 dark:bg-slate-800/50 p-4 rounded-lg border dark:border-slate-700">
+                        <div className="flex items-center justify-between gap-2">
+                          <button
+                            onClick={() => handleSelectAccount(account.id)}
+                            disabled={isLoadingDatabases}
+                            className="flex-grow text-left font-bold text-slate-800 dark:text-slate-100 flex items-center gap-2 disabled:cursor-not-allowed disabled:opacity-60 hover:text-blue-600 dark:hover:text-blue-400 transition-colors"
+                            title={`Load databases for ${account.name}`}
+                          >
+                            <CloudIcon className="w-5 h-5 text-slate-500" />
+                            {account.name}
+                          </button>
+                          <button
+                            onClick={(e) => { e.stopPropagation(); handleQuickExploreAccount(account); }}
+                            className="flex items-center gap-1.5 px-2 py-1.5 text-xs font-medium text-blue-600 dark:text-blue-400 bg-white dark:bg-slate-700 hover:bg-blue-50 dark:hover:bg-slate-600 border border-slate-200 dark:border-slate-600 rounded-md transition-colors shadow-sm ml-2"
+                            title={`Quickly open Data Explorer for ${account.name} (first database)`}
+                          >
+                            <DataGridIcon className="w-4 h-4" />
+                            <span className="hidden sm:inline">Explorer</span>
+                          </button>
+                        </div>
+                        {selectedAccountId === account.id && (
+                          <div className="mt-3 pl-7 animate-fade-in">
+                            {isLoadingDatabases ? (
+                              <div className="text-sm text-slate-500 dark:text-slate-400 py-2">Loading databases...</div>
+                            ) : dbError ? (
+                              <div className="text-red-600 bg-red-50 border border-red-200 text-sm mt-2 p-3 rounded-md dark:bg-red-900/30 dark:border-red-500/50 dark:text-red-300">
+                                {dbError}
+                              </div>
+                            ) : accountDatabases.length > 0 ? (
+                              <div className="flex flex-wrap gap-2">
+                                {accountDatabases.map(db => (
+                                  <button
+                                    key={db.name}
+                                    type="button"
+                                    onClick={() => handleConnectDatabase(db)}
+                                    className="px-4 py-2 border border-blue-600 dark:border-blue-500 text-sm font-medium rounded-md text-white bg-blue-600 hover:bg-blue-700 dark:bg-blue-600 dark:hover:bg-blue-500 focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-offset-slate-50 dark:focus:ring-offset-slate-800 focus:ring-blue-500 transition-colors"
+                                    title={`Connect to the ${db.name} database`}
+                                  >
+                                    {db.name}
+                                  </button>
+                                ))}
+                              </div>
+                            ) : (
+                              <p className="text-xs text-slate-500 dark:text-slate-400">No databases found in this account.</p>
+                            )}
+                          </div>
+                        )}
+                      </div>
+                    ))}
+                  </div>
                 )}
               </div>
             )}
@@ -1209,25 +1435,25 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
           {/* Query Generator */}
           <div id="tutorial-prompt-section" className={`bg-white dark:bg-slate-800 rounded-xl shadow-md p-6 transition-opacity duration-500 ${isQuerySectionDisabled ? 'opacity-40 pointer-events-none' : 'opacity-100'}`}>
             {showContextBanner && contextForRender && (
-                <div id="tutorial-context-banner" className="relative bg-blue-50 dark:bg-slate-800/60 border border-blue-200 dark:border-blue-500/30 rounded-lg p-4 mb-6 animate-fade-in">
-                    <div className="flex items-start justify-between">
-                        <div className="flex items-center gap-3">
-                             <PinIcon className="w-6 h-6 text-blue-500 flex-shrink-0" />
-                            <div>
-                                <h4 className="font-bold text-blue-800 dark:text-blue-300">Query Context Active</h4>
-                                <p className="text-sm text-blue-700 dark:text-blue-200/80">
-                                    Using results from <strong className="font-mono">{contextForRender.source}</strong> ({Array.isArray(contextForRender.data) ? contextForRender.data.length : 1} items) as context for the next query.
-                                </p>
-                                <button onClick={() => setIsContextViewerOpen(true)} className="text-sm text-blue-600 dark:text-blue-400 hover:underline font-semibold mt-1" title="View the full data being used as context">View Data</button>
-                            </div>
-                        </div>
-                        <button onClick={() => setIntermediateContext(null)} className="p-1.5 rounded-full text-blue-600 dark:text-blue-400 hover:bg-blue-200/50 dark:hover:bg-blue-900/40" title="Remove the active query context">
-                            <XIcon className="w-4 h-4" />
-                        </button>
+              <div id="tutorial-context-banner" className="relative bg-blue-50 dark:bg-slate-800/60 border border-blue-200 dark:border-blue-500/30 rounded-lg p-4 mb-6 animate-fade-in">
+                <div className="flex items-start justify-between">
+                  <div className="flex items-center gap-3">
+                    <PinIcon className="w-6 h-6 text-blue-500 flex-shrink-0" />
+                    <div>
+                      <h4 className="font-bold text-blue-800 dark:text-blue-300">Query Context Active</h4>
+                      <p className="text-sm text-blue-700 dark:text-blue-200/80">
+                        Using results from <strong className="font-mono">{contextForRender.source}</strong> ({Array.isArray(contextForRender.data) ? contextForRender.data.length : 1} items) as context for the next query.
+                      </p>
+                      <button onClick={() => setIsContextViewerOpen(true)} className="text-sm text-blue-600 dark:text-blue-400 hover:underline font-semibold mt-1" title="View the full data being used as context">View Data</button>
                     </div>
+                  </div>
+                  <button onClick={() => setIntermediateContext(null)} className="p-1.5 rounded-full text-blue-600 dark:text-blue-400 hover:bg-blue-200/50 dark:hover:bg-blue-900/40" title="Remove the active query context">
+                    <XIcon className="w-4 h-4" />
+                  </button>
                 </div>
+              </div>
             )}
-            
+
             <div className="space-y-4">
               <label htmlFor="userInput" className="block text-lg font-medium text-slate-700 dark:text-slate-300">
                 Enter your command:
@@ -1239,48 +1465,48 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
                 onKeyDown={(e) => {
                   if ((e.metaKey || e.ctrlKey) && e.key === 'g') {
                     e.preventDefault();
-                    if (!isLoading && userInput.trim() && !isQuerySectionDisabled) {
+                    if (!isLoading && userInput.trim() && !isQuerySectionDisabled && !isPromptUnchanged) {
                       handleGenerateQueryClick();
                     }
                   }
                 }}
-                placeholder={isQuerySectionDisabled ? "Connect to a database to begin..." : (selectedCollection ? `Querying '${selectedCollection}'... e.g., 'Find all users from Canada'` : "e.g., 'Find all users from Canada and sort them by name'")}
+                placeholder={isQuerySectionDisabled ? "Connect to a database to begin..." : (selectedCollections.length > 0 ? `Querying ${selectedCollections.length > 1 ? `${selectedCollections.length} collections` : `'${selectedCollections[0]}'`}... e.g., 'Find all users from Canada'` : "e.g., 'Find all users from Canada and sort them by name'")}
                 className="w-full h-28 p-4 bg-slate-50 dark:bg-slate-700/50 border border-slate-300 dark:border-slate-600 rounded-lg text-slate-900 dark:text-slate-100 focus:ring-2 focus:ring-blue-500 focus:border-blue-500 transition-colors duration-200 placeholder-slate-400 dark:placeholder-slate-500 resize-none"
                 disabled={isLoading || isQuerySectionDisabled}
               />
               <button
                 onClick={handleGenerateQueryClick}
-                disabled={isLoading || !userInput.trim() || isQuerySectionDisabled}
+                disabled={isLoading || !userInput.trim() || isQuerySectionDisabled || isPromptUnchanged}
                 className="w-full flex justify-center items-center px-6 py-3 border border-transparent text-base font-medium rounded-md text-white bg-blue-600 hover:bg-blue-700 disabled:bg-slate-300 dark:disabled:bg-slate-600 disabled:text-slate-500 dark:disabled:text-slate-400 disabled:cursor-not-allowed focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-blue-500 transition-all duration-200 transform hover:scale-[1.02] active:scale-[0.99]"
-                title="Generate MongoDB code from your natural language command (⌘ + G)"
+                title={isPromptUnchanged ? "The current query matches the prompt. Modify the prompt to generate a new query." : "Generate MongoDB code from your natural language command (⌘ + G)"}
               >
                 {generateButtonText}
               </button>
             </div>
 
             <div id="tutorial-results-area" className="mt-8">
-               <div className={`flex flex-col sm:flex-row items-start sm:items-center justify-between gap-3 mb-4 ${isQuerySectionDisabled ? 'hidden' : ''}`}>
-                    <h3 className="text-lg font-medium text-slate-700 dark:text-slate-300">
-                        Query Output
-                    </h3>
-                    <button
-                        id="tutorial-notebook-button"
-                        onClick={() => setIsNotebookPanelOpen(true)}
-                        className="flex items-center gap-2 px-4 py-2 border border-slate-300 dark:border-slate-600 text-sm font-medium rounded-md text-slate-600 dark:text-slate-300 bg-white dark:bg-slate-800 hover:bg-slate-100 dark:hover:bg-slate-700 hover:text-slate-800 dark:hover:text-slate-100 focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-offset-white dark:focus:ring-offset-slate-800 focus:ring-blue-500 transition-colors"
-                        title="View your session as a reproducible Jupyter Notebook"
-                    >
-                        <NotebookIcon className="w-5 h-5" />
-                        <span>View Notebook</span>
-                    </button>
-                </div>
+              <div className={`flex flex-col sm:flex-row items-start sm:items-center justify-between gap-3 mb-4 ${isQuerySectionDisabled ? 'hidden' : ''}`}>
+                <h3 className="text-lg font-medium text-slate-700 dark:text-slate-300">
+                  Query Output
+                </h3>
+                <button
+                  id="tutorial-notebook-button"
+                  onClick={() => setIsNotebookPanelOpen(true)}
+                  className="flex items-center gap-2 px-4 py-2 border border-slate-300 dark:border-slate-600 text-sm font-medium rounded-md text-slate-600 dark:text-slate-300 bg-white dark:bg-slate-800 hover:bg-slate-100 dark:hover:bg-slate-700 hover:text-slate-800 dark:hover:text-slate-100 focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-offset-white dark:focus:ring-offset-slate-800 focus:ring-blue-500 transition-colors"
+                  title="View your session as a reproducible Jupyter Notebook"
+                >
+                  <NotebookIcon className="w-5 h-5" />
+                  <span>View Notebook</span>
+                </button>
+              </div>
 
               {isLoading && !isDemoModeForDebugStep && !isDemoModeForResultsStep && !isDemoModeForRunStep && <Loader message="Generating query..." />}
-              
+
               {error && !isDemoModeForDebugStep && !isDemoModeForResultsStep && !isDemoModeForRunStep && (
-                  <div className="bg-red-100 border border-red-400 text-red-700 px-4 py-3 rounded-lg animate-fade-in dark:bg-red-900/50 dark:border-red-500/50 dark:text-red-300" role="alert">
-                      <strong className="font-bold">Error: </strong>
-                      <span className="block sm:inline">{error}</span>
-                  </div>
+                <div className="bg-red-100 border border-red-400 text-red-700 px-4 py-3 rounded-lg animate-fade-in dark:bg-red-900/50 dark:border-red-500/50 dark:text-red-300" role="alert">
+                  <strong className="font-bold">Error: </strong>
+                  <span className="block sm:inline">{error}</span>
+                </div>
               )}
 
               {/* Tutorial Demo for Run, Edit & Save Steps */}
@@ -1288,13 +1514,13 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
                 <div className="space-y-8 animate-fade-in">
                   <QueryDisplay
                     code={mockFindUsersQuery.generated_code}
-                    onCodeChange={() => {}}
-                    onRunQuery={() => {}}
-                    onSaveQuery={() => {}}
+                    onCodeChange={() => { }}
+                    onRunQuery={() => { }}
+                    onSaveQuery={() => { }}
                     isExecuting={false}
                     historyCount={1}
                     historyIndex={0}
-                    onNavigateHistory={() => {}}
+                    onNavigateHistory={() => { }}
                   />
                 </div>
               )}
@@ -1304,26 +1530,26 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
                 <div className="space-y-8 animate-fade-in">
                   <QueryDisplay
                     code={mockFindUsersQuery.generated_code}
-                    onCodeChange={() => {}}
-                    onRunQuery={() => {}}
-                    onSaveQuery={() => {}}
+                    onCodeChange={() => { }}
+                    onRunQuery={() => { }}
+                    onSaveQuery={() => { }}
                     isExecuting={false}
                     historyCount={1}
                     historyIndex={0}
-                    onNavigateHistory={() => {}}
+                    onNavigateHistory={() => { }}
                   />
                   <QueryResult
                     isExecuting={false}
                     executionError={null}
                     executionResult={mockUserFindResult}
-                    onDebug={() => {}}
+                    onDebug={() => { }}
                     isDebugging={false}
                     debuggingResult={null}
                     debugError={null}
                     sourceCollection={'users'}
-                    onSetIntermediateContext={() => {}}
+                    onSetIntermediateContext={() => { }}
                     intermediateContext={null}
-                    onAnalyze={() => {}}
+                    onAnalyze={() => { }}
                     isAnalyzing={false}
                     analysisResult={null}
                     analysisError={null}
@@ -1332,32 +1558,32 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
                   />
                 </div>
               )}
-              
+
               {/* Tutorial Demo for Debug View */}
               {isDemoModeForDebugStep && (
                 <div className="space-y-8 animate-fade-in">
                   <QueryDisplay
                     code={"db['users'].find({}).sor([('name', 1)])"}
-                    onCodeChange={() => {}}
-                    onRunQuery={() => {}}
-                    onSaveQuery={() => {}}
+                    onCodeChange={() => { }}
+                    onRunQuery={() => { }}
+                    onSaveQuery={() => { }}
                     isExecuting={false}
                     historyCount={1}
                     historyIndex={0}
-                    onNavigateHistory={() => {}}
+                    onNavigateHistory={() => { }}
                   />
                   <QueryResult
                     isExecuting={false}
                     executionError={"MongoDB query error: unknown operator: $sor (MongoServerError)"}
                     executionResult={null}
-                    onDebug={() => {}}
+                    onDebug={() => { }}
                     isDebugging={false}
                     debuggingResult={null}
                     debugError={null}
                     sourceCollection={'users'}
-                    onSetIntermediateContext={() => {}}
+                    onSetIntermediateContext={() => { }}
                     intermediateContext={null}
-                    onAnalyze={() => {}}
+                    onAnalyze={() => { }}
                     isAnalyzing={false}
                     analysisResult={null}
                     analysisError={null}
@@ -1370,46 +1596,46 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
               {/* Real results */}
               {(!isLoading && !error && !isDemoModeForResultsStep && !isDemoModeForDebugStep && !isDemoModeForContextActiveStep && !isDemoModeForRunStep && !isDemoModeForSaveStep) && (
                 <div className="space-y-8">
-                    {editableCode ? (
-                        <>
-                            <QueryDisplay
-                                code={editableCode}
-                                onCodeChange={setEditableCode}
-                                onRunQuery={handleRunQuery}
-                                onSaveQuery={handleOpenSaveDialog}
-                                isExecuting={isExecuting}
-                                historyCount={codeHistory.length}
-                                historyIndex={historyIndex}
-                                onNavigateHistory={handleNavigateHistory}
-                            />
-                            <QueryResult
-                                isExecuting={isExecuting}
-                                executionError={executionError}
-                                executionResult={executionResult}
-                                onDebug={handleDebugQuery}
-                                isDebugging={isDebugging}
-                                debuggingResult={debuggingResult}
-                                debugError={debugError}
-                                sourceCollection={querySourceCollection}
-                                onSetIntermediateContext={handleSetIntermediateContext}
-                                intermediateContext={intermediateContext}
-                                onAnalyze={handleAnalyzeQuery}
-                                isAnalyzing={isAnalyzing}
-                                analysisResult={analysisResult}
-                                analysisError={analysisError}
-                            />
-                        </>
-                    ) : (
-                        <div className="text-center text-slate-500 dark:text-slate-400 py-10 border-2 border-dashed border-slate-300 dark:border-slate-600 rounded-lg">
-                            <p>{isQuerySectionDisabled ? 'Connect to a database to generate queries.' : 'Your generated query will appear here.'}</p>
-                        </div>
-                    )}
+                  {editableCode ? (
+                    <>
+                      <QueryDisplay
+                        code={editableCode}
+                        onCodeChange={setEditableCode}
+                        onRunQuery={handleRunQuery}
+                        onSaveQuery={handleOpenSaveDialog}
+                        isExecuting={isExecuting}
+                        historyCount={codeHistory.length}
+                        historyIndex={historyIndex}
+                        onNavigateHistory={handleNavigateHistory}
+                      />
+                      <QueryResult
+                        isExecuting={isExecuting}
+                        executionError={executionError}
+                        executionResult={executionResult}
+                        onDebug={handleDebugQuery}
+                        isDebugging={isDebugging}
+                        debuggingResult={debuggingResult}
+                        debugError={debugError}
+                        sourceCollection={querySourceCollection}
+                        onSetIntermediateContext={handleSetIntermediateContext}
+                        intermediateContext={intermediateContext}
+                        onAnalyze={handleAnalyzeQuery}
+                        isAnalyzing={isAnalyzing}
+                        analysisResult={analysisResult}
+                        analysisError={analysisError}
+                      />
+                    </>
+                  ) : (
+                    <div className="text-center text-slate-500 dark:text-slate-400 py-10 border-2 border-dashed border-slate-300 dark:border-slate-600 rounded-lg">
+                      <p>{isQuerySectionDisabled ? 'Connect to a database to generate queries.' : 'Your generated query will appear here.'}</p>
+                    </div>
+                  )}
                 </div>
               )}
             </div>
           </div>
         </main>
-        
+
         <footer className="text-center mt-8 text-slate-500 dark:text-slate-400 text-sm space-y-1">
           <p>Powered by Microsoft Azure and Google Gemini. For internal use only.</p>
           <p className="text-xs">
@@ -1417,15 +1643,15 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
           </p>
         </footer>
       </div>
-      
+
       <Tutorial
         isActive={isTutorialActive}
         onStepChange={setTutorialStepIndex}
         onClose={() => {
-            setIsTutorialActive(false);
-            localStorage.setItem('hasSeenTutorial', 'true');
+          setIsTutorialActive(false);
+          localStorage.setItem('hasSeenTutorial', 'true');
         }}
-       />
+      />
 
       {contextViewerDrawer}
       {notebookPanelDrawer}
@@ -1434,7 +1660,7 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
       {shareQueryDialog}
       {shortcutCheatsheet}
 
-       <style>{`
+      <style>{`
           @keyframes fade-in {
             from { opacity: 0; transform: translateY(-10px); }
             to { opacity: 1; transform: translateY(0); }
diff --git a/frontend/services/geminiService.ts b/frontend/services/geminiService.ts
index 9f6a1e9..0ca8c9b 100644
--- a/frontend/services/geminiService.ts
+++ b/frontend/services/geminiService.ts
@@ -1,12 +1,14 @@
-
-import { QueryResultData, DbInfo, CollectionInfo, DebuggingResult, AnalysisResult } from '../types';
+import { QueryResultData, DbInfo, CollectionInfo, DebuggingResult, AnalysisResult, SchemaRelationshipsResponse } from '../types';
 import { USE_MSAL_AUTH, API_BASE_URL } from '../app.config';
 import { mockDelay, mockFindUsersQuery, mockUpdateProductsQuery, mockDefaultQuery, mockDebuggingResult, mockAnalysisResult } from './mockData';
+import { msalInstance, loginRequest } from '../authConfig';
+import { getAuthErrorMessage, isAuthenticationExpiredError } from '../utils/authErrorHandler';
 
 /**
  * Sends the user's natural language prompt to the backend for processing by the Gemini API.
  * The backend is responsible for securely calling the AI model and returning the structured result.
  * @param userInput The natural language query from the user.
+ * @param accountId The Azure Cosmos DB account ID to fetch schema context for.
  * @param dbInfo Optional information about the connected database to provide context to the AI.
  * @param collectionContext Optional information about a specific collection to provide even more detailed context.
  * @param intermediateContext Optional data from a previous query result to be used as context.
@@ -14,15 +16,17 @@ import { mockDelay, mockFindUsersQuery, mockUpdateProductsQuery, mockDefaultQuer
  */
 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) {
         console.log("DEV MODE: Returning mock AI-generated query.");
         await mockDelay(1500); // Simulate AI thinking time
-        
+
         const lowerInput = userInput.toLowerCase();
         if (lowerInput.includes('user')) {
             return Promise.resolve(mockFindUsersQuery);
@@ -33,18 +37,41 @@ export const generateMongoQuery = async (
         return Promise.resolve(mockDefaultQuery);
     }
     // --- END DEVELOPMENT MOCK ---
-    
+
     console.log("Sending prompt to backend for query generation:", userInput);
 
+    // Acquire Token
+    let accessToken = "";
+    try {
+        const account = msalInstance.getAllAccounts()[0]; // Assume first account
+        if (!account) {
+            throw new Error("No active account found. Please sign in.");
+        }
+        const response = await msalInstance.acquireTokenSilent({
+            ...loginRequest,
+            account: account
+        });
+        accessToken = response.accessToken;
+    } catch (error) {
+        console.error("Failed to acquire token silently:", getAuthErrorMessage(error));
+        throw new Error("Authentication failed. Please refresh the page and sign in again.");
+    }
+
     const response = await fetch(`${API_BASE_URL}/query/nl2query`, {
         method: 'POST',
         headers: {
             'Content-Type': 'application/json',
+            'Authorization': `Bearer ${accessToken}`
         },
         body: JSON.stringify({
             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
         }),
     });
@@ -73,7 +100,7 @@ export const debugMongoQuery = async (query: string, errorMessage: string): Prom
         return Promise.resolve(mockDebuggingResult);
     }
     // --- END DEVELOPMENT MOCK ---
-    
+
     console.log("Sending failed query to backend for debugging...");
 
     const response = await fetch(`${API_BASE_URL}/query/debug`, {
@@ -116,7 +143,7 @@ export const analyzeQueryResult = async (queryResult: any): Promise<AnalysisResu
         throw new Error("No mock analysis available for this data. Please implement a new mock in services/mockData.ts");
     }
     // --- END DEVELOPMENT MOCK ---
-    
+
     console.log("Sending query result to backend for analysis...");
 
     const response = await fetch(`${API_BASE_URL}/query/analyze`, {
@@ -133,4 +160,78 @@ export const analyzeQueryResult = async (queryResult: any): Promise<AnalysisResu
 
     const result: AnalysisResult = await response.json();
     return result;
+};
+
+
+/**
+ * Infers relationships between selected collections using AI.
+ * @param accountId The Azure Cosmos DB account ID.
+ * @param databaseName The name of the database.
+ * @param collectionNames The list of collections to analyze.
+ * @returns A promise that resolves with the inferred relationships.
+ */
+export const inferSchemaRelationships = async (
+    accountId: string,
+    databaseName: string,
+    collectionNames: string[]
+): Promise<SchemaRelationshipsResponse> => {
+    // --- DEVELOPMENT MOCK ---
+    if (!USE_MSAL_AUTH) {
+        console.log("DEV MODE: Returning mock AI relationship inference.");
+        await mockDelay(2000);
+        return Promise.resolve({
+            relationships: [
+                {
+                    source_collection: collectionNames[0] || "orders",
+                    source_field: "userId",
+                    target_collection: collectionNames[1] || "users",
+                    target_field: "_id",
+                    description: "Inferred foreign key relationship based on field name similarity.",
+                    confidence: 0.95
+                }
+            ]
+        });
+    }
+    // --- END DEVELOPMENT MOCK ---
+
+    console.log("Sending schema inference request to backend...");
+
+    // Acquire Token
+    let accessToken = "";
+    try {
+        const account = msalInstance.getAllAccounts()[0];
+        if (!account) {
+            throw new Error("No active account found. Please sign in.");
+        }
+        const response = await msalInstance.acquireTokenSilent({
+            ...loginRequest,
+            account: account
+        });
+        accessToken = response.accessToken;
+    } catch (error) {
+        console.error("Failed to acquire token silently:", getAuthErrorMessage(error));
+        throw new Error("Authentication failed. Please refresh the page and sign in again.");
+    }
+
+    const response = await fetch(`${API_BASE_URL}/query/infer-relationships`, {
+        method: 'POST',
+        headers: {
+            'Content-Type': 'application/json',
+            'Authorization': `Bearer ${accessToken}`
+        },
+        body: JSON.stringify({
+            account_id: accountId,
+            database_name: databaseName,
+            collection_names: collectionNames
+        }),
+    });
+
+    if (!response.ok) {
+        const errorData = await response.json().catch(() => ({}));
+        const errorMessage = errorData.detail || errorData.message || 'The AI model failed to infer relationships.';
+        throw new Error(errorMessage);
+    }
+
+    const result: SchemaRelationshipsResponse = await response.json();
+    return result;
 };
\ No newline at end of file
diff --git a/frontend/types.ts b/frontend/types.ts
index 30bd9ec..c89fad6 100644
--- a/frontend/types.ts
+++ b/frontend/types.ts
@@ -17,7 +17,7 @@ export interface DbConfig {
 }
 
 export interface QueryResultData {
-  generated_code: string;
+    generated_code: string;
 }
 
 /**
@@ -47,10 +47,10 @@ export type SelectedResource = {
  * Represents the detailed information of a successfully connected database.
  */
 export interface DbInfo {
-  name:string;
-  collections: CollectionSummary[];
-  totalDocuments: number;
-  size: string | null;
+    name: string;
+    collections: CollectionSummary[];
+    totalDocuments: number;
+    size: string | null;
 }
 
 /**
@@ -75,7 +75,7 @@ export interface ChartJSDataset {
     borderColor?: string | string[];
     borderWidth?: number;
     // Allow other chart.js dataset properties
-    [key: string]: any; 
+    [key: string]: any;
 }
 
 export interface ChartJSData {
@@ -100,13 +100,13 @@ export interface AnalysisResult {
  * Represents a single step in the user's query workflow for notebook generation.
  */
 export interface NotebookStep {
-  id: string;
-  type: 'query' | 'note'; // Add type to distinguish steps
-  prompt: string; // For query: NL prompt. For note: markdown content.
-  query?: string; // Query code is optional now (for notes)
-  resultSample?: any; // Result sample is optional (for notes)
-  contextSource?: string;
-  isEditing?: boolean; // Transient state for UI editing of notes
+    id: string;
+    type: 'query' | 'note'; // Add type to distinguish steps
+    prompt: string; // For query: NL prompt. For note: markdown content.
+    query?: string; // Query code is optional now (for notes)
+    resultSample?: any; // Result sample is optional (for notes)
+    contextSource?: string;
+    isEditing?: boolean; // Transient state for UI editing of notes
 }
 
 /**
@@ -170,4 +170,23 @@ export interface DocumentHistoryResponse {
     document_id: string;
     history_entries: DocumentHistoryEntry[];
     total_entries: number;
+}
+
+/**
+ * Represents an inferred relationship between two collections.
+ */
+export interface Relationship {
+    source_collection: string;
+    source_field: string;
+    target_collection: string;
+    target_field: string;
+    description: string;
+    confidence: number;
+}
+
+/**
+ * Represents the response from the schema relationship inference endpoint.
+ */
+export interface SchemaRelationshipsResponse {
+    relationships: Relationship[];
 }
\ No newline at end of file