enable image attachments to the chatHEADmain

5 files changed, 340 insertions, 4 deletions

diff --git a/src/chat-view.ts b/src/chat-view.ts
index 55b730d..9261254 100644
--- a/src/chat-view.ts
+++ b/src/chat-view.ts
@@ -7,9 +7,17 @@ import { ToolModal } from "./tool-modal";
 import { TOOL_REGISTRY } from "./tools";
 import type { OllamaToolDefinition } from "./tools";
 import { collectVaultContext, formatVaultContext } from "./vault-context";
+import { setCurrentAttachments, clearCurrentAttachments } from "./image-attachments";
+import type { ImageAttachment } from "./image-attachments";
 
 export const VIEW_TYPE_CHAT = "ai-pulse-chat";
 
+interface PendingAttachment {
+	file: File;
+	dataUrl: string;
+	mimeType: string;
+}
+
 export class ChatView extends ItemView {
 	private plugin: AIPulse;
 	private messages: ChatMessage[] = [];
@@ -21,6 +29,10 @@ export class ChatView extends ItemView {
 	private scrollDebounceTimer: ReturnType<typeof setTimeout> | null = null;
 	private bubbleContent: Map<HTMLDivElement, string> = new Map();
 	private modelBadge: HTMLDivElement | null = null;
+	private pendingAttachments: PendingAttachment[] = [];
+	private attachmentStrip: HTMLDivElement | null = null;
+	private attachButton: HTMLButtonElement | null = null;
+	private fileInput: HTMLInputElement | null = null;
 
 	constructor(leaf: WorkspaceLeaf, plugin: AIPulse) {
 		super(leaf);
@@ -94,6 +106,7 @@ export class ChatView extends ItemView {
 			const modal = new ToolModal(this.plugin);
 			modal.onClose = () => {
 				this.updateToolsButtonState();
+				this.updateAttachButtonVisibility();
 			};
 			modal.open();
 			(document.activeElement as HTMLElement)?.blur();
@@ -110,6 +123,9 @@ export class ChatView extends ItemView {
 		clearBtn.addEventListener("click", () => {
 			this.messages = [];
 			this.bubbleContent.clear();
+			this.pendingAttachments = [];
+			clearCurrentAttachments();
+			this.renderAttachmentStrip();
 			if (this.messageContainer !== null) {
 				this.messageContainer.empty();
 			}
@@ -117,6 +133,39 @@ export class ChatView extends ItemView {
 		});
 
 		const inputRow = messagesArea.createDiv({ cls: "ai-pulse-input-row" });
+
+		// --- Attachment preview strip (above input controls) ---
+		this.attachmentStrip = inputRow.createDiv({ cls: "ai-pulse-attachment-strip" });
+		this.attachmentStrip.style.display = "none";
+
+		// --- Attach button (left of textarea) ---
+		this.attachButton = inputRow.createEl("button", {
+			cls: "ai-pulse-attach-btn",
+			attr: { "aria-label": "Attach image" },
+		});
+		setIcon(this.attachButton, "image-plus");
+		this.updateAttachButtonVisibility();
+
+		// Hidden file input
+		this.fileInput = inputRow.createEl("input", {
+			type: "file",
+			attr: {
+				accept: "image/jpeg,image/png,image/gif,image/webp,image/bmp,image/svg+xml",
+				multiple: "",
+				style: "display:none",
+			},
+		});
+
+		this.attachButton.addEventListener("click", () => {
+			this.fileInput?.click();
+		});
+
+		this.fileInput.addEventListener("change", () => {
+			if (this.fileInput === null || this.fileInput.files === null) return;
+			void this.handleFileSelection(this.fileInput.files);
+			this.fileInput.value = ""; // Reset so same file can be re-selected
+		});
+
 		this.textarea = inputRow.createEl("textarea", {
 			attr: { placeholder: "Type a message...", rows: "2" },
 		});
@@ -151,12 +200,17 @@ export class ChatView extends ItemView {
 		this.contentEl.empty();
 		this.messages = [];
 		this.bubbleContent.clear();
+		this.pendingAttachments = [];
+		clearCurrentAttachments();
 		this.messageContainer = null;
 		this.textarea = null;
 		this.sendButton = null;
 		this.toolsButton = null;
 		this.modelBadge = null;
 		this.abortController = null;
+		this.attachmentStrip = null;
+		this.attachButton = null;
+		this.fileInput = null;
 	}
 
 	private getEnabledTools(): OllamaToolDefinition[] {
@@ -185,6 +239,12 @@ export class ChatView extends ItemView {
 		this.toolsButton.toggleClass("ai-pulse-tools-active", this.hasAnyToolEnabled());
 	}
 
+	private updateAttachButtonVisibility(): void {
+		if (this.attachButton === null) return;
+		const visible = this.plugin.settings.enabledTools["save_image"] === true;
+		this.attachButton.style.display = visible ? "" : "none";
+	}
+
 	private updateModelBadge(): void {
 		if (this.modelBadge === null) return;
 		const model = this.plugin.settings.model;
@@ -203,7 +263,7 @@ export class ChatView extends ItemView {
 		}
 
 		const text = this.textarea.value.trim();
-		if (text === "") {
+		if (text === "" && this.pendingAttachments.length === 0) {
 			return;
 		}
 
@@ -212,13 +272,44 @@ export class ChatView extends ItemView {
 			return;
 		}
 
-		// Append user message
+		// Convert pending attachments to ImageAttachment format for the tool
+		let messageContent = text;
+		if (this.pendingAttachments.length > 0) {
+			const imageAttachments: ImageAttachment[] = [];
+			for (const pa of this.pendingAttachments) {
+				const arrayBuffer = await pa.file.arrayBuffer();
+				const bytes = new Uint8Array(arrayBuffer);
+				let binary = "";
+				for (const b of bytes) binary += String.fromCharCode(b);
+				const base64 = btoa(binary);
+
+				imageAttachments.push({
+					base64,
+					mimeType: pa.mimeType,
+					originalName: pa.file.name,
+					arrayBuffer,
+				});
+			}
+
+			// Set the module-level attachments for the save_image tool to access
+			setCurrentAttachments(imageAttachments);
+
+			// Prepend context note to the message for the LLM
+			const count = this.pendingAttachments.length;
+			messageContent = `[${count} image(s) are attached to this message. You MUST use the save_image tool to save them to the vault. Infer from the user's message how and where these images should be saved and embedded. Assume the user wants the images attached to whatever note they are asking you to create or edit.]\n\n${text}`;
+
+			// Clear the UI attachments
+			this.pendingAttachments = [];
+			this.renderAttachmentStrip();
+		}
+
+		// Append user message (show original text in UI)
 		this.appendMessage("user", text);
 		this.textarea.value = "";
 		this.scrollToBottom();
 
-		// Track in message history
-		this.messages.push({ role: "user", content: text });
+		// Track the augmented message in history for the LLM
+		this.messages.push({ role: "user", content: messageContent });
 
 		// Switch to streaming state
 		this.abortController = new AbortController();
@@ -324,6 +415,9 @@ export class ChatView extends ItemView {
 		} catch (err: unknown) {
 			const isAbort = err instanceof DOMException && err.name === "AbortError";
 
+			// Clear stale attachments on error
+			clearCurrentAttachments();
+
 			// Clean up the streaming bubble
 			// eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion
 			const bubble = currentBubble as HTMLDivElement | null;
@@ -737,6 +831,71 @@ export class ChatView extends ItemView {
 		}
 	}
 
+	private async handleFileSelection(files: FileList): Promise<void> {
+		for (let i = 0; i < files.length; i++) {
+			const file = files[i];
+			if (file === undefined) continue;
+			if (!file.type.startsWith("image/")) continue;
+
+			const dataUrl = await this.readFileAsDataUrl(file);
+			this.pendingAttachments.push({
+				file,
+				dataUrl,
+				mimeType: file.type,
+			});
+		}
+		this.renderAttachmentStrip();
+	}
+
+	private readFileAsDataUrl(file: File): Promise<string> {
+		return new Promise<string>((resolve, reject) => {
+			const reader = new FileReader();
+			reader.onload = () => {
+				if (typeof reader.result === "string") {
+					resolve(reader.result);
+				} else {
+					reject(new Error("Failed to read file as data URL."));
+				}
+			};
+			reader.onerror = () => reject(new Error("File read error."));
+			reader.readAsDataURL(file);
+		});
+	}
+
+	private renderAttachmentStrip(): void {
+		if (this.attachmentStrip === null) return;
+		this.attachmentStrip.empty();
+
+		if (this.pendingAttachments.length === 0) {
+			this.attachmentStrip.style.display = "none";
+			return;
+		}
+
+		this.attachmentStrip.style.display = "flex";
+
+		for (let i = 0; i < this.pendingAttachments.length; i++) {
+			const attachment = this.pendingAttachments[i];
+			if (attachment === undefined) continue;
+
+			const thumb = this.attachmentStrip.createDiv({ cls: "ai-pulse-attachment-thumb" });
+			thumb.createEl("img", {
+				attr: { src: attachment.dataUrl, alt: attachment.file.name },
+			});
+
+			const removeBtn = thumb.createEl("button", {
+				cls: "ai-pulse-attachment-remove",
+				attr: { "aria-label": "Remove" },
+			});
+			setIcon(removeBtn, "x");
+
+			const index = i;
+			removeBtn.addEventListener("click", () => {
+				this.pendingAttachments.splice(index, 1);
+				this.renderAttachmentStrip();
+			});
+		}
+	}
+
 	private scrollToBottom(): void {
 		if (this.messageContainer === null) return;
 		const lastChild = this.messageContainer.lastElementChild;
@@ -769,5 +928,8 @@ export class ChatView extends ItemView {
 			this.sendButton.textContent = streaming ? "Stop" : "Send";
 			this.sendButton.toggleClass("ai-pulse-stop-btn", streaming);
 		}
+		if (this.attachButton !== null) {
+			this.attachButton.disabled = streaming;
+		}
 	}
 }
diff --git a/src/context/tools/save-image.json b/src/context/tools/save-image.json
new file mode 100644
index 0000000..6b8bb47
--- /dev/null
+++ b/src/context/tools/save-image.json
@@ -0,0 +1,26 @@
+{
+	"id": "save_image",
+	"label": "Save Image",
+	"description": "Save attached image(s) to the vault at a specified path.",
+	"friendlyName": "Save Image",
+	"requiresApproval": true,
+	"definition": {
+		"type": "function",
+		"function": {
+			"name": "save_image",
+			"description": "Save image(s) attached to the current chat message into the vault. The user has attached image(s) to their message — this tool writes them as files. You provide the vault-relative path WITHOUT the file extension (the correct extension is detected automatically from the image type). If multiple images are attached and you provide a single path, they will be saved as path_1.ext, path_2.ext, etc. If no images are attached, this tool returns an error. This action requires user approval.",
+			"parameters": {
+				"type": "object",
+				"required": [
+					"file_path"
+				],
+				"properties": {
+					"file_path": {
+						"type": "string",
+						"description": "The vault-relative path for the image WITHOUT the file extension. The extension is added automatically based on the image type (e.g., .jpg, .png). Example: 'attachments/cool-keyboard' will become 'attachments/cool-keyboard.jpg'. For multiple images with a single path, they are numbered: 'attachments/cool-keyboard_1.jpg', 'attachments/cool-keyboard_2.jpg', etc."
+					}
+				}
+			}
+		}
+	}
+}
diff --git a/src/image-attachments.ts b/src/image-attachments.ts
new file mode 100644
index 0000000..add2afd
--- /dev/null
+++ b/src/image-attachments.ts
@@ -0,0 +1,24 @@
+export interface ImageAttachment {
+	base64: string;
+	mimeType: string;
+	originalName: string;
+	arrayBuffer: ArrayBuffer;
+}
+
+let currentAttachments: ImageAttachment[] = [];
+
+export function setCurrentAttachments(attachments: ImageAttachment[]): void {
+	currentAttachments = attachments;
+}
+
+export function getCurrentAttachments(): ImageAttachment[] {
+	return currentAttachments;
+}
+
+export function clearCurrentAttachments(): void {
+	currentAttachments = [];
+}
+
+export function hasCurrentAttachments(): boolean {
+	return currentAttachments.length > 0;
+}
diff --git a/src/ollama-client.ts b/src/ollama-client.ts
index b778798..f86e5e7 100644
--- a/src/ollama-client.ts
+++ b/src/ollama-client.ts
@@ -2,6 +2,7 @@ import { Platform, requestUrl, TFile } from "obsidian";
 import type { App } from "obsidian";
 import type { OllamaToolDefinition } from "./tools";
 import { findToolByName } from "./tools";
+import { hasCurrentAttachments } from "./image-attachments";
 import systemPromptData from "./context/system-prompt.json";
 import markdownRulesData from "./context/obsidian-markdown-rules.json";
 
@@ -293,6 +294,8 @@ function preValidateTool(
 			return preValidateBatchMoveFile(app, args);
 		case "batch_set_frontmatter":
 			return preValidateBatchSetFrontmatter(app, args);
+		case "save_image":
+			return preValidateSaveImage(args);
 		default:
 			return null;
 	}
@@ -401,6 +404,19 @@ function preValidateSetFrontmatter(app: App, args: Record<string, unknown>): str
 	return null;
 }
 
+function preValidateSaveImage(args: Record<string, unknown>): string | null {
+	const filePath = typeof args["file_path"] === "string" ? args["file_path"] : "";
+	if (filePath === "") {
+		return "Error: file_path parameter is required.";
+	}
+
+	if (!hasCurrentAttachments()) {
+		return "Error: No images are attached to the current message. The user must attach images before you can save them.";
+	}
+
+	return null;
+}
+
 // -- Batch tool validators -----------------------------------------------------
 
 function preValidateBatchEditFile(app: App, args: Record<string, unknown>): string | null {
diff --git a/src/tools.ts b/src/tools.ts
index df45b96..0cc5b51 100644
--- a/src/tools.ts
+++ b/src/tools.ts
@@ -17,6 +17,8 @@ import batchDeleteFileCtx from "./context/tools/batch-delete-file.json";
 import batchMoveFileCtx from "./context/tools/batch-move-file.json";
 import batchSetFrontmatterCtx from "./context/tools/batch-set-frontmatter.json";
 import batchEditFileCtx from "./context/tools/batch-edit-file.json";
+import saveImageCtx from "./context/tools/save-image.json";
+import { getCurrentAttachments, clearCurrentAttachments } from "./image-attachments";
 
 /**
  * Schema for an Ollama tool definition (function calling).
@@ -459,6 +461,91 @@ async function executeSetFrontmatter(app: App, args: Record<string, unknown>): P
 }
 
 // ---------------------------------------------------------------------------
+// Save image tool
+// ---------------------------------------------------------------------------
+
+/**
+ * Map MIME types to file extensions.
+ */
+function mimeToExtension(mimeType: string): string {
+	const map: Record<string, string> = {
+		"image/jpeg": ".jpg",
+		"image/png": ".png",
+		"image/gif": ".gif",
+		"image/webp": ".webp",
+		"image/bmp": ".bmp",
+		"image/svg+xml": ".svg",
+	};
+	return map[mimeType] ?? ".png";
+}
+
+/**
+ * Execute the "save_image" tool.
+ * Saves attached image(s) to the vault at the specified path.
+ */
+async function executeSaveImage(app: App, args: Record<string, unknown>): Promise<string> {
+	const filePath = typeof args["file_path"] === "string" ? args["file_path"] : "";
+	if (filePath === "") {
+		return "Error: file_path parameter is required.";
+	}
+
+	const attachments = getCurrentAttachments();
+	if (attachments.length === 0) {
+		return "Error: No images are attached to the current message.";
+	}
+
+	const savedPaths: string[] = [];
+	const errors: string[] = [];
+
+	for (let i = 0; i < attachments.length; i++) {
+		const attachment = attachments[i];
+		if (attachment === undefined) continue;
+
+		const ext = mimeToExtension(attachment.mimeType);
+		const fullPath = attachments.length === 1
+			? `${filePath}${ext}`
+			: `${filePath}_${i + 1}${ext}`;
+
+		// Check if file already exists
+		const existing = app.vault.getAbstractFileByPath(fullPath);
+		if (existing !== null) {
+			errors.push(`"${fullPath}" already exists — skipped.`);
+			continue;
+		}
+
+		// Ensure parent folder exists
+		const lastSlash = fullPath.lastIndexOf("/");
+		if (lastSlash > 0) {
+			const folderPath = fullPath.substring(0, lastSlash);
+			const folder = app.vault.getFolderByPath(folderPath);
+			if (folder === null) {
+				await app.vault.createFolder(folderPath);
+			}
+		}
+
+		try {
+			await app.vault.createBinary(fullPath, attachment.arrayBuffer);
+			savedPaths.push(fullPath);
+		} catch (err: unknown) {
+			const msg = err instanceof Error ? err.message : "Unknown error";
+			errors.push(`"${fullPath}": ${msg}`);
+		}
+	}
+
+	clearCurrentAttachments();
+
+	const parts: string[] = [];
+	if (savedPaths.length > 0) {
+		parts.push(`Saved ${savedPaths.length} image(s):\n${savedPaths.map(p => `- ${p}`).join("\n")}`);
+	}
+	if (errors.length > 0) {
+		parts.push(`Errors:\n${errors.map(e => `- ${e}`).join("\n")}`);
+	}
+
+	return parts.join("\n\n");
+}
+
+// ---------------------------------------------------------------------------
 // Batch tool execute functions
 // ---------------------------------------------------------------------------
 
@@ -858,6 +945,27 @@ export const TOOL_REGISTRY: ToolEntry[] = [
 		},
 		execute: executeSetFrontmatter,
 	},
+	{
+		...asToolContext(saveImageCtx as Record<string, unknown>),
+		approvalMessage: (args) => {
+			const filePath = typeof args["file_path"] === "string" ? args["file_path"] : "unknown";
+			const count = getCurrentAttachments().length;
+			return `Save ${count} image(s) to "${filePath}"?`;
+		},
+		summarize: (args) => {
+			const filePath = typeof args["file_path"] === "string" ? args["file_path"] : "";
+			const count = getCurrentAttachments().length;
+			return `${count} image(s) \u2192 "/${filePath}"`;
+		},
+		summarizeResult: (result) => {
+			if (result.startsWith("Error")) return result;
+			if (result.includes("declined")) return "Declined by user";
+			const match = result.match(/Saved (\d+) image/);
+			if (match !== null) return `${match[1]} image(s) saved`;
+			return "Images saved";
+		},
+		execute: executeSaveImage,
+	},
 	// --- Batch tools ---
 	{
 		...asToolContext(batchSearchFilesCtx as Record<string, unknown>),