1 files changed, 139 insertions, 0 deletions
diff --git a/frontend/src/components/nodes/LLMNode.tsx b/frontend/src/components/nodes/LLMNode.tsx
new file mode 100644
index 0000000..cdd402c
--- /dev/null
+++ b/frontend/src/components/nodes/LLMNode.tsx
@@ -0,0 +1,139 @@
+import { useEffect } from 'react';
+import { Handle, Position, type NodeProps, useUpdateNodeInternals, useEdges } from 'reactflow';
+import type { NodeData } from '../../store/flowStore';
+import { Loader2, MessageSquare } from 'lucide-react';
+import useFlowStore from '../../store/flowStore';
+
+const LLMNode = ({ id, data, selected }: NodeProps<NodeData>) => {
+  const { updateNodeData } = useFlowStore();
+  const updateNodeInternals = useUpdateNodeInternals();
+  const edges = useEdges();
+
+  // Force update handles when traces change
+  useEffect(() => {
+    updateNodeInternals(id);
+  }, [id, data.outgoingTraces, data.inputs, updateNodeInternals]);
+
+  // Determine how many input handles to show
+  // We want to ensure there is always at least one empty handle at the bottom
+  // plus all currently connected handles.
+  
+  // Find all edges connected to this node's inputs
+  const connectedHandles = new Set(
+    edges
+      .filter(e => e.target === id)
+      .map(e => e.targetHandle)
+  );
+  
+  // Logic: 
+  // If input-0 is connected, show input-1.
+  // If input-1 is connected, show input-2.
+  // We can just iterate until we find an unconnected one.
+  
+  let handleCount = 1;
+  while (connectedHandles.has(`input-${handleCount - 1}`)) {
+    handleCount++;
+  }
+  
+  // But wait, if we delete an edge to input-0, we still want input-1 to exist if it's connected?
+  // No, usually in this designs, we just render up to max(connected_index) + 1.
+  
+  // Let's get the max index connected
+  let maxConnectedIndex = -1;
+  edges.filter(e => e.target === id).forEach(e => {
+    const idx = parseInt(e.targetHandle?.replace('input-', '') || '0');
+    if (!isNaN(idx) && idx > maxConnectedIndex) {
+      maxConnectedIndex = idx;
+    }
+  });
+  
+  const inputsToShow = Math.max(maxConnectedIndex + 2, 1);
+
+  return (
+    <div className={`px-4 py-2 shadow-md rounded-md bg-white border-2 min-w-[200px] ${selected ? 'border-blue-500' : 'border-gray-200'}`}>
+      <div className="flex items-center mb-2">
+        <div className="rounded-full w-8 h-8 flex justify-center items-center bg-gray-100">
+          {data.status === 'loading' ? (
+            <Loader2 className="w-4 h-4 animate-spin text-blue-500" />
+          ) : (
+            <MessageSquare className="w-4 h-4 text-gray-600" />
+          )}
+        </div>
+        <div className="ml-2">
+          <div className="text-sm font-bold truncate max-w-[150px]">{data.label}</div>
+          <div className="text-xs text-gray-500">{data.model}</div>
+        </div>
+      </div>
+
+      {/* Dynamic Inputs */}
+      <div className="absolute left-0 top-0 bottom-0 flex flex-col justify-center w-4">
+         {Array.from({ length: inputsToShow }).map((_, i) => {
+           // Find the connected edge to get color
+           const connectedEdge = edges.find(e => e.target === id && e.targetHandle === `input-${i}`);
+           const edgeColor = connectedEdge?.style?.stroke as string;
+           
+           return (
+             <div key={i} className="relative h-4 w-4 my-1">
+               <Handle 
+                  type="target" 
+                  position={Position.Left} 
+                  id={`input-${i}`}
+                  className="!w-3 !h-3 !left-[-6px]"
+                  style={{ 
+                    top: '50%', 
+                    transform: 'translateY(-50%)',
+                    backgroundColor: edgeColor || '#3b82f6', // Default blue if not connected
+                    border: edgeColor ? 'none' : undefined 
+                  }}
+               />
+               <span className="absolute left-4 top-[-2px] text-[9px] text-gray-400 pointer-events-none">
+                 {i}
+               </span>
+             </div>
+           );
+         })}
+      </div>
+
+      {/* Dynamic Outputs (Traces) */}
+      <div className="absolute right-0 top-0 bottom-0 flex flex-col justify-center w-4">
+          {/* 1. Outgoing Traces (Pass-through + Self) */}
+          {data.outgoingTraces && data.outgoingTraces.map((trace, i) => (
+             <div key={trace.id} className="relative h-4 w-4 my-1" title={`Trace: ${trace.id}`}>
+                <Handle 
+                   type="source" 
+                   position={Position.Right} 
+                   id={`trace-${trace.id}`}
+                   className="!w-3 !h-3 !right-[-6px]"
+                   style={{ 
+                     backgroundColor: trace.color,
+                     top: '50%', 
+                     transform: 'translateY(-50%)' 
+                   }}
+                />
+             </div>
+          ))}
+          
+          {/* 2. New Branch Generator Handle (Always visible) */}
+          <div className="relative h-4 w-4 my-1" title="Create New Branch">
+             <Handle 
+                type="source" 
+                position={Position.Right} 
+                id="new-trace"
+                className="!w-3 !h-3 !bg-gray-400 !right-[-6px]"
+                style={{ top: '50%', transform: 'translateY(-50%)' }}
+             />
+             <span className="absolute right-4 top-[-2px] text-[9px] text-gray-400 pointer-events-none w-max">
+               + New
+             </span>
+          </div>
+      </div>
+      
+      {data.status === 'error' && (
+        <div className="text-xs text-red-500 mt-2">Error</div>
+      )}
+    </div>
+  );
+};
+
+export default LLMNode;
+