From 67d2e063e021ab6ace61a14f683f510050b549e4 Mon Sep 17 00:00:00 2001
From: Amin Khorramii <khorramiamindev@gmail.com>
Date: Mon, 4 Aug 2025 12:19:50 +0200
Subject: [PATCH] python prompt notebook in progress

---
 .../components/tabs/notebooks/CellDivider.tsx |  14 +-
 .../tabs/notebooks/NotebooksWorkspace.tsx     |  35 +-
 .../components/tabs/notebooks/PromptCell.tsx  | 813 ++++++++++++++++++
 frontend/src/hooks/notebooks/useNotebookAI.ts | 274 ++++++
 .../src/hooks/notebooks/useNotebookContext.ts | 346 ++++++++
 frontend/src/lib/python/types.ts              |   8 +-
 6 files changed, 1479 insertions(+), 11 deletions(-)
 create mode 100644 frontend/src/components/tabs/notebooks/PromptCell.tsx
 create mode 100644 frontend/src/hooks/notebooks/useNotebookAI.ts
 create mode 100644 frontend/src/hooks/notebooks/useNotebookContext.ts

diff --git a/frontend/src/components/tabs/notebooks/CellDivider.tsx b/frontend/src/components/tabs/notebooks/CellDivider.tsx
index 581474d..3bb876b 100644
--- a/frontend/src/components/tabs/notebooks/CellDivider.tsx
+++ b/frontend/src/components/tabs/notebooks/CellDivider.tsx
@@ -1,5 +1,5 @@
 import React, { useState } from "react";
-import { Code2, Type } from "lucide-react";
+import { Code2, Type, Sparkles } from "lucide-react";
 import { usePythonStore } from "@/store/pythonStore";
 
 interface CellDividerProps {
@@ -14,7 +14,7 @@ const CellDivider: React.FC<CellDividerProps> = ({ insertIndex, isLastCell = fal
   const { createCell } = usePythonStore();
   const [isHovered, setIsHovered] = useState(false);
 
-  const handleCreateCell = (type: 'code' | 'markdown') => {
+  const handleCreateCell = (type: 'code' | 'markdown' | 'prompt') => {
     createCell(type, "", insertIndex);
   };
 
@@ -45,12 +45,20 @@ const CellDivider: React.FC<CellDividerProps> = ({ insertIndex, isLastCell = fal
         </button>
         <button
           onClick={() => handleCreateCell('markdown')}
-          className="flex items-center gap-2 px-3 py-2 text-xs text-white/70 hover:text-white hover:bg-white/10 transition-colors rounded-r-md"
+          className="flex items-center gap-2 px-3 py-2 text-xs text-white/70 hover:text-white hover:bg-white/10 transition-colors border-r border-white/10"
           title="Add Text Cell"
         >
           <Type size={12} />
           <span>Text</span>
         </button>
+        <button
+          onClick={() => handleCreateCell('prompt')}
+          className="flex items-center gap-2 px-3 py-2 text-xs text-purple-400/70 hover:text-purple-300 hover:bg-purple-500/10 transition-colors rounded-r-md"
+          title="Add AI Prompt Cell"
+        >
+          <Sparkles size={12} />
+          <span>Prompt</span>
+        </button>
       </div>
     </div>
   );
diff --git a/frontend/src/components/tabs/notebooks/NotebooksWorkspace.tsx b/frontend/src/components/tabs/notebooks/NotebooksWorkspace.tsx
index 5054999..13ff40d 100644
--- a/frontend/src/components/tabs/notebooks/NotebooksWorkspace.tsx
+++ b/frontend/src/components/tabs/notebooks/NotebooksWorkspace.tsx
@@ -16,6 +16,7 @@ import {
   Clock,
   AlertCircle,
   RefreshCw,
+  Sparkles,
 } from 'lucide-react';
 
 import { usePythonStore } from '@/store/pythonStore';
@@ -35,6 +36,7 @@ import {
 } from '@/utils/notebookExport';
 
 import PythonCell from './PythonCell';
+import PromptCell from './PromptCell';
 import CellDivider from './CellDivider';
 import ScriptHistory from './ScriptHistory';
 import PackageManager from './PackageManager';
@@ -617,6 +619,18 @@ const NotebooksWorkspace: React.FC = () => {
               </Button>
             </Tooltip>
 
+            <Tooltip content="Add New AI Prompt Cell" placement="bottom">
+              <Button
+                variant="ghost"
+                size="sm"
+                onClick={() => createCell('prompt')}
+                className="h-8 bg-gradient-to-r from-purple-500/10 to-pink-500/10 hover:from-purple-500/20 hover:to-pink-500/20 border border-purple-500/20"
+              >
+                <Sparkles size={14} className="mr-1 text-purple-400" />
+                <span>Prompt</span>
+              </Button>
+            </Tooltip>
+
             <Tooltip
               content="Execute All Cells (⌘+Shift+Enter)"
               placement="bottom"
@@ -696,12 +710,21 @@ const NotebooksWorkspace: React.FC = () => {
           <div className="flex-1 overflow-y-auto overflow-x-visible p-8 space-y-0">
             {cells.map((cell, index) => (
               <React.Fragment key={cell.id}>
-                <PythonCell
-                  cell={cell}
-                  isActive={cell.id === activeCellId}
-                  onActivate={() => setActiveCellId(cell.id)}
-                  cellNumber={index + 1}
-                />
+                {cell.type === 'prompt' ? (
+                  <PromptCell
+                    cell={cell}
+                    isActive={cell.id === activeCellId}
+                    onActivate={() => setActiveCellId(cell.id)}
+                    cellNumber={index + 1}
+                  />
+                ) : (
+                  <PythonCell
+                    cell={cell}
+                    isActive={cell.id === activeCellId}
+                    onActivate={() => setActiveCellId(cell.id)}
+                    cellNumber={index + 1}
+                  />
+                )}
                 {/* Cell Divider - always show after each cell */}
                 <CellDivider
                   insertIndex={index + 1}
diff --git a/frontend/src/components/tabs/notebooks/PromptCell.tsx b/frontend/src/components/tabs/notebooks/PromptCell.tsx
new file mode 100644
index 0000000..bd7e492
--- /dev/null
+++ b/frontend/src/components/tabs/notebooks/PromptCell.tsx
@@ -0,0 +1,813 @@
+import React, { useState, useRef, useCallback } from 'react';
+import {
+  Sparkles,
+  Play,
+  Square,
+  Trash2,
+  MoreVertical,
+  ChevronUp,
+  ChevronDown,
+  Copy,
+  Eye,
+  ChevronRight,
+  Wand2,
+  Brain,
+  Lock,
+  Check,
+  X,
+  Code2,
+  ArrowDown,
+  Zap,
+} from 'lucide-react';
+
+import { usePythonStore } from '@/store/pythonStore';
+import { useAuth } from '@/hooks/auth/useAuth';
+import { useNotebookContext } from '@/hooks/notebooks/useNotebookContext';
+import { useNotebookAI } from '@/hooks/notebooks/useNotebookAI';
+import { Button } from '@/components/ui/Button';
+import AuthModal from '@/components/auth/AuthModal';
+import MonacoEditor from '../query/MonacoEditor';
+import MonacoErrorBoundary from './MonacoErrorBoundary';
+import type { PythonCell as PythonCellType } from '@/lib/python/types';
+import Tooltip from '@/components/ui/Tooltip';
+
+interface PromptCellProps {
+  cell: PythonCellType;
+  isActive: boolean;
+  onActivate: () => void;
+  cellNumber: number;
+}
+
+const PromptCell: React.FC<PromptCellProps> = ({
+  cell,
+  isActive,
+  onActivate,
+  cellNumber,
+}) => {
+  const {
+    updateCell,
+    deleteCell,
+    moveCell,
+    toggleCellInputCollapse,
+    createCell,
+    executeCell,
+    cells,
+  } = usePythonStore();
+
+  // Helper function to update prompt cell properties
+  const updatePromptCell = (cellId: string, code: string, additionalProps?: Partial<PythonCellType>) => {
+    const state = usePythonStore.getState();
+    const updatedCells = state.cells.map((c) =>
+      c.id === cellId 
+        ? { ...c, code, updatedAt: new Date(), ...additionalProps }
+        : c
+    );
+    usePythonStore.setState({ cells: updatedCells });
+    // Mark as unsaved
+    state.markAsUnsaved();
+  };
+
+  // Helper function to split generated code into executable blocks
+  const splitCodeIntoBlocks = (code: string): string[] => {
+    if (!code) return [];
+    
+    // Split by comment separator or natural code blocks
+    const blocks = code.split(/\n\n# --- Next Code Block ---\n\n/);
+    
+    // Further split by logical breaks (imports, function definitions, etc.)
+    const splitBlocks: string[] = [];
+    
+    blocks.forEach(block => {
+      const lines = block.split('\n');
+      let currentBlock: string[] = [];
+      
+      for (let i = 0; i < lines.length; i++) {
+        const line = lines[i];
+        const trimmedLine = line.trim();
+        
+        // Check if this line should start a new block
+        const isNewBlockTrigger = 
+          (trimmedLine.startsWith('import ') || trimmedLine.startsWith('from ')) && currentBlock.length > 0 ||
+          (trimmedLine.startsWith('def ') || trimmedLine.startsWith('class ')) && currentBlock.length > 0 ||
+          (trimmedLine.startsWith('# ') && trimmedLine.length > 10) && currentBlock.length > 0;
+        
+        if (isNewBlockTrigger) {
+          // Save current block and start new one
+          if (currentBlock.length > 0) {
+            splitBlocks.push(currentBlock.join('\n').trim());
+            currentBlock = [];
+          }
+        }
+        
+        currentBlock.push(line);
+      }
+      
+      // Add remaining block
+      if (currentBlock.length > 0) {
+        splitBlocks.push(currentBlock.join('\n').trim());
+      }
+    });
+    
+    return splitBlocks.filter(block => block.trim().length > 0);
+  };
+
+  // Create executable code cells from generated code
+  const createCodeCells = () => {
+    if (!cell.generatedCode) return;
+    
+    const codeBlocks = splitCodeIntoBlocks(cell.generatedCode);
+    if (codeBlocks.length === 0) return;
+    
+    // Find the current cell index
+    const currentIndex = cells.findIndex(c => c.id === cell.id);
+    if (currentIndex === -1) return;
+    
+    // Create cells after the current prompt cell
+    codeBlocks.forEach((code, index) => {
+      const insertIndex = currentIndex + 1 + index;
+      createCell('code', code, insertIndex);
+    });
+  };
+
+  // Execute all generated code cells
+  const executeGeneratedCode = async () => {
+    if (!cell.generatedCode) return;
+    
+    // First create the code cells
+    createCodeCells();
+    
+    // Wait a bit for cells to be created, then execute them
+    setTimeout(async () => {
+      const currentIndex = cells.findIndex(c => c.id === cell.id);
+      if (currentIndex === -1) return;
+      
+      const codeBlocks = splitCodeIntoBlocks(cell.generatedCode || '');
+      
+      // Execute each created cell sequentially
+      for (let i = 0; i < codeBlocks.length; i++) {
+        const cellIndex = currentIndex + 1 + i;
+        const updatedCells = usePythonStore.getState().cells;
+        
+        if (cellIndex < updatedCells.length) {
+          const cellToExecute = updatedCells[cellIndex];
+          try {
+            await executeCell(cellToExecute.id);
+            // Small delay between executions
+            await new Promise(resolve => setTimeout(resolve, 500));
+          } catch (error) {
+            console.error(`Error executing cell ${i + 1}:`, error);
+            // Continue with next cell even if one fails
+          }
+        }
+      }
+    }, 100);
+  };
+
+  const { isAuthenticated } = useAuth();
+  const { 
+    context, 
+    getAIContext, 
+    refreshContext, 
+    toggleTableSelection, 
+    toggleVariableSelection,
+    toggleAllTables,
+    toggleAllVariables 
+  } = useNotebookContext();
+  
+  const { status: aiStatus, processAIRequest, resetStatus } = useNotebookAI();
+  const [showMenu, setShowMenu] = useState(false);
+  const [showAuthModal, setShowAuthModal] = useState(false);
+  const [showContext, setShowContext] = useState(false);
+  const menuRef = useRef<HTMLDivElement>(null);
+  const cellRef = useRef<HTMLDivElement>(null);
+  const textareaRef = useRef<HTMLTextAreaElement>(null);
+
+  // Auto-resize textarea
+  const adjustTextareaHeight = useCallback(() => {
+    const textarea = textareaRef.current;
+    if (textarea) {
+      textarea.style.height = 'auto';
+      textarea.style.height = `${Math.max(120, textarea.scrollHeight)}px`;
+    }
+  }, []);
+
+  const handleInputChange = (e: React.ChangeEvent<HTMLTextAreaElement>) => {
+    updateCell(cell.id, e.target.value);
+    adjustTextareaHeight();
+  };
+
+  const handleCodeChange = (value: string | undefined) => {
+    updateCell(cell.id, value || '');
+  };
+
+  // Parse AI response to show code blocks and explanations separately
+  const parseAIOutput = (response: string) => {
+    if (!response) return [];
+    
+    const parts = [];
+    const lines = response.split('\n');
+    let currentPart = { type: 'text', content: '' };
+    let inCodeBlock = false;
+    
+    for (const line of lines) {
+      if (line.trim().startsWith('```python')) {
+        // Starting a code block
+        if (currentPart.content.trim()) {
+          parts.push(currentPart);
+        }
+        currentPart = { type: 'code', content: '' };
+        inCodeBlock = true;
+      } else if (line.trim() === '```' && inCodeBlock) {
+        // Ending a code block
+        parts.push(currentPart);
+        currentPart = { type: 'text', content: '' };
+        inCodeBlock = false;
+      } else {
+        // Regular line
+        if (currentPart.content) {
+          currentPart.content += '\n';
+        }
+        currentPart.content += line;
+      }
+    }
+    
+    if (currentPart.content.trim()) {
+      parts.push(currentPart);
+    }
+    
+    return parts.filter(part => part.content.trim());
+  };
+
+  const handleExecutePrompt = async () => {
+    if (!cell.code.trim()) return;
+    
+    // Phase 2: Authentication check
+    if (!isAuthenticated) {
+      setShowAuthModal(true);
+      return;
+    }
+    
+    // Phase 3: Get current context
+    const aiContext = getAIContext();
+    
+    // Mark cell as processing
+    updatePromptCell(cell.id, cell.code, {
+      isProcessing: true,
+      generatedCode: undefined,
+      aiResponse: undefined,
+    });
+    
+    try {
+      // Phase 4: Real AI processing with DataKit
+      console.log('Processing AI request with context:', aiContext);
+      
+      const aiResponse = await processAIRequest(
+        {
+          prompt: cell.code,
+          context: aiContext,
+        },
+        // Streaming callback for real-time updates
+        (chunk) => {
+          updatePromptCell(cell.id, cell.code, {
+            isProcessing: true,
+            aiResponse: aiStatus.currentResponse,
+          });
+        }
+      );
+      
+      if (aiResponse.error) {
+        // Handle AI errors
+        updatePromptCell(cell.id, cell.code, {
+          isProcessing: false,
+          aiResponse: `❌ AI Error: ${aiResponse.error}\n\nPlease try again or check your connection.`,
+        });
+        return;
+      }
+      
+      // Update cell with successful AI response - put code directly in the cell
+      updatePromptCell(cell.id, aiResponse.generatedCode || cell.code, {
+        isProcessing: false,
+        generatedCode: aiResponse.generatedCode,
+        aiResponse: aiResponse.explanation || 'AI processing completed.',
+      });
+      
+      // Log usage information
+      if (aiResponse.tokensUsed || aiResponse.creditsUsed) {
+        console.log('AI Usage:', {
+          tokens: aiResponse.tokensUsed,
+          credits: aiResponse.creditsUsed,
+        });
+      }
+      
+    } catch (error) {
+      console.error('Error executing AI prompt:', error);
+      updatePromptCell(cell.id, cell.code, {
+        isProcessing: false,
+        aiResponse: `❌ Error: ${error instanceof Error ? error.message : 'Unknown error'}\n\nPlease try again.`,
+      });
+    }
+  };
+
+  const handleAuthSuccess = () => {
+    setShowAuthModal(false);
+    // Reset AI status after auth
+    resetStatus();
+    // After successful auth, automatically execute the prompt
+    setTimeout(() => {
+      handleExecutePrompt();
+    }, 100);
+  };
+
+  const handleCopyPrompt = () => {
+    navigator.clipboard.writeText(cell.code);
+    setShowMenu(false);
+  };
+
+  // Auto-adjust height on mount and content change
+  React.useEffect(() => {
+    adjustTextareaHeight();
+  }, [cell.code, adjustTextareaHeight]);
+
+  // Close menu when clicking outside
+  React.useEffect(() => {
+    const handleClickOutside = (event: MouseEvent) => {
+      if (menuRef.current && !menuRef.current.contains(event.target as Node)) {
+        setShowMenu(false);
+      }
+    };
+
+    document.addEventListener('mousedown', handleClickOutside);
+    return () => document.removeEventListener('mousedown', handleClickOutside);
+  }, []);
+
+  return (
+    <div
+      ref={cellRef}
+      className={`border rounded-lg overflow-visible transition-colors ${
+        isActive
+          ? 'border-primary/50 bg-primary/5'
+          : 'border-white/10 bg-black/20'
+      }`}
+      onClick={onActivate}
+    >
+      {/* Cell Header */}
+      <div className="flex items-center justify-between px-2.5 py-1.5 bg-gradient-to-r from-purple-500/10 to-pink-500/10 border-b border-white/10 relative">
+        <div className="flex items-center gap-3">
+          {/* Collapse/Expand Input Button */}
+          <button
+            onClick={(e) => {
+              e.stopPropagation();
+              toggleCellInputCollapse(cell.id);
+            }}
+            className="p-1 hover:bg-white/10 rounded transition-colors"
+            title={cell.isInputCollapsed ? 'Expand input' : 'Collapse input'}
+          >
+            <ChevronRight
+              size={14}
+              className={`text-white/50 transition-transform ${
+                cell.isInputCollapsed ? '' : 'rotate-90'
+              }`}
+            />
+          </button>
+
+          {/* Cell Number with AI Icon */}
+          <div className="flex items-center gap-2">
+            <Sparkles className="w-4 h-4 text-purple-400" />
+            <span className="text-xs text-white/50">Prompt</span>
+            <span className="text-xs font-mono bg-gradient-to-r from-purple-500/20 to-pink-500/20 px-2 py-1 rounded">
+              {cellNumber}
+            </span>
+          </div>
+
+          {/* Processing Status */}
+          {cell.isProcessing && (
+            <div className="flex items-center gap-2 text-xs text-purple-400">
+              <div className="w-2 h-2 bg-purple-400 rounded-full animate-pulse" />
+              <span>{aiStatus.currentResponse ? 'AI responding...' : 'AI thinking...'}</span>
+            </div>
+          )}
+          
+          {/* Context Status */}
+          {isAuthenticated && !cell.isProcessing && (
+            <div className="flex items-center gap-3">
+              <button
+                onClick={() => setShowContext(!showContext)}
+                className="flex items-center gap-2 text-xs text-blue-400 hover:text-blue-300 transition-colors"
+                title="Toggle context information"
+              >
+                <span>📊</span>
+                <span>{context.selectedTables.size + context.selectedVariables.size}/{context.tables.length + context.variables.length} selected</span>
+              </button>
+              
+              {/* Code Ready Indicator */}
+              {cell.generatedCode && (
+                <div className="flex items-center gap-2 text-xs text-green-400">
+                  <Code2 size={12} />
+                  <span>Code generated</span>
+                </div>
+              )}
+            </div>
+          )}
+          
+          {/* Authentication Status */}
+          {!isAuthenticated && (
+            <button
+              onClick={() => setShowAuthModal(true)}
+              className="flex items-center gap-2 text-xs text-yellow-400 hover:text-yellow-300 transition-colors"
+            >
+              <Lock size={12} />
+              <span>Sign in required</span>
+            </button>
+          )}
+        </div>
+
+        <div className="flex items-center gap-2">
+          {/* Execute Button */}
+          <Tooltip 
+            content={isAuthenticated ? "Execute Prompt (AI)" : "Sign in required to use AI features"} 
+            placement="bottom"
+          >
+            <Button
+              variant="ghost"
+              size="icon"
+              className={`h-8 w-8 bg-gradient-to-r from-purple-500/20 to-pink-500/20 hover:from-purple-500/30 hover:to-pink-500/30 ${
+                !isAuthenticated ? 'opacity-75' : ''
+              }`}
+              onClick={(e) => {
+                e.stopPropagation();
+                handleExecutePrompt();
+              }}
+              disabled={cell.isProcessing || (!isAuthenticated && !cell.code.trim())}
+            >
+              {cell.isProcessing ? (
+                <Square size={14} />
+              ) : !isAuthenticated ? (
+                <Lock size={14} className="text-purple-400" />
+              ) : (
+                <Brain size={14} className="text-purple-400" />
+              )}
+            </Button>
+          </Tooltip>
+
+          {/* Delete Button */}
+          <Button
+            variant="ghost"
+            size="icon"
+            className="h-8 w-8 text-red-400 hover:text-red-300 hover:bg-red-500/10"
+            onClick={(e) => {
+              e.stopPropagation();
+              deleteCell(cell.id);
+            }}
+            title="Delete Cell"
+          >
+            <Trash2 size={14} />
+          </Button>
+
+          {/* Menu Button */}
+          <div className="relative" ref={menuRef}>
+            <Button
+              variant="ghost"
+              size="icon"
+              className="h-8 w-8"
+              onClick={(e) => {
+                e.stopPropagation();
+                setShowMenu(!showMenu);
+              }}
+            >
+              <MoreVertical size={14} />
+            </Button>
+
+            {/* Dropdown Menu */}
+            {showMenu && (
+              <div className="absolute right-0 top-full mt-1 bg-black border border-white/10 rounded shadow-xl z-50 min-w-40">
+                <button
+                  className="w-full px-3 py-2 text-left text-sm text-white/80 hover:bg-white/10 flex items-center gap-2"
+                  onClick={(e) => {
+                    e.preventDefault();
+                    e.stopPropagation();
+                    moveCell(cell.id, 'up');
+                    setTimeout(() => setShowMenu(false), 100);
+                  }}
+                >
+                  <ChevronUp size={14} />
+                  Move Up
+                </button>
+                <button
+                  className="w-full px-3 py-2 text-left text-sm text-white/80 hover:bg-white/10 flex items-center gap-2"
+                  onClick={(e) => {
+                    e.preventDefault();
+                    e.stopPropagation();
+                    moveCell(cell.id, 'down');
+                    setTimeout(() => setShowMenu(false), 100);
+                  }}
+                >
+                  <ChevronDown size={14} />
+                  Move Down
+                </button>
+                <div className="border-t border-white/10" />
+                <button
+                  className="w-full px-3 py-2 text-left text-sm text-white/80 hover:bg-white/10 flex items-center gap-2"
+                  onClick={handleCopyPrompt}
+                >
+                  <Copy size={14} />
+                  Copy Prompt
+                </button>
+                <button
+                  className="w-full px-3 py-2 text-left text-sm text-white/80 hover:bg-white/10 flex items-center gap-2"
+                  onClick={() => {
+                    toggleCellInputCollapse(cell.id);
+                    setShowMenu(false);
+                  }}
+                >
+                  <Eye size={14} />
+                  {cell.isInputCollapsed ? 'Show' : 'Hide'} Input
+                </button>
+              </div>
+            )}
+          </div>
+        </div>
+      </div>
+
+      {/* Cell Input Area */}
+      {!cell.isInputCollapsed && (
+        <div className="border-b border-white/10">
+          {/* Show original prompt as textarea when no code is generated, otherwise show Monaco editor */}
+          {!cell.generatedCode ? (
+            <div className="p-4">
+              <textarea
+                ref={textareaRef}
+                value={cell.code}
+                onChange={handleInputChange}
+                placeholder={isAuthenticated 
+                  ? "What do you want DataKit to do?\n\nExamples:\n• Analyze the sales data and create a chart showing trends\n• Find customers with orders > $1000 and their contact info  \n• Install matplotlib and create a scatter plot of price vs quantity\n• Show me the top 10 products by revenue this quarter"
+                  : "Sign in to use AI-powered data analysis and code generation.\n\nOnce authenticated, you can:\n• Ask questions about your data in natural language\n• Generate Python code automatically\n• Create visualizations with simple prompts\n• Install packages and run complex analysis"
+                }
+                className={`w-full min-h-[120px] p-3 bg-black/20 border border-white/10 rounded-md text-white placeholder-white/40 focus:outline-none focus:border-purple-400/50 focus:ring-1 focus:ring-purple-400/50 transition-colors resize-none font-mono text-sm leading-relaxed ${
+                  !isAuthenticated ? 'opacity-75' : ''
+                }`}
+                style={{ height: 'auto' }}
+                disabled={!isAuthenticated}
+              />
+            </div>
+          ) : (
+            <MonacoErrorBoundary cellId={cell.id}>
+              <MonacoEditor
+                key={`monaco-${cell.id}`}
+                value={cell.code}
+                onChange={handleCodeChange}
+                onExecute={() => {/* TODO: Add execution for prompt cells */}}
+                language="python"
+                height="auto"
+                minHeight={80}
+                maxHeight={2000}
+                options={{
+                  minimap: { enabled: false },
+                  scrollBeyondLastLine: false,
+                  lineNumbers: 'on',
+                  folding: true,
+                  wordWrap: 'on',
+                  automaticLayout: true,
+                  tabSize: 4,
+                  insertSpaces: true,
+                  renderWhitespace: 'boundary',
+                  bracketPairColorization: { enabled: true },
+                }}
+              />
+            </MonacoErrorBoundary>
+          )}
+
+          {/* Context Information */}
+          {showContext && isAuthenticated && (
+            <div className="mt-4 p-3 bg-black/30 border border-blue-500/20 rounded-md">
+              <div className="flex items-center justify-between mb-3">
+                <div className="flex items-center gap-2">
+                  <span className="text-blue-400">📊</span>
+                  <span className="text-xs text-blue-400 font-medium">Context Selection</span>
+                </div>
+                <button
+                  onClick={refreshContext}
+                  className="text-xs text-blue-400/70 hover:text-blue-400 transition-colors"
+                  disabled={context.isLoading}
+                >
+                  {context.isLoading ? 'Refreshing...' : 'Refresh'}
+                </button>
+              </div>
+              
+              <div className="space-y-3">
+                {/* Tables Section */}
+                {context.tables.length > 0 && (
+                  <div>
+                    <div className="flex items-center justify-between mb-2">
+                      <span className="text-xs text-green-400 font-medium">DuckDB Tables ({context.tables.length})</span>
+                      <div className="flex gap-1">
+                        <button
+                          onClick={() => toggleAllTables(true)}
+                          className="text-xs text-green-400/70 hover:text-green-400 transition-colors px-1"
+                          title="Select all tables"
+                        >
+                          All
+                        </button>
+                        <span className="text-white/30">|</span>
+                        <button
+                          onClick={() => toggleAllTables(false)}
+                          className="text-xs text-red-400/70 hover:text-red-400 transition-colors px-1"
+                          title="Deselect all tables"
+                        >
+                          None
+                        </button>
+                      </div>
+                    </div>
+                    <div className="space-y-1">
+                      {context.tables.map(table => (
+                        <div key={table.name} className="flex items-center gap-2">
+                          <button
+                            onClick={() => toggleTableSelection(table.name)}
+                            className={`w-4 h-4 rounded border flex items-center justify-center transition-colors ${
+                              context.selectedTables.has(table.name)
+                                ? 'bg-green-500 border-green-500 text-white'
+                                : 'border-white/30 hover:border-green-400'
+                            }`}
+                          >
+                            {context.selectedTables.has(table.name) && <Check size={10} />}
+                          </button>
+                          <span className="text-xs text-white/80 flex-1">
+                            {table.name} <span className="text-white/50">({table.schema.length} cols)</span>
+                          </span>
+                        </div>
+                      ))}
+                    </div>
+                  </div>
+                )}
+                
+                {/* DataFrames Section */}
+                {context.variables.filter(v => v.type === 'DataFrame').length > 0 && (
+                  <div>
+                    <div className="flex items-center justify-between mb-2">
+                      <span className="text-xs text-blue-400 font-medium">DataFrames ({context.variables.filter(v => v.type === 'DataFrame').length})</span>
+                      <div className="flex gap-1">
+                        <button
+                          onClick={() => {
+                            context.variables.filter(v => v.type === 'DataFrame').forEach(v => {
+                              if (!context.selectedVariables.has(v.name)) {
+                                toggleVariableSelection(v.name);
+                              }
+                            });
+                          }}
+                          className="text-xs text-blue-400/70 hover:text-blue-400 transition-colors px-1"
+                          title="Select all DataFrames"
+                        >
+                          All
+                        </button>
+                        <span className="text-white/30">|</span>
+                        <button
+                          onClick={() => {
+                            context.variables.filter(v => v.type === 'DataFrame').forEach(v => {
+                              if (context.selectedVariables.has(v.name)) {
+                                toggleVariableSelection(v.name);
+                              }
+                            });
+                          }}
+                          className="text-xs text-red-400/70 hover:text-red-400 transition-colors px-1"
+                          title="Deselect all DataFrames"
+                        >
+                          None
+                        </button>
+                      </div>
+                    </div>
+                    <div className="space-y-1">
+                      {context.variables.filter(v => v.type === 'DataFrame').map(variable => (
+                        <div key={variable.name} className="flex items-center gap-2">
+                          <button
+                            onClick={() => toggleVariableSelection(variable.name)}
+                            className={`w-4 h-4 rounded border flex items-center justify-center transition-colors ${
+                              context.selectedVariables.has(variable.name)
+                                ? 'bg-blue-500 border-blue-500 text-white'
+                                : 'border-white/30 hover:border-blue-400'
+                            }`}
+                          >
+                            {context.selectedVariables.has(variable.name) && <Check size={10} />}
+                          </button>
+                          <span className="text-xs text-white/80 flex-1">
+                            {variable.name} <span className="text-white/50">({variable.shape?.[0]}×{variable.shape?.[1]})</span>
+                          </span>
+                        </div>
+                      ))}
+                    </div>
+                  </div>
+                )}
+                
+                {/* Other Variables Section */}
+                {context.variables.filter(v => v.type !== 'DataFrame').length > 0 && (
+                  <div>
+                    <div className="flex items-center justify-between mb-2">
+                      <span className="text-xs text-yellow-400 font-medium">Variables ({context.variables.filter(v => v.type !== 'DataFrame').length})</span>
+                      <div className="flex gap-1">
+                        <button
+                          onClick={() => {
+                            context.variables.filter(v => v.type !== 'DataFrame').forEach(v => {
+                              if (!context.selectedVariables.has(v.name)) {
+                                toggleVariableSelection(v.name);
+                              }
+                            });
+                          }}
+                          className="text-xs text-yellow-400/70 hover:text-yellow-400 transition-colors px-1"
+                          title="Select all variables"
+                        >
+                          All
+                        </button>
+                        <span className="text-white/30">|</span>
+                        <button
+                          onClick={() => {
+                            context.variables.filter(v => v.type !== 'DataFrame').forEach(v => {
+                              if (context.selectedVariables.has(v.name)) {
+                                toggleVariableSelection(v.name);
+                              }
+                            });
+                          }}
+                          className="text-xs text-red-400/70 hover:text-red-400 transition-colors px-1"
+                          title="Deselect all variables"
+                        >
+                          None
+                        </button>
+                      </div>
+                    </div>
+                    <div className="space-y-1">
+                      {context.variables.filter(v => v.type !== 'DataFrame').map(variable => (
+                        <div key={variable.name} className="flex items-center gap-2">
+                          <button
+                            onClick={() => toggleVariableSelection(variable.name)}
+                            className={`w-4 h-4 rounded border flex items-center justify-center transition-colors ${
+                              context.selectedVariables.has(variable.name)
+                                ? 'bg-yellow-500 border-yellow-500 text-white'
+                                : 'border-white/30 hover:border-yellow-400'
+                            }`}
+                          >
+                            {context.selectedVariables.has(variable.name) && <Check size={10} />}
+                          </button>
+                          <span className="text-xs text-white/80 flex-1">
+                            {variable.name} <span className="text-white/50">({variable.type})</span>
+                          </span>
+                        </div>
+                      ))}
+                    </div>
+                  </div>
+                )}
+                
+                {context.tables.length === 0 && context.variables.length === 0 && (
+                  <div className="text-white/50 italic text-xs">No context detected</div>
+                )}
+              </div>
+            </div>
+          )}
+          
+
+        </div>
+      )}
+      
+      {/* AI Output Area - shows parsed response with code blocks and explanations */}
+      {cell.aiResponse && (
+        <div className="border-t border-white/5">
+          <div className="p-4 space-y-4">
+            {parseAIOutput(cell.aiResponse).map((part, index) => (
+              <div key={index}>
+                {part.type === 'code' ? (
+                  <div className="bg-black/30 border border-green-500/20 rounded-md overflow-hidden">
+                    <div className="px-3 py-2 bg-green-500/10 border-b border-green-500/20 flex items-center gap-2">
+                      <Code2 size={12} className="text-green-400" />
+                      <span className="text-xs text-green-400 font-medium">Generated Code</span>
+                    </div>
+                    <div className="p-3">
+                      <pre className="text-sm text-white/90 whitespace-pre-wrap font-mono">
+                        {part.content.trim()}
+                      </pre>
+                    </div>
+                  </div>
+                ) : (
+                  <div className="bg-black/20 border border-blue-500/20 rounded-md p-3">
+                    <div className="flex items-center gap-2 mb-2">
+                      <Brain size={12} className="text-blue-400" />
+                      <span className="text-xs text-blue-400 font-medium">AI Explanation</span>
+                    </div>
+                    <div className="text-sm text-white/80 leading-relaxed whitespace-pre-wrap">
+                      {part.content.trim()}
+                    </div>
+                  </div>
+                )}
+              </div>
+            ))}
+          </div>
+        </div>
+      )}
+      
+      {/* Authentication Modal */}
+      <AuthModal
+        isOpen={showAuthModal}
+        onClose={() => setShowAuthModal(false)}
+        onLoginSuccess={handleAuthSuccess}
+        defaultMode="login"
+      />
+    </div>
+  );
+};
+
+export default PromptCell;
\ No newline at end of file
diff --git a/frontend/src/hooks/notebooks/useNotebookAI.ts b/frontend/src/hooks/notebooks/useNotebookAI.ts
new file mode 100644
index 0000000..cadd46a
--- /dev/null
+++ b/frontend/src/hooks/notebooks/useNotebookAI.ts
@@ -0,0 +1,274 @@
+import { useState, useCallback } from 'react';
+import { useAuth } from '@/hooks/auth/useAuth';
+import { DataKitProvider } from '@/lib/ai/providers/datakit';
+import { aiService } from '@/lib/ai/aiService';
+import type { AIMessage } from '@/lib/ai/types';
+import type { NotebookContext } from './useNotebookContext';
+
+export interface NotebookAIRequest {
+  prompt: string;
+  context: {
+    hasContext: boolean;
+    contextText: string;
+    tables: any[];
+    variables: any[];
+    summary?: any;
+  };
+}
+
+export interface NotebookAIResponse {
+  generatedCode?: string;
+  explanation?: string;
+  suggestions?: string[];
+  error?: string;
+  tokensUsed?: {
+    input: number;
+    output: number;
+  };
+  creditsUsed?: number;
+}
+
+export interface NotebookAIStatus {
+  isProcessing: boolean;
+  currentResponse: string;
+  error: string | null;
+}
+
+/**
+ * Hook for AI-powered notebook interactions using DataKit models
+ * Integrates with context detection and provides streaming responses
+ */
+export const useNotebookAI = () => {
+  const { isAuthenticated } = useAuth();
+  const [status, setStatus] = useState<NotebookAIStatus>({
+    isProcessing: false,
+    currentResponse: '',
+    error: null,
+  });
+
+  // Initialize DataKit provider
+  const initializeAI = useCallback(async () => {
+    if (!isAuthenticated) {
+      throw new Error('Authentication required for AI features');
+    }
+
+    try {
+      // Initialize DataKit provider with datakit-fast model for cost efficiency
+      const datakitProvider = new DataKitProvider('datakit-fast');
+      
+      // Validate the provider
+      const isValid = await datakitProvider.validateApiKey();
+      if (!isValid) {
+        throw new Error('DataKit API validation failed. Please check your authentication.');
+      }
+
+      // Set the provider in the AI service
+      aiService.setProvider('datakit', datakitProvider);
+      
+      return datakitProvider;
+    } catch (error) {
+      console.error('Failed to initialize DataKit AI:', error);
+      throw error;
+    }
+  }, [isAuthenticated]);
+
+  // Create system prompt for notebook context
+  const createNotebookSystemPrompt = useCallback((context: NotebookAIRequest['context']) => {
+    let systemPrompt = `You are an expert data analyst and Python programmer working in a Jupyter-style notebook environment. Your goal is to help users analyze data and write Python code.
+
+Environment:
+- Python environment with pandas, numpy, matplotlib, plotly, and other common data science libraries
+- DuckDB integration available for SQL queries
+- All code should be production-ready and well-documented
+
+Guidelines:
+1. Generate clean, efficient Python code in code blocks
+2. Provide clear explanations between code blocks
+3. Use appropriate libraries for the task
+4. Handle potential errors gracefully
+5. Follow Python best practices and PEP 8 style
+6. Structure response as: explanation → code blocks → explanation → etc.
+7. Keep explanations concise and focused on what the code does
+
+`;
+
+    if (context.hasContext) {
+      systemPrompt += `Current Data Context:
+${context.contextText}
+
+When generating code:
+- Reference the available tables and variables by their exact names
+- Consider the data types and schema when writing queries
+- Suggest appropriate analysis based on the available data
+- Use DuckDB SQL syntax when querying tables
+- Use pandas operations when working with DataFrames
+
+`;
+    } else {
+      systemPrompt += `No specific data context available. Generate general-purpose code and suggest loading sample data when appropriate.
+
+`;
+    }
+
+    return systemPrompt;
+  }, []);
+
+  // Process AI request with streaming
+  const processAIRequest = useCallback(async (
+    request: NotebookAIRequest,
+    onChunk?: (chunk: string) => void
+  ): Promise<NotebookAIResponse> => {
+    if (!isAuthenticated) {
+      throw new Error('Authentication required');
+    }
+
+    setStatus({ isProcessing: true, currentResponse: '', error: null });
+
+    try {
+      // Initialize AI if needed
+      const provider = await initializeAI();
+
+      // Create messages for the AI
+      const systemPrompt = createNotebookSystemPrompt(request.context);
+      const messages: AIMessage[] = [
+        { role: 'system', content: systemPrompt },
+        { role: 'user', content: request.prompt }
+      ];
+
+      let fullResponse = '';
+      let tokensUsed = { input: 0, output: 0 };
+      let creditsUsed = 0;
+
+      // Use streaming for better UX
+      await provider.generateStreamCompletion(
+        messages,
+        (chunk) => {
+          if (chunk.content) {
+            fullResponse += chunk.content;
+            setStatus(prev => ({ 
+              ...prev, 
+              currentResponse: fullResponse 
+            }));
+            onChunk?.(chunk.content);
+          }
+          if (chunk.done && chunk.usage) {
+            tokensUsed = {
+              input: chunk.usage.promptTokens || 0,
+              output: chunk.usage.completionTokens || 0,
+            };
+          }
+          if (chunk._datakit?.creditsUsed) {
+            creditsUsed = chunk._datakit.creditsUsed;
+          }
+        },
+        {
+          temperature: 0.1, // Low temperature for more deterministic code generation
+          maxTokens: 2000,  // Reasonable limit for code generation
+        }
+      );
+
+      // Parse the response to extract code and explanation
+      const parsed = parseAIResponse(fullResponse);
+
+      setStatus({ isProcessing: false, currentResponse: fullResponse, error: null });
+
+      return {
+        ...parsed,
+        tokensUsed,
+        creditsUsed,
+      };
+
+    } catch (error) {
+      const errorMessage = error instanceof Error ? error.message : 'AI processing failed';
+      setStatus({ 
+        isProcessing: false, 
+        currentResponse: '', 
+        error: errorMessage 
+      });
+      
+      return {
+        error: errorMessage,
+      };
+    }
+  }, [isAuthenticated, initializeAI, createNotebookSystemPrompt]);
+
+  // Parse AI response to extract code blocks and explanations
+  const parseAIResponse = useCallback((response: string): Omit<NotebookAIResponse, 'tokensUsed' | 'creditsUsed'> => {
+    // Extract Python code blocks
+    const codeMatches = response.match(/```python\n([\s\S]*?)\n```/g);
+    let generatedCode = '';
+    
+    if (codeMatches && codeMatches.length > 0) {
+      // Combine all code blocks
+      generatedCode = codeMatches
+        .map(match => match.replace(/```python\n|\n```/g, '').trim())
+        .join('\n\n# --- Next Code Block ---\n\n');
+    } else {
+      // Fallback: look for any code-like content
+      const lines = response.split('\n');
+      const codeLines = lines.filter(line => 
+        line.trim().startsWith('import ') ||
+        line.trim().startsWith('from ') ||
+        line.trim().startsWith('df') ||
+        line.trim().startsWith('pd.') ||
+        line.trim().startsWith('plt.') ||
+        line.trim().startsWith('print(') ||
+        /^\s*[a-zA-Z_][a-zA-Z0-9_]*\s*=/.test(line)
+      );
+      
+      if (codeLines.length > 0) {
+        generatedCode = codeLines.join('\n');
+      }
+    }
+
+    // Extract explanation (everything that's not code)
+    let explanation = response;
+    if (codeMatches) {
+      codeMatches.forEach(match => {
+        explanation = explanation.replace(match, '');
+      });
+    }
+    explanation = explanation.trim();
+
+    // Extract suggestions (look for bullet points or numbered lists)
+    const suggestions: string[] = [];
+    const suggestionPatterns = [
+      /(?:^|\n)[-*•]\s+(.+?)(?=\n|$)/g,
+      /(?:^|\n)\d+\.\s+(.+?)(?=\n|$)/g,
+    ];
+    
+    suggestionPatterns.forEach(pattern => {
+      let match;
+      while ((match = pattern.exec(response)) !== null) {
+        const suggestion = match[1].trim();
+        if (suggestion.length > 10 && !suggestions.includes(suggestion)) {
+          suggestions.push(suggestion);
+        }
+      }
+    });
+
+    return {
+      generatedCode: generatedCode || undefined,
+      explanation: explanation || undefined,
+      suggestions: suggestions.length > 0 ? suggestions : undefined,
+    };
+  }, []);
+
+  // Reset status
+  const resetStatus = useCallback(() => {
+    setStatus({
+      isProcessing: false,
+      currentResponse: '',
+      error: null,
+    });
+  }, []);
+
+  return {
+    status,
+    processAIRequest,
+    resetStatus,
+    isReady: isAuthenticated,
+  };
+};
+
+export default useNotebookAI;
\ No newline at end of file
diff --git a/frontend/src/hooks/notebooks/useNotebookContext.ts b/frontend/src/hooks/notebooks/useNotebookContext.ts
new file mode 100644
index 0000000..3c5cac3
--- /dev/null
+++ b/frontend/src/hooks/notebooks/useNotebookContext.ts
@@ -0,0 +1,346 @@
+import { useState, useEffect, useCallback } from 'react';
+import { useDuckDBStore } from '@/store/duckDBStore';
+import { usePythonStore } from '@/store/pythonStore';
+import { getPythonVariables } from '@/lib/python/executor';
+
+export interface DuckDBTable {
+  name: string;
+  schema: Array<{ name: string; type: string }>;
+  rowCount?: number;
+  objectType: 'table' | 'view';
+}
+
+export interface PythonVariable {
+  name: string;
+  type: string;
+  shape?: number[];
+  columns?: string[];
+  length?: number;
+  value?: string;
+  dtype?: string;
+}
+
+export interface NotebookContext {
+  tables: DuckDBTable[];
+  variables: PythonVariable[];
+  selectedTables: Set<string>;
+  selectedVariables: Set<string>;
+  isLoading: boolean;
+  error: string | null;
+  lastUpdated: Date | null;
+}
+
+export interface ContextSummary {
+  tableCount: number;
+  variableCount: number;
+  availableDataFrames: string[];
+  availableTables: string[];
+  contextDescription: string;
+}
+
+/**
+ * Hook for detecting and managing notebook context (tables, variables, etc.)
+ * This provides AI with contextual information about the user's data and environment
+ */
+export const useNotebookContext = () => {
+  const [context, setContext] = useState<NotebookContext>({
+    tables: [],
+    variables: [],
+    selectedTables: new Set(),
+    selectedVariables: new Set(),
+    isLoading: false,
+    error: null,
+    lastUpdated: null,
+  });
+
+  const { 
+    isInitialized: isDuckDBInitialized, 
+    getAvailableTables, 
+    getTableSchema,
+    getObjectType,
+    registeredTables 
+  } = useDuckDBStore();
+
+  const { 
+    pyodide: { isInitialized: isPyodideInitialized },
+    globalVariables 
+  } = usePythonStore();
+
+  // Detect DuckDB tables and their schemas
+  const detectTables = useCallback(async (): Promise<DuckDBTable[]> => {
+    if (!isDuckDBInitialized) {
+      return [];
+    }
+
+    try {
+      const tableNames = getAvailableTables();
+      const tables: DuckDBTable[] = [];
+
+      for (const tableName of tableNames) {
+        try {
+          const [schema, objectType] = await Promise.all([
+            getTableSchema(tableName),
+            getObjectType(tableName),
+          ]);
+
+          if (schema && objectType) {
+            tables.push({
+              name: tableName,
+              schema,
+              objectType,
+            });
+          }
+        } catch (error) {
+          console.warn(`Failed to get schema for table ${tableName}:`, error);
+        }
+      }
+
+      return tables;
+    } catch (error) {
+      console.error('Failed to detect tables:', error);
+      throw error;
+    }
+  }, [isDuckDBInitialized, getAvailableTables, getTableSchema, getObjectType]);
+
+  // Detect Python variables
+  const detectVariables = useCallback(async (): Promise<PythonVariable[]> => {
+    if (!isPyodideInitialized) {
+      return [];
+    }
+
+    try {
+      const variables = await getPythonVariables();
+      return Object.entries(variables).map(([name, info]) => ({
+        name,
+        type: info.type,
+        shape: info.shape,
+        columns: info.columns,
+        length: info.length,
+        value: info.value,
+        dtype: info.dtype,
+      }));
+    } catch (error) {
+      console.error('Failed to detect variables:', error);
+      throw error;
+    }
+  }, [isPyodideInitialized]);
+
+  // Refresh context by detecting both tables and variables
+  const refreshContext = useCallback(async () => {
+    setContext(prev => ({ ...prev, isLoading: true, error: null }));
+
+    try {
+      const [tables, variables] = await Promise.all([
+        detectTables(),
+        detectVariables(),
+      ]);
+
+      setContext(prev => ({
+        tables,
+        variables,
+        // Auto-select new tables and variables, keep existing selections
+        selectedTables: new Set([
+          ...Array.from(prev.selectedTables).filter(name => tables.some(t => t.name === name)),
+          ...tables.filter(t => !prev.selectedTables.has(t.name)).map(t => t.name)
+        ]),
+        selectedVariables: new Set([
+          ...Array.from(prev.selectedVariables).filter(name => variables.some(v => v.name === name)),
+          ...variables.filter(v => !prev.selectedVariables.has(v.name)).map(v => v.name)
+        ]),
+        isLoading: false,
+        error: null,
+        lastUpdated: new Date(),
+      }));
+    } catch (error) {
+      const errorMessage = error instanceof Error ? error.message : 'Unknown error';
+      setContext(prev => ({
+        ...prev,
+        isLoading: false,
+        error: errorMessage,
+      }));
+    }
+  }, [detectTables, detectVariables]);
+
+  // Auto-refresh context when environments are ready
+  useEffect(() => {
+    if (isDuckDBInitialized || isPyodideInitialized) {
+      refreshContext();
+    }
+  }, [isDuckDBInitialized, isPyodideInitialized, refreshContext]);
+
+  // Auto-refresh when DuckDB tables change
+  useEffect(() => {
+    if (isDuckDBInitialized && registeredTables.size > 0) {
+      refreshContext();
+    }
+  }, [isDuckDBInitialized, registeredTables, refreshContext]);
+
+  // Auto-refresh when Python variables change (debounced)
+  useEffect(() => {
+    if (isPyodideInitialized && Object.keys(globalVariables).length > 0) {
+      const timeoutId = setTimeout(() => {
+        refreshContext();
+      }, 1000); // Debounce to avoid excessive refreshes
+
+      return () => clearTimeout(timeoutId);
+    }
+  }, [isPyodideInitialized, globalVariables, refreshContext]);
+
+  // Toggle table selection
+  const toggleTableSelection = useCallback((tableName: string) => {
+    setContext(prev => {
+      const newSelectedTables = new Set(prev.selectedTables);
+      if (newSelectedTables.has(tableName)) {
+        newSelectedTables.delete(tableName);
+      } else {
+        newSelectedTables.add(tableName);
+      }
+      return { ...prev, selectedTables: newSelectedTables };
+    });
+  }, []);
+
+  // Toggle variable selection
+  const toggleVariableSelection = useCallback((variableName: string) => {
+    setContext(prev => {
+      const newSelectedVariables = new Set(prev.selectedVariables);
+      if (newSelectedVariables.has(variableName)) {
+        newSelectedVariables.delete(variableName);
+      } else {
+        newSelectedVariables.add(variableName);
+      }
+      return { ...prev, selectedVariables: newSelectedVariables };
+    });
+  }, []);
+
+  // Select/deselect all tables
+  const toggleAllTables = useCallback((selectAll: boolean) => {
+    setContext(prev => ({
+      ...prev,
+      selectedTables: selectAll ? new Set(prev.tables.map(t => t.name)) : new Set(),
+    }));
+  }, []);
+
+  // Select/deselect all variables
+  const toggleAllVariables = useCallback((selectAll: boolean) => {
+    setContext(prev => ({
+      ...prev,
+      selectedVariables: selectAll ? new Set(prev.variables.map(v => v.name)) : new Set(),
+    }));
+  }, []);
+
+  // Generate context summary for AI (only selected items)
+  const getContextSummary = useCallback((): ContextSummary => {
+    const { tables, variables, selectedTables, selectedVariables } = context;
+    
+    const selectedTablesList = tables.filter(t => selectedTables.has(t.name));
+    const selectedVariablesList = variables.filter(v => selectedVariables.has(v.name));
+    
+    const availableDataFrames = selectedVariablesList
+      .filter(v => v.type === 'DataFrame')
+      .map(v => v.name);
+    
+    const availableTables = selectedTablesList.map(t => t.name);
+    
+    let contextDescription = '';
+    
+    if (selectedTablesList.length > 0) {
+      contextDescription += `Available DuckDB tables: ${selectedTablesList.map(t => 
+        `${t.name} (${t.schema.length} columns: ${t.schema.map(c => `${c.name}:${c.type}`).join(', ')})`
+      ).join('; ')}. `;
+    }
+    
+    if (availableDataFrames.length > 0) {
+      contextDescription += `Available DataFrames: ${selectedVariablesList
+        .filter(v => v.type === 'DataFrame')
+        .map(v => `${v.name} (${v.shape?.[0]}x${v.shape?.[1]}, columns: ${v.columns?.join(', ')})`)
+        .join('; ')}. `;
+    }
+    
+    const otherVariables = selectedVariablesList.filter(v => v.type !== 'DataFrame');
+    if (otherVariables.length > 0) {
+      contextDescription += `Other variables: ${otherVariables
+        .map(v => `${v.name} (${v.type})`)
+        .join(', ')}. `;
+    }
+
+    return {
+      tableCount: selectedTablesList.length,
+      variableCount: selectedVariablesList.length,
+      availableDataFrames,
+      availableTables,
+      contextDescription: contextDescription.trim(),
+    };
+  }, [context]);
+
+  // Generate formatted context for AI prompts (only selected items)
+  const getAIContext = useCallback(() => {
+    const summary = getContextSummary();
+    
+    if (summary.tableCount === 0 && summary.variableCount === 0) {
+      return {
+        hasContext: false,
+        contextText: 'No data tables or variables selected for context.',
+        tables: [],
+        variables: [],
+      };
+    }
+
+    let contextText = 'Selected notebook context:\n\n';
+    
+    const selectedTables = context.tables.filter(t => context.selectedTables.has(t.name));
+    const selectedVariables = context.variables.filter(v => context.selectedVariables.has(v.name));
+    
+    if (selectedTables.length > 0) {
+      contextText += 'DuckDB Tables:\n';
+      selectedTables.forEach(table => {
+        contextText += `- ${table.name} (${table.objectType}): ${table.schema.length} columns\n`;
+        table.schema.forEach(col => {
+          contextText += `  * ${col.name}: ${col.type}\n`;
+        });
+      });
+      contextText += '\n';
+    }
+    
+    if (selectedVariables.length > 0) {
+      contextText += 'Python Variables:\n';
+      selectedVariables.forEach(variable => {
+        let varDesc = `- ${variable.name}: ${variable.type}`;
+        if (variable.shape) {
+          varDesc += ` (shape: ${variable.shape.join('x')})`;
+        }
+        if (variable.columns) {
+          varDesc += ` (columns: ${variable.columns.join(', ')})`;
+        }
+        if (variable.length !== undefined) {
+          varDesc += ` (length: ${variable.length})`;
+        }
+        if (variable.value) {
+          varDesc += ` (value: ${variable.value})`;
+        }
+        contextText += varDesc + '\n';
+      });
+    }
+
+    return {
+      hasContext: true,
+      contextText,
+      tables: selectedTables,
+      variables: selectedVariables,
+      summary,
+    };
+  }, [context, getContextSummary]);
+
+  return {
+    context,
+    refreshContext,
+    getContextSummary,
+    getAIContext,
+    toggleTableSelection,
+    toggleVariableSelection,
+    toggleAllTables,
+    toggleAllVariables,
+    isReady: isDuckDBInitialized || isPyodideInitialized,
+  };
+};
+
+export default useNotebookContext;
\ No newline at end of file
diff --git a/frontend/src/lib/python/types.ts b/frontend/src/lib/python/types.ts
index 07e7d8d..6b08470 100644
--- a/frontend/src/lib/python/types.ts
+++ b/frontend/src/lib/python/types.ts
@@ -1,17 +1,21 @@
 import type { PyodideInterface } from "pyodide";
 
-export type CellType = 'code' | 'markdown';
+export type CellType = 'code' | 'markdown' | 'prompt';
 
 export interface PythonCell {
   id: string;
   type: CellType;
-  code: string; // For code cells, this is Python code; for markdown cells, this is markdown content
+  code: string; // For code cells, this is Python code; for markdown cells, this is markdown content; for prompt cells, this is the user's natural language prompt
   output: CellOutput[];
   executionCount: number | null;
   isExecuting: boolean;
   isEditing?: boolean; // For markdown cells, tracks if currently editing
   isInputCollapsed?: boolean; // Whether the cell input is collapsed
   isOutputCollapsed?: boolean; // Whether the cell output is collapsed
+  // Prompt-specific properties
+  isProcessing?: boolean; // For prompt cells, tracks if AI is processing
+  generatedCode?: string; // For prompt cells, stores the AI-generated code
+  aiResponse?: string; // For prompt cells, stores the full AI response
   createdAt: Date;
   updatedAt: Date;
 }