add streaming of ai text

2 files changed, 230 insertions, 18 deletions

diff --git a/src/chat-view.ts b/src/chat-view.ts
index 16b9676..9263c57 100644
--- a/src/chat-view.ts
+++ b/src/chat-view.ts
@@ -1,7 +1,7 @@
 import { ItemView, Notice, WorkspaceLeaf, setIcon } from "obsidian";
 import type AIOrganizer from "./main";
 import type { ChatMessage, ToolCallEvent } from "./ollama-client";
-import { sendChatMessage } from "./ollama-client";
+import { sendChatMessageStreaming } from "./ollama-client";
 import { SettingsModal } from "./settings-modal";
 import { ToolModal } from "./tool-modal";
 import { TOOL_REGISTRY } from "./tools";
@@ -16,6 +16,8 @@ export class ChatView extends ItemView {
 	private textarea: HTMLTextAreaElement | null = null;
 	private sendButton: HTMLButtonElement | null = null;
 	private toolsButton: HTMLButtonElement | null = null;
+	private abortController: AbortController | null = null;
+	private scrollDebounceTimer: ReturnType<typeof setTimeout> | null = null;
 
 	constructor(leaf: WorkspaceLeaf, plugin: AIOrganizer) {
 		super(leaf);
@@ -86,6 +88,11 @@ export class ChatView extends ItemView {
 		});
 
 		this.sendButton.addEventListener("click", () => {
+			if (this.abortController !== null) {
+				// Currently streaming — abort
+				this.abortController.abort();
+				return;
+			}
 			void this.handleSend();
 		});
 
@@ -94,12 +101,16 @@ export class ChatView extends ItemView {
 	}
 
 	async onClose(): Promise<void> {
+		if (this.abortController !== null) {
+			this.abortController.abort();
+		}
 		this.contentEl.empty();
 		this.messages = [];
 		this.messageContainer = null;
 		this.textarea = null;
 		this.sendButton = null;
 		this.toolsButton = null;
+		this.abortController = null;
 	}
 
 	private getEnabledTools(): OllamaToolDefinition[] {
@@ -146,8 +157,12 @@ export class ChatView extends ItemView {
 		// Track in message history
 		this.messages.push({ role: "user", content: text });
 
-		// Disable input
-		this.setInputEnabled(false);
+		// Switch to streaming state
+		this.abortController = new AbortController();
+		this.setStreamingState(true);
+
+		// Create the assistant bubble for streaming into
+		const streamingBubble = this.createStreamingBubble();
 
 		try {
 			const enabledTools = this.getEnabledTools();
@@ -158,30 +173,52 @@ export class ChatView extends ItemView {
 				this.scrollToBottom();
 			};
 
-			const response = await sendChatMessage(
-				this.plugin.settings.ollamaUrl,
-				this.plugin.settings.model,
-				this.messages,
-				hasTools ? enabledTools : undefined,
-				hasTools ? this.plugin.app : undefined,
-				hasTools ? onToolCall : undefined,
-			);
+			const onChunk = (chunk: string): void => {
+				streamingBubble.textContent += chunk;
+				this.debouncedScrollToBottom();
+			};
 
+			const response = await sendChatMessageStreaming({
+				ollamaUrl: this.plugin.settings.ollamaUrl,
+				model: this.plugin.settings.model,
+				messages: this.messages,
+				tools: hasTools ? enabledTools : undefined,
+				app: hasTools ? this.plugin.app : undefined,
+				onChunk,
+				onToolCall: hasTools ? onToolCall : undefined,
+				abortSignal: this.abortController.signal,
+			});
+
+			// Finalize the streaming bubble
+			streamingBubble.removeClass("ai-organizer-streaming");
 			this.messages.push({ role: "assistant", content: response });
-			this.appendMessage("assistant", response);
 			this.scrollToBottom();
 		} catch (err: unknown) {
+			// Finalize bubble even on error
+			streamingBubble.removeClass("ai-organizer-streaming");
+
 			const errMsg = err instanceof Error ? err.message : "Unknown error.";
 			new Notice(errMsg);
 			this.appendMessage("error", `Error: ${errMsg}`);
 			this.scrollToBottom();
 		}
 
-		// Re-enable input
-		this.setInputEnabled(true);
+		// Restore normal state
+		this.abortController = null;
+		this.setStreamingState(false);
 		this.textarea.focus();
 	}
 
+	private createStreamingBubble(): HTMLDivElement {
+		if (this.messageContainer === null) {
+			// Should not happen, but satisfy TS
+			throw new Error("Message container not initialized.");
+		}
+		return this.messageContainer.createDiv({
+			cls: "ai-organizer-message assistant ai-organizer-streaming",
+		});
+	}
+
 	private appendMessage(role: "user" | "assistant" | "error", content: string): void {
 		if (this.messageContainer === null) {
 			return;
@@ -227,13 +264,21 @@ export class ChatView extends ItemView {
 		}
 	}
 
-	private setInputEnabled(enabled: boolean): void {
+	private debouncedScrollToBottom(): void {
+		if (this.scrollDebounceTimer !== null) return;
+		this.scrollDebounceTimer = setTimeout(() => {
+			this.scrollDebounceTimer = null;
+			this.scrollToBottom();
+		}, 50);
+	}
+
+	private setStreamingState(streaming: boolean): void {
 		if (this.textarea !== null) {
-			this.textarea.disabled = !enabled;
+			this.textarea.disabled = streaming;
 		}
 		if (this.sendButton !== null) {
-			this.sendButton.disabled = !enabled;
-			this.sendButton.textContent = enabled ? "Send" : "...";
+			this.sendButton.textContent = streaming ? "Stop" : "Send";
+			this.sendButton.toggleClass("ai-organizer-stop-btn", streaming);
 		}
 	}
 }
diff --git a/src/ollama-client.ts b/src/ollama-client.ts
index 91bd40c..66badc6 100644
--- a/src/ollama-client.ts
+++ b/src/ollama-client.ts
@@ -190,3 +190,170 @@ export async function sendChatMessage(
 
 	throw new Error("Tool calling loop exceeded maximum iterations.");
 }
+
+/**
+ * Streaming chat options.
+ */
+export interface StreamingChatOptions {
+	ollamaUrl: string;
+	model: string;
+	messages: ChatMessage[];
+	tools?: OllamaToolDefinition[];
+	app?: App;
+	onChunk: (text: string) => void;
+	onToolCall?: (event: ToolCallEvent) => void;
+	abortSignal?: AbortSignal;
+}
+
+/**
+ * Parse ndjson lines from a streamed response body.
+ * Handles partial lines that may span across chunks from the reader.
+ */
+async function* readNdjsonStream(
+	reader: ReadableStreamDefaultReader<Uint8Array>,
+	decoder: TextDecoder,
+): AsyncGenerator<Record<string, unknown>> {
+	let buffer = "";
+
+	while (true) {
+		const { done, value } = await reader.read();
+		if (done) break;
+
+		buffer += decoder.decode(value, { stream: true });
+		const lines = buffer.split("\n");
+		// Last element may be incomplete — keep it in buffer
+		buffer = lines.pop() ?? "";
+
+		for (const line of lines) {
+			const trimmed = line.trim();
+			if (trimmed === "") continue;
+			yield JSON.parse(trimmed) as Record<string, unknown>;
+		}
+	}
+
+	// Process any remaining data in buffer
+	const trimmed = buffer.trim();
+	if (trimmed !== "") {
+		yield JSON.parse(trimmed) as Record<string, unknown>;
+	}
+}
+
+/**
+ * Send a chat message with streaming.
+ * Streams text chunks via onChunk callback. Supports tool-calling agent loop:
+ * tool execution rounds are non-streamed, only the final text response streams.
+ * Returns the full accumulated response text.
+ */
+export async function sendChatMessageStreaming(
+	opts: StreamingChatOptions,
+): Promise<string> {
+	const { ollamaUrl, model, messages, tools, app, onChunk, onToolCall, abortSignal } = opts;
+	const maxIterations = 10;
+	let iterations = 0;
+
+	const workingMessages = messages.map((m) => ({ ...m }));
+
+	while (iterations < maxIterations) {
+		iterations++;
+
+		const body: Record<string, unknown> = {
+			model,
+			messages: workingMessages,
+			stream: true,
+		};
+
+		if (tools !== undefined && tools.length > 0) {
+			body.tools = tools;
+		}
+
+		const response = await fetch(`${ollamaUrl}/api/chat`, {
+			method: "POST",
+			headers: { "Content-Type": "application/json" },
+			body: JSON.stringify(body),
+			signal: abortSignal,
+		});
+
+		if (!response.ok) {
+			throw new Error(`Ollama returned status ${response.status}.`);
+		}
+
+		if (response.body === null) {
+			throw new Error("Response body is null — streaming not supported.");
+		}
+
+		const reader = response.body.getReader();
+		const decoder = new TextDecoder();
+
+		let content = "";
+		const toolCalls: ToolCallResponse[] = [];
+
+		try {
+			for await (const chunk of readNdjsonStream(reader, decoder)) {
+				const msg = chunk.message as Record<string, unknown> | undefined;
+				if (msg !== undefined && msg !== null) {
+					if (typeof msg.content === "string" && msg.content !== "") {
+						content += msg.content;
+						onChunk(msg.content);
+					}
+					if (Array.isArray(msg.tool_calls)) {
+						toolCalls.push(...(msg.tool_calls as ToolCallResponse[]));
+					}
+				}
+			}
+		} catch (err: unknown) {
+			if (err instanceof DOMException && err.name === "AbortError") {
+				// User cancelled — return whatever we accumulated
+				return content;
+			}
+			throw err;
+		}
+
+		// If no tool calls, we're done
+		if (toolCalls.length === 0) {
+			return content;
+		}
+
+		// Tool calling: append assistant message and execute tools
+		const assistantMsg: ChatMessage = {
+			role: "assistant",
+			content,
+			tool_calls: toolCalls,
+		};
+		workingMessages.push(assistantMsg);
+
+		if (app === undefined) {
+			throw new Error("App reference required for tool execution.");
+		}
+
+		for (const tc of toolCalls) {
+			const fnName = tc.function.name;
+			const fnArgs = tc.function.arguments;
+			const toolEntry = findToolByName(fnName);
+
+			let result: string;
+			if (toolEntry === undefined) {
+				result = `Error: Unknown tool "${fnName}".`;
+			} else {
+				result = await toolEntry.execute(app, fnArgs);
+			}
+
+			if (onToolCall !== undefined) {
+				const friendlyName = toolEntry !== undefined ? toolEntry.friendlyName : fnName;
+				const summary = toolEntry !== undefined ? toolEntry.summarize(fnArgs) : `Called ${fnName}`;
+				const resultSummary = toolEntry !== undefined ? toolEntry.summarizeResult(result) : "";
+				onToolCall({ toolName: fnName, friendlyName, summary, resultSummary, args: fnArgs, result });
+			}
+
+			workingMessages.push({
+				role: "tool",
+				tool_name: fnName,
+				content: result,
+			});
+		}
+
+		// Reset content for next streaming round
+		// (tool call content was intermediate, next round streams the final answer)
+	}
+
+	throw new Error("Tool calling loop exceeded maximum iterations.");
+}