From f1afba70ff23ff3a665ee91c5a9e386e5490f8dc Mon Sep 17 00:00:00 2001 From: Florian Egger Date: Wed, 20 May 2026 22:46:42 +0200 Subject: [PATCH] Add structured memory injection and telemetry tracking - Introduce buildMessagesWithMemory() to prepend memory context as a system message before LLM calls - Record LLM call telemetry (tokens, duration) for follow-up requests in both streaming and non-streaming paths - Add telemetry coverage for tool execution (success/failure, args, duration) - Update tool-executor tests to verify telemetry integration with TelemetryManager --- src/chat-view.ts | 45 +++++++++++++++++- tests/tool-executor.test.ts | 95 +++++++++++++++++++++++++++++++++++-- 2 files changed, 135 insertions(+), 5 deletions(-) diff --git a/src/chat-view.ts b/src/chat-view.ts index 0890e43..02c2792 100755 --- a/src/chat-view.ts +++ b/src/chat-view.ts @@ -607,6 +607,18 @@ export class ChatView extends ItemView { return filterToolsForMode(allTools, this.currentAgentMode); } + /** + * Build messages for the LLM, injecting structured memory context if available. + * This is the canonical message builder used for all LLM calls in this view. + */ + private buildMessagesWithMemory(baseMessages: OllamaMessage[]): OllamaMessage[] { + const memoryContext = this.structuredMemoryManager?.buildMemoryContext(); + if (memoryContext) { + return [{ role: 'system', content: memoryContext }, ...baseMessages]; + } + return baseMessages; + } + buildMessages(userMessageContent: string, tools?: OllamaTool[]): OllamaMessage[] { const systemContent = getSystemPromptForMode(this.currentAgentMode); const messages: OllamaMessage[] = []; @@ -730,12 +742,25 @@ export class ChatView extends ItemView { if (followUpMessages.length > 0) { const finalMessages = [...messages, followUp, ...followUpMessages]; + const followUpStartTime = Date.now(); const response = await this.ollamaClient.chat(finalMessages, tools); + const followUpDurationMs = Date.now() - followUpStartTime; const finalResponse = response.content || fullResponse; this.updateMessageById(assistantMessageId, { content: finalResponse, isStreaming: false, }); + + // Record follow-up LLM call telemetry + this.telemetryManager?.recordLlmCall({ + model: this.settings.model, + promptTokens: Math.round(finalMessages.reduce((sum, m) => sum + m.content.length, 0) / 4), + completionTokens: Math.round(finalResponse.length / 4), + totalTokens: Math.round( + (finalMessages.reduce((sum, m) => sum + m.content.length, 0) + finalResponse.length) / 4 + ), + durationMs: followUpDurationMs, + }); } } @@ -777,12 +802,25 @@ export class ChatView extends ItemView { if (followUpMessages.length > 0) { const finalMessages = [...messages, followUp, ...followUpMessages]; + const followUpStartTime = Date.now(); const response = await this.ollamaClient.chat(finalMessages, tools); + const followUpDurationMs = Date.now() - followUpStartTime; const finalResponse = response.content || 'Actions applied successfully.'; this.updateMessageById(assistantMessageId, { content: finalResponse, isStreaming: false, }); + + // Record follow-up LLM call telemetry + this.telemetryManager?.recordLlmCall({ + model: this.settings.model, + promptTokens: Math.round(finalMessages.reduce((sum, m) => sum + m.content.length, 0) / 4), + completionTokens: Math.round(finalResponse.length / 4), + totalTokens: Math.round( + (finalMessages.reduce((sum, m) => sum + m.content.length, 0) + finalResponse.length) / 4 + ), + durationMs: followUpDurationMs, + }); } else { this.updateMessageById(assistantMessageId, { content: 'Actions applied successfully.', @@ -1035,7 +1073,10 @@ export class ChatView extends ItemView { const completeMessages = this.conversationStateManager.getCompleteMessages(userMessageWithContext); - const stream = this.ollamaClient.streamChat(completeMessages, tools); + // Prepend structured memory as a system message if available + const messagesWithMemory = this.buildMessagesWithMemory(completeMessages); + + const stream = this.ollamaClient.streamChat(messagesWithMemory, tools); let fullResponse = ''; let toolCalls: OllamaToolCall[] = []; @@ -1087,7 +1128,7 @@ export class ChatView extends ItemView { if (toolCalls.length > 0) { await this.processToolCalls( toolCalls, - completeMessages, + messagesWithMemory, tools, fullResponse, assistantMessageId diff --git a/tests/tool-executor.test.ts b/tests/tool-executor.test.ts index 9c207ec..9c3ef4a 100755 --- a/tests/tool-executor.test.ts +++ b/tests/tool-executor.test.ts @@ -2,6 +2,7 @@ import { ToolExecutor } from '../src/tool-executor'; import { TFile } from 'obsidian'; import { ToolCall, ToolResult } from '../src/types'; import { ErrorHandler } from '../src/error-handler'; +import { TelemetryManager } from '../src/tool-telemetry'; // Mock Obsidian types interface MockVault { @@ -14,9 +15,9 @@ interface MockVault { delete: (file: any) => Promise; } interface MockApp { - metadataCache: { - getFileCache: jest.Mock; - }; + metadataCache: { + getFileCache: jest.Mock; + }; // Mock app properties if needed } interface MockNotice { @@ -1242,4 +1243,92 @@ describe('ToolExecutor', () => { }); }); }); + + describe('telemetry integration', () => { + let telemetryManager: TelemetryManager; + let telemetryExecutor: ToolExecutor; + + beforeEach(() => { + telemetryManager = new TelemetryManager({ enabled: true, maxEntries: 100 }); + telemetryExecutor = new ToolExecutor( + mockVault as unknown as any, + mockApp as unknown as any, + telemetryManager + ); + }); + + it('should record successful tool calls in telemetry', async () => { + mockVault.create = jest.fn().mockResolvedValue(null); + + const call: ToolCall = { + id: 'call_t1', + type: 'function', + function: { + name: 'create_file', + arguments: JSON.stringify({ path: 'test.md', content: 'hello' }), + }, + }; + + await telemetryExecutor.handleToolCall(call); + const entries = telemetryManager.getEntriesByType('tool_call'); + expect(entries).toHaveLength(1); + expect((entries[0] as any).toolName).toBe('create_file'); + expect((entries[0] as any).success).toBe(true); + expect((entries[0] as any).durationMs).toBeGreaterThanOrEqual(0); + }); + + it('should record failed tool calls in telemetry', async () => { + const call: ToolCall = { + id: 'call_t2', + type: 'function', + function: { + name: 'create_file', + arguments: JSON.stringify({ path: '/invalid/path.md', content: 'hello' }), + }, + }; + + await expect(telemetryExecutor.handleToolCall(call)).rejects.toThrow(); + const entries = telemetryManager.getEntriesByType('tool_call'); + expect(entries).toHaveLength(1); + expect((entries[0] as any).toolName).toBe('create_file'); + expect((entries[0] as any).success).toBe(false); + }); + + it('should include parsed args in telemetry', async () => { + mockVault.create = jest.fn().mockResolvedValue(null); + + const call: ToolCall = { + id: 'call_t3', + type: 'function', + function: { + name: 'create_file', + arguments: JSON.stringify({ path: 'note.md', content: 'data' }), + }, + }; + + await telemetryExecutor.handleToolCall(call); + const entries = telemetryManager.getEntriesByType('tool_call'); + expect((entries[0] as any).args).toEqual({ path: 'note.md', content: 'data' }); + }); + + it('should not record telemetry when telemetry manager is undefined', async () => { + const noTelemetryExecutor = new ToolExecutor( + mockVault as unknown as any, + mockApp as unknown as any + ); + mockVault.create = jest.fn().mockResolvedValue(null); + + const call: ToolCall = { + id: 'call_t4', + type: 'function', + function: { + name: 'create_file', + arguments: JSON.stringify({ path: 'x.md', content: 'y' }), + }, + }; + + // Should not throw + await noTelemetryExecutor.handleToolCall(call); + }); + }); });