fix: make ops draggable / reorderable

docetl/api.py

@@ -203,7 +203,7 @@ def optimize(
             yaml_file_suffix=self.name,
             max_threads=max_threads,
         )
-        optimized_config = runner.optimize(return_pipeline=False)
+        optimized_config, _ = runner.optimize(return_pipeline=False)
 
         updated_pipeline = Pipeline(
             name=self.name,

docetl/builder.py

@@ -474,7 +474,7 @@ def _load_optimized_ops(self):
         else:
             self.console.log("[yellow]No optimized operations found[/yellow]")
 
-    def optimize(self):
+    def optimize(self) -> float:
         """
         Optimize the entire pipeline defined in the configuration.
 
@@ -602,6 +602,8 @@ def optimize(self):
             f"[bold]Total cost: ${self.llm_client.total_cost + self.operations_cost:.2f}[/bold]"
         )
 
+        return self.llm_client.total_cost + self.operations_cost
+
     def _run_partial_step(
         self,
         step: Dict[str, Any],

docetl/operations/resolve.py

@@ -416,7 +416,7 @@ def merge_clusters(item1, item2):
             f"({(comparisons_saved / total_possible_comparisons) * 100:.2f}%)[/green]"
         )
         self.console.log(
-            f"[blue]Number of pairs to compare: {len(filtered_pairs)}[/blue]"
+            f"[blue]Number of pairs to compare: {len(blocked_pairs)}[/blue]"
         )
 
         # Compare pairs and update clusters in real-time

docetl/runner.py

@@ -476,22 +476,25 @@ def _save_checkpoint(self, step_name: str, operation_name: str, data: List[Dict]
         )
 
     def optimize(
-        self, save: bool = False, return_pipeline: bool = True, **kwargs
-    ) -> Union[Dict, "DSLRunner"]:
+        self,
+        save: bool = False,
+        return_pipeline: bool = True,
+        **kwargs,
+    ) -> Tuple[Union[Dict, "DSLRunner"], float]:
         builder = Optimizer(
             self,
             max_threads=self.max_threads,
             **kwargs,
         )
-        builder.optimize()
+        cost = builder.optimize()
         if save:
             builder.save_optimized_config(f"{self.base_name}_opt.yaml")
             self.optimized_config_path = f"{self.base_name}_opt.yaml"
 
         if return_pipeline:
-            return DSLRunner(builder.clean_optimized_config(), self.max_threads)
+            return DSLRunner(builder.clean_optimized_config(), self.max_threads), cost
 
-        return builder.clean_optimized_config()
+        return builder.clean_optimized_config(), cost
 
 
 if __name__ == "__main__":

server/app/routes/pipeline.py

@@ -1,4 +1,3 @@
-from docetl.builder import Optimizer
 from fastapi import APIRouter, HTTPException, WebSocket, WebSocketDisconnect
 from server.app.models import PipelineRequest
 from docetl.runner import DSLRunner
@@ -12,7 +11,7 @@
 def run_pipeline(request: PipelineRequest):
     try:
         runner = DSLRunner.from_yaml(request.yaml_config)
-        cost = runner.run()
+        cost = runner.load_run_save()
         return {"cost": cost, "message": "Pipeline executed successfully"}
     except Exception as e:
         print(e)
@@ -24,24 +23,20 @@ async def websocket_run_pipeline(websocket: WebSocket):
     await websocket.accept()
     try:
         config = await websocket.receive_json()
-        runner = (
-            DSLRunner.from_yaml(config["yaml_config"])
-            if not config.get("optimize", False)
-            else Optimizer.from_yaml(config["yaml_config"])
-        )
+        runner = DSLRunner.from_yaml(config["yaml_config"])
 
         if config.get("clear_intermediate", False):
             runner.clear_intermediate()
 
         if config.get("optimize", False):
 
             async def run_pipeline():
-                return await asyncio.to_thread(runner.optimize)
+                return await asyncio.to_thread(runner.optimize, return_pipeline=False)
 
         else:
 
             async def run_pipeline():
-                return await asyncio.to_thread(runner.run)
+                return await asyncio.to_thread(runner.load_run_save)
 
         pipeline_task = asyncio.create_task(run_pipeline())
 
@@ -62,7 +57,7 @@ async def run_pipeline():
             await asyncio.sleep(0.5)
 
         # Final check to send any remaining output
-        cost = await pipeline_task
+        result = await pipeline_task
 
         console_output = runner.console.file.getvalue()
         if console_output:
@@ -73,7 +68,7 @@ async def run_pipeline():
 
         # If optimize is true, send back the optimized operations
         if config.get("optimize", False):
-            optimized_config = runner.clean_optimized_config()
+            optimized_config, cost = result
             # find the operation that has optimize = true
             optimized_op = None
             for op in optimized_config["operations"]:
@@ -102,7 +97,7 @@ async def run_pipeline():
                     "type": "result",
                     "data": {
                         "message": "Pipeline executed successfully",
-                        "cost": cost,
+                        "cost": result,
                     },
                 }
             )

tests/test_synth_resolve.py

@@ -80,7 +80,7 @@ def test_synth_resolve(config_yaml):
     runner = DSLRunner.from_yaml(config_yaml)
 
     # Run the optimization
-    optimized_config = runner.optimize(save=True, return_pipeline=False)
+    optimized_config, _ = runner.optimize(save=True, return_pipeline=False)
 
     # Check if a resolve operation was synthesized
     synthesized_resolve_found = False

website/src/app/playground/page.tsx

@@ -89,9 +89,12 @@ const RightPanelIcon: React.FC<{ isActive: boolean }> = ({ isActive }) => (
 
 const CodeEditorPipelineApp: React.FC = () => {
   const [isLocalhost, setIsLocalhost] = useState(true);
+  // Add client-side only rendering for the cost display
+  const [isMounted, setIsMounted] = useState(false);
 
   useEffect(() => {
     setIsLocalhost(window.location.hostname === 'localhost' || window.location.hostname === '127.0.0.1');
+    setIsMounted(true);
   }, []);
 
   if (!isLocalhost) {
@@ -133,22 +136,12 @@ const CodeEditorPipelineApp: React.FC = () => {
     const newOperation: Operation = {
       id: String(Date.now()),
       llmType: llmType as 'LLM' | 'non-LLM',
-      type: type as 'map' | 'reduce' | 'filter' | 'resolve' | 'parallel_map' | 'unnest' | 'split' | 'gather',
+      type: type as 'map' | 'reduce' | 'filter' | 'resolve' | 'parallel_map' | 'unnest' | 'split' | 'gather' ,
       name: name,
     };
     setOperations([...operations, newOperation]);
   };
 
-  const handleDragEnd = (result: DropResult) => {
-    if (!result.destination) return;
-
-    const items = Array.from(operations);
-    const [reorderedItem] = items.splice(result.source.index, 1);
-    items.splice(result.destination.index, 0, reorderedItem);
-
-    setOperations(items);
-  };
-
   return (
     <BookmarkProvider>
     <div className="h-screen flex flex-col bg-gray-50">
@@ -185,7 +178,12 @@ const CodeEditorPipelineApp: React.FC = () => {
               </TooltipContent>
             </Tooltip>
           </TooltipProvider>
-          <span className="text-sm font-medium text-gray-600">Cost: ${cost.toFixed(2)}</span>
+          {/* Only render the cost when client-side */}
+          {isMounted && (
+            <span className="text-sm font-medium text-gray-600">
+              Cost: ${cost.toFixed(2)}
+            </span>
+          )}
           <TooltipProvider>
             <Tooltip>
               <TooltipTrigger asChild>
@@ -269,9 +267,7 @@ const CodeEditorPipelineApp: React.FC = () => {
           <ResizablePanel defaultSize={60} minSize={30}>
             <ResizablePanelGroup direction="vertical">
               <ResizablePanel defaultSize={70} minSize={5}>
-                <PipelineGUI 
-                  onDragEnd={handleDragEnd}
-                />
+                <PipelineGUI />
               </ResizablePanel>
               {showOutput && <ResizableHandle withHandle className="h-2 bg-gray-200 hover:bg-gray-300 transition-colors duration-200" />}
               {showOutput && (

website/src/components/OperationCard.tsx

@@ -501,28 +501,34 @@ export const OperationCard: React.FC<{ index: number }> = ({ index }) => {
           {...provided.draggableProps} 
           className={`mb-2 relative rounded-sm shadow-sm w-full ${pipelineOutput?.operationId === operation.id ? 'bg-white border-blue-500 border-2' : 'bg-white'}`}
         >
-          <div {...provided.dragHandleProps} className="absolute left-0 top-0 bottom-0 w-5 flex items-center justify-center cursor-move hover:bg-gray-100">
-            <GripVertical size={14} />
+          {/* Move the drag handle div outside of the ml-5 container */}
+          <div 
+            {...provided.dragHandleProps} 
+            className="absolute left-0 top-0 bottom-0 w-6 flex items-center justify-center cursor-move hover:bg-gray-100 border-r border-gray-100"
+          >
+            <GripVertical size={14} className="text-gray-400" />
           </div>
-          <div className="ml-5">
-          <OperationHeader
-            name={operation.name}
-            type={operation.type}
-            llmType={operation.llmType}
-            disabled={isLoadingOutputs || pipelineOutput === undefined}
-            currOp={operation.id === pipelineOutput?.operationId}
-            onEdit={(name) => {
-              dispatch({ type: 'UPDATE_NAME', payload: name });
-              debouncedUpdate();
-            }}
-            onDelete={() => setOperations(prev => prev.filter(op => op.id !== operation.id))}
-            onRunOperation={handleRunOperation}
-            onToggleSettings={() => dispatch({ type: 'TOGGLE_SETTINGS' })}
-            onShowOutput={onShowOutput}
-            onOptimize={onOptimize}
-          />
+
+          {/* Adjust the left margin to accommodate the drag handle */}
+          <div className="ml-6">
+            <OperationHeader
+              name={operation.name}
+              type={operation.type}
+              llmType={operation.llmType}
+              disabled={isLoadingOutputs || pipelineOutput === undefined}
+              currOp={operation.id === pipelineOutput?.operationId}
+              onEdit={(name) => {
+                dispatch({ type: 'UPDATE_NAME', payload: name });
+                debouncedUpdate();
+              }}
+              onDelete={() => setOperations(prev => prev.filter(op => op.id !== operation.id))}
+              onRunOperation={handleRunOperation}
+              onToggleSettings={() => dispatch({ type: 'TOGGLE_SETTINGS' })}
+              onShowOutput={onShowOutput}
+              onOptimize={onOptimize}
+            />
             <CardContent className="py-2 px-3">
-              {createOperationComponent(operation, handleOperationUpdate,isSchemaExpanded, () => dispatch({ type: 'TOGGLE_SCHEMA' }))}
+              {createOperationComponent(operation, handleOperationUpdate, isSchemaExpanded, () => dispatch({ type: 'TOGGLE_SCHEMA' }))}
             </CardContent>
             {operation.llmType === 'LLM' && (
               <Guardrails
@@ -566,4 +572,4 @@ const SkeletonCard: React.FC = () => (
   </div>
 );
 
-export default OperationCard;
+export default OperationCard;

website/src/components/PipelineGui.tsx

@@ -30,9 +30,7 @@ import {
 } from "@/components/ui/alert-dialog"
 
 
-const PipelineGUI: React.FC<{ 
-  onDragEnd: (result: DropResult) => void;
-}> = ({ onDragEnd }) => {
+const PipelineGUI: React.FC = () => {
   const fileInputRef = useRef<HTMLInputElement>(null);
   const { operations, setOperations, pipelineName, setPipelineName, sampleSize, setSampleSize, numOpRun, setNumOpRun, currentFile, setCurrentFile, output, unsavedChanges, setFiles, setOutput, isLoadingOutputs, setIsLoadingOutputs, files, setCost, defaultModel, setDefaultModel, setTerminalOutput, saveProgress, clearPipelineState } = usePipelineContext();
   const [isSettingsOpen, setIsSettingsOpen] = useState(false);
@@ -365,6 +363,21 @@ const PipelineGUI: React.FC<{
     setIsSettingsOpen(false);
   };
 
+  const handleDragEnd = (result: DropResult) => {
+    if (!result.destination) return;
+
+    const { source, destination } = result;
+
+    if (source.droppableId === 'operations' && destination.droppableId === 'operations') {
+      setOperations(prevOperations => {
+        const newOperations = Array.from(prevOperations);
+        const [removed] = newOperations.splice(source.index, 1);
+        newOperations.splice(destination.index, 0, removed);
+        return newOperations;
+      });
+    }
+  };
+
   return (
     <div className="h-full overflow-auto">
       <div className="sticky top-0 z-10 p-2 bg-white">
@@ -535,21 +548,25 @@ const PipelineGUI: React.FC<{
         </div>
       </div>
       <div className="p-2">
-      <DragDropContext onDragEnd={onDragEnd}>
-        <Droppable droppableId="operations">
-          {(provided) => (
-            <div {...provided.droppableProps} ref={provided.innerRef}>
-              {operations.map((op, index) => (
-                <OperationCard 
-                  key={op.id} 
-                  index={index} 
-                />
-              ))}
-              {provided.placeholder}
-            </div>
-          )}
-        </Droppable>
-      </DragDropContext>
+        <DragDropContext onDragEnd={handleDragEnd}>
+          <Droppable droppableId="operations" type="operation">
+            {(provided, snapshot) => (
+              <div
+                {...provided.droppableProps}
+                ref={provided.innerRef}
+                className={`space-y-2 ${snapshot.isDraggingOver ? 'bg-gray-50' : ''}`}
+              >
+                {operations.map((op, index) => (
+                  <OperationCard 
+                    key={op.id} 
+                    index={index} 
+                  />
+                ))}
+                {provided.placeholder}
+              </div>
+            )}
+          </Droppable>
+        </DragDropContext>
       </div>
       <Dialog open={isSettingsOpen} onOpenChange={setIsSettingsOpen}>
         <DialogContent>
@@ -622,4 +639,4 @@ const PipelineGUI: React.FC<{
   );
 };
 
-export default PipelineGUI;
+export default PipelineGUI;

website/src/contexts/PipelineContext.tsx

@@ -56,7 +56,7 @@ export const PipelineProvider: React.FC<{ children: React.ReactNode }> = ({ chil
     sampleSize: loadFromLocalStorage(localStorageKeys.SAMPLE_SIZE_KEY, mockSampleSize),
     files: loadFromLocalStorage(localStorageKeys.FILES_KEY, mockFiles),
     cost: loadFromLocalStorage(localStorageKeys.COST_KEY, 0),
-    defaultModel: loadFromLocalStorage(localStorageKeys.DEFAULT_MODEL_KEY, "gpt-4-mini"),
+    defaultModel: loadFromLocalStorage(localStorageKeys.DEFAULT_MODEL_KEY, "gpt-4o-mini"),
   }));
 
   const [unsavedChanges, setUnsavedChanges] = useState(false);
@@ -103,7 +103,7 @@ export const PipelineProvider: React.FC<{ children: React.ReactNode }> = ({ chil
       sampleSize: mockSampleSize,
       files: mockFiles,
       cost: 0,
-      defaultModel: "gpt-4-mini",
+      defaultModel: "gpt-4o-mini",
     });
     setUnsavedChanges(false);
     console.log('Pipeline state cleared!');