feat: add code operations to the ui (#169)

website/package.json

@@ -12,6 +12,7 @@
     "@agbishop/react-ansi-18": "^4.0.6",
     "@ai-sdk/openai": "^0.0.70",
     "@hookform/resolvers": "^3.9.0",
+    "@monaco-editor/react": "^4.6.0",
     "@next/third-parties": "^14.2.11",
     "@radix-ui/react-accordion": "^1.2.0",
     "@radix-ui/react-alert-dialog": "^1.1.2",

website/src/app/types.ts

@@ -15,7 +15,10 @@ export type Operation = {
     | "unnest"
     | "split"
     | "gather"
-    | "sample";
+    | "sample"
+    | "code_map"
+    | "code_reduce"
+    | "code_filter";
   name: string;
   prompt?: string;
   output?: { schema: SchemaItem[] };

website/src/components/AIChatPanel.tsx

@@ -106,7 +106,7 @@ const AIChatPanel: React.FC<AIChatPanelProps> = ({ onClose }) => {
 Core Capabilities:
 - DocETL enables users to create sophisticated data processing workflows with LLM calls, like crowdsourcing pipelines
 - Each pipeline processes documents through a sequence of operations
-- Operations can be LLM-based (map, reduce, resolve, filter) or utility-based (unnest, split, gather, sample)
+- Operations can be LLM-based (map, reduce, resolve, filter) or utility-based (unnest, split, gather, sample) or code-based (python for map, reduce, and filter)
 
 Operation Details:
 - Every LLM operation has:
@@ -116,6 +116,12 @@ Operation Details:
   - Map/Filter: Access current doc with '{{ input.keyname }}'
   - Reduce: Loop through docs with '{% for doc in inputs %}...{% endfor %}'
   - Resolve: Compare docs with '{{ input1 }}/{{ input2 }}' and canonicalize with '{{ inputs }}'
+- Code-based operations:
+  - Map: Define a transform function (def transform(doc: dict) -> dict), where the returned dict will have key-value pairs that will be added to the output document
+  - Filter: Define a transform function (def transform(doc: dict) -> bool), where the function should return true if the document should be included in the output
+  - Reduce: Define a transform function (def transform(docs: list[dict]) -> dict), where the returned dict will have key-value pairs that will *be* the output document (unless "pass_through" is set to true, then the first original doc for every group will also be returned)
+  - Only do imports of common libraries, inside the function definition
+  - Only suggest code-based operations if the task is one that is easily expressed in code, and LLMs or crowd workers are incapable of doing it correctly (e.g., word count, simple regex, etc.)
 
 Your Role:
 - Help users optimize pipelines and overcome challenges

website/src/components/Output.tsx

@@ -2,7 +2,7 @@ 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 } from "lucide-react";
+import { Loader2, Download, ChevronDown, Circle } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { Progress } from "@/components/ui/progress";
 import BookmarkableText from "@/components/BookmarkableText";
@@ -19,6 +19,47 @@ import { useWebSocket } from "@/contexts/WebSocketContext";
 import AnsiRenderer from "./AnsiRenderer";
 import clsx from "clsx";
 
+const TinyPieChart: React.FC<{ percentage: number }> = ({ percentage }) => {
+  const size = 16;
+  const radius = 6;
+  const strokeWidth = 2;
+  const center = size / 2;
+  const circumference = 2 * Math.PI * radius;
+  const strokeDasharray = `${
+    (percentage * circumference) / 100
+  } ${circumference}`;
+
+  return (
+    <div className="relative w-4 h-4 flex items-center justify-center">
+      <svg
+        className="w-4 h-4 -rotate-90"
+        viewBox={`0 0 ${size} ${size}`}
+        width={size}
+        height={size}
+      >
+        <circle
+          cx={center}
+          cy={center}
+          r={radius}
+          className="fill-none stroke-gray-200"
+          strokeWidth={strokeWidth}
+        />
+        <circle
+          cx={center}
+          cy={center}
+          r={radius}
+          className={clsx(
+            "fill-none transition-all duration-500 ease-out",
+            percentage > 100 ? "stroke-emerald-500" : "stroke-rose-500"
+          )}
+          strokeWidth={strokeWidth}
+          strokeDasharray={strokeDasharray}
+        />
+      </svg>
+    </div>
+  );
+};
+
 export const ConsoleContent: React.FC = () => {
   const { terminalOutput, setTerminalOutput, optimizerProgress } =
     usePipelineContext();
@@ -442,7 +483,7 @@ export const Output: React.FC = () => {
               <TooltipProvider>
                 <Tooltip>
                   <TooltipTrigger>
-                    <div className="flex items-center px-2 py-1 border border-gray-200 rounded-md">
+                    <div className="flex items-center px-2 py-1 border border-gray-200 rounded-md cursor-help">
                       <span
                         className={clsx(
                           "font-medium",
@@ -457,6 +498,14 @@ export const Output: React.FC = () => {
                       >
                         {selectivityFactor}×
                       </span>
+                      {selectivityFactor !== "N/A" &&
+                        Number(selectivityFactor) < 1 && (
+                          <div className="ml-1">
+                            <TinyPieChart
+                              percentage={Number(selectivityFactor) * 100}
+                            />
+                          </div>
+                        )}
                     </div>
                   </TooltipTrigger>
                   <TooltipContent>

website/src/components/PipelineGui.tsx

@@ -735,6 +735,41 @@ const PipelineGUI: React.FC = () => {
                 >
                   Sample
                 </DropdownMenuItem>
+                <DropdownMenuSeparator />
+                <DropdownMenuLabel>Code Operations</DropdownMenuLabel>
+                <DropdownMenuItem
+                  onClick={() =>
+                    handleAddOperation(
+                      "non-LLM",
+                      "code_map",
+                      "Untitled Code Map"
+                    )
+                  }
+                >
+                  Code Map
+                </DropdownMenuItem>
+                <DropdownMenuItem
+                  onClick={() =>
+                    handleAddOperation(
+                      "non-LLM",
+                      "code_reduce",
+                      "Untitled Code Reduce"
+                    )
+                  }
+                >
+                  Code Reduce
+                </DropdownMenuItem>
+                <DropdownMenuItem
+                  onClick={() =>
+                    handleAddOperation(
+                      "non-LLM",
+                      "code_filter",
+                      "Untitled Code Filter"
+                    )
+                  }
+                >
+                  Code Filter
+                </DropdownMenuItem>
               </DropdownMenuContent>
             </DropdownMenu>
             <div className="flex space-x-2">