src/routes/canvas/chat-node.svelte

<script lang="ts">
	import { previewImage } from "$lib/components/inference-playground/img-preview.svelte";
	import { TextareaAutosize } from "$lib/spells/textarea-autosize.svelte";
	import { images } from "$lib/state/images.svelte.js";
	import { models } from "$lib/state/models.svelte";
	import { token } from "$lib/state/token.svelte";
	import { PipelineTag, type Model } from "$lib/types.js";
	import { AsyncQueue } from "$lib/utils/queue.js";
	import { InferenceClient } from "@huggingface/inference";
	import type { ChatCompletionInputMessage } from "@huggingface/tasks";
	import { Handle, Position, useSvelteFlow, type Edge, type Node, type NodeProps } from "@xyflow/svelte";
	import { marked } from "marked";
	import { FileUpload } from "melt/builders";
	import { ElementSize } from "runed";
	import { onMount } from "svelte";
	import { fade } from "svelte/transition";
	import IconImage from "~icons/carbon/image-reference";
	import IconMaximize from "~icons/carbon/maximize";
	import IconCode from "~icons/lucide/code";
	import IconCopy from "~icons/lucide/copy";
	import IconLoading from "~icons/lucide/loader-2";
	import IconAttachment from "~icons/lucide/paperclip";
	import IconAdd from "~icons/lucide/plus";
	import IconSend from "~icons/lucide/send";
	import IconStop from "~icons/lucide/square";
	import IconX from "~icons/lucide/x";
	import ModelPicker from "./model-picker.svelte";
	import ProviderPicker from "./provider-picker.svelte";
	import { edges, nodes } from "./state.js";

	type Props = Omit<NodeProps, "data"> & {
		data: { query: string; response: string; modelId?: Model["id"]; provider?: string; imageIds?: string[] };
	};
	let { id, data }: Props = $props();

	let { updateNodeData, updateNode, getNode, getViewport } = useSvelteFlow();
	onMount(() => {
		if (!data.modelId) data.modelId = models.trending[0]?.id;
		if (!data.provider) data.provider = "auto";
		updateNode(id, { height: undefined });
	});

	const autosized = new TextareaAutosize();

	let isLoading = $state(false);
	let abortController = $state<AbortController | null>(null);

	const history = $derived.by(function getNodeHistory() {
		const node = nodes.current.find(n => n.id === id);
		if (!node) return [];

		let history: Array<Omit<Node, "data"> & { data: Props["data"] }> = [
			node as Omit<Node, "data"> & { data: Props["data"] },
		];
		let target = node.id;

		while (true) {
			const parentEdge = edges.current.find(edge => edge.target === target);
			if (!parentEdge) break; // No more parents found

			const parentNode = nodes.current.find(n => n.id === parentEdge.source);
			if (!parentNode) {
				// Optional: clean up broken edges
				// edges.current = edges.current.filter(e => e.id !== parentEdge.id);
				break;
			}

			history.unshift(parentNode as Omit<Node, "data"> & { data: Props["data"] });
			target = parentNode.id; // Move up the chain
		}

		return history;
	});

	async function handleSubmit(e: SubmitEvent) {
		e.preventDefault();
		isLoading = true;
		abortController = new AbortController();
		updateNodeData(id, { response: "" });

		try {
			const client = new InferenceClient(token.value);

			const messages: ChatCompletionInputMessage[] = await Promise.all(
				history.flatMap(async n => {
					const res: ChatCompletionInputMessage[] = [];
					if (n.data.query) {
						let content: string | Array<{ type: string; text?: string; image_url?: { url: string } }> = n.data.query;

						// If node has images, convert to multimodal format
						if (n.data.imageIds && n.data.imageIds.length > 0) {
							const urls = await Promise.all(n.data.imageIds.map(k => images.get(k)));
							content = [
								{
									type: "text",
									text: n.data.query,
								},
								...n.data.imageIds.map((_imgKey, i) => ({
									type: "image_url",
									image_url: { url: urls[i] as string },
								})),
							];
						}

						res.push({
							role: "user",
							content,
						});
					}
					if (n.data.response) {
						res.push({
							role: "assistant",
							content: n.data.response,
						});
					}

					return res;
				}),
			).then(arr => arr.flat());

			const stream = client.chatCompletionStream(
				{
					// eslint-disable-next-line @typescript-eslint/no-explicit-any
					provider: (data.provider || "auto") as any,
					model: data.modelId,
					messages,
					temperature: 0.5,
					top_p: 0.7,
				},
				{
					signal: abortController.signal,
				},
			);

			for await (const chunk of stream) {
				if (abortController.signal.aborted) {
					break;
				}
				if (chunk.choices && chunk.choices.length > 0) {
					const newContent = chunk.choices[0]?.delta.content ?? "";
					updateNodeData(id, { response: data.response + newContent });
				}
			}
		} catch (error) {
			if (error instanceof Error && error.name === "AbortError") {
				// Generation was aborted, this is expected
				return;
			}
			// Re-throw other errors
			throw error;
		} finally {
			isLoading = false;
			abortController = null;
		}
	}

	function stopGeneration(e: MouseEvent) {
		e.stopPropagation();
		e.preventDefault();
		if (abortController) {
			abortController.abort();
		}
	}

	let node = $state<HTMLElement>();
	const size = new ElementSize(() => node);

	const parsedResponse = $derived.by(() => {
		if (!data.response) return "";
		return marked(data.response);
	});

	const characterCount = $derived(data.query?.length || 0);

	function insertCodeSnippet() {
		const currentQuery = data.query || "";
		const codeTemplate = "```\n\n```";
		const newQuery = currentQuery + (currentQuery ? "\n\n" : "") + codeTemplate;
		updateNodeData(id, { query: newQuery });

		setTimeout(() => {
			const textarea = document.getElementById(`message-${id}`) as HTMLTextAreaElement;
			if (textarea) {
				const cursorPosition = newQuery.length - 4;
				textarea.focus();
				textarea.setSelectionRange(cursorPosition, cursorPosition);
			}
		}, 0);
	}

	// Helper function to get actual node dimensions from DOM, converted to flow coordinates
	function getNodeDimensions(nodeId: string) {
		const nodeElement = document.querySelector(`[data-id="${nodeId}"]`);
		if (nodeElement) {
			const rect = nodeElement.getBoundingClientRect();
			const viewport = getViewport();
			// Convert from screen coordinates to flow coordinates
			const flowWidth = rect.width / viewport.zoom;
			const flowHeight = rect.height / viewport.zoom;
			return { width: flowWidth, height: flowHeight };
		}
		// Fallback to default size
		return { width: 500, height: 200 };
	}

	// Helper function to find a non-overlapping position
	function findAvailablePosition(
		startX: number,
		startY: number,
		newNodeWidth = 500,
		newNodeHeight = 200,
		constrainBelow = false,
	) {
		const spacing = 40;
		let x = startX;
		let y = startY;

		// Check if position overlaps with any existing node
		const isOverlapping = (testX: number, testY: number) => {
			return nodes.current.some(node => {
				if (node.id === id) return false; // Don't check against self

				const existingDims = getNodeDimensions(node.id);
				const nodeWidth = existingDims.width;
				const nodeHeight = existingDims.height;

				// Check for overlap with proper spacing
				return !(
					testX >= node.position.x + nodeWidth + spacing ||
					testX + newNodeWidth + spacing <= node.position.x ||
					testY >= node.position.y + nodeHeight + spacing ||
					testY + newNodeHeight + spacing <= node.position.y
				);
			});
		};

		// For add node (constrainBelow = true), maintain same Y level and search horizontally
		if (constrainBelow) {
			const fixedY = y; // Always use the same Y distance from parent

			// Try the preferred X position first
			if (!isOverlapping(x, fixedY)) return { x, y: fixedY };

			// Search horizontally at the fixed Y level, checking each position carefully
			let testX = x;
			let attempts = 0;
			while (attempts < 50) {
				// Try moving right by small increments
				testX += 50;
				if (!isOverlapping(testX, fixedY)) {
					return { x: testX, y: fixedY };
				}
				attempts++;
			}

			// Try going left from original position
			testX = x;
			attempts = 0;
			while (attempts < 50) {
				testX -= 50;
				if (!isOverlapping(testX, fixedY)) {
					return { x: testX, y: fixedY };
				}
				attempts++;
			}

			// Fallback: force position far to the right
			return { x: x + 2000, y: fixedY };
		}

		// For duplicate (constrainBelow = false), use spiral pattern
		let offset = 0;
		while (offset < 1000) {
			// Try right
			if (!isOverlapping(x + offset, y)) return { x: x + offset, y };
			// Try left
			if (!isOverlapping(x - offset, y)) return { x: x - offset, y };
			// Try down
			if (!isOverlapping(x, y + offset)) return { x, y: y + offset };
			// Try up
			if (!isOverlapping(x, y - offset)) return { x, y: y - offset };
			// Try diagonal combinations
			if (!isOverlapping(x + offset, y + offset)) return { x: x + offset, y: y + offset };
			if (!isOverlapping(x - offset, y + offset)) return { x: x - offset, y: y + offset };
			if (!isOverlapping(x + offset, y - offset)) return { x: x + offset, y: y - offset };
			if (!isOverlapping(x - offset, y - offset)) return { x: x - offset, y: y - offset };

			offset += 50;
		}

		// Fallback: return original position with larger offset
		return { x: x + 150, y: y + 150 };
	}

	const currentModel = $derived(models.all.find(m => m.id === data.modelId));
	const supportsImgUpload = $derived(currentModel?.pipeline_tag === PipelineTag.ImageTextToText);

	const fileQueue = new AsyncQueue();
	const fileUpload = new FileUpload({
		accept: "image/*",
		multiple: true,
		disabled: () => !supportsImgUpload,
		async onAccept(file) {
			fileQueue.add(async () => {
				const key = await images.upload(file);

				const prev = data.imageIds ?? [];
				await updateNodeData(id, { imageIds: [...prev, key] });

				// We're dealing with files ourselves, so we don't want fileUpload to have any internal state,
				// to avoid conflicts
				if (fileQueue.queue.length <= 1) fileUpload.clear();
			});
		},
	});
</script>

<div
	class="chat-node group relative flex h-full min-h-[200px] w-full max-w-[1000px]
	min-w-[700px] flex-col items-stretch rounded-2xl border border-gray-200 bg-white p-6 shadow-sm"
	bind:this={node}
>
	<!-- Model and Provider selectors -->
	<div class="mb-4 flex gap-3">
		<div class="flex-[3]">
			<ModelPicker modelId={data.modelId} onModelSelect={modelId => updateNodeData(id, { modelId })} />
		</div>
		<div class="flex-[2]">
			<ProviderPicker
				provider={data.provider}
				modelId={data.modelId}
				onProviderSelect={provider => updateNodeData(id, { provider })}
			/>
		</div>
	</div>

	<form class="flex flex-col" onsubmit={handleSubmit}>
		<div class="relative">
			<label for={`message-${id}`} class="mb-1.5 block text-xs font-medium text-gray-600">Message</label>

			<!-- Message container with top bar, textarea, and bottom bar -->
			<div
				class="rounded-lg border border-gray-200 bg-gray-50 focus-within:border-gray-900 focus-within:ring-2 focus-within:ring-gray-900/10"
				{...fileUpload.dropzone}
				onclick={undefined}
			>
				{#if fileUpload.isDragging}
					<div
						class="absolute inset-2 z-10 flex flex-col items-center justify-center rounded-xl bg-gray-800/50 backdrop-blur-md"
						transition:fade={{ duration: 100 }}
					>
						<IconImage />
						<p>Drop the image here to upload</p>
					</div>
				{/if}
				<!-- Top bar with buttons and character count -->
				<div class="flex items-center justify-between border-b border-gray-200 px-4 py-2">
					<div class="flex items-center gap-2">
						<button
							type="button"
							class="flex items-center gap-1.5 rounded-md px-2 py-1 text-xs text-gray-600 transition-colors"
							class:opacity-50={!supportsImgUpload}
							class:cursor-not-allowed={!supportsImgUpload}
							class:hover:bg-gray-200={supportsImgUpload}
							class:hover:text-gray-900={supportsImgUpload}
							title={supportsImgUpload ? "Attach file" : "Model doesn't support images"}
							{...fileUpload.trigger}
						>
							<IconAttachment class="h-4 w-4" />
							<input {...fileUpload.input} />
						</button>
						<button
							type="button"
							class="flex items-center gap-1.5 rounded-md px-2 py-1 text-xs text-gray-600 transition-colors hover:bg-gray-200 hover:text-gray-900"
							title="Code snippet"
							onclick={insertCodeSnippet}
						>
							<IconCode class="h-4 w-4" />
						</button>
					</div>
					<div class="text-xs text-gray-500">
						{characterCount}
					</div>
				</div>

				<!-- Textarea -->
				<textarea
					id={`message-${id}`}
					class="nodrag min-h-[80px] w-full resize-none border-0 bg-transparent px-4
					py-3 text-sm text-gray-900 placeholder-gray-500 focus:ring-0 focus:outline-none"
					placeholder="Type your message..."
					value={data.query}
					oninput={evt => {
						updateNodeData(id, { query: evt.currentTarget.value });
					}}
					{@attach autosized.attachment}
				></textarea>
			</div>
		</div>

		<div class="mt-2">
			<div class="flex items-center gap-2">
				{#each data.imageIds ?? [] as imgKey (imgKey)}
					{#await images.get(imgKey)}
						<!-- nothing -->
					{:then imgSrc}
						<div class="group/img relative">
							<button
								aria-label="expand"
								class="absolute inset-0 z-10 grid place-items-center rounded-md bg-gray-800/70 text-white opacity-0 group-hover/img:opacity-100"
								onclick={() => previewImage(imgSrc)}
							>
								<IconMaximize />
							</button>
							<img src={imgSrc} alt="uploaded" class="size-12 rounded-md object-cover" />
							<button
								aria-label="remove"
								type="button"
								onclick={async e => {
									e.stopPropagation();
									updateNodeData(id, { imageIds: data.imageIds?.filter(i => i !== imgKey) });
									images.delete(imgKey);
								}}
								class="invisible absolute -top-1 -right-1 z-20 grid size-5 place-items-center rounded-full bg-gray-800 text-xs text-white group-hover/img:visible hover:bg-gray-700"
							>
								✕
							</button>
						</div>
					{/await}
				{/each}
			</div>
		</div>

		<!-- Bottom bar with hint and send button (outside textarea) -->
		<div class="mt-2 flex items-center justify-between">
			<div class="text-xs text-gray-500">Shift + Enter for newline</div>
			<button
				type={isLoading ? "button" : "submit"}
				onclick={isLoading ? stopGeneration : undefined}
				class="flex items-center gap-2 rounded-xl
				bg-black px-4 py-2 text-sm font-medium
				text-white transition-all hover:scale-[1.02] hover:bg-gray-900
				focus:ring-2 focus:ring-gray-900/20 focus:outline-none
				active:scale-[0.98]"
			>
				{#if isLoading}
					<IconStop class="h-4 w-4" />
					Stop
				{:else}
					<IconSend class="h-4 w-4" />
					Send
				{/if}
			</button>
		</div>
	</form>

	{#if data.response || isLoading}
		<div class="mt-4">
			<div class="mb-2 flex items-center gap-2 text-xs font-medium text-gray-600">
				Response
				{#if isLoading}
					<IconLoading class="h-3 w-3 animate-spin" />
				{/if}
			</div>
			{#if data.response}
				<div class="prose prose-sm max-w-none text-gray-800">
					{@html parsedResponse}
				</div>
			{:else if isLoading}
				<div class="text-sm text-gray-500">Generating response...</div>
			{/if}
		</div>
	{/if}

	<!-- Action buttons -->
	<div class="abs-x-center absolute -bottom-4 z-10 flex gap-2 opacity-0 transition-all group-hover:opacity-100">
		<!-- Duplicate button -->
		<button
			class="flex items-center gap-1.5 rounded-full bg-gray-700
			px-4 py-2 text-xs font-medium text-white
			shadow-sm transition-all hover:scale-[1.02]
			hover:bg-gray-600 focus:ring-2 focus:ring-gray-600/20 focus:outline-none active:scale-[0.98]"
			onclick={() => {
				const curr = getNode(id);
				const newNodeId = crypto.randomUUID();
				const currentDims = getNodeDimensions(id);
				const preferredPos = findAvailablePosition(
					(curr?.position.x ?? 100) + 50,
					(curr?.position.y ?? 0) + 50,
					currentDims.width,
					currentDims.height,
				);
				const newNode: Node = {
					id: newNodeId,
					position: preferredPos,
					data: {
						query: data.query,
						response: data.response,
						modelId: data.modelId,
						provider: data.provider,
					},
					type: "chat",
					width: undefined,
					height: undefined,
				};
				nodes.current.push(newNode);

				// Copy only incoming edges (parent connections)
				const incomingEdges = edges.current.filter(edge => edge.target === id);
				for (const edge of incomingEdges) {
					const newEdge: Edge = {
						id: crypto.randomUUID(),
						source: edge.source,
						target: newNodeId,
						animated: edge.animated,
						label: edge.label,
						data: edge.data,
					};
					edges.current.push(newEdge);
				}
			}}
		>
			<IconCopy class="h-3 w-3" />
			Duplicate
		</button>

		<!-- Add node button -->
		<button
			class="flex items-center gap-1.5 rounded-full bg-black
			px-4 py-2 text-xs font-medium text-white
			shadow-sm transition-all hover:scale-[1.02]
			hover:bg-gray-900 focus:ring-2 focus:ring-gray-900/20 focus:outline-none active:scale-[0.98]"
			onclick={() => {
				const curr = getNode(id);
				const currentDims = getNodeDimensions(id);
				const preferredPos = findAvailablePosition(
					curr?.position.x ?? 100,
					(curr?.position.y ?? 0) + size.height + 40,
					currentDims.width,
					currentDims.height,
					true, // constrainBelow = true for Add Node
				);
				const newNode: Node = {
					id: crypto.randomUUID(),
					position: preferredPos,
					data: { query: "", response: "", modelId: data.modelId, provider: data.provider },
					type: "chat",
					width: undefined,
					height: undefined,
				};
				nodes.current.push(newNode);
				const edge: Edge = {
					id: crypto.randomUUID(),
					source: curr!.id,
					target: newNode.id,
					animated: true,
					label: "",
					data: {},
				};
				edges.current.push(edge);
			}}
		>
			<IconAdd class="h-3 w-3" />
			Add Node
		</button>
	</div>

	<!-- Close button -->
	<button
		class="absolute top-3 right-3 rounded-lg p-1.5 text-gray-400 transition-colors
		hover:bg-red-50 hover:text-red-500 focus:ring-2 focus:ring-red-500/20 focus:outline-none"
		onclick={() => (nodes.current = nodes.current.filter(n => n.id !== id))}
	>
		<IconX class="h-4 w-4" />
	</button>
</div>

<Handle type="target" position={Position.Top} class="h-3 w-3 border-2 border-white bg-gray-500 shadow-sm" />
<Handle
	type="source"
	position={Position.Bottom}
	class="h-3 w-3 border-2 border-white bg-gray-500 opacity-0 shadow-sm"
/>

<!-- <NodeResizeControl minWidth={200} minHeight={150}> -->
<!-- 	<IconResize class="absolute right-2 bottom-2" /> -->
<!-- </NodeResizeControl> -->