From 67a7f646a8fa198162ec73baee2f3a1b28e173dc Mon Sep 17 00:00:00 2001
From: Shreya Shankar <ss.shankar505@gmail.com>
Date: Fri, 29 Nov 2024 00:48:15 -0600
Subject: [PATCH] feat: add column view dialog (#214)

* feat: add column view dialog

* feat: add column view dialog

* increase docker time

* feat: add feedback indicator in the row view

* feat: add prompt engineering flow for notes

* feat: add prompt engineering flow for notes

* feat: support resolve prompts in the prompt editor

* fix: fix build errors

* docs: update playground

* docs: update playground

* docs: update playground

* docs: update playground
---
 .github/workflows/docker-ci.yml               |    2 +-
 docs/playground/index.md                      |   23 +-
 website/package-lock.json                     |   17 +
 website/package.json                          |    2 +
 website/src/app/api/chat/route.ts             |    2 +-
 website/src/app/globals.css                   |    6 +-
 website/src/app/page.tsx                      |    8 +-
 website/src/app/playground/page.tsx           |    6 +-
 website/src/app/types.ts                      |   17 +-
 website/src/components/BookmarkableText.tsx   |  290 -----
 website/src/components/BookmarksPanel.tsx     |  139 ++-
 website/src/components/ColumnDialog.tsx       |  545 +++++++++
 website/src/components/CommandMenu.tsx        |   61 -
 website/src/components/DatasetView.tsx        |   33 +-
 website/src/components/MarkdownCell.tsx       |   93 ++
 website/src/components/OperationCard.tsx      |   84 +-
 website/src/components/Output.tsx             |   28 +-
 website/src/components/PipelineGui.tsx        |   41 +-
 website/src/components/PrettyJSON.tsx         |  145 +++
 .../components/PromptImprovementDialog.tsx    | 1013 +++++++++++++++++
 website/src/components/ResizableDataTable.tsx |  369 +++---
 website/src/components/RowNavigator.tsx       |  128 +++
 website/src/components/SearchableCell.tsx     |  170 +++
 website/src/contexts/BookmarkContext.tsx      |   35 +-
 website/src/contexts/PipelineContext.tsx      |   33 +-
 website/src/hooks/use-toast.ts                |   15 +-
 26 files changed, 2569 insertions(+), 736 deletions(-)
 delete mode 100644 website/src/components/BookmarkableText.tsx
 create mode 100644 website/src/components/ColumnDialog.tsx
 delete mode 100644 website/src/components/CommandMenu.tsx
 create mode 100644 website/src/components/MarkdownCell.tsx
 create mode 100644 website/src/components/PrettyJSON.tsx
 create mode 100644 website/src/components/PromptImprovementDialog.tsx
 create mode 100644 website/src/components/RowNavigator.tsx
 create mode 100644 website/src/components/SearchableCell.tsx

diff --git a/.github/workflows/docker-ci.yml b/.github/workflows/docker-ci.yml
index 0dcebf87..376cbc6a 100644
--- a/.github/workflows/docker-ci.yml
+++ b/.github/workflows/docker-ci.yml
@@ -40,7 +40,7 @@ jobs:
           docetl
           
         # Wait for container to start up
-        sleep 120
+        sleep 240
         
         # Check if container is still running
         if [ "$(docker ps -q -f name=docetl-test)" ]; then
diff --git a/docs/playground/index.md b/docs/playground/index.md
index 89f2eaee..c4192e4a 100644
--- a/docs/playground/index.md
+++ b/docs/playground/index.md
@@ -1,12 +1,15 @@
 # Playground
 
-The playground is a web app that allows you to interactively build DocETL pipelines. The playground is built with Next.js and TypeScript. We use the `docetl` Python package (built from this source code) to process the data with a FastAPI server. We stream out the logs from the FastAPI server to the frontend so you can see the pipeline execution progress and outputs in real time.
+The DocETL Playground is an integrated development environment (IDE) for building and testing document processing pipelines. Built with Next.js and TypeScript, it provides a real-time interface to develop, test and refine your pipelines through a FastAPI backend.
 
-## Why an interactive playground?
+## Why a Playground? 🤔
 
-Often, unstructured data analysis tasks are fuzzy and require iteration. You might start with a prompt, see the outputs for a sample, then realize you need to tweak the prompt or change the definition of the task you want the LLM to do. Or, you might want to create a complex pipeline that involves multiple operations, but you are unsure of what prompts you want to use for each step, so you want to build your pipeline one operation at a time.
+This **interactive playground** streamlines development from prototype to production! **Our (in-progress) user studies show 100% of developers** found building pipelines significantly faster and easier with our playground vs traditional approaches.
 
-The playground allows you to do just that.
+Building complex LLM pipelines for your data often requires experimentation and iteration. The IDE lets you:
+- 🚀 Test prompts and see results instantly
+- ✨ Refine operations based on sample outputs  
+- 🔄 Build complex pipelines step-by-step
 
 ## Installation
 
@@ -18,9 +21,9 @@ The easiest way to get started is using Docker:
 
 1. Create the required environment files:
 
-Create `.env` in the root directory:
+Create `.env` in the root directory (for the FastAPI backend):
 ```bash
-OPENAI_API_KEY=your_api_key_here
+OPENAI_API_KEY=your_api_key_here # Or your LLM provider's API key
 BACKEND_ALLOW_ORIGINS=
 BACKEND_HOST=localhost
 BACKEND_PORT=8000
@@ -29,11 +32,11 @@ FRONTEND_HOST=localhost
 FRONTEND_PORT=3000
 ```
 
-Create `.env.local` in the `website` directory:
+Create `.env.local` in the `website` directory (for the frontend) **note that this must be in the `website` directory**:
 ```bash
-OPENAI_API_KEY=sk-xxx
-OPENAI_API_BASE=https://api.openai.com/v1
-MODEL_NAME=gpt-4o-mini
+OPENAI_API_KEY=sk-xxx # For the AI assistant in the interface
+OPENAI_API_BASE=https://api.openai.com/v1 # For the AI assistant in the interface
+MODEL_NAME=gpt-4o-mini # For the AI assistant in the interface
 
 NEXT_PUBLIC_BACKEND_HOST=localhost
 NEXT_PUBLIC_BACKEND_PORT=8000
diff --git a/website/package-lock.json b/website/package-lock.json
index e4240d57..b56e3f00 100644
--- a/website/package-lock.json
+++ b/website/package-lock.json
@@ -37,6 +37,7 @@
         "@tanstack/react-query": "^5.59.15",
         "@tanstack/react-table": "^8.20.5",
         "@tanstack/react-virtual": "^3.10.9",
+        "@types/diff": "^6.0.0",
         "@types/mime-types": "^2.1.4",
         "@types/react-resizable": "^3.0.8",
         "ai": "^3.4.29",
@@ -47,6 +48,7 @@
         "clsx": "^2.1.1",
         "cmdk": "^1.0.0",
         "css-loader": "^7.1.2",
+        "diff": "^7.0.0",
         "framer-motion": "^11.5.4",
         "gray-matter": "^4.0.3",
         "js-yaml": "^4.1.0",
@@ -3888,6 +3890,12 @@
         "@types/ms": "*"
       }
     },
+    "node_modules/@types/diff": {
+      "version": "6.0.0",
+      "resolved": "https://registry.npmjs.org/@types/diff/-/diff-6.0.0.tgz",
+      "integrity": "sha512-dhVCYGv3ZSbzmQaBSagrv1WJ6rXCdkyTcDyoNu1MD8JohI7pR7k8wdZEm+mvdxRKXyHVwckFzWU1vJc+Z29MlA==",
+      "license": "MIT"
+    },
     "node_modules/@types/diff-match-patch": {
       "version": "1.0.36",
       "resolved": "https://registry.npmjs.org/@types/diff-match-patch/-/diff-match-patch-1.0.36.tgz",
@@ -6462,6 +6470,15 @@
       "resolved": "https://registry.npmjs.org/didyoumean/-/didyoumean-1.2.2.tgz",
       "integrity": "sha512-gxtyfqMg7GKyhQmb056K7M3xszy/myH8w+B4RT+QXBQsvAOdc3XymqDDPHx1BgPgsdAA5SIifona89YtRATDzw=="
     },
+    "node_modules/diff": {
+      "version": "7.0.0",
+      "resolved": "https://registry.npmjs.org/diff/-/diff-7.0.0.tgz",
+      "integrity": "sha512-PJWHUb1RFevKCwaFA9RlG5tCd+FO5iRh9A8HEtkmBH2Li03iJriB6m6JIN4rGz3K3JLawI7/veA1xzRKP6ISBw==",
+      "license": "BSD-3-Clause",
+      "engines": {
+        "node": ">=0.3.1"
+      }
+    },
     "node_modules/diff-match-patch": {
       "version": "1.0.5",
       "resolved": "https://registry.npmjs.org/diff-match-patch/-/diff-match-patch-1.0.5.tgz",
diff --git a/website/package.json b/website/package.json
index 26152b6f..a0c4dc69 100644
--- a/website/package.json
+++ b/website/package.json
@@ -38,6 +38,7 @@
     "@tanstack/react-query": "^5.59.15",
     "@tanstack/react-table": "^8.20.5",
     "@tanstack/react-virtual": "^3.10.9",
+    "@types/diff": "^6.0.0",
     "@types/mime-types": "^2.1.4",
     "@types/react-resizable": "^3.0.8",
     "ai": "^3.4.29",
@@ -48,6 +49,7 @@
     "clsx": "^2.1.1",
     "cmdk": "^1.0.0",
     "css-loader": "^7.1.2",
+    "diff": "^7.0.0",
     "framer-motion": "^11.5.4",
     "gray-matter": "^4.0.3",
     "js-yaml": "^4.1.0",
diff --git a/website/src/app/api/chat/route.ts b/website/src/app/api/chat/route.ts
index 6447b7b0..cac7e49b 100644
--- a/website/src/app/api/chat/route.ts
+++ b/website/src/app/api/chat/route.ts
@@ -6,7 +6,7 @@ export const maxDuration = 60;
 const openai = createOpenAI({
   apiKey: process.env.OPENAI_API_KEY,
   baseURL: process.env.OPENAI_API_BASE,
-  compatibility: 'strict', // strict mode, enable when using the OpenAI API
+  compatibility: "strict", // strict mode, enable when using the OpenAI API
 });
 
 export async function POST(req: Request) {
diff --git a/website/src/app/globals.css b/website/src/app/globals.css
index 6418a0e3..09cd89e5 100644
--- a/website/src/app/globals.css
+++ b/website/src/app/globals.css
@@ -13,11 +13,11 @@
 
 @layer base {
   :root {
-    --background: 211 100% 95%;
+    --background: 211 100% 98%;
     --foreground: 211 5% 0%;
-    --card: 211 50% 90%;
+    --card: 211 50% 95%;
     --card-foreground: 211 5% 10%;
-    --popover: 211 100% 95%;
+    --popover: 211 100% 98%;
     --popover-foreground: 211 100% 0%;
     --primary: 211 100% 50%;
     --primary-foreground: 0 0% 100%;
diff --git a/website/src/app/page.tsx b/website/src/app/page.tsx
index 71d1c0d6..576ca3de 100644
--- a/website/src/app/page.tsx
+++ b/website/src/app/page.tsx
@@ -27,16 +27,16 @@ export default function Home() {
 
         <div className="max-w-lg mx-auto">
           <p className="text-sm sm:text-md mb-1 text-gray-600">
-            <em>New NotebookLM Podcast!</em>{" "}
+            <em>New IDE Released!</em>{" "}
             <a
-              href="https://notebooklm.google.com/notebook/ef73248b-5a43-49cd-9976-432d20f9fa4f/audio?pli=1"
+              href="https://ucbepic.github.io/docetl/playground/"
               target="_blank"
               rel="noopener noreferrer"
               className="text-blue-500 underline"
             >
-              Sept 28, 2024
+              Dec 2, 2024
             </a>
-            . Thanks to Shabie from our Discord community!
+            ! Try out our new web-based IDE.
           </p>
           <p className="text-sm sm:text-md mb-6 text-gray-600">
             <em>New blog post!</em>{" "}
diff --git a/website/src/app/playground/page.tsx b/website/src/app/playground/page.tsx
index 69da2708..5d73b37c 100644
--- a/website/src/app/playground/page.tsx
+++ b/website/src/app/playground/page.tsx
@@ -256,7 +256,7 @@ const CodeEditorPipelineApp: React.FC = () => {
     <BookmarkProvider>
       <div className="h-screen flex flex-col bg-gray-50">
         <div className="p-1 flex justify-between items-center border-b">
-          <div className="flex-1 flex">
+          <div className="flex-1 flex items-center">
             <Menubar className="border-none bg-transparent shadow-none">
               <MenubarMenu>
                 <MenubarTrigger>File</MenubarTrigger>
@@ -316,7 +316,7 @@ const CodeEditorPipelineApp: React.FC = () => {
                 <TooltipTrigger asChild>
                   <Button
                     variant="ghost"
-                    size="icon"
+                    size="sm"
                     onClick={() => {
                       saveProgress();
                       toast({
@@ -325,7 +325,7 @@ const CodeEditorPipelineApp: React.FC = () => {
                         duration: 3000,
                       });
                     }}
-                    className={`relative ${
+                    className={`relative h-8 px-2 ${
                       unsavedChanges ? "border-orange-500" : ""
                     }`}
                   >
diff --git a/website/src/app/types.ts b/website/src/app/types.ts
index f3873789..9d86ade0 100644
--- a/website/src/app/types.ts
+++ b/website/src/app/types.ts
@@ -52,25 +52,26 @@ export interface SchemaItem {
 export interface UserNote {
   id: string;
   note: string;
+  metadata: {
+    columnId?: string;
+    rowIndex?: number;
+    mainColumnValue?: unknown;
+    rowContent?: Record<string, unknown>;
+    operationName?: string;
+  };
 }
 
 export interface Bookmark {
   id: string;
-  text: string;
-  source: string;
   color: string;
   notes: UserNote[];
 }
 
 export interface BookmarkContextType {
   bookmarks: Bookmark[];
-  addBookmark: (
-    text: string,
-    source: string,
-    color: string,
-    notes: UserNote[]
-  ) => void;
+  addBookmark: (color: string, notes: UserNote[]) => void;
   removeBookmark: (id: string) => void;
+  getNotesForRowAndColumn: (rowIndex: number, columnId: string) => UserNote[];
 }
 
 export interface OutputType {
diff --git a/website/src/components/BookmarkableText.tsx b/website/src/components/BookmarkableText.tsx
deleted file mode 100644
index e87c8e93..00000000
--- a/website/src/components/BookmarkableText.tsx
+++ /dev/null
@@ -1,290 +0,0 @@
-import React, { useState, useRef, useEffect } from "react";
-import { Button } from "@/components/ui/button";
-import {
-  Popover,
-  PopoverContent,
-  PopoverTrigger,
-} from "@/components/ui/popover";
-import { Bookmark, BookmarkPlus, X } from "lucide-react";
-import { useBookmarkContext } from "@/contexts/BookmarkContext";
-import { Textarea } from "@/components/ui/textarea";
-import {
-  Select,
-  SelectContent,
-  SelectItem,
-  SelectTrigger,
-  SelectValue,
-} from "@/components/ui/select";
-import { UserNote } from "@/app/types";
-import {
-  Form,
-  FormControl,
-  FormField,
-  FormItem,
-  FormLabel,
-} from "@/components/ui/form";
-import { useForm } from "react-hook-form";
-import { zodResolver } from "@hookform/resolvers/zod";
-import * as z from "zod";
-import { useToast } from "@/hooks/use-toast";
-
-interface BookmarkableTextProps {
-  children: React.ReactNode;
-  source: string;
-  className?: string;
-}
-
-const formSchema = z.object({
-  editedText: z.string().min(1, "Edited text is required"),
-  color: z.string(),
-  note: z.string(),
-});
-
-const BookmarkableText: React.FC<BookmarkableTextProps> = ({
-  children,
-  source,
-  className = "overflow-y-auto",
-}) => {
-  const [buttonPosition, setButtonPosition] = useState({ x: 0, y: 0 });
-  const [showButton, setShowButton] = useState(false);
-  const [isPopoverOpen, setIsPopoverOpen] = useState(false);
-  const textRef = useRef<HTMLDivElement>(null);
-  const buttonRef = useRef<HTMLButtonElement>(null);
-  const popoverRef = useRef<HTMLDivElement>(null);
-  const { addBookmark } = useBookmarkContext();
-  const { toast } = useToast();
-
-  const form = useForm<z.infer<typeof formSchema>>({
-    resolver: zodResolver(formSchema),
-    defaultValues: {
-      editedText: "",
-      color: "#FF0000",
-      note: "",
-    },
-  });
-
-  const handleBookmark = (values: z.infer<typeof formSchema>) => {
-    const userNotes: UserNote[] = [{ id: "default", note: values.note }];
-    addBookmark(values.editedText, source, values.color, userNotes);
-    setShowButton(false);
-    setIsPopoverOpen(false);
-    toast({
-      title: "Bookmark Added",
-      description: "Your bookmark has been successfully added.",
-    });
-  };
-
-  // Listen for selection changes
-  useEffect(() => {
-    const handleSelectionChange = () => {
-      if (isPopoverOpen) return;
-
-      const selection = window.getSelection();
-      if (!selection || selection.isCollapsed || !selection.toString().trim()) {
-        setShowButton(false);
-      }
-    };
-
-    document.addEventListener("selectionchange", handleSelectionChange);
-    document.addEventListener("mousedown", handleSelectionChange);
-
-    return () => {
-      document.removeEventListener("selectionchange", handleSelectionChange);
-      document.removeEventListener("mousedown", handleSelectionChange);
-    };
-  }, [isPopoverOpen]);
-
-  const handleMultiElementSelection = (
-    event: React.MouseEvent | React.TouchEvent
-  ) => {
-    if (isPopoverOpen) return;
-
-    const selection = window.getSelection();
-    const text = selection?.toString().trim();
-
-    if (!selection || !text) {
-      setShowButton(false);
-      return;
-    }
-
-    const range = selection.getRangeAt(0);
-    const rect = range.getBoundingClientRect();
-
-    form.setValue("editedText", text);
-    setButtonPosition({
-      x: rect.left + rect.width / 2,
-      y: rect.top,
-    });
-    setShowButton(true);
-  };
-
-  const handlePopoverOpenChange = (open: boolean) => {
-    setIsPopoverOpen(open);
-    if (!open) {
-      const selection = window.getSelection();
-      if (!selection || selection.isCollapsed) {
-        setShowButton(false);
-      }
-    }
-  };
-
-  const handleClosePopover = () => {
-    setShowButton(false);
-    setIsPopoverOpen(false);
-  };
-
-  return (
-    <div
-      ref={textRef}
-      onMouseUp={handleMultiElementSelection}
-      onTouchEnd={handleMultiElementSelection}
-      className={className}
-    >
-      {children}
-      {showButton && (
-        <Popover
-          open={isPopoverOpen}
-          onOpenChange={handlePopoverOpenChange}
-          modal={true}
-        >
-          <PopoverTrigger asChild>
-            <Button
-              size="icon"
-              ref={buttonRef}
-              variant="default"
-              aria-label="Bookmark"
-              className="shadow-md"
-              style={{
-                position: "absolute",
-                left: `${buttonPosition.x}px`,
-                top: `${buttonPosition.y}px`,
-                transform: "translate(-50%, -100%)",
-              }}
-              onClick={(e) => {
-                e.stopPropagation();
-                setIsPopoverOpen(true);
-              }}
-            >
-              <BookmarkPlus />
-            </Button>
-          </PopoverTrigger>
-          <PopoverContent
-            className="w-[300px]"
-            ref={popoverRef}
-            onInteractOutside={(e) => {
-              if (!buttonRef.current?.contains(e.target as Node)) {
-                e.preventDefault();
-              }
-            }}
-            side="top"
-            align="start"
-            sideOffset={5}
-          >
-            <div className="flex justify-between items-center mb-2">
-              <div className="flex items-center">
-                <Bookmark className="h-5 w-5 text-primary mr-2" />
-                <h4 className="font-medium">Edit Bookmark</h4>
-              </div>
-              <Button
-                variant="ghost"
-                size="icon"
-                onClick={handleClosePopover}
-                className="h-6 w-6"
-              >
-                <X className="h-4 w-4" />
-              </Button>
-            </div>
-            <Form {...form}>
-              <form onSubmit={form.handleSubmit(handleBookmark)}>
-                <FormField
-                  control={form.control}
-                  name="editedText"
-                  render={({ field }) => (
-                    <FormItem>
-                      <FormControl>
-                        <Textarea {...field} className="min-h-[100px] mb-2" />
-                      </FormControl>
-                    </FormItem>
-                  )}
-                />
-                <FormField
-                  control={form.control}
-                  name="color"
-                  render={({ field }) => (
-                    <FormItem>
-                      <Select
-                        onValueChange={field.onChange}
-                        defaultValue={field.value}
-                      >
-                        <FormControl>
-                          <SelectTrigger className="w-full">
-                            <SelectValue placeholder="Select a color" />
-                          </SelectTrigger>
-                        </FormControl>
-                        <SelectContent>
-                          <SelectItem value="#FF0000">
-                            <div className="flex items-center">
-                              <div className="w-4 h-4 rounded-full bg-[#FF0000] mr-2"></div>
-                              Red
-                            </div>
-                          </SelectItem>
-                          <SelectItem value="#00FF00">
-                            <div className="flex items-center">
-                              <div className="w-4 h-4 rounded-full bg-[#00FF00] mr-2"></div>
-                              Green
-                            </div>
-                          </SelectItem>
-                          <SelectItem value="#0000FF">
-                            <div className="flex items-center">
-                              <div className="w-4 h-4 rounded-full bg-[#0000FF] mr-2"></div>
-                              Blue
-                            </div>
-                          </SelectItem>
-                          <SelectItem value="#FFFF00">
-                            <div className="flex items-center">
-                              <div className="w-4 h-4 rounded-full bg-[#FFFF00] mr-2"></div>
-                              Yellow
-                            </div>
-                          </SelectItem>
-                          <SelectItem value="#FF00FF">
-                            <div className="flex items-center">
-                              <div className="w-4 h-4 rounded-full bg-[#FF00FF] mr-2"></div>
-                              Magenta
-                            </div>
-                          </SelectItem>
-                          <SelectItem value="#00FFFF">
-                            <div className="flex items-center">
-                              <div className="w-4 h-4 rounded-full bg-[#00FFFF] mr-2"></div>
-                              Cyan
-                            </div>
-                          </SelectItem>
-                        </SelectContent>
-                      </Select>
-                    </FormItem>
-                  )}
-                />
-                <FormField
-                  control={form.control}
-                  name="note"
-                  render={({ field }) => (
-                    <FormItem>
-                      <FormLabel>Notes:</FormLabel>
-                      <FormControl>
-                        <Textarea {...field} className="min-h-[50px] mb-1" />
-                      </FormControl>
-                    </FormItem>
-                  )}
-                />
-                <Button type="submit" className="w-full mt-2">
-                  Add Bookmark
-                </Button>
-              </form>
-            </Form>
-          </PopoverContent>
-        </Popover>
-      )}
-    </div>
-  );
-};
-
-export default BookmarkableText;
diff --git a/website/src/components/BookmarksPanel.tsx b/website/src/components/BookmarksPanel.tsx
index d6c7992b..4c5e63c1 100644
--- a/website/src/components/BookmarksPanel.tsx
+++ b/website/src/components/BookmarksPanel.tsx
@@ -8,6 +8,8 @@ import {
   Filter,
   Trash2,
   X,
+  MessageSquare,
+  Maximize2,
 } from "lucide-react";
 import {
   Select,
@@ -22,23 +24,26 @@ import {
   PopoverContent,
   PopoverTrigger,
 } from "@/components/ui/popover";
+import { UserNote } from "@/app/types";
 
 const BookmarksPanel: React.FC = () => {
   const { bookmarks, removeBookmark } = useBookmarkContext();
   const [searchTerm, setSearchTerm] = useState("");
   const [expandedBookmarkId, setExpandedBookmarkId] = useState<string | null>(
-    null,
+    null
   );
   const [selectedColor, setSelectedColor] = useState<string | "all">("all");
 
   const uniqueColors = Array.from(
-    new Set(bookmarks.map((bookmark) => bookmark.color)),
+    new Set(bookmarks.map((bookmark) => bookmark.color))
   );
 
   const filteredBookmarks = bookmarks.filter(
     (bookmark) =>
-      bookmark.text.toLowerCase().includes(searchTerm.toLowerCase()) &&
-      (selectedColor === "all" || bookmark.color === selectedColor),
+      bookmark.notes.some((note) =>
+        note.note.toLowerCase().includes(searchTerm.toLowerCase())
+      ) &&
+      (selectedColor === "all" || bookmark.color === selectedColor)
   );
 
   const toggleBookmarkExpansion = (id: string) => {
@@ -59,12 +64,36 @@ const BookmarksPanel: React.FC = () => {
     setSelectedColor("all");
   };
 
+  const renderNoteContent = (note: UserNote) => {
+    return (
+      <div className="mt-2 space-y-1">
+        {note.metadata?.columnId && (
+          <div className="flex items-center gap-2 text-xs text-muted-foreground">
+            <MessageSquare size={12} />
+            <span>Column: {note.metadata.columnId}</span>
+            {note.metadata.rowIndex !== undefined && (
+              <span>• Row: {note.metadata.rowIndex}</span>
+            )}
+          </div>
+        )}
+        {note.metadata?.rowContent && (
+          <div className="text-xs bg-muted/50 p-2 rounded-md mt-1">
+            <div className="font-semibold mb-1">Row Context:</div>
+            <pre className="overflow-x-auto">
+              {JSON.stringify(note.metadata.rowContent, null, 2)}
+            </pre>
+          </div>
+        )}
+      </div>
+    );
+  };
+
   return (
     <div className="h-full p-4 flex flex-col">
       <div className="flex justify-between items-center mb-2">
         <h2 className="text-sm font-bold flex items-center uppercase">
           <Bookmark className="mr-2" size={16} />
-          Notes
+          Notes & Feedback
         </h2>
         <Button
           variant="ghost"
@@ -76,6 +105,14 @@ const BookmarksPanel: React.FC = () => {
           Clear All
         </Button>
       </div>
+      <div className="text-xs mb-4 bg-muted/50 p-2 rounded-md flex items-center gap-2">
+        <span className="text-muted-foreground">Tip:</span>
+        <span>
+          Click{" "}
+          <Maximize2 className="h-3 w-3 inline-block mx-0.5 text-primary" /> in
+          any output column to leave feedback
+        </span>
+      </div>
       <div className="flex mb-2">
         <Input
           type="text"
@@ -90,69 +127,81 @@ const BookmarksPanel: React.FC = () => {
               <Filter className="h-4 w-4" />
               {selectedColor !== "all" && (
                 <span className="absolute -top-2 -right-2 bg-primary text-white text-xs rounded-full w-5 h-5 flex items-center justify-center">
-                  1
+                  {1}
                 </span>
               )}
             </Button>
           </PopoverTrigger>
           <PopoverContent className="w-48">
-            <Select
-              value={selectedColor}
-              onValueChange={(value) => setSelectedColor(value as string)}
-            >
-              <SelectTrigger className="border-none shadow-none">
-                <SelectValue placeholder="Filter by color" />
-              </SelectTrigger>
-              <SelectContent>
-                <SelectItem value="all">All colors</SelectItem>
-                {uniqueColors.map((color) => (
-                  <SelectItem key={color} value={color}>
-                    <div className="flex items-center">
-                      <div
-                        className="w-3 h-3 rounded-full mr-2"
-                        style={{ backgroundColor: color }}
-                      />
-                    </div>
-                  </SelectItem>
-                ))}
-              </SelectContent>
-            </Select>
+            <div className="space-y-2">
+              <Select
+                value={selectedColor}
+                onValueChange={(value) => setSelectedColor(value as string)}
+              >
+                <SelectTrigger className="border-none shadow-none">
+                  <SelectValue placeholder="Filter by color" />
+                </SelectTrigger>
+                <SelectContent>
+                  <SelectItem value="all">All colors</SelectItem>
+                  {uniqueColors.map((color) => (
+                    <SelectItem key={color} value={color}>
+                      <div className="flex items-center">
+                        <div
+                          className="w-3 h-3 rounded-full mr-2"
+                          style={{ backgroundColor: color }}
+                        />
+                      </div>
+                    </SelectItem>
+                  ))}
+                </SelectContent>
+              </Select>
+            </div>
           </PopoverContent>
         </Popover>
       </div>
       <div className="overflow-y-auto flex-grow">
         {filteredBookmarks.map((bookmark) => (
-          <div key={bookmark.id} className="mb-2">
+          <div
+            key={bookmark.id}
+            className="mb-1 rounded hover:bg-muted/50 transition-colors"
+          >
             <div
-              className="flex items-start cursor-pointer"
+              className="flex items-start cursor-pointer px-1.5 py-1"
               onClick={() => toggleBookmarkExpansion(bookmark.id)}
             >
               <div
-                className="w-3 h-3 rounded-full mr-2 mt-1"
+                className="w-3 h-3 rounded-full mr-2 mt-1 flex-shrink-0"
                 style={{
                   backgroundColor: bookmark.color,
                   minWidth: "0.75rem",
                   minHeight: "0.75rem",
                 }}
               />
-              <span
-                className={`flex-grow ${expandedBookmarkId === bookmark.id ? "whitespace-normal" : "whitespace-nowrap overflow-hidden text-ellipsis"}`}
-              >
-                {bookmark.text}
-              </span>
-              {expandedBookmarkId === bookmark.id ? (
-                <ChevronUp className="ml-2 min-w-4 min-h-4 mt-1" size={16} />
-              ) : (
-                <ChevronDown className="ml-2 min-w-4 min-h-4 mt-1" size={16} />
-              )}
+              <div className="flex-grow min-w-0 mr-1">
+                <div
+                  className={`text-sm italic ${
+                    expandedBookmarkId === bookmark.id
+                      ? "whitespace-normal"
+                      : "whitespace-nowrap overflow-hidden text-ellipsis"
+                  }`}
+                >
+                  &quot;{bookmark.notes[0]?.note || "No notes"}&quot;
+                </div>
+              </div>
+              <div className="flex-shrink-0">
+                {expandedBookmarkId === bookmark.id ? (
+                  <ChevronUp className="text-muted-foreground" size={16} />
+                ) : (
+                  <ChevronDown className="text-muted-foreground" size={16} />
+                )}
+              </div>
             </div>
             {expandedBookmarkId === bookmark.id && (
-              <div className="mt-2 ml-5 text-sm text-gray-600">
-                <p>Source: {bookmark.source}</p>
+              <div className="mt-2 ml-5">
                 {bookmark.notes.map((note, index) => (
-                  <p key={index}>
-                    Note {index + 1}: {note.note}
-                  </p>
+                  <div key={note.id || index} className="border-l-2 pl-2 mb-2">
+                    {renderNoteContent(note)}
+                  </div>
                 ))}
                 <Button
                   variant="destructive"
diff --git a/website/src/components/ColumnDialog.tsx b/website/src/components/ColumnDialog.tsx
new file mode 100644
index 00000000..e26d278e
--- /dev/null
+++ b/website/src/components/ColumnDialog.tsx
@@ -0,0 +1,545 @@
+import React, { useState, useEffect, useCallback } from "react";
+import { Dialog, DialogContent } from "@/components/ui/dialog";
+import { Button } from "@/components/ui/button";
+import {
+  ResizableHandle,
+  ResizablePanel,
+  ResizablePanelGroup,
+} from "@/components/ui/resizable";
+import { SearchableCell } from "@/components/SearchableCell";
+import { PrettyJSON } from "@/components/PrettyJSON";
+import { RowNavigator } from "@/components/RowNavigator";
+import { ChevronDown, Eye, ChevronLeft, ChevronRight } from "lucide-react";
+import {
+  HoverCard,
+  HoverCardContent,
+  HoverCardTrigger,
+} from "@/components/ui/hover-card";
+import { useBookmarkContext } from "@/contexts/BookmarkContext";
+import { Textarea } from "@/components/ui/textarea";
+import { UserNote } from "@/app/types";
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from "@/components/ui/select";
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipTrigger,
+  TooltipProvider,
+} from "@/components/ui/tooltip";
+
+interface ObservabilityIndicatorProps {
+  row: Record<string, unknown>;
+  currentOperation: string;
+}
+
+const ObservabilityIndicator = React.memo(
+  ({ row, currentOperation }: ObservabilityIndicatorProps) => {
+    const observabilityEntries = Object.entries(row).filter(
+      ([key]) => key === `_observability_${currentOperation}`
+    );
+
+    if (observabilityEntries.length === 0) return null;
+
+    return (
+      <HoverCard>
+        <HoverCardTrigger asChild>
+          <div className="cursor-help">
+            <Eye className="h-4 w-4 text-muted-foreground hover:text-primary" />
+          </div>
+        </HoverCardTrigger>
+        <HoverCardContent
+          className="w-[800px] max-h-[600px] overflow-auto"
+          side="right"
+          align="start"
+        >
+          <div className="space-y-4">
+            <h3 className="text-lg font-semibold border-b pb-2">
+              LLM Call(s) for {currentOperation}
+            </h3>
+            <div className="space-y-2">
+              {observabilityEntries.map(([key, value]) => (
+                <div key={key} className="flex flex-col gap-1">
+                  <div className="text-sm text-muted-foreground">
+                    {typeof value === "object"
+                      ? JSON.stringify(value, null, 2)
+                      : String(value)}
+                  </div>
+                </div>
+              ))}
+            </div>
+          </div>
+        </HoverCardContent>
+      </HoverCard>
+    );
+  }
+);
+ObservabilityIndicator.displayName = "ObservabilityIndicator";
+
+export interface ColumnDialogProps<T extends Record<string, unknown>> {
+  isOpen: boolean;
+  onClose: () => void;
+  columnId: string;
+  columnHeader: string;
+  data: T[];
+  currentIndex: number;
+  onNavigate: (direction: "prev" | "next") => void;
+  onJumpToRow: (index: number) => void;
+  currentOperation: string;
+}
+
+export function ColumnDialog<T extends Record<string, unknown>>({
+  isOpen,
+  onClose,
+  columnId,
+  columnHeader,
+  data,
+  currentIndex,
+  onNavigate,
+  onJumpToRow,
+  currentOperation,
+}: ColumnDialogProps<T>) {
+  const [splitView, setSplitView] = useState(false);
+  const [compareIndex, setCompareIndex] = useState<number | null>(null);
+  const [expandedFields, setExpandedFields] = useState<string[]>([]);
+  const [isAllExpanded, setIsAllExpanded] = useState(false);
+  const [feedbackColor, setFeedbackColor] = useState("#FF0000");
+  const [showPreviousNotes, setShowPreviousNotes] = useState(false);
+
+  const currentRow = data[currentIndex];
+  const compareRow = compareIndex !== null ? data[compareIndex] : null;
+  const currentValue = currentRow[columnId];
+  const compareValue = compareRow?.[columnId];
+
+  const otherFields = Object.entries(currentRow)
+    .filter(([key]) => key !== columnId && !key.startsWith("_"))
+    .sort((a, b) => a[0].localeCompare(b[0]));
+
+  const toggleField = (fieldKey: string) => {
+    setExpandedFields((prev) =>
+      prev.includes(fieldKey)
+        ? prev.filter((k) => k !== fieldKey)
+        : [...prev, fieldKey]
+    );
+  };
+
+  const toggleAllFields = () => {
+    if (isAllExpanded) {
+      setExpandedFields([]);
+    } else {
+      setExpandedFields(otherFields.map(([key]) => key));
+    }
+    setIsAllExpanded(!isAllExpanded);
+  };
+
+  const renderContent = (value: unknown) => {
+    if (value === null || value === undefined) {
+      return <span className="text-muted-foreground">No value</span>;
+    }
+
+    if (typeof value === "object") {
+      return (
+        <SearchableCell
+          content={JSON.stringify(value, null, 2)}
+          isResizing={false}
+        >
+          {(searchTerm) => (searchTerm ? null : <PrettyJSON data={value} />)}
+        </SearchableCell>
+      );
+    }
+
+    if (typeof value === "string") {
+      return <SearchableCell content={value} isResizing={false} />;
+    }
+
+    return String(value);
+  };
+
+  const handleKeyDown = useCallback(
+    (e: KeyboardEvent) => {
+      if (
+        e.target instanceof HTMLTextAreaElement ||
+        e.target instanceof HTMLInputElement
+      ) {
+        return; // Don't handle shortcuts when typing in inputs
+      }
+
+      if (e.key === "ArrowLeft") {
+        onNavigate("prev");
+      } else if (e.key === "ArrowRight") {
+        onNavigate("next");
+      }
+    },
+    [onNavigate]
+  );
+
+  useEffect(() => {
+    window.addEventListener("keydown", handleKeyDown);
+    return () => window.removeEventListener("keydown", handleKeyDown);
+  }, [handleKeyDown]);
+
+  const renderRowContent = (row: T | null, value: unknown) => {
+    if (!row) return null;
+    const { addBookmark, getNotesForRowAndColumn } = useBookmarkContext();
+
+    const handleSubmitFeedback = (feedbackText: string) => {
+      if (!feedbackText.trim()) return;
+
+      const filteredRowContent = Object.fromEntries(
+        Object.entries(row).filter(([key]) => !key.startsWith("_observability"))
+      );
+
+      const feedback: UserNote[] = [
+        {
+          id: Date.now().toString(),
+          note: feedbackText,
+          metadata: {
+            columnId,
+            rowIndex: currentIndex,
+            mainColumnValue: row[columnId],
+            rowContent: filteredRowContent,
+            operationName: currentOperation,
+          },
+        },
+      ];
+
+      addBookmark(feedbackColor, feedback);
+    };
+
+    const otherFields = Object.entries(row)
+      .filter(([key]) => key !== columnId && !key.startsWith("_"))
+      .sort((a, b) => a[0].localeCompare(b[0]));
+
+    const existingNotes = getNotesForRowAndColumn(currentIndex, columnId);
+
+    return (
+      <ResizablePanelGroup direction="horizontal" className="h-full">
+        <ResizablePanel defaultSize={75} minSize={30}>
+          <div className="h-full flex">
+            <TooltipProvider>
+              <Tooltip>
+                <TooltipTrigger asChild>
+                  <div className="flex-none flex items-stretch w-8">
+                    <Button
+                      variant="ghost"
+                      onClick={() => onNavigate("prev")}
+                      className="h-full w-8 rounded-none hover:bg-muted/20 flex items-center justify-center bg-muted/5"
+                      disabled={currentIndex === 0}
+                      aria-label="Previous example (Left arrow key)"
+                    >
+                      <ChevronLeft className="h-12 w-12 absolute" />
+                    </Button>
+                  </div>
+                </TooltipTrigger>
+                <TooltipContent side="right">
+                  <p>Previous example</p>
+                  <p className="text-xs text-muted-foreground">
+                    Left arrow key
+                  </p>
+                </TooltipContent>
+              </Tooltip>
+
+              <div className="flex-1 overflow-auto">
+                <div className="px-4 py-2">{renderContent(value)}</div>
+              </div>
+
+              <Tooltip>
+                <TooltipTrigger asChild>
+                  <div className="flex-none flex items-stretch w-8">
+                    <Button
+                      variant="ghost"
+                      onClick={() => onNavigate("next")}
+                      className="h-full w-8 rounded-none hover:bg-muted/20 flex items-center justify-center bg-muted/5"
+                      disabled={currentIndex === data.length - 1}
+                      aria-label="Next example (Right arrow key)"
+                    >
+                      <ChevronRight className="h-12 w-12 absolute" />
+                    </Button>
+                  </div>
+                </TooltipTrigger>
+                <TooltipContent side="left">
+                  <p>Next example</p>
+                  <p className="text-xs text-muted-foreground">
+                    Right arrow key
+                  </p>
+                </TooltipContent>
+              </Tooltip>
+            </TooltipProvider>
+          </div>
+        </ResizablePanel>
+        <ResizableHandle withHandle />
+        <ResizablePanel defaultSize={25} minSize={20}>
+          <ResizablePanelGroup direction="vertical">
+            <ResizablePanel defaultSize={60} minSize={30}>
+              <div className="h-full overflow-auto bg-muted/10">
+                <div className="sticky top-0 bg-background/95 backdrop-blur-sm z-10 p-2 border-b">
+                  <div className="flex items-center justify-between">
+                    <span className="font-medium text-sm">Other Keys</span>
+                    <Button
+                      variant="ghost"
+                      size="sm"
+                      onClick={toggleAllFields}
+                      className="h-7 px-2"
+                    >
+                      <ChevronDown
+                        className={`h-4 w-4 transition-transform ${
+                          isAllExpanded ? "rotate-180" : ""
+                        }`}
+                      />
+                      <span className="ml-1 text-xs">
+                        {isAllExpanded ? "Collapse" : "Expand"}
+                      </span>
+                    </Button>
+                  </div>
+                </div>
+                <div className="divide-y divide-border">
+                  {otherFields.map(([key, value]) => (
+                    <div key={key} className="bg-background/50">
+                      <button
+                        onClick={() => toggleField(key)}
+                        className="w-full text-left p-3 hover:bg-muted/30 flex items-center justify-between"
+                      >
+                        <span className="font-medium text-sm truncate">
+                          {key}
+                        </span>
+                        <ChevronDown
+                          className={`h-4 w-4 flex-shrink-0 transition-transform ${
+                            expandedFields.includes(key) ? "rotate-180" : ""
+                          }`}
+                        />
+                      </button>
+                      {expandedFields.includes(key) && (
+                        <div className="p-3 bg-muted/10">
+                          {renderContent(value)}
+                        </div>
+                      )}
+                    </div>
+                  ))}
+                </div>
+              </div>
+            </ResizablePanel>
+            <ResizableHandle withHandle />
+            <ResizablePanel defaultSize={40} minSize={20}>
+              <div className="h-full bg-muted/10 flex flex-col">
+                <div className="flex-none p-2 border-b">
+                  <h3 className="text-sm font-medium">Feedback</h3>
+                  <p className="text-xs text-muted-foreground">
+                    Share your thoughts on this output
+                  </p>
+                </div>
+                <div className="flex-1 overflow-auto p-2">
+                  {existingNotes.length > 0 && (
+                    <div className="mb-2">
+                      <button
+                        onClick={() => setShowPreviousNotes(!showPreviousNotes)}
+                        className="flex items-center justify-between w-full text-sm font-medium text-muted-foreground hover:text-foreground mb-2"
+                      >
+                        <span>Previous Notes ({existingNotes.length})</span>
+                        <ChevronDown
+                          className={`h-4 w-4 transition-transform ${
+                            showPreviousNotes ? "rotate-180" : ""
+                          }`}
+                        />
+                      </button>
+                      {showPreviousNotes && (
+                        <div className="space-y-2 mb-2">
+                          {existingNotes.map((note) => (
+                            <div
+                              key={note.id}
+                              className="p-2 rounded-lg bg-muted/20 text-sm"
+                            >
+                              <span className="italic">"{note.note}"</span>
+                            </div>
+                          ))}
+                        </div>
+                      )}
+                    </div>
+                  )}
+                  <div className="flex-none">
+                    <Textarea
+                      placeholder="e.g., I don't like how short these outputs are"
+                      className="min-h-[80px] bg-transparent resize-none"
+                    />
+                  </div>
+                  <div className="flex-none mt-2">
+                    <Select
+                      value={feedbackColor}
+                      onValueChange={setFeedbackColor}
+                    >
+                      <SelectTrigger className="w-full">
+                        <SelectValue>
+                          <div className="flex items-center">
+                            <div
+                              className="w-4 h-4 rounded-full mr-2"
+                              style={{ backgroundColor: feedbackColor }}
+                            />
+                            {feedbackColor === "#FF0000"
+                              ? "Red"
+                              : feedbackColor === "#00FF00"
+                              ? "Green"
+                              : feedbackColor === "#0000FF"
+                              ? "Blue"
+                              : feedbackColor === "#FFFF00"
+                              ? "Yellow"
+                              : feedbackColor === "#FF00FF"
+                              ? "Magenta"
+                              : "Cyan"}
+                          </div>
+                        </SelectValue>
+                      </SelectTrigger>
+                      <SelectContent>
+                        <SelectItem value="#FF0000">
+                          <div className="flex items-center">
+                            <div className="w-4 h-4 rounded-full bg-[#FF0000] mr-2" />
+                            Red
+                          </div>
+                        </SelectItem>
+                        <SelectItem value="#00FF00">
+                          <div className="flex items-center">
+                            <div className="w-4 h-4 rounded-full bg-[#00FF00] mr-2" />
+                            Green
+                          </div>
+                        </SelectItem>
+                        <SelectItem value="#0000FF">
+                          <div className="flex items-center">
+                            <div className="w-4 h-4 rounded-full bg-[#0000FF] mr-2" />
+                            Blue
+                          </div>
+                        </SelectItem>
+                        <SelectItem value="#FFFF00">
+                          <div className="flex items-center">
+                            <div className="w-4 h-4 rounded-full bg-[#FFFF00] mr-2" />
+                            Yellow
+                          </div>
+                        </SelectItem>
+                        <SelectItem value="#FF00FF">
+                          <div className="flex items-center">
+                            <div className="w-4 h-4 rounded-full bg-[#FF00FF] mr-2" />
+                            Magenta
+                          </div>
+                        </SelectItem>
+                        <SelectItem value="#00FFFF">
+                          <div className="flex items-center">
+                            <div className="w-4 h-4 rounded-full bg-[#00FFFF] mr-2" />
+                            Cyan
+                          </div>
+                        </SelectItem>
+                      </SelectContent>
+                    </Select>
+                  </div>
+                  <div className="flex-none mt-2">
+                    <Button
+                      className="w-full"
+                      onClick={(e) => {
+                        const textarea = (e.target as HTMLElement)
+                          .closest(".flex-col")
+                          ?.querySelector("textarea");
+                        if (textarea) {
+                          handleSubmitFeedback(textarea.value);
+                          textarea.value = "";
+                        }
+                      }}
+                    >
+                      Submit Feedback
+                    </Button>
+                  </div>
+                </div>
+              </div>
+            </ResizablePanel>
+          </ResizablePanelGroup>
+        </ResizablePanel>
+      </ResizablePanelGroup>
+    );
+  };
+
+  return (
+    <Dialog open={isOpen} onOpenChange={() => onClose()}>
+      <DialogContent className="max-w-[95vw] max-h-[95vh] w-full h-full flex flex-col p-0 bg-background rounded-lg overflow-hidden">
+        <div className="flex flex-col h-full">
+          <div className="flex-none flex items-center justify-between px-6 py-3 border-b bg-background">
+            <div className="flex items-center gap-4 flex-1">
+              <h2 className="text-lg font-semibold">{columnHeader}</h2>
+              <div className="flex items-center gap-4">
+                <ObservabilityIndicator
+                  row={currentRow}
+                  currentOperation={currentOperation}
+                />
+                <RowNavigator
+                  currentRow={currentIndex}
+                  totalRows={data.length}
+                  onNavigate={onNavigate}
+                  onJumpToRow={onJumpToRow}
+                />
+                <span className="text-sm text-muted-foreground">
+                  Use ← → arrow keys to navigate
+                </span>
+                <Button
+                  variant="outline"
+                  size="sm"
+                  onClick={() => {
+                    setSplitView(!splitView);
+                    if (!splitView && compareIndex === null) {
+                      setCompareIndex(
+                        Math.min(currentIndex + 1, data.length - 1)
+                      );
+                    }
+                  }}
+                >
+                  {splitView ? "Single View" : "Split View"}
+                </Button>
+              </div>
+            </div>
+          </div>
+
+          <div className="flex-1 min-h-0">
+            {splitView ? (
+              <ResizablePanelGroup direction="vertical" className="h-full">
+                <ResizablePanel defaultSize={50} minSize={20}>
+                  {renderRowContent(currentRow, currentValue)}
+                </ResizablePanel>
+                <ResizableHandle withHandle />
+                <ResizablePanel defaultSize={50} minSize={20}>
+                  <div className="flex flex-col h-full">
+                    <div className="flex-none flex items-center px-6 py-1.5 border-b bg-muted/30">
+                      <div className="flex items-center gap-4">
+                        <ObservabilityIndicator
+                          row={compareRow ?? {}}
+                          currentOperation={currentOperation}
+                        />
+                        <RowNavigator
+                          currentRow={compareIndex ?? 0}
+                          totalRows={data.length}
+                          onNavigate={(direction) => {
+                            setCompareIndex((prev) =>
+                              direction === "next"
+                                ? Math.min((prev ?? 0) + 1, data.length - 1)
+                                : Math.max((prev ?? 0) - 1, 0)
+                            );
+                          }}
+                          onJumpToRow={(index) => setCompareIndex(index)}
+                          label="Compare Row"
+                          compact={true}
+                        />
+                      </div>
+                    </div>
+                    <div className="flex-1 min-h-0">
+                      {renderRowContent(compareRow, compareValue)}
+                    </div>
+                  </div>
+                </ResizablePanel>
+              </ResizablePanelGroup>
+            ) : (
+              <div className="h-full">
+                {renderRowContent(currentRow, currentValue)}
+              </div>
+            )}
+          </div>
+        </div>
+      </DialogContent>
+    </Dialog>
+  );
+}
diff --git a/website/src/components/CommandMenu.tsx b/website/src/components/CommandMenu.tsx
deleted file mode 100644
index 1f8789b9..00000000
--- a/website/src/components/CommandMenu.tsx
+++ /dev/null
@@ -1,61 +0,0 @@
-import React from "react";
-import { Dialog, DialogContent } from "@/components/ui/dialog";
-import {
-  Command,
-  CommandInput,
-  CommandList,
-  CommandEmpty,
-  CommandGroup,
-  CommandItem,
-} from "@/components/ui/command";
-import { Bookmark, Search } from "lucide-react";
-import { useBookmarkContext } from "@/contexts/BookmarkContext";
-
-interface CommandMenuProps {
-  isOpen: boolean;
-  onClose: () => void;
-  onBookmarkCurrent: () => void;
-}
-
-const CommandMenu: React.FC<CommandMenuProps> = ({
-  isOpen,
-  onClose,
-  onBookmarkCurrent,
-}) => {
-  const { bookmarks } = useBookmarkContext();
-
-  return (
-    <Dialog open={isOpen} onOpenChange={onClose}>
-      <DialogContent className="p-0">
-        <Command className="rounded-lg border shadow-md">
-          <CommandInput placeholder="Type a command or search..." />
-          <CommandList>
-            <CommandEmpty>No results found.</CommandEmpty>
-            <CommandGroup heading="Actions">
-              <CommandItem onSelect={onBookmarkCurrent}>
-                <Bookmark className="mr-2 h-4 w-4" />
-                <span>Bookmark current selection</span>
-              </CommandItem>
-              {/* Add more actions here */}
-            </CommandGroup>
-            <CommandGroup heading="Bookmarks">
-              {bookmarks.map((bookmark) => (
-                <CommandItem
-                  key={bookmark.id}
-                  onSelect={() =>
-                    console.log("Navigate to bookmark", bookmark.id)
-                  }
-                >
-                  <Search className="mr-2 h-4 w-4" />
-                  <span>{bookmark.text.substring(0, 30)}...</span>
-                </CommandItem>
-              ))}
-            </CommandGroup>
-          </CommandList>
-        </Command>
-      </DialogContent>
-    </Dialog>
-  );
-};
-
-export default CommandMenu;
diff --git a/website/src/components/DatasetView.tsx b/website/src/components/DatasetView.tsx
index 33659a62..d74e4b78 100644
--- a/website/src/components/DatasetView.tsx
+++ b/website/src/components/DatasetView.tsx
@@ -5,7 +5,6 @@ import { useInfiniteQuery } from "@tanstack/react-query";
 import { Input } from "@/components/ui/input";
 import { Button } from "@/components/ui/button";
 import { ChevronUp, ChevronDown, Search } from "lucide-react";
-import BookmarkableText from "@/components/BookmarkableText";
 import { Skeleton } from "@/components/ui/skeleton";
 import { Loader2 } from "lucide-react";
 import {
@@ -507,25 +506,21 @@ const DatasetView: React.FC<{ file: File | null }> = ({ file }) => {
             : "No matches"}
         </span>
       </form>
-      <BookmarkableText source="dataset">
-        <div ref={parentRef} className="flex-grow overflow-y-auto">
-          {lines.map((lineContent, index) => (
-            <div key={index} className="flex">
-              <span className="inline-block w-12 flex-shrink-0 text-gray-500 select-none text-right pr-2">
-                {index + 1}
-              </span>
-              <div className="flex-grow">
-                <pre className="whitespace-pre-wrap break-words font-mono text-sm">
-                  {highlightMatches(lineContent, index)}
-                </pre>
-              </div>
+      <div ref={parentRef} className="flex-grow overflow-y-auto">
+        {lines.map((lineContent, index) => (
+          <div key={index} className="flex">
+            <span className="inline-block w-12 flex-shrink-0 text-gray-500 select-none text-right pr-2">
+              {index + 1}
+            </span>
+            <div className="flex-grow">
+              <pre className="whitespace-pre-wrap break-words font-mono text-sm">
+                {highlightMatches(lineContent, index)}
+              </pre>
             </div>
-          ))}
-          {isFetching && (
-            <div className="text-center py-4">Loading more...</div>
-          )}
-        </div>
-      </BookmarkableText>
+          </div>
+        ))}
+        {isFetching && <div className="text-center py-4">Loading more...</div>}
+      </div>
     </div>
   );
 };
diff --git a/website/src/components/MarkdownCell.tsx b/website/src/components/MarkdownCell.tsx
new file mode 100644
index 00000000..557041d3
--- /dev/null
+++ b/website/src/components/MarkdownCell.tsx
@@ -0,0 +1,93 @@
+import React from "react";
+import ReactMarkdown from "react-markdown";
+
+interface MarkdownCellProps {
+  content: string;
+}
+
+export const MarkdownCell = React.memo(({ content }: MarkdownCellProps) => {
+  return (
+    <ReactMarkdown
+      components={{
+        h1: ({ children }) => (
+          <div style={{ fontWeight: "bold", fontSize: "1.5rem" }}>
+            {children}
+          </div>
+        ),
+        h2: ({ children }) => (
+          <div style={{ fontWeight: "bold", fontSize: "1.25rem" }}>
+            {children}
+          </div>
+        ),
+        h3: ({ children }) => (
+          <div style={{ fontWeight: "bold", fontSize: "1.1rem" }}>
+            {children}
+          </div>
+        ),
+        h4: ({ children }) => (
+          <div style={{ fontWeight: "bold" }}>{children}</div>
+        ),
+        ul: ({ children }) => (
+          <ul
+            style={{
+              listStyleType: "•",
+              paddingLeft: "1rem",
+              margin: "0.25rem 0",
+            }}
+          >
+            {children}
+          </ul>
+        ),
+        ol: ({ children }) => (
+          <ol
+            style={{
+              listStyleType: "decimal",
+              paddingLeft: "1rem",
+              margin: "0.25rem 0",
+            }}
+          >
+            {children}
+          </ol>
+        ),
+        li: ({ children }) => (
+          <li style={{ marginBottom: "0.125rem" }}>{children}</li>
+        ),
+        code: ({
+          className,
+          children,
+          inline,
+          ...props
+        }: {
+          className?: string;
+          children: React.ReactNode;
+          inline?: boolean;
+        }) => {
+          const match = /language-(\w+)/.exec(className || "");
+          return !inline && match ? (
+            <pre className="bg-slate-100 p-2 rounded">
+              <code className={className} {...props}>
+                {children}
+              </code>
+            </pre>
+          ) : (
+            <code className="bg-slate-100 px-1 rounded" {...props}>
+              {children}
+            </code>
+          );
+        },
+        pre: ({ children }) => (
+          <pre className="bg-slate-100 p-2 rounded">{children}</pre>
+        ),
+        blockquote: ({ children }) => (
+          <blockquote className="border-l-4 border-slate-300 pl-4 my-2 italic">
+            {children}
+          </blockquote>
+        ),
+      }}
+    >
+      {content}
+    </ReactMarkdown>
+  );
+});
+
+MarkdownCell.displayName = "MarkdownCell";
diff --git a/website/src/components/OperationCard.tsx b/website/src/components/OperationCard.tsx
index 18f3ab29..99e1a612 100644
--- a/website/src/components/OperationCard.tsx
+++ b/website/src/components/OperationCard.tsx
@@ -55,6 +55,7 @@ import {
   TooltipProvider,
   TooltipTrigger,
 } from "./ui/tooltip";
+import { PromptImprovementDialog } from "@/components/PromptImprovementDialog";
 
 // Separate components
 const OperationHeader: React.FC<{
@@ -75,6 +76,7 @@ const OperationHeader: React.FC<{
   onAIEdit: (instruction: string) => void;
   onToggleExpand: () => void;
   onToggleVisibility: () => void;
+  onImprovePrompt: () => void;
 }> = React.memo(
   ({
     name,
@@ -94,6 +96,7 @@ const OperationHeader: React.FC<{
     onAIEdit,
     onToggleExpand,
     onToggleVisibility,
+    onImprovePrompt,
   }) => {
     const [isEditing, setIsEditing] = useState(false);
     const [editedName, setEditedName] = useState(name);
@@ -197,11 +200,23 @@ const OperationHeader: React.FC<{
             <ListCollapse size={14} className="text-primary" />
             <span className="text-xs text-primary">Show outputs</span>
           </Button>
+          {llmType === "LLM" && (
+            <Button
+              variant="ghost"
+              size="sm"
+              className="flex items-center gap-1 px-2 h-6"
+              disabled={disabled}
+              onClick={onImprovePrompt}
+            >
+              <Wand2 size={14} className="text-primary" />
+              <span className="text-xs text-primary">Improve prompt</span>
+            </Button>
+          )}
         </div>
 
         {/* Centered title */}
         <div
-          className={`flex-grow flex justify-center ${
+          className={`flex-grow flex justify-center mx-20 ${
             !visibility ? "opacity-50" : ""
           }`}
         >
@@ -211,12 +226,12 @@ const OperationHeader: React.FC<{
               onChange={(e) => setEditedName(e.target.value)}
               onBlur={handleEditComplete}
               onKeyPress={(e) => e.key === "Enter" && handleEditComplete()}
-              className="text-sm font-medium w-1/2 font-mono text-center"
+              className="text-sm font-medium max-w-[200px] font-mono text-center"
               autoFocus
             />
           ) : (
             <span
-              className={`text-sm font-medium cursor-pointer ${
+              className={`text-sm font-medium cursor-pointer truncate max-w-[200px] ${
                 llmType === "LLM"
                   ? "bg-gradient-to-r from-blue-500 to-purple-500 text-transparent bg-clip-text"
                   : ""
@@ -844,6 +859,60 @@ export const OperationCard: React.FC<{ index: number }> = ({ index }) => {
     handleOperationUpdate(updatedOperation);
   }, [operation, handleOperationUpdate]);
 
+  const [showPromptImprovement, setShowPromptImprovement] = useState(false);
+
+  const handlePromptSave = (
+    newPrompt:
+      | string
+      | { comparison_prompt: string; resolution_prompt: string },
+    schemaChanges?: Array<[string, string]>
+  ) => {
+    if (!operation) return;
+
+    let updatedOperation = { ...operation };
+
+    if (operation.type === "resolve") {
+      if (typeof newPrompt === "object") {
+        updatedOperation = {
+          ...updatedOperation,
+          otherKwargs: {
+            ...operation.otherKwargs,
+            comparison_prompt: newPrompt.comparison_prompt,
+            resolution_prompt: newPrompt.resolution_prompt,
+          },
+        };
+      }
+    } else {
+      if (typeof newPrompt === "string") {
+        updatedOperation.prompt = newPrompt;
+      }
+    }
+
+    // Handle schema changes
+    if (schemaChanges?.length && operation.output?.schema) {
+      const updatedSchema = operation.output.schema.map((item) => {
+        const change = schemaChanges.find(([oldKey]) => oldKey === item.key);
+        if (change) {
+          return { ...item, key: change[1] };
+        }
+        return item;
+      });
+
+      updatedOperation.output = {
+        ...operation.output,
+        schema: updatedSchema,
+      };
+    }
+
+    handleOperationUpdate(updatedOperation);
+    toast({
+      title: "Success",
+      description: `Prompt${
+        operation.type === "resolve" ? "s" : ""
+      } and schema updated successfully`,
+    });
+  };
+
   if (!operation) {
     return <SkeletonCard />;
   }
@@ -905,6 +974,7 @@ export const OperationCard: React.FC<{ index: number }> = ({ index }) => {
                 onAIEdit={handleAIEdit}
                 onToggleExpand={() => dispatch({ type: "TOGGLE_EXPAND" })}
                 onToggleVisibility={handleVisibilityToggle}
+                onImprovePrompt={() => setShowPromptImprovement(true)}
               />
               {isExpanded && operation.visibility !== false && (
                 <>
@@ -946,6 +1016,14 @@ export const OperationCard: React.FC<{ index: number }> = ({ index }) => {
                 otherKwargs={operation.otherKwargs || {}}
                 onSettingsSave={handleSettingsSave}
               />
+              {operation.llmType === "LLM" && (
+                <PromptImprovementDialog
+                  open={showPromptImprovement}
+                  onOpenChange={setShowPromptImprovement}
+                  currentOperation={operation}
+                  onSave={handlePromptSave}
+                />
+              )}
             </div>
           </Card>
         )}
diff --git a/website/src/components/Output.tsx b/website/src/components/Output.tsx
index 7e5889b8..b2e7609c 100644
--- a/website/src/components/Output.tsx
+++ b/website/src/components/Output.tsx
@@ -2,10 +2,8 @@ import React, { useState, useEffect, useMemo } from "react";
 import { ColumnType } from "@/components/ResizableDataTable";
 import ResizableDataTable from "@/components/ResizableDataTable";
 import { usePipelineContext } from "@/contexts/PipelineContext";
-import { Loader2, Download, ChevronDown, Circle } from "lucide-react";
+import { Loader2, Download, ChevronDown } from "lucide-react";
 import { Button } from "@/components/ui/button";
-import { Progress } from "@/components/ui/progress";
-import BookmarkableText from "@/components/BookmarkableText";
 import { Operation, OutputRow } from "@/app/types";
 import { Parser } from "json2csv";
 import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
@@ -212,8 +210,8 @@ export const Output: React.FC = () => {
           if (parsedOutputs.length > 0) {
             if ("date" in parsedOutputs[0]) {
               parsedOutputs.sort((a, b) => {
-                const dateA = (a as any).date;
-                const dateB = (b as any).date;
+                const dateA = (a as OutputRow & { date?: string }).date;
+                const dateB = (b as OutputRow & { date?: string }).date;
                 if (dateA && dateB) {
                   return new Date(dateB).getTime() - new Date(dateA).getTime();
                 }
@@ -266,7 +264,7 @@ export const Output: React.FC = () => {
     fetchData();
   }, [output, isLoadingOutputs]);
 
-  const columns: ColumnType<any>[] = React.useMemo(() => {
+  const columns: ColumnType<OutputRow>[] = React.useMemo(() => {
     const importantColumns = operation?.output?.schema
       ? operation.output.schema.map((field) => field.key)
       : [];
@@ -275,7 +273,7 @@ export const Output: React.FC = () => {
       ? Object.keys(outputs[0]).map((key) => ({
           accessorKey: key,
           header: key,
-          cell: ({ getValue }: { getValue: () => any }) => {
+          cell: ({ getValue }: { getValue: () => unknown }) => {
             const value = getValue();
             const stringValue =
               typeof value === "object" && value !== null
@@ -294,13 +292,17 @@ export const Output: React.FC = () => {
 
   const TableContent = () => (
     <div className="flex-1 min-h-0">
-      {isLoadingOutputs ? (
+      {!opName ? (
+        <div className="flex items-center justify-center h-full">
+          <p className="text-muted-foreground">No operation selected.</p>
+        </div>
+      ) : isLoadingOutputs ? (
         <div className="flex items-center justify-center h-full">
           <Loader2 className="w-8 h-8 animate-spin text-blue-500" />
           <span className="ml-2 text-muted-foreground">Loading outputs...</span>
         </div>
       ) : outputs.length > 0 ? (
-        <BookmarkableText source="output" className="h-full">
+        <div className="h-full">
           <ResizableDataTable
             data={outputs}
             columns={columns}
@@ -312,7 +314,7 @@ export const Output: React.FC = () => {
             startingRowHeight={180}
             currentOperation={opName}
           />
-        </BookmarkableText>
+        </div>
       ) : (
         <div className="flex items-center justify-center h-full">
           <p className="text-muted-foreground">No outputs available.</p>
@@ -402,7 +404,7 @@ export const Output: React.FC = () => {
       const groupedData = useMemo(() => {
         const intersectionKeys = new Set(
           outputs.flatMap((row) => {
-            // @ts-ignore
+            // @ts-expect-error Record type needs refinement
             const kvPairs = row[visualizationColumn.name] as Record<
               string,
               unknown
@@ -419,7 +421,7 @@ export const Output: React.FC = () => {
         } = {};
 
         outputs.forEach((row) => {
-          // @ts-ignore
+          // @ts-expect-error Record type needs refinement
           const kvPairs = row[visualizationColumn.name] as Record<
             string,
             unknown
@@ -480,7 +482,7 @@ export const Output: React.FC = () => {
                           JSON.stringify(value, null, 2)
                         )
                       )
-                      // @ts-ignore
+                      // @ts-expect-error
                     ).map((str) => JSON.parse(str));
 
                     // Calculate percentage of total documents
diff --git a/website/src/components/PipelineGui.tsx b/website/src/components/PipelineGui.tsx
index 17d612bb..25d452cf 100644
--- a/website/src/components/PipelineGui.tsx
+++ b/website/src/components/PipelineGui.tsx
@@ -161,10 +161,6 @@ const PipelineGUI: React.FC = () => {
     prompt: undefined,
     operationName: undefined,
   });
-  const [tempSystemPrompt, setTempSystemPrompt] = useState({
-    datasetDescription: systemPrompt.datasetDescription || "",
-    persona: systemPrompt.persona || "",
-  });
 
   const { submitTask } = useOptimizeCheck({
     onComplete: (result) => {
@@ -634,7 +630,6 @@ const PipelineGUI: React.FC = () => {
     setIsSettingsOpen(false);
     setOptimizerModel(tempOptimizerModel);
     setAutoOptimizeCheck(tempAutoOptimizeCheck);
-    setSystemPrompt(tempSystemPrompt);
   };
 
   const handleDragEnd = (result: DropResult) => {
@@ -778,14 +773,16 @@ const PipelineGUI: React.FC = () => {
                         <Textarea
                           id="datasetDescription"
                           placeholder="a collection of documents"
-                          value={tempSystemPrompt.datasetDescription}
-                          onChange={(e) =>
-                            setTempSystemPrompt((prev) => ({
-                              ...prev,
-                              datasetDescription: e.target.value,
-                            }))
-                          }
-                          onBlur={() => setSystemPrompt(tempSystemPrompt)}
+                          defaultValue={systemPrompt.datasetDescription}
+                          onBlur={(e) => {
+                            const value = e.target.value;
+                            setTimeout(() => {
+                              setSystemPrompt((prev) => ({
+                                ...prev,
+                                datasetDescription: value,
+                              }));
+                            }, 0);
+                          }}
                           className="h-[3.5rem]"
                         />
                       </div>
@@ -799,14 +796,16 @@ const PipelineGUI: React.FC = () => {
                         <Textarea
                           id="persona"
                           placeholder="a helpful assistant"
-                          value={tempSystemPrompt.persona}
-                          onChange={(e) =>
-                            setTempSystemPrompt((prev) => ({
-                              ...prev,
-                              persona: e.target.value,
-                            }))
-                          }
-                          onBlur={() => setSystemPrompt(tempSystemPrompt)}
+                          defaultValue={systemPrompt.persona}
+                          onBlur={(e) => {
+                            const value = e.target.value;
+                            setTimeout(() => {
+                              setSystemPrompt((prev) => ({
+                                ...prev,
+                                persona: value,
+                              }));
+                            }, 0);
+                          }}
                           className="h-[3.5rem]"
                         />
                       </div>
diff --git a/website/src/components/PrettyJSON.tsx b/website/src/components/PrettyJSON.tsx
new file mode 100644
index 00000000..42ab124b
--- /dev/null
+++ b/website/src/components/PrettyJSON.tsx
@@ -0,0 +1,145 @@
+import React, { useState } from "react";
+import { ChevronDown } from "lucide-react";
+
+interface PrettyJSONProps {
+  data: unknown;
+}
+
+export const PrettyJSON = React.memo(({ data }: PrettyJSONProps) => {
+  const [expandedPaths, setExpandedPaths] = useState<string[]>(() => {
+    const paths: string[] = [];
+
+    const collectPaths = (value: unknown, path: string = "") => {
+      if (Array.isArray(value)) {
+        paths.push(path);
+        value.forEach((item, index) => {
+          if (typeof item === "object" && item !== null) {
+            collectPaths(item, `${path}.${index}`);
+          }
+        });
+      } else if (typeof value === "object" && value !== null) {
+        paths.push(path);
+        Object.entries(value).forEach(([key, val]) => {
+          if (typeof val === "object" && val !== null) {
+            collectPaths(val, `${path}.${key}`);
+          }
+        });
+      }
+    };
+
+    collectPaths(data);
+    return paths;
+  });
+
+  const renderValue = (
+    value: unknown,
+    path: string = "",
+    level: number = 0
+  ): React.ReactNode => {
+    if (value === null)
+      return <span className="text-red-600 font-bold">null</span>;
+    if (value === undefined)
+      return <span className="text-red-600 font-bold">undefined</span>;
+
+    if (Array.isArray(value)) {
+      if (value.length === 0)
+        return <span className="text-slate-700 font-bold">[]</span>;
+
+      const isExpanded = expandedPaths.includes(path);
+      return (
+        <div className="relative">
+          <button
+            onClick={() => {
+              setExpandedPaths((prev) =>
+                prev.includes(path)
+                  ? prev.filter((p) => p !== path)
+                  : [...prev, path]
+              );
+            }}
+            className="text-left hover:text-primary transition-colors inline-flex items-center gap-1 font-bold"
+          >
+            <ChevronDown
+              className={`h-4 w-4 transition-transform ${
+                isExpanded ? "rotate-180" : ""
+              }`}
+            />
+            <span className="text-slate-900">Array ({value.length} items)</span>
+          </button>
+          {isExpanded && (
+            <div className="ml-4 border-l pl-4 mt-2 space-y-2">
+              {value.map((item, index) => (
+                <div key={index} className="relative">
+                  <span className="text-slate-600 font-bold">{index}: </span>
+                  {renderValue(item, `${path}.${index}`, level + 1)}
+                </div>
+              ))}
+            </div>
+          )}
+        </div>
+      );
+    }
+
+    if (typeof value === "object") {
+      const entries = Object.entries(value as Record<string, unknown>);
+      if (entries.length === 0)
+        return <span className="text-slate-700 font-bold">{}</span>;
+
+      const isExpanded = expandedPaths.includes(path);
+      return (
+        <div className="relative">
+          <button
+            onClick={() => {
+              setExpandedPaths((prev) =>
+                prev.includes(path)
+                  ? prev.filter((p) => p !== path)
+                  : [...prev, path]
+              );
+            }}
+            className="text-left hover:text-primary transition-colors inline-flex items-center gap-1 font-bold"
+          >
+            <ChevronDown
+              className={`h-4 w-4 transition-transform ${
+                isExpanded ? "rotate-180" : ""
+              }`}
+            />
+            <span className="text-slate-900">
+              Object ({entries.length} properties)
+            </span>
+          </button>
+          {isExpanded && (
+            <div className="ml-4 border-l pl-4 mt-2 space-y-2">
+              {entries.map(([key, val]) => (
+                <div key={key} className="relative">
+                  <span className="text-slate-900 font-bold">{key}: </span>
+                  {renderValue(val, `${path}.${key}`, level + 1)}
+                </div>
+              ))}
+            </div>
+          )}
+        </div>
+      );
+    }
+
+    if (typeof value === "string") {
+      return <span className="text-emerald-600 font-bold">"{value}"</span>;
+    }
+
+    if (typeof value === "number") {
+      return <span className="text-blue-700 font-bold">{value}</span>;
+    }
+
+    if (typeof value === "boolean") {
+      return (
+        <span className="text-purple-700 font-bold">{value.toString()}</span>
+      );
+    }
+
+    return <span className="text-slate-900 font-bold">{String(value)}</span>;
+  };
+
+  return (
+    <div className="font-mono text-sm space-y-1.5">{renderValue(data)}</div>
+  );
+});
+
+PrettyJSON.displayName = "PrettyJSON";
diff --git a/website/src/components/PromptImprovementDialog.tsx b/website/src/components/PromptImprovementDialog.tsx
new file mode 100644
index 00000000..529c1a36
--- /dev/null
+++ b/website/src/components/PromptImprovementDialog.tsx
@@ -0,0 +1,1013 @@
+import React, { useState, useCallback, useEffect, useRef } from "react";
+import { useChat } from "ai/react";
+import { Operation } from "@/app/types";
+import { Button } from "@/components/ui/button";
+import {
+  Dialog,
+  DialogContent,
+  DialogHeader,
+  DialogTitle,
+  DialogDescription,
+} from "@/components/ui/dialog";
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from "@/components/ui/select";
+import { ScrollArea } from "@/components/ui/scroll-area";
+import { Loader2, ArrowLeft } from "lucide-react";
+import ReactMarkdown from "react-markdown";
+import { usePipelineContext } from "@/contexts/PipelineContext";
+import { useBookmarkContext } from "@/contexts/BookmarkContext";
+import { diffLines } from "diff";
+import {
+  Popover,
+  PopoverContent,
+  PopoverTrigger,
+} from "@/components/ui/popover";
+import { Textarea } from "@/components/ui/textarea";
+
+type Step = "select" | "analyze";
+
+interface PromptImprovementDialogProps {
+  open: boolean;
+  onOpenChange: (open: boolean) => void;
+  currentOperation: Operation;
+  onSave: (
+    newPrompt:
+      | string
+      | { comparison_prompt: string; resolution_prompt: string },
+    schemaChanges?: Array<[string, string]>
+  ) => void;
+}
+
+interface Revision {
+  messages: {
+    id: string;
+    role: "system" | "user" | "assistant";
+    content: string;
+  }[];
+  prompt: string | null;
+  timestamp: number;
+}
+
+interface FeedbackConnection {
+  fromRevision: number;
+  toRevision: number;
+  feedback: string;
+}
+
+function extractTagContent(text: string, tag: string): string | null {
+  const regex = new RegExp(`<${tag}>(.*?)</${tag}>`, "s");
+  const match = text.match(regex);
+  return match ? match[1].trim() : null;
+}
+
+type PromptType = "comparison" | "resolve";
+
+function DiffView({
+  oldText,
+  newText,
+  type,
+}: {
+  oldText: string;
+  newText: string;
+  type?: PromptType;
+}) {
+  const diff = diffLines(oldText, newText);
+
+  return (
+    <div>
+      {type && (
+        <div className="text-sm font-medium mb-2">
+          {type === "comparison" ? "Comparison Prompt:" : "Resolution Prompt:"}
+        </div>
+      )}
+      <div className="font-mono text-sm whitespace-pre-wrap">
+        {diff.map((part, index) => (
+          <span
+            key={index}
+            className={
+              part.added
+                ? "bg-green-100 dark:bg-green-900/30"
+                : part.removed
+                ? "bg-red-100 dark:bg-red-900/30"
+                : ""
+            }
+          >
+            {part.value}
+          </span>
+        ))}
+      </div>
+    </div>
+  );
+}
+
+function removePromptAndSchemaTags(text: string): string {
+  return text
+    .replace(/<prompt>[\s\S]*?<\/prompt>/g, "")
+    .replace(/<schema>[\s\S]*?<\/schema>/g, "")
+    .replace(/<comparison_prompt>[\s\S]*?<\/comparison_prompt>/g, "")
+    .replace(/<resolution_prompt>[\s\S]*?<\/resolution_prompt>/g, "");
+}
+
+const getSystemContent = (
+  pipelineState: string,
+  selectedOperation: Operation
+) => `You are a prompt engineering expert. Analyze the current operation's prompt${
+  selectedOperation.type === "resolve" ? "s" : ""
+} and suggest improvements based on the pipeline state.
+
+Current pipeline state:
+${pipelineState}
+
+Focus on the operation named "${selectedOperation.name}" ${
+  selectedOperation.type === "resolve"
+    ? `with comparison prompt:
+${selectedOperation.otherKwargs?.comparison_prompt || ""}
+
+and resolution prompt:
+${selectedOperation.otherKwargs?.resolution_prompt || ""}`
+    : `with prompt:
+${selectedOperation.prompt}`
+}
+
+${
+  selectedOperation.output?.schema
+    ? `
+Current output schema keys:
+${selectedOperation.output.schema.map((item) => `- ${item.key}`).join("\n")}
+`
+    : ""
+}
+
+IMPORTANT: 
+1. ${
+  selectedOperation.type === "resolve"
+    ? "You must ALWAYS include complete revised prompts wrapped in <comparison_prompt></comparison_prompt> AND <resolution_prompt></resolution_prompt> tags in your response"
+    : "You must ALWAYS include a complete revised prompt wrapped in <prompt></prompt> tags in your response"
+}, even if you're just responding to feedback.
+
+2. Only suggest schema key changes if absolutely necessary - when the current keys are misleading, incorrect, or ambiguous. If the schema keys are fine, don't suggest changes. Include changes in <schema> tags as a list of "oldkey,newkey" pairs, one per line. Example:
+<schema>
+misleading_key,accurate_key
+ambiguous_name,specific_name
+</schema>
+
+When responding:
+1. Briefly acknowledge/analyze any feedback (1-2 sentences)
+2. ALWAYS provide ${
+  selectedOperation.type === "resolve"
+    ? "complete revised prompts wrapped in <comparison_prompt></comparison_prompt> AND <resolution_prompt></resolution_prompt> tags"
+    : "a complete revised prompt wrapped in <prompt></prompt> tags"
+}
+3. The prompt${
+  selectedOperation.type === "resolve" ? "s" : ""
+} should include all previous improvements plus any new changes
+4. Make prompts specific and concise:
+   - For subjective terms like "detailed" or "comprehensive", provide examples or metrics (e.g. "include 3-5 key points per section")
+   - For qualitative instructions like "long output", specify length (e.g. "200-300 words") based on my feedback or provide examples
+   - When using adjectives, include a reference point (e.g. "technical like API documentation" vs "simple like a blog post")`;
+
+function extractSchemaChanges(text: string): Array<[string, string]> {
+  const schemaContent = extractTagContent(text, "schema");
+  if (!schemaContent) return [];
+
+  return schemaContent
+    .split("\n")
+    .filter((line) => line.trim())
+    .map((line) => {
+      const [oldKey, newKey] = line.split(",").map((k) => k.trim());
+      return [oldKey, newKey] as [string, string];
+    });
+}
+
+function AutosizeTextarea({
+  value,
+  onChange,
+  onBlur,
+  type,
+}: {
+  value: string;
+  onChange: (e: React.ChangeEvent<HTMLTextAreaElement>) => void;
+  onBlur?: () => void;
+  type?: PromptType;
+}) {
+  const [height, setHeight] = useState("auto");
+  const hiddenRef = useRef<HTMLDivElement>(null);
+
+  useEffect(() => {
+    if (hiddenRef.current) {
+      setHeight(`${hiddenRef.current.scrollHeight}px`);
+    }
+  }, [value]);
+
+  return (
+    <div className="relative">
+      {type && (
+        <div className="text-sm font-medium mb-2">
+          {type === "comparison" ? "Comparison Prompt:" : "Resolution Prompt:"}
+        </div>
+      )}
+      <div
+        ref={hiddenRef}
+        className="invisible absolute whitespace-pre-wrap break-words font-mono text-sm p-2"
+        style={{ width: "calc(100% - 4px)" }}
+      >
+        {value + "\n"}
+      </div>
+      <Textarea
+        value={value}
+        onChange={onChange}
+        onBlur={onBlur}
+        className="font-mono text-sm resize-none"
+        style={{ height, minHeight: height }}
+        autoFocus={!type || type === "comparison"}
+      />
+    </div>
+  );
+}
+
+function buildRevisionTree(
+  revisions: Revision[],
+  connections: FeedbackConnection[]
+) {
+  type TreeNode = {
+    index: number;
+    revision: Revision;
+    children: TreeNode[];
+    feedback?: string;
+  };
+
+  const nodes: TreeNode[] = revisions.map((revision, index) => ({
+    index,
+    revision,
+    children: [],
+  }));
+
+  // Build the tree structure
+  connections.forEach((conn) => {
+    const parentNode = nodes[conn.fromRevision];
+    const childNode = nodes[conn.toRevision];
+    if (parentNode && childNode) {
+      childNode.feedback = conn.feedback;
+      parentNode.children.push(childNode);
+    }
+  });
+
+  // Return only root nodes (nodes without parents)
+  return nodes.filter(
+    (node) => !connections.some((conn) => conn.toRevision === node.index)
+  );
+}
+
+function RevisionTreeNode({
+  node,
+  depth = 0,
+  selectedIndex,
+  onSelect,
+}: {
+  node: ReturnType<typeof buildRevisionTree>[0];
+  depth?: number;
+  selectedIndex: number | null;
+  onSelect: (index: number) => void;
+}) {
+  return (
+    <div className="flex flex-col">
+      <div
+        className="flex items-center relative"
+        style={{ paddingLeft: `${depth * 20}px` }}
+      >
+        {depth > 0 && (
+          <div
+            className="absolute left-0 top-1/2 w-4 h-px border-t-2 border-muted-foreground/20"
+            style={{ left: `${(depth - 1) * 20 + 10}px` }}
+          />
+        )}
+        <button
+          onClick={() => onSelect(node.index)}
+          className={`flex items-center gap-2 w-full text-left py-1.5 px-2 text-sm hover:bg-muted rounded-sm transition-colors ${
+            selectedIndex === node.index ? "bg-muted" : ""
+          }`}
+        >
+          <div className="w-1.5 h-1.5 rounded-full bg-primary shrink-0" />
+          <div className="min-w-0">
+            <div className="font-medium truncate">
+              {node.index === 0
+                ? "Initial version"
+                : node.feedback || `Revision ${node.index}`}
+            </div>
+            <div className="text-xs text-muted-foreground">
+              {new Date(node.revision.timestamp).toLocaleTimeString()}
+            </div>
+          </div>
+        </button>
+      </div>
+      {node.children.length > 0 && (
+        <div className="relative">
+          <div
+            className="absolute left-0 top-0 bottom-4 border-l-2 border-muted-foreground/20"
+            style={{ left: `${depth * 20 + 10}px` }}
+          />
+          <div className="flex flex-col">
+            {node.children.map((child, i) => (
+              <RevisionTreeNode
+                key={i}
+                node={child}
+                depth={depth + 1}
+                selectedIndex={selectedIndex}
+                onSelect={onSelect}
+              />
+            ))}
+          </div>
+        </div>
+      )}
+    </div>
+  );
+}
+
+// Add new component for schema changes diff view
+function SchemaChangesDiff({ changes }: { changes: Array<[string, string]> }) {
+  if (changes.length === 0) return null;
+
+  return (
+    <div className="border rounded-md p-4">
+      <h3 className="font-medium mb-2">Schema Key Changes:</h3>
+      <div className="font-mono text-sm space-y-1">
+        {changes.map(([oldKey, newKey], index) => (
+          <div key={index}>
+            <span className="line-through text-red-500">{oldKey}</span>
+            <span className="mx-2">→</span>
+            <span className="text-green-500">{newKey}</span>
+          </div>
+        ))}
+      </div>
+    </div>
+  );
+}
+
+// Add helper function to extract both prompts for resolve operations
+function extractPrompts(text: string): {
+  comparisonPrompt?: string;
+  resolvePrompt?: string;
+  prompt?: string;
+} {
+  const comparisonPrompt = extractTagContent(text, "comparison_prompt");
+  const resolvePrompt = extractTagContent(text, "resolution_prompt");
+  const prompt = extractTagContent(text, "prompt");
+
+  return {
+    comparisonPrompt,
+    resolvePrompt,
+    prompt,
+  };
+}
+
+// Update the state to use a single editedPrompt
+type EditedPrompt =
+  | string
+  | { comparison_prompt: string; resolution_prompt: string };
+
+export function PromptImprovementDialog({
+  open,
+  onOpenChange,
+  currentOperation,
+  onSave,
+}: PromptImprovementDialogProps) {
+  const { operations, serializeState } = usePipelineContext();
+  const { bookmarks } = useBookmarkContext();
+  const [step, setStep] = useState<Step>("select");
+  const [selectedOperationId, setSelectedOperationId] = useState<string>(
+    currentOperation.id
+  );
+  const [editedPrompt, setEditedPrompt] = useState<EditedPrompt | null>(null);
+  const [feedbackText, setFeedbackText] = useState("");
+  const [popoverOpen, setPopoverOpen] = useState(false);
+  const [isDirectEditing, setIsDirectEditing] = useState(false);
+  const [revisions, setRevisions] = useState<Revision[]>([]);
+  const [connections, setConnections] = useState<FeedbackConnection[]>([]);
+  const [selectedRevisionIndex, setSelectedRevisionIndex] = useState<
+    number | null
+  >(null);
+  const [expectingNewRevision, setExpectingNewRevision] = useState(false);
+  const [chatKey, setChatKey] = useState(0);
+
+  const selectedOperation = operations.find(
+    (op) => op.id === selectedOperationId
+  );
+
+  const relevantBookmarks = selectedOperation
+    ? bookmarks.flatMap((bookmark) =>
+        bookmark.notes.filter(
+          (note) => note.metadata?.operationName === selectedOperation.name
+        )
+      )
+    : [];
+
+  const { messages, isLoading, append, setMessages } = useChat({
+    api: "/api/chat",
+    id: `prompt-improvement-${chatKey}`,
+    onFinish: () => {
+      // Optional: handle completion
+    },
+  });
+
+  // Update the effect that handles new messages
+  useEffect(() => {
+    if (!isLoading && messages.length > 0 && expectingNewRevision) {
+      const lastMessage = messages[messages.length - 1];
+      if (lastMessage.role === "assistant") {
+        if (selectedOperation?.type === "resolve") {
+          const extractedPrompts = extractPrompts(lastMessage.content);
+          if (
+            extractedPrompts.comparisonPrompt &&
+            extractedPrompts.resolvePrompt
+          ) {
+            // Skip if the prompts are the same as what we have in editedPrompt
+            if (
+              typeof editedPrompt === "object" &&
+              editedPrompt?.comparison_prompt ===
+                extractedPrompts.comparisonPrompt &&
+              editedPrompt?.resolution_prompt === extractedPrompts.resolvePrompt
+            ) {
+              return;
+            }
+
+            const newPrompt = {
+              comparison_prompt: extractedPrompts.comparisonPrompt,
+              resolution_prompt: extractedPrompts.resolvePrompt,
+            };
+            setEditedPrompt(newPrompt);
+
+            // Create new revision
+            const newRevision = {
+              messages: JSON.parse(JSON.stringify(messages)),
+              prompt: newPrompt,
+              timestamp: Date.now(),
+            };
+
+            // Add revision and update selected index
+            setRevisions((prev) => {
+              const newRevisions = [...prev, newRevision] as Revision[];
+              setSelectedRevisionIndex(newRevisions.length - 1);
+              return newRevisions;
+            });
+          }
+        } else {
+          const extractedPrompt = extractTagContent(
+            lastMessage.content,
+            "prompt"
+          );
+          if (extractedPrompt) {
+            // Skip if the prompt is the same as what we have in editedPrompt
+            if (
+              typeof editedPrompt === "string" &&
+              editedPrompt === extractedPrompt
+            ) {
+              return;
+            }
+
+            setEditedPrompt(extractedPrompt);
+
+            // Create new revision
+            const newRevision = {
+              messages: JSON.parse(JSON.stringify(messages)),
+              prompt: extractedPrompt,
+              timestamp: Date.now(),
+            };
+
+            // Add revision and update selected index
+            setRevisions((prev) => {
+              const newRevisions = [...prev, newRevision];
+              setSelectedRevisionIndex(newRevisions.length - 1);
+              return newRevisions;
+            });
+          }
+        }
+
+        // If this isn't the first revision, create a connection
+        if (revisions.length > 0) {
+          const parentIndex = selectedRevisionIndex ?? revisions.length - 1;
+          setConnections((prev) => [
+            ...prev,
+            {
+              fromRevision: parentIndex,
+              toRevision: revisions.length,
+              feedback: feedbackText,
+            },
+          ]);
+        }
+
+        setExpectingNewRevision(false);
+      }
+    }
+  }, [messages, isLoading, expectingNewRevision, selectedOperation?.type]);
+
+  useEffect(() => {
+    if (selectedRevisionIndex !== null && revisions[selectedRevisionIndex]) {
+      const revision = revisions[selectedRevisionIndex];
+      // Set messages to a deep copy of the revision's messages
+      setMessages(JSON.parse(JSON.stringify(revision.messages)));
+
+      // Update the prompt
+      if (revision.prompt) {
+        setEditedPrompt(revision.prompt);
+      }
+    }
+  }, [selectedRevisionIndex, revisions]);
+
+  const handleImprove = useCallback(async () => {
+    setExpectingNewRevision(true); // Set flag before getting first response
+    setEditedPrompt(null);
+    setStep("analyze");
+    setRevisions([]);
+    setConnections([]);
+    setSelectedRevisionIndex(null);
+
+    const selectedOperation = operations.find(
+      (op) => op.id === selectedOperationId
+    );
+    if (!selectedOperation) return;
+
+    const bookmarksSection =
+      relevantBookmarks.length > 0
+        ? `\nMake sure to reflect my Feedback for this operation:\n${relevantBookmarks
+            .map((note) => `- ${note.note}`)
+            .join("\n")}`
+        : "\nNo feedback found for this operation.";
+
+    const pipelineState = await serializeState();
+    const systemContent = getSystemContent(pipelineState, selectedOperation);
+
+    // First set the system message
+    setMessages([
+      {
+        role: "system",
+        content: systemContent,
+        id: "system-1",
+      },
+    ]);
+
+    // Then append the user message with the appropriate prompt(s)
+    await append({
+      role: "user",
+      content: `Please analyze and improve my prompt${
+        selectedOperation.type === "resolve" ? "s" : ""
+      }:${
+        selectedOperation.type === "resolve"
+          ? `\nComparison prompt:\n${
+              selectedOperation.otherKwargs?.comparison_prompt || ""
+            }\n\nResolution prompt:\n${
+              selectedOperation.otherKwargs?.resolution_prompt || ""
+            }`
+          : `\n${selectedOperation.prompt}`
+      }${bookmarksSection}`,
+      id: "user-1",
+    });
+
+    // Create first revision after we get the response
+    // This will happen in the useEffect that watches messages
+  }, [
+    selectedOperationId,
+    operations,
+    serializeState,
+    append,
+    setMessages,
+    relevantBookmarks,
+  ]);
+
+  const handleBack = () => {
+    setStep("select");
+    // Clear messages when going back
+    setMessages([]);
+  };
+
+  const handleClose = () => {
+    onOpenChange(false);
+    // Reset state when dialog closes
+    setTimeout(() => {
+      setStep("select");
+      setMessages([]);
+    }, 200);
+  };
+
+  const handleFeedbackSubmit = useCallback(async () => {
+    if (!feedbackText.trim()) return;
+
+    setExpectingNewRevision(true); // Set flag before getting response
+
+    const sourceRevisionIndex = selectedRevisionIndex ?? revisions.length - 1;
+    const sourceRevision = revisions[sourceRevisionIndex];
+
+    // Create a new array with the source revision's messages
+    const newMessages = [...sourceRevision.messages];
+
+    // Add the feedback message with appropriate prompt instructions based on operation type
+    const feedbackMessage = {
+      role: "user" as const,
+      content: `Consider this feedback and provide ${
+        selectedOperation?.type === "resolve"
+          ? "updated prompts"
+          : "an updated prompt"
+      }: ${feedbackText}
+
+Remember to ${
+        selectedOperation?.type === "resolve"
+          ? "include the complete revised prompts wrapped in <comparison_prompt></comparison_prompt> AND <resolution_prompt></resolution_prompt> tags"
+          : "include the complete revised prompt wrapped in <prompt></prompt> tags"
+      }, incorporating all previous improvements plus any new changes based on this feedback.`,
+      id: `user-feedback-${newMessages.length}`,
+    };
+
+    // Update messages with the feedback before sending to API
+    setMessages([...newMessages]);
+
+    // Send the feedback
+    await append(feedbackMessage);
+
+    setFeedbackText("");
+    setPopoverOpen(false);
+  }, [
+    feedbackText,
+    selectedRevisionIndex,
+    revisions,
+    messages,
+    append,
+    selectedOperation,
+  ]);
+
+  // Update the direct edit handlers
+  const handleDirectEditStart = () => {
+    setIsDirectEditing(true);
+  };
+
+  const handleDirectEditComplete = () => {
+    setIsDirectEditing(false);
+  };
+
+  const handleRevisionSelect = (index: number) => {
+    setSelectedRevisionIndex(index);
+    const revision = revisions[index];
+
+    // Increment the chat key to force a new instance
+    setChatKey((prev) => prev + 1);
+
+    // In the next tick, set up the new chat state
+    setTimeout(() => {
+      setMessages(JSON.parse(JSON.stringify(revision.messages)));
+      setEditedPrompt(revision.prompt || "");
+    }, 0);
+
+    // Reset direct edit state if it was active
+    if (isDirectEditing) {
+      setIsDirectEditing(false);
+    }
+  };
+
+  // Update the save handler
+  const handleSave = () => {
+    if (editedPrompt && selectedOperation) {
+      const lastMessage = messages[messages.length - 1];
+      const schemaChanges = lastMessage
+        ? extractSchemaChanges(lastMessage.content)
+        : [];
+      onSave(editedPrompt, schemaChanges);
+      setMessages([]);
+      setRevisions([]);
+      onOpenChange(false);
+    }
+  };
+
+  return (
+    <Dialog open={open} onOpenChange={handleClose}>
+      <DialogContent className="max-w-7xl h-[90vh] flex flex-col gap-4">
+        <DialogHeader className="flex-none">
+          <div className="flex items-center gap-4 mb-2">
+            {step === "analyze" && (
+              <Button
+                variant="ghost"
+                size="icon"
+                onClick={handleBack}
+                className="h-6 w-6"
+              >
+                <ArrowLeft className="h-4 w-4" />
+              </Button>
+            )}
+            <DialogTitle>Rewrite Prompt</DialogTitle>
+          </div>
+          <div className="flex items-center gap-2 mt-2">
+            <div
+              className={`h-2 flex-1 rounded-full ${
+                step === "select" ? "bg-primary" : "bg-muted"
+              }`}
+            />
+            <div
+              className={`h-2 flex-1 rounded-full ${
+                step === "analyze" ? "bg-primary" : "bg-muted"
+              }`}
+            />
+          </div>
+          <DialogDescription className="mt-2">
+            {step === "select"
+              ? "Select an operation to improve its prompt"
+              : "DocETL is analyzing and suggesting improvements"}
+          </DialogDescription>
+        </DialogHeader>
+
+        <ScrollArea className="flex-1 w-full h-[calc(80vh-8rem)] overflow-y-auto">
+          {step === "select" ? (
+            <div className="flex flex-col gap-4 pr-4 pb-4">
+              <Select
+                value={selectedOperationId}
+                onValueChange={setSelectedOperationId}
+              >
+                <SelectTrigger>
+                  <SelectValue placeholder="Select operation" />
+                </SelectTrigger>
+                <SelectContent>
+                  {operations
+                    .filter((op) => op.llmType === "LLM")
+                    .map((op) => (
+                      <SelectItem key={op.id} value={op.id}>
+                        {op.name}
+                      </SelectItem>
+                    ))}
+                </SelectContent>
+              </Select>
+
+              {selectedOperation && (
+                <>
+                  <div className="text-sm">
+                    <div className="font-medium mb-2">
+                      Current Prompt
+                      {selectedOperation.type === "resolve" ? "s" : ""}:
+                    </div>
+                    {selectedOperation.type === "resolve" ? (
+                      <div className="space-y-4">
+                        <div>
+                          <div className="font-medium text-sm mb-1">
+                            Comparison Prompt:
+                          </div>
+                          <pre className="bg-muted p-2 rounded-md whitespace-pre-wrap">
+                            {selectedOperation.otherKwargs?.comparison_prompt ||
+                              ""}
+                          </pre>
+                        </div>
+                        <div>
+                          <div className="font-medium text-sm mb-1">
+                            Resolution Prompt:
+                          </div>
+                          <pre className="bg-muted p-2 rounded-md whitespace-pre-wrap">
+                            {selectedOperation.otherKwargs?.resolution_prompt ||
+                              ""}
+                          </pre>
+                        </div>
+                      </div>
+                    ) : (
+                      <pre className="bg-muted p-2 rounded-md whitespace-pre-wrap">
+                        {selectedOperation.prompt}
+                      </pre>
+                    )}
+                  </div>
+
+                  <div className="text-sm">
+                    <div className="font-medium mb-2">Feedback:</div>
+                    <div className="bg-muted p-2 rounded-md">
+                      {relevantBookmarks.length > 0 ? (
+                        <ul className="list-disc list-inside space-y-1">
+                          {relevantBookmarks.map((note, index) => (
+                            <li key={index}>{note.note}</li>
+                          ))}
+                        </ul>
+                      ) : (
+                        <p className="text-muted-foreground">
+                          No feedback or bookmarks found for this operation.
+                        </p>
+                      )}
+                    </div>
+                  </div>
+                </>
+              )}
+
+              <Button
+                onClick={handleImprove}
+                disabled={isLoading || !selectedOperation}
+                className="mt-4"
+              >
+                Continue to Analysis
+              </Button>
+            </div>
+          ) : (
+            <div className="flex flex-col gap-4 pr-4 pb-4">
+              {messages.length === 0 ? (
+                <div className="flex flex-col items-center justify-center py-8">
+                  <Loader2 className="h-8 w-8 animate-spin mb-4" />
+                  <p className="text-sm text-muted-foreground">
+                    Starting analysis...
+                  </p>
+                </div>
+              ) : (
+                <>
+                  <div className="border rounded-md p-4">
+                    {messages
+                      .filter((m) => m.role === "assistant")
+                      .slice(-1)
+                      .map((message, index) => (
+                        <div
+                          key={`${selectedRevisionIndex}-${index}`}
+                          className="prose prose-sm max-w-none relative"
+                        >
+                          <ReactMarkdown>
+                            {isLoading
+                              ? message.content
+                              : removePromptAndSchemaTags(message.content)}
+                          </ReactMarkdown>
+                          {isLoading && (
+                            <div className="absolute -right-2 -bottom-2">
+                              <Loader2 className="h-4 w-4 animate-spin" />
+                            </div>
+                          )}
+                        </div>
+                      ))}
+                  </div>
+
+                  {!isLoading && editedPrompt && selectedOperation && (
+                    <>
+                      <div className="border rounded-md p-4">
+                        <h3 className="font-medium mb-2">Prompt Changes:</h3>
+                        {isDirectEditing ? (
+                          selectedOperation.type === "resolve" ? (
+                            <div className="space-y-4">
+                              <AutosizeTextarea
+                                value={
+                                  typeof editedPrompt === "object"
+                                    ? editedPrompt.comparison_prompt
+                                    : ""
+                                }
+                                onChange={(e) =>
+                                  setEditedPrompt((prev) =>
+                                    typeof prev === "object"
+                                      ? {
+                                          ...prev,
+                                          comparison_prompt: e.target.value,
+                                        }
+                                      : prev
+                                  )
+                                }
+                                type="comparison"
+                              />
+                              <AutosizeTextarea
+                                value={
+                                  typeof editedPrompt === "object"
+                                    ? editedPrompt.resolution_prompt
+                                    : ""
+                                }
+                                onChange={(e) =>
+                                  setEditedPrompt((prev) =>
+                                    typeof prev === "object"
+                                      ? {
+                                          ...prev,
+                                          resolution_prompt: e.target.value,
+                                        }
+                                      : prev
+                                  )
+                                }
+                                type="resolve"
+                              />
+                            </div>
+                          ) : (
+                            <AutosizeTextarea
+                              value={
+                                typeof editedPrompt === "string"
+                                  ? editedPrompt
+                                  : ""
+                              }
+                              onChange={(e) => setEditedPrompt(e.target.value)}
+                            />
+                          )
+                        ) : selectedOperation.type === "resolve" &&
+                          typeof editedPrompt === "object" ? (
+                          <div className="space-y-4">
+                            <DiffView
+                              oldText={
+                                selectedOperation.otherKwargs
+                                  ?.comparison_prompt || ""
+                              }
+                              newText={editedPrompt.comparison_prompt}
+                              type="comparison"
+                            />
+                            <DiffView
+                              oldText={
+                                selectedOperation.otherKwargs
+                                  ?.resolution_prompt || ""
+                              }
+                              newText={editedPrompt.resolution_prompt}
+                              type="resolve"
+                            />
+                          </div>
+                        ) : (
+                          <DiffView
+                            oldText={selectedOperation.prompt || ""}
+                            newText={
+                              typeof editedPrompt === "string"
+                                ? editedPrompt
+                                : ""
+                            }
+                          />
+                        )}
+                      </div>
+
+                      {messages.length > 0 && (
+                        <SchemaChangesDiff
+                          changes={extractSchemaChanges(
+                            messages[messages.length - 1].content
+                          )}
+                        />
+                      )}
+
+                      <div className="flex gap-2 justify-end">
+                        <Button
+                          variant="secondary"
+                          onClick={() => {
+                            if (isDirectEditing) {
+                              // When switching to diff view, save the changes first
+                              handleDirectEditComplete();
+                            } else {
+                              handleDirectEditStart();
+                            }
+                          }}
+                        >
+                          {isDirectEditing ? "See diff" : "Directly edit"}
+                        </Button>
+
+                        <Popover
+                          open={popoverOpen}
+                          onOpenChange={setPopoverOpen}
+                          modal
+                        >
+                          <PopoverTrigger asChild>
+                            <Button variant="secondary">Add feedback</Button>
+                          </PopoverTrigger>
+                          <PopoverContent
+                            className="w-[500px] max-h-[80vh]"
+                            side="left"
+                            align="start"
+                          >
+                            <div className="flex flex-col gap-2 h-full">
+                              <div className="flex-1 overflow-y-auto border-b max-h-[500px]">
+                                <div className="font-medium text-sm mb-1">
+                                  Revision History:
+                                </div>
+                                <div>
+                                  {buildRevisionTree(
+                                    revisions,
+                                    connections
+                                  ).map((node, index) => (
+                                    <RevisionTreeNode
+                                      key={index}
+                                      node={node}
+                                      selectedIndex={selectedRevisionIndex}
+                                      onSelect={handleRevisionSelect}
+                                    />
+                                  ))}
+                                </div>
+                              </div>
+                              <div className="flex-none pt-2">
+                                <Textarea
+                                  placeholder="What would you like to improve?"
+                                  value={feedbackText}
+                                  onChange={(e) =>
+                                    setFeedbackText(e.target.value)
+                                  }
+                                  className="min-h-[100px] mb-2"
+                                />
+                                <div className="flex justify-end">
+                                  <Button
+                                    onClick={handleFeedbackSubmit}
+                                    disabled={!feedbackText.trim()}
+                                  >
+                                    Submit Feedback
+                                  </Button>
+                                </div>
+                              </div>
+                            </div>
+                          </PopoverContent>
+                        </Popover>
+
+                        <Button onClick={handleSave} disabled={!editedPrompt}>
+                          Save and Overwrite
+                        </Button>
+                      </div>
+                    </>
+                  )}
+                </>
+              )}
+            </div>
+          )}
+        </ScrollArea>
+      </DialogContent>
+    </Dialog>
+  );
+}
diff --git a/website/src/components/ResizableDataTable.tsx b/website/src/components/ResizableDataTable.tsx
index 93866b91..35441e9d 100644
--- a/website/src/components/ResizableDataTable.tsx
+++ b/website/src/components/ResizableDataTable.tsx
@@ -37,6 +37,7 @@ import {
   ChevronDown,
   Search,
   Eye,
+  Maximize2,
 } from "lucide-react";
 import {
   DropdownMenu,
@@ -54,7 +55,9 @@ import {
   HoverCardContent,
   HoverCardTrigger,
 } from "@/components/ui/hover-card";
-
+import { ColumnDialog } from "@/components/ColumnDialog";
+import { SearchableCell } from "@/components/SearchableCell";
+import { PrettyJSON } from "@/components/PrettyJSON";
 export type DataType = Record<string, unknown>;
 export type ColumnType<T extends DataType> = ColumnDef<T> & {
   initialWidth?: number;
@@ -360,6 +363,7 @@ interface ColumnHeaderProps {
   filterValue: string;
   onSort: () => void;
   sortDirection: false | "asc" | "desc";
+  onExpand: () => void;
 }
 
 const ColumnHeader = React.memo(
@@ -371,6 +375,7 @@ const ColumnHeader = React.memo(
     filterValue,
     onSort,
     sortDirection,
+    onExpand,
   }: ColumnHeaderProps) => {
     const histogramData = useMemo(() => {
       if (!stats) return [];
@@ -435,75 +440,87 @@ const ColumnHeader = React.memo(
         <div
           className={`${
             isBold ? "font-bold" : ""
-          } text-sm px-1 flex items-center gap-2`}
+          } text-sm px-1 flex items-center`}
         >
-          <Button
-            variant="ghost"
-            size="sm"
-            className="h-6 w-6 p-0"
-            onClick={onSort}
-          >
-            {sortDirection === false && (
-              <svg
-                xmlns="http://www.w3.org/2000/svg"
-                width="14"
-                height="14"
-                viewBox="0 0 24 24"
-                fill="none"
-                stroke="currentColor"
-                strokeWidth="2"
-                strokeLinecap="round"
-                strokeLinejoin="round"
-                className="text-muted-foreground"
-              >
-                <path d="m3 16 4 4 4-4" />
-                <path d="M7 20V4" />
-                <path d="m21 8-4-4-4 4" />
-                <path d="M17 4v16" />
-              </svg>
-            )}
-            {sortDirection === "asc" && (
-              <svg
-                xmlns="http://www.w3.org/2000/svg"
-                width="14"
-                height="14"
-                viewBox="0 0 24 24"
-                fill="none"
-                stroke="currentColor"
-                strokeWidth="2"
-                strokeLinecap="round"
-                strokeLinejoin="round"
-              >
-                <path d="m3 8 4-4 4 4" />
-                <path d="M7 4v16" />
-              </svg>
-            )}
-            {sortDirection === "desc" && (
-              <svg
-                xmlns="http://www.w3.org/2000/svg"
-                width="14"
-                height="14"
-                viewBox="0 0 24 24"
-                fill="none"
-                stroke="currentColor"
-                strokeWidth="2"
-                strokeLinecap="round"
-                strokeLinejoin="round"
-              >
-                <path d="m3 16 4 4 4-4" />
-                <path d="M7 20V4" />
-              </svg>
-            )}
-          </Button>
-          {header}
+          <div className="flex items-center">
+            <Button
+              variant="ghost"
+              size="sm"
+              className="h-6 w-6 p-0"
+              onClick={onSort}
+            >
+              {sortDirection === false && (
+                <svg
+                  xmlns="http://www.w3.org/2000/svg"
+                  width="14"
+                  height="14"
+                  viewBox="0 0 24 24"
+                  fill="none"
+                  stroke="currentColor"
+                  strokeWidth="2"
+                  strokeLinecap="round"
+                  strokeLinejoin="round"
+                  className="text-muted-foreground"
+                >
+                  <path d="m3 16 4 4 4-4" />
+                  <path d="M7 20V4" />
+                  <path d="m21 8-4-4-4 4" />
+                  <path d="M17 4v16" />
+                </svg>
+              )}
+              {sortDirection === "asc" && (
+                <svg
+                  xmlns="http://www.w3.org/2000/svg"
+                  width="14"
+                  height="14"
+                  viewBox="0 0 24 24"
+                  fill="none"
+                  stroke="currentColor"
+                  strokeWidth="2"
+                  strokeLinecap="round"
+                  strokeLinejoin="round"
+                >
+                  <path d="m3 8 4-4 4 4" />
+                  <path d="M7 4v16" />
+                </svg>
+              )}
+              {sortDirection === "desc" && (
+                <svg
+                  xmlns="http://www.w3.org/2000/svg"
+                  width="14"
+                  height="14"
+                  viewBox="0 0 24 24"
+                  fill="none"
+                  stroke="currentColor"
+                  strokeWidth="2"
+                  strokeLinecap="round"
+                  strokeLinejoin="round"
+                >
+                  <path d="m3 16 4 4 4-4" />
+                  <path d="M7 20V4" />
+                </svg>
+              )}
+            </Button>
+            <Button
+              variant="ghost"
+              size="sm"
+              className="h-6 w-6 p-0"
+              onClick={onExpand}
+            >
+              <Maximize2 className="h-3 w-3 text-muted-foreground" />
+            </Button>
+          </div>
+          <span className="ml-2">{header}</span>
         </div>
         <div
           className={`${isBold ? "font-bold" : ""} space-y-2 ${
-            filterValue ? "bg-primary/5 rounded-md p-1" : ""
+            filterValue ? "bg-primary/5 rounded-md" : ""
           }`}
         >
-          <div className="flex items-center gap-2">
-            <Search className="h-3 w-3 text-muted-foreground" />
+          <div className="flex items-center h-6">
+            <div className="flex items-center w-6">
+              <Search className="h-3 w-3 text-muted-foreground ml-1.5" />
+            </div>
             <Input
               placeholder="Filter..."
               value={filterValue}
@@ -717,163 +734,6 @@ const MarkdownCell = React.memo(({ content }: MarkdownCellProps) => {
 });
 MarkdownCell.displayName = "MarkdownCell";
 
-interface SearchableCellProps {
-  content: string;
-  isResizing: boolean;
-}
-
-const SearchableCell = React.memo(
-  ({ content, isResizing }: SearchableCellProps) => {
-    const [searchTerm, setSearchTerm] = useState("");
-    const [highlightedContent, setHighlightedContent] = useState(content);
-    const [currentMatchIndex, setCurrentMatchIndex] = useState(0);
-    const [matchCount, setMatchCount] = useState(0);
-    const containerRef = useRef<HTMLDivElement>(null);
-
-    // Search functionality
-    useEffect(() => {
-      if (!searchTerm) {
-        setHighlightedContent(content);
-        setMatchCount(0);
-        setCurrentMatchIndex(0);
-        return;
-      }
-
-      try {
-        const regex = new RegExp(`(${searchTerm})`, "gi");
-        const matches = content.match(regex);
-        const matchesCount = matches ? matches.length : 0;
-        setMatchCount(matchesCount);
-
-        if (matchesCount > 0) {
-          const highlighted = content.replace(
-            regex,
-            (match) => `<mark class="search-match">${match}</mark>`
-          );
-          setHighlightedContent(highlighted);
-
-          // Scroll to current match
-          setTimeout(() => {
-            if (containerRef.current) {
-              const marks =
-                containerRef.current.getElementsByClassName("search-match");
-              if (marks.length > 0 && currentMatchIndex < marks.length) {
-                marks[currentMatchIndex].scrollIntoView({
-                  behavior: "smooth",
-                  block: "center",
-                });
-              }
-            }
-          }, 100);
-        } else {
-          setHighlightedContent(content);
-        }
-      } catch {
-        setHighlightedContent(content);
-        setMatchCount(0);
-      }
-    }, [searchTerm, content, currentMatchIndex]);
-
-    const navigateMatches = useCallback(
-      (direction: "next" | "prev") => {
-        if (matchCount === 0) return;
-
-        if (direction === "next") {
-          setCurrentMatchIndex((prev) => (prev + 1) % matchCount);
-        } else {
-          setCurrentMatchIndex((prev) => (prev - 1 + matchCount) % matchCount);
-        }
-      },
-      [matchCount]
-    );
-
-    // Style for search matches
-    useEffect(() => {
-      const styleId = "search-match-style";
-      let styleElement = document.getElementById(styleId) as HTMLStyleElement;
-
-      if (!styleElement) {
-        styleElement = document.createElement("style");
-        styleElement.id = styleId;
-        document.head.appendChild(styleElement);
-      }
-
-      styleElement.textContent = `
-        mark.search-match {
-          background-color: hsl(var(--primary) / 0.2);
-          color: inherit;
-          padding: 0;
-          border-radius: 2px;
-        }
-        mark.search-match:nth-of-type(${currentMatchIndex + 1}) {
-          background-color: hsl(var(--primary) / 0.5);
-        }
-      `;
-
-      return () => {
-        if (styleElement && styleElement.parentNode) {
-          styleElement.parentNode.removeChild(styleElement);
-        }
-      };
-    }, [currentMatchIndex]);
-
-    return (
-      <div className="relative">
-        <div className="sticky top-0 z-10 bg-background/80 backdrop-blur-sm rounded-md">
-          <div className="flex items-center gap-2 p-1">
-            <Search className="h-3 w-3 text-muted-foreground" />
-            <Input
-              placeholder="Search in cell..."
-              value={searchTerm}
-              onChange={(e) => {
-                setSearchTerm(e.target.value);
-                setCurrentMatchIndex(0);
-              }}
-              className="h-6 text-xs border-none shadow-none focus-visible:ring-0"
-            />
-            {matchCount > 0 && (
-              <div className="flex items-center gap-1">
-                <Button
-                  variant="ghost"
-                  size="sm"
-                  className="h-5 w-5 p-0"
-                  onClick={() => navigateMatches("prev")}
-                >
-                  <ChevronLeft className="h-3 w-3" />
-                </Button>
-                <span className="text-xs text-muted-foreground whitespace-nowrap">
-                  {currentMatchIndex + 1}/{matchCount}
-                </span>
-                <Button
-                  variant="ghost"
-                  size="sm"
-                  className="h-5 w-5 p-0"
-                  onClick={() => navigateMatches("next")}
-                >
-                  <ChevronRight className="h-3 w-3" />
-                </Button>
-              </div>
-            )}
-          </div>
-        </div>
-        <div ref={containerRef}>
-          {isResizing ? (
-            <div>{content}</div>
-          ) : searchTerm ? (
-            <div
-              dangerouslySetInnerHTML={{ __html: highlightedContent }}
-              className="prose prose-sm max-w-none"
-            />
-          ) : (
-            <MarkdownCell content={content} />
-          )}
-        </div>
-      </div>
-    );
-  }
-);
-SearchableCell.displayName = "SearchableCell";
-
 interface ObservabilityIndicatorProps {
   row: Record<string, unknown>;
   currentOperation: string;
@@ -1028,18 +888,17 @@ function ResizableDataTable<T extends DataType>({
 
     if (cellValue == null) return false;
 
-    // Handle different types of values
-    if (typeof cellValue === "number") {
-      return cellValue.toString().includes(searchValue);
-    }
+    // Convert the cell value to a searchable string based on its type
+    let searchableString = "";
 
-    if (Array.isArray(cellValue)) {
-      return cellValue.some((item) =>
-        String(item).toLowerCase().includes(searchValue)
-      );
+    if (typeof cellValue === "object") {
+      // Handle both arrays and objects by converting to JSON string
+      searchableString = JSON.stringify(cellValue);
+    } else {
+      searchableString = String(cellValue);
     }
 
-    return String(cellValue).toLowerCase().includes(searchValue);
+    return searchableString.toLowerCase().includes(searchValue);
   };
 
   // Add this state to store original row indices
@@ -1161,6 +1020,16 @@ function ResizableDataTable<T extends DataType>({
     saveSettings();
   }, [columns, data, saveSettings, table]);
 
+  const [dialogOpen, setDialogOpen] = useState(false);
+  const [activeColumn, setActiveColumn] = useState<string | null>(null);
+  const [currentValueIndex, setCurrentValueIndex] = useState(0);
+
+  const handleColumnExpand = (columnId: string) => {
+    setActiveColumn(columnId);
+    setCurrentValueIndex(0);
+    setDialogOpen(true);
+  };
+
   return (
     <div className="w-full overflow-auto">
       <div className="mb-2 flex justify-between items-center">
@@ -1285,6 +1154,7 @@ function ResizableDataTable<T extends DataType>({
                           }
                         }}
                         sortDirection={header.column.getIsSorted()}
+                        onExpand={() => handleColumnExpand(header.column.id)}
                       />
                     )}
                     <ColumnResizer header={header} />
@@ -1340,6 +1210,17 @@ function ResizableDataTable<T extends DataType>({
                             content={cell.getValue() as string}
                             isResizing={isResizing}
                           />
+                        ) : typeof cell.getValue() === "object" ? (
+                          <SearchableCell
+                            content={JSON.stringify(cell.getValue(), null, 2)}
+                            isResizing={isResizing}
+                          >
+                            {(searchTerm) =>
+                              searchTerm ? null : (
+                                <PrettyJSON data={cell.getValue()} />
+                              )
+                            }
+                          </SearchableCell>
                         ) : (
                           flexRender(
                             cell.column.columnDef.cell,
@@ -1393,6 +1274,34 @@ function ResizableDataTable<T extends DataType>({
           </Button>
         </div>
       )}
+
+      {activeColumn && (
+        <ColumnDialog
+          isOpen={dialogOpen}
+          onClose={() => {
+            setDialogOpen(false);
+            setActiveColumn(null);
+          }}
+          columnId={activeColumn}
+          columnHeader={
+            columns.find((col) => (col.accessorKey || col.id) === activeColumn)
+              ?.header as string
+          }
+          data={data}
+          currentIndex={currentValueIndex}
+          onNavigate={(direction) => {
+            if (direction === "next") {
+              setCurrentValueIndex((prev) =>
+                Math.min(prev + 1, data.length - 1)
+              );
+            } else {
+              setCurrentValueIndex((prev) => Math.max(prev - 1, 0));
+            }
+          }}
+          onJumpToRow={(index) => setCurrentValueIndex(index)}
+          currentOperation={currentOperation}
+        />
+      )}
     </div>
   );
 }
diff --git a/website/src/components/RowNavigator.tsx b/website/src/components/RowNavigator.tsx
new file mode 100644
index 00000000..fd9e784b
--- /dev/null
+++ b/website/src/components/RowNavigator.tsx
@@ -0,0 +1,128 @@
+import React, { useState, useRef } from "react";
+import { Button } from "@/components/ui/button";
+import { Input } from "@/components/ui/input";
+import { ChevronLeft, ChevronRight } from "lucide-react";
+
+interface RowNavigatorProps {
+  currentRow: number;
+  totalRows: number;
+  onNavigate: (direction: "prev" | "next") => void;
+  onJumpToRow: (index: number) => void;
+  disabled?: boolean;
+  label?: string;
+  compact?: boolean;
+}
+
+export const RowNavigator = React.memo(
+  ({
+    currentRow,
+    totalRows,
+    onNavigate,
+    onJumpToRow,
+    disabled = false,
+    label = "Row",
+    compact = false,
+  }: RowNavigatorProps) => {
+    const [isEditing, setIsEditing] = useState(false);
+    const [inputValue, setInputValue] = useState("");
+    const inputRef = useRef<HTMLInputElement>(null);
+
+    const handleSubmit = () => {
+      const rowNum = parseInt(inputValue);
+      if (!isNaN(rowNum) && rowNum >= 1 && rowNum <= totalRows) {
+        onJumpToRow(rowNum - 1);
+        setIsEditing(false);
+      }
+      setInputValue("");
+    };
+
+    return (
+      <div className="flex items-center gap-2">
+        <Button
+          variant="ghost"
+          size="sm"
+          onClick={() => onNavigate("prev")}
+          disabled={currentRow === 0 || disabled}
+          className={`${compact ? "h-6 w-6" : "h-8 w-8"} p-0`}
+        >
+          <ChevronLeft className={`${compact ? "h-3 w-3" : "h-4 w-4"}`} />
+        </Button>
+
+        <div className="flex items-center gap-1.5">
+          <span
+            className={`${
+              compact ? "text-xs" : "text-sm"
+            } text-muted-foreground`}
+          >
+            {label}
+          </span>
+          {isEditing ? (
+            <form
+              onSubmit={(e) => {
+                e.preventDefault();
+                handleSubmit();
+              }}
+              className="relative"
+            >
+              <Input
+                ref={inputRef}
+                type="number"
+                min={1}
+                max={totalRows}
+                value={inputValue}
+                onChange={(e) => setInputValue(e.target.value)}
+                onBlur={() => {
+                  if (inputValue) handleSubmit();
+                  else setIsEditing(false);
+                }}
+                className={`${
+                  compact ? "w-16 h-6 text-xs" : "w-20 h-8 text-sm"
+                } pr-8`}
+                autoFocus
+              />
+              <span className="absolute right-2 top-1/2 -translate-y-1/2 text-xs text-muted-foreground">
+                ↵
+              </span>
+            </form>
+          ) : (
+            <Button
+              variant="ghost"
+              size="sm"
+              className={`${
+                compact
+                  ? "h-6 min-w-[3rem] px-1.5 text-xs"
+                  : "h-8 min-w-[4rem] px-2 text-sm"
+              } font-mono`}
+              onClick={() => {
+                setIsEditing(true);
+                setTimeout(() => inputRef.current?.focus(), 0);
+              }}
+              disabled={disabled}
+            >
+              {currentRow + 1}
+            </Button>
+          )}
+          <span
+            className={`${
+              compact ? "text-xs" : "text-sm"
+            } text-muted-foreground`}
+          >
+            of {totalRows}
+          </span>
+        </div>
+
+        <Button
+          variant="ghost"
+          size="sm"
+          onClick={() => onNavigate("next")}
+          disabled={currentRow === totalRows - 1 || disabled}
+          className={`${compact ? "h-6 w-6" : "h-8 w-8"} p-0`}
+        >
+          <ChevronRight className={`${compact ? "h-3 w-3" : "h-4 w-4"}`} />
+        </Button>
+      </div>
+    );
+  }
+);
+
+RowNavigator.displayName = "RowNavigator";
diff --git a/website/src/components/SearchableCell.tsx b/website/src/components/SearchableCell.tsx
new file mode 100644
index 00000000..5f65f7d7
--- /dev/null
+++ b/website/src/components/SearchableCell.tsx
@@ -0,0 +1,170 @@
+import React, { useState, useEffect, useRef } from "react";
+import { Search } from "lucide-react";
+import { Input } from "@/components/ui/input";
+import { Button } from "@/components/ui/button";
+import { ChevronLeft, ChevronRight } from "lucide-react";
+import { MarkdownCell } from "@/components/MarkdownCell";
+
+interface SearchableCellProps {
+  content: string;
+  isResizing: boolean;
+  children?: (searchTerm: string) => React.ReactNode;
+}
+
+export const SearchableCell = React.memo(
+  ({ content, isResizing, children }: SearchableCellProps) => {
+    const [searchTerm, setSearchTerm] = useState("");
+    const [highlightedContent, setHighlightedContent] = useState(content);
+    const [currentMatchIndex, setCurrentMatchIndex] = useState(0);
+    const [matchCount, setMatchCount] = useState(0);
+    const containerRef = useRef<HTMLDivElement>(null);
+
+    useEffect(() => {
+      if (!searchTerm) {
+        setHighlightedContent(content);
+        setMatchCount(0);
+        setCurrentMatchIndex(0);
+        return;
+      }
+
+      try {
+        const regex = new RegExp(`(${searchTerm})`, "gi");
+        const matches = content.match(regex);
+        const matchesCount = matches ? matches.length : 0;
+        setMatchCount(matchesCount);
+
+        if (matchesCount > 0) {
+          const highlighted = content.replace(
+            regex,
+            (match) => `<mark class="search-match">${match}</mark>`
+          );
+          setHighlightedContent(highlighted);
+
+          setTimeout(() => {
+            if (containerRef.current) {
+              const marks =
+                containerRef.current.getElementsByClassName("search-match");
+              if (marks.length > 0 && currentMatchIndex < marks.length) {
+                marks[currentMatchIndex].scrollIntoView({
+                  behavior: "smooth",
+                  block: "center",
+                });
+              }
+            }
+          }, 100);
+        } else {
+          setHighlightedContent(content);
+        }
+      } catch {
+        setHighlightedContent(content);
+        setMatchCount(0);
+      }
+    }, [searchTerm, content, currentMatchIndex]);
+
+    const navigateMatches = (direction: "next" | "prev") => {
+      if (matchCount === 0) return;
+
+      if (direction === "next") {
+        setCurrentMatchIndex((prev) => (prev + 1) % matchCount);
+      } else {
+        setCurrentMatchIndex((prev) => (prev - 1 + matchCount) % matchCount);
+      }
+    };
+
+    useEffect(() => {
+      const styleId = "search-match-style";
+      let styleElement = document.getElementById(styleId) as HTMLStyleElement;
+
+      if (!styleElement) {
+        styleElement = document.createElement("style");
+        styleElement.id = styleId;
+        document.head.appendChild(styleElement);
+      }
+
+      styleElement.textContent = `
+        mark.search-match {
+          background-color: hsl(var(--primary) / 0.2);
+          color: inherit;
+          padding: 0;
+          border-radius: 2px;
+        }
+        mark.search-match:nth-of-type(${currentMatchIndex + 1}) {
+          background-color: hsl(var(--primary) / 0.5);
+        }
+      `;
+
+      return () => {
+        if (styleElement && styleElement.parentNode) {
+          styleElement.parentNode.removeChild(styleElement);
+        }
+      };
+    }, [currentMatchIndex]);
+
+    return (
+      <div className="relative">
+        <div className="sticky top-0 z-10 bg-background/80 backdrop-blur-sm rounded-md">
+          <div className="flex items-center gap-2 p-1">
+            <Search className="h-3 w-3 text-muted-foreground" />
+            <Input
+              placeholder="Search in cell..."
+              value={searchTerm}
+              onChange={(e) => {
+                setSearchTerm(e.target.value);
+                setCurrentMatchIndex(0);
+              }}
+              className="h-6 text-xs border-none shadow-none focus-visible:ring-0"
+            />
+            {matchCount > 0 && (
+              <div className="flex items-center gap-1">
+                <Button
+                  variant="ghost"
+                  size="sm"
+                  onClick={() => navigateMatches("prev")}
+                  disabled={currentMatchIndex === 0}
+                  className="h-5 w-5 p-0"
+                >
+                  <ChevronLeft className="h-3 w-3" />
+                </Button>
+                <span className="text-xs text-muted-foreground whitespace-nowrap">
+                  {currentMatchIndex + 1}/{matchCount}
+                </span>
+                <Button
+                  variant="ghost"
+                  size="sm"
+                  onClick={() => navigateMatches("next")}
+                  disabled={currentMatchIndex === matchCount - 1}
+                  className="h-5 w-5 p-0"
+                >
+                  <ChevronRight className="h-3 w-3" />
+                </Button>
+              </div>
+            )}
+          </div>
+        </div>
+        <div ref={containerRef}>
+          {children ? (
+            searchTerm ? (
+              <div
+                dangerouslySetInnerHTML={{ __html: highlightedContent }}
+                className="prose prose-sm max-w-none"
+              />
+            ) : (
+              children(searchTerm)
+            )
+          ) : isResizing ? (
+            <div>{content}</div>
+          ) : searchTerm ? (
+            <div
+              dangerouslySetInnerHTML={{ __html: highlightedContent }}
+              className="prose prose-sm max-w-none"
+            />
+          ) : (
+            <MarkdownCell content={content} />
+          )}
+        </div>
+      </div>
+    );
+  }
+);
+
+SearchableCell.displayName = "SearchableCell";
diff --git a/website/src/contexts/BookmarkContext.tsx b/website/src/contexts/BookmarkContext.tsx
index 2e714c72..08c226f8 100644
--- a/website/src/contexts/BookmarkContext.tsx
+++ b/website/src/contexts/BookmarkContext.tsx
@@ -9,14 +9,14 @@ import { Bookmark, BookmarkContextType, UserNote } from "@/app/types";
 import { BOOKMARKS_STORAGE_KEY } from "@/app/localStorageKeys";
 
 const BookmarkContext = createContext<BookmarkContextType | undefined>(
-  undefined,
+  undefined
 );
 
 export const useBookmarkContext = () => {
   const context = useContext(BookmarkContext);
   if (!context) {
     throw new Error(
-      "useBookmarkContext must be used within a BookmarkProvider",
+      "useBookmarkContext must be used within a BookmarkProvider"
     );
   }
   return context;
@@ -37,16 +37,9 @@ export const BookmarkProvider: React.FC<{ children: ReactNode }> = ({
     localStorage.setItem(BOOKMARKS_STORAGE_KEY, JSON.stringify(bookmarks));
   }, [bookmarks]);
 
-  const addBookmark = (
-    text: string,
-    source: string,
-    color: string,
-    notes: UserNote[],
-  ) => {
+  const addBookmark = (color: string, notes: UserNote[]) => {
     const newBookmark: Bookmark = {
       id: Date.now().toString(),
-      text,
-      source,
       color,
       notes,
     };
@@ -55,13 +48,31 @@ export const BookmarkProvider: React.FC<{ children: ReactNode }> = ({
 
   const removeBookmark = (id: string) => {
     setBookmarks((prevBookmarks) =>
-      prevBookmarks.filter((bookmark) => bookmark.id !== id),
+      prevBookmarks.filter((bookmark) => bookmark.id !== id)
+    );
+  };
+
+  const getNotesForRowAndColumn = (
+    rowIndex: number,
+    columnId: string
+  ): UserNote[] => {
+    return bookmarks.flatMap((bookmark) =>
+      bookmark.notes.filter(
+        (note) =>
+          note.metadata?.rowIndex === rowIndex &&
+          note.metadata?.columnId === columnId
+      )
     );
   };
 
   return (
     <BookmarkContext.Provider
-      value={{ bookmarks, addBookmark, removeBookmark }}
+      value={{
+        bookmarks,
+        addBookmark,
+        removeBookmark,
+        getNotesForRowAndColumn,
+      }}
     >
       {children}
     </BookmarkContext.Provider>
diff --git a/website/src/contexts/PipelineContext.tsx b/website/src/contexts/PipelineContext.tsx
index b5c13dfc..c0e690bb 100644
--- a/website/src/contexts/PipelineContext.tsx
+++ b/website/src/contexts/PipelineContext.tsx
@@ -212,9 +212,28 @@ const serializeState = async (state: PipelineState): Promise<string> => {
   const bookmarksDetails = bookmarks
     .map((bookmark: Bookmark) => {
       return `
-- Text: "${bookmark.text}"
-  Source: ${bookmark.source}
-  Context: ${bookmark.notes[0].note || "None"}`;
+- Color: ${bookmark.color}
+  Notes: ${bookmark.notes
+    .map(
+      (note) => `
+    "${note.note}"${
+        note.metadata?.columnId
+          ? `
+    Column: ${note.metadata.columnId}${
+              note.metadata.rowIndex !== undefined
+                ? `
+    Row: ${note.metadata.rowIndex}`
+                : ""
+            }`
+          : ""
+      }${
+        note.metadata?.operationName
+          ? `
+    Operation: ${note.metadata.operationName}`
+          : ""
+      }`
+    )
+    .join("\n")}`;
     })
     .join("\n");
 
@@ -227,7 +246,9 @@ Input Dataset File: ${
 
 Pipeline operations:${operationsDetails}
 
-My notes:${bookmarks.length > 0 ? bookmarksDetails : "\nNo notes added yet"}
+My feedback:${
+    bookmarks.length > 0 ? bookmarksDetails : "\nNo feedback added yet"
+  }
 ${
   currentOperationName && outputSample
     ? `
@@ -358,6 +379,10 @@ export const PipelineProvider: React.FC<{ children: React.ReactNode }> = ({
       localStorageKeys.HIGH_LEVEL_GOAL_KEY,
       JSON.stringify(stateRef.current.highLevelGoal)
     );
+    localStorage.setItem(
+      localStorageKeys.SYSTEM_PROMPT_KEY,
+      JSON.stringify(stateRef.current.systemPrompt)
+    );
     setUnsavedChanges(false);
   }, []);
 
diff --git a/website/src/hooks/use-toast.ts b/website/src/hooks/use-toast.ts
index 2e9d0205..8636f2b0 100644
--- a/website/src/hooks/use-toast.ts
+++ b/website/src/hooks/use-toast.ts
@@ -22,13 +22,6 @@ type ActionTypes = {
   REMOVE_TOAST: "REMOVE_TOAST";
 };
 
-const actionTypes: ActionTypes = {
-  ADD_TOAST: "ADD_TOAST",
-  UPDATE_TOAST: "UPDATE_TOAST",
-  DISMISS_TOAST: "DISMISS_TOAST",
-  REMOVE_TOAST: "REMOVE_TOAST",
-} as const;
-
 let count = 0;
 
 function genId() {
@@ -160,7 +153,13 @@ function toast({ ...props }: Toast) {
       ...props,
       id,
       open: true,
-      style: { wordWrap: "break-word", maxWidth: "100%" },
+      style: {
+        wordWrap: "break-word",
+        maxWidth: "100%",
+        maxHeight: "80vh",
+        overflow: "auto",
+        padding: "8px",
+      },
       onOpenChange: (open) => {
         if (!open) dismiss();
       },