diff --git a/AIClasses/AIFunctionCall.ts b/AIClasses/AIFunctionCall.ts index 7042f6a..bc499af 100644 --- a/AIClasses/AIFunctionCall.ts +++ b/AIClasses/AIFunctionCall.ts @@ -1,10 +1,12 @@ +import { AIFunction } from "Enums/AIFunction"; + // platform agnostic function call class used to execute the requested function export class AIFunctionCall { - public readonly name: string; + public readonly name: AIFunction; public readonly arguments: Record; public readonly toolId?: string; - constructor(name: string, args: Record, toolId?: string) { + constructor(name: AIFunction, args: Record, toolId?: string) { this.name = name; this.arguments = args; this.toolId = toolId; diff --git a/AIClasses/Claude/Claude.ts b/AIClasses/Claude/Claude.ts index 204ce58..9acc56a 100644 --- a/AIClasses/Claude/Claude.ts +++ b/AIClasses/Claude/Claude.ts @@ -6,6 +6,7 @@ import { StreamingService, type IStreamChunk } from "Services/StreamingService"; import type { Conversation } from "Conversations/Conversation"; import { AIProvider, AIProviderURL } from "Enums/ApiProvider"; import { AIFunctionCall } from "AIClasses/AIFunctionCall"; +import { fromString as aiFunctionFromString } from "Enums/AIFunction"; import type { IAIFunctionDefinition } from "AIClasses/FunctionDefinitions/IAIFunctionDefinition"; import type { AIFunctionDefinitions } from "AIClasses/FunctionDefinitions/AIFunctionDefinitions"; import { isValidJson } from "Helpers/Helpers"; @@ -114,7 +115,7 @@ export class Claude implements IAIClass { try { const args = JSON.parse(this.accumulatedFunctionArgs) as Record; functionCall = new AIFunctionCall( - this.accumulatedFunctionName, + aiFunctionFromString(this.accumulatedFunctionName), args as Record, this.accumulatedFunctionId || undefined ); diff --git a/AIClasses/Gemini/Gemini.ts b/AIClasses/Gemini/Gemini.ts index 6e03988..ff83508 100644 --- a/AIClasses/Gemini/Gemini.ts +++ b/AIClasses/Gemini/Gemini.ts @@ -7,6 +7,7 @@ import type { Conversation } from "Conversations/Conversation"; import { Role } from "Enums/Role"; import { AIProvider, AIProviderURL } from "Enums/ApiProvider"; import { AIFunctionCall } from "AIClasses/AIFunctionCall"; +import { fromString as aiFunctionFromString } from "Enums/AIFunction"; import type { IAIFunctionDefinition } from "AIClasses/FunctionDefinitions/IAIFunctionDefinition"; import type VaultkeeperAIPlugin from "main"; import type { AIFunctionDefinitions } from "AIClasses/FunctionDefinitions/AIFunctionDefinitions"; @@ -139,7 +140,7 @@ export class Gemini implements IAIClass { // If streaming is complete and we have accumulated a function call, return it if (isComplete && this.accumulatedFunctionName) { functionCall = new AIFunctionCall( - this.accumulatedFunctionName, + aiFunctionFromString(this.accumulatedFunctionName), this.accumulatedFunctionArgs as Record ); } diff --git a/AIClasses/OpenAI/OpenAI.ts b/AIClasses/OpenAI/OpenAI.ts index 7b51ef9..783273a 100644 --- a/AIClasses/OpenAI/OpenAI.ts +++ b/AIClasses/OpenAI/OpenAI.ts @@ -6,6 +6,7 @@ import { StreamingService, type IStreamChunk } from "Services/StreamingService"; import type { Conversation } from "Conversations/Conversation"; import { AIProvider, AIProviderURL } from "Enums/ApiProvider"; import { AIFunctionCall } from "AIClasses/AIFunctionCall"; +import { fromString as aiFunctionFromString } from "Enums/AIFunction"; import type { IAIFunctionDefinition } from "AIClasses/FunctionDefinitions/IAIFunctionDefinition"; import type VaultkeeperAIPlugin from "main"; import type { AIFunctionDefinitions } from "AIClasses/FunctionDefinitions/AIFunctionDefinitions"; @@ -225,7 +226,7 @@ export class OpenAI implements IAIClass { try { const args = JSON.parse(firstToolCall.arguments) as Record; functionCall = new AIFunctionCall( - firstToolCall.name, + aiFunctionFromString(firstToolCall.name), args as Record, firstToolCall.id || undefined ); diff --git a/Enums/AIFunction.ts b/Enums/AIFunction.ts index 818f476..ebd5d86 100644 --- a/Enums/AIFunction.ts +++ b/Enums/AIFunction.ts @@ -8,4 +8,12 @@ export enum AIFunction { // only used by gemini RequestWebSearch = "request_web_search", +} + +export function fromString(functionName: string): AIFunction { + const enumValue = Object.values(AIFunction).find((value: string) => value === functionName); + if (enumValue) { + return enumValue as AIFunction; + } + throw new Error(`Unknown function name: ${functionName}`); } \ No newline at end of file diff --git a/Services/AIFunctionService.ts b/Services/AIFunctionService.ts index 34a40af..13dd482 100644 --- a/Services/AIFunctionService.ts +++ b/Services/AIFunctionService.ts @@ -97,7 +97,7 @@ export class AIFunctionService { case AIFunction.RequestWebSearch: return new AIFunctionResponse(functionCall.name, {}, functionCall.toolId) - default: + default: { const error = `Unknown function request ${functionCall.name}` console.error(error); return new AIFunctionResponse( @@ -105,6 +105,7 @@ export class AIFunctionService { { error: error }, functionCall.toolId ); + } } } @@ -142,7 +143,7 @@ export class AIFunctionService { } private async writeVaultFile(filePath: string, content: string): Promise { - const result: any = await this.fileSystemService.writeFile(normalizePath(filePath), content); + const result: boolean | unknown = await this.fileSystemService.writeFile(normalizePath(filePath), content); return typeof result === "boolean" ? { success: result } : { success: false, error: result }; } diff --git a/__tests__/AIClasses/Claude.test.ts b/__tests__/AIClasses/Claude.test.ts index 0c5846b..3bb9c86 100644 --- a/__tests__/AIClasses/Claude.test.ts +++ b/__tests__/AIClasses/Claude.test.ts @@ -62,7 +62,7 @@ describe('Claude', () => { mockFunctionDefinitions = { getQueryActions: vi.fn().mockReturnValue([ { - name: 'test_function', + name: 'search_vault_filestion', description: 'Test function', parameters: { type: 'object', @@ -130,7 +130,7 @@ describe('Claude', () => { type: 'content_block_start', content_block: { type: 'tool_use', - name: 'search_files', + name: 'search_vault_files', id: 'tool_123' } }); @@ -148,7 +148,7 @@ describe('Claude', () => { type: 'content_block_start', content_block: { type: 'tool_use', - name: 'search_files', + name: 'search_vault_files', id: 'tool_123' } })); @@ -176,14 +176,14 @@ describe('Claude', () => { })); expect(result.functionCall).toBeDefined(); - expect(result.functionCall?.name).toBe('search_files'); + expect(result.functionCall?.name).toBe('search_vault_files'); expect(result.functionCall?.arguments).toEqual({ query: 'test' }); expect(result.functionCall?.toolId).toBe('tool_123'); }); it('should handle content_block_stop and finalize function call', () => { // Setup - (claude as any).accumulatedFunctionName = 'test_func'; + (claude as any).accumulatedFunctionName = 'search_vault_files'; (claude as any).accumulatedFunctionArgs = '{"param":"value"}'; (claude as any).accumulatedFunctionId = 'func_456'; @@ -194,7 +194,7 @@ describe('Claude', () => { const result = (claude as any).parseStreamChunk(chunk); expect(result.functionCall).toBeDefined(); - expect(result.functionCall?.name).toBe('test_func'); + expect(result.functionCall?.name).toBe('search_vault_files'); expect(result.functionCall?.arguments).toEqual({ param: 'value' }); expect(result.functionCall?.toolId).toBe('func_456'); @@ -240,7 +240,7 @@ describe('Claude', () => { it('should handle invalid JSON in function arguments gracefully', () => { // Setup - (claude as any).accumulatedFunctionName = 'test_func'; + (claude as any).accumulatedFunctionName = 'search_vault_files'; (claude as any).accumulatedFunctionArgs = 'invalid json {'; (claude as any).accumulatedFunctionId = 'func_789'; @@ -300,7 +300,7 @@ describe('Claude', () => { JSON.stringify({ functionCall: { id: 'call_123', - name: 'search_files', + name: 'search_vault_files', args: { query: 'test' } } }), @@ -315,7 +315,7 @@ describe('Claude', () => { expect(result[0].content[0]).toEqual({ type: 'tool_use', id: 'call_123', - name: 'search_files', + name: 'search_vault_files', input: { query: 'test' } }); }); @@ -412,7 +412,7 @@ describe('Claude', () => { JSON.stringify({ functionCall: { id: 'call_456', - name: 'search_files', + name: 'search_vault_files', args: { query: 'test' } } }), @@ -433,7 +433,7 @@ describe('Claude', () => { it('should map function definitions to Claude tool format', () => { const definitions = [ { - name: 'search_files', + name: 'search_vault_files', description: 'Search for files', parameters: { type: 'object', @@ -459,7 +459,7 @@ describe('Claude', () => { expect(result).toHaveLength(2); expect(result[0]).toEqual({ - name: 'search_files', + name: 'search_vault_files', description: 'Search for files', input_schema: definitions[0].parameters }); diff --git a/__tests__/AIClasses/Gemini.test.ts b/__tests__/AIClasses/Gemini.test.ts index 0cc3347..555585d 100644 --- a/__tests__/AIClasses/Gemini.test.ts +++ b/__tests__/AIClasses/Gemini.test.ts @@ -62,7 +62,7 @@ describe('Gemini', () => { mockFunctionDefinitions = { getQueryActions: vi.fn().mockReturnValue([ { - name: 'test_function', + name: 'search_vault_filestion', description: 'Test function', parameters: { type: 'object', @@ -145,7 +145,7 @@ describe('Gemini', () => { content: { parts: [{ functionCall: { - name: 'search_files', + name: 'search_vault_files', args: { query: 'test' } @@ -157,7 +157,7 @@ describe('Gemini', () => { (gemini as any).parseStreamChunk(chunk); - expect((gemini as any).accumulatedFunctionName).toBe('search_files'); + expect((gemini as any).accumulatedFunctionName).toBe('search_vault_files'); expect((gemini as any).accumulatedFunctionArgs).toEqual({ query: 'test' }); }); @@ -168,7 +168,7 @@ describe('Gemini', () => { content: { parts: [{ functionCall: { - name: 'test_func', + name: 'search_vault_files', args: { param1: 'value1' } @@ -201,7 +201,7 @@ describe('Gemini', () => { it('should finalize function call on completion', () => { // Setup accumulated state - (gemini as any).accumulatedFunctionName = 'search_files'; + (gemini as any).accumulatedFunctionName = 'search_vault_files'; (gemini as any).accumulatedFunctionArgs = { query: 'test' }; const chunk = JSON.stringify({ @@ -215,7 +215,7 @@ describe('Gemini', () => { expect(result.isComplete).toBe(true); expect(result.shouldContinue).toBe(true); expect(result.functionCall).toBeDefined(); - expect(result.functionCall?.name).toBe('search_files'); + expect(result.functionCall?.name).toBe('search_vault_files'); expect(result.functionCall?.arguments).toEqual({ query: 'test' }); }); @@ -357,7 +357,7 @@ describe('Gemini', () => { '', JSON.stringify({ functionCall: { - name: 'search_files', + name: 'search_vault_files', args: { query: 'test' } } }), @@ -372,7 +372,7 @@ describe('Gemini', () => { expect(result[0].parts).toHaveLength(1); expect(result[0].parts[0]).toEqual({ functionCall: { - name: 'search_files', + name: 'search_vault_files', args: { query: 'test' } } }); @@ -381,7 +381,7 @@ describe('Gemini', () => { it('should convert function response to Gemini format', () => { const responseContent = JSON.stringify({ functionResponse: { - name: 'search_files', + name: 'search_vault_files', response: ['file1.txt', 'file2.txt'] } }); @@ -399,7 +399,7 @@ describe('Gemini', () => { // Gemini API requires both 'name' and 'response' fields expect(result[0].parts[0]).toEqual({ functionResponse: { - name: 'search_files', + name: 'search_vault_files', response: ['file1.txt', 'file2.txt'] } }); @@ -466,7 +466,7 @@ describe('Gemini', () => { it('should map function definitions to Gemini format', () => { const definitions = [ { - name: 'search_files', + name: 'search_vault_files', description: 'Search for files', parameters: { type: 'object', @@ -481,7 +481,7 @@ describe('Gemini', () => { expect(result).toHaveLength(1); expect(result[0]).toEqual({ - name: 'search_files', + name: 'search_vault_files', description: 'Search for files', parameters: definitions[0].parameters }); diff --git a/__tests__/AIClasses/OpenAI.test.ts b/__tests__/AIClasses/OpenAI.test.ts index e7af1b2..4f726f8 100644 --- a/__tests__/AIClasses/OpenAI.test.ts +++ b/__tests__/AIClasses/OpenAI.test.ts @@ -62,7 +62,7 @@ describe('OpenAI', () => { mockFunctionDefinitions = { getQueryActions: vi.fn().mockReturnValue([ { - name: 'test_function', + name: 'search_vault_filestion', description: 'Test function', parameters: { type: 'object', @@ -139,7 +139,7 @@ describe('OpenAI', () => { index: 0, id: 'call_123', function: { - name: 'search_files', + name: 'search_vault_files', arguments: '{"qu' } }] @@ -178,7 +178,7 @@ describe('OpenAI', () => { expect(result.isComplete).toBe(true); expect(result.shouldContinue).toBe(true); expect(result.functionCall).toBeDefined(); - expect(result.functionCall?.name).toBe('search_files'); + expect(result.functionCall?.name).toBe('search_vault_files'); expect(result.functionCall?.arguments).toEqual({ query: 'test' }); expect(result.functionCall?.toolId).toBe('call_123'); }); @@ -193,7 +193,7 @@ describe('OpenAI', () => { index: 0, id: 'call_1', function: { - name: 'first_func', + name: 'search_vault_files', arguments: '{"a":1}' } }, @@ -201,7 +201,7 @@ describe('OpenAI', () => { index: 1, id: 'call_2', function: { - name: 'second_func', + name: 'read_vault_files', arguments: '{"b":2}' } } @@ -224,7 +224,7 @@ describe('OpenAI', () => { // Should only return the first tool call (index 0) expect(result.functionCall).toBeDefined(); - expect(result.functionCall?.name).toBe('first_func'); + expect(result.functionCall?.name).toBe('search_vault_files'); expect((openai as any).accumulatedToolCalls.size).toBe(2); }); @@ -262,7 +262,7 @@ describe('OpenAI', () => { // Setup invalid arguments (openai as any).accumulatedToolCalls.set(0, { id: 'call_123', - name: 'test_func', + name: 'search_vault_files', arguments: 'invalid json {' }); @@ -344,7 +344,7 @@ describe('OpenAI', () => { JSON.stringify({ functionCall: { id: 'call_123', - name: 'search_files', + name: 'search_vault_files', args: { query: 'test' } } }), @@ -370,7 +370,7 @@ describe('OpenAI', () => { id: 'call_123', type: 'function', function: { - name: 'search_files', + name: 'search_vault_files', arguments: '{"query":"test"}' } }); @@ -496,7 +496,7 @@ describe('OpenAI', () => { it('should map function definitions to OpenAI tool format', () => { const definitions = [ { - name: 'search_files', + name: 'search_vault_files', description: 'Search for files', parameters: { type: 'object', @@ -524,7 +524,7 @@ describe('OpenAI', () => { expect(result[0]).toEqual({ type: 'function', function: { - name: 'search_files', + name: 'search_vault_files', description: 'Search for files', parameters: definitions[0].parameters } diff --git a/__tests__/Services/ChatService.test.ts b/__tests__/Services/ChatService.test.ts index c5ed6fe..0eb63f6 100644 --- a/__tests__/Services/ChatService.test.ts +++ b/__tests__/Services/ChatService.test.ts @@ -6,6 +6,7 @@ import { Conversation } from '../../Conversations/Conversation'; import { ConversationContent } from '../../Conversations/ConversationContent'; import { Role } from '../../Enums/Role'; import { AIFunctionCall } from '../../AIClasses/AIFunctionCall'; +import { AIFunction, fromString } from '../../Enums/AIFunction'; /** * INTEGRATION TESTS - Simplified @@ -390,7 +391,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should return content unchanged when content is empty', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('test_function', { arg1: 'value1' }); + const functionCall = new AIFunctionCall(AIFunction.SearchVaultFiles, { arg1: 'value1' }); const result = sanitize('', functionCall); @@ -399,7 +400,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should return content unchanged when content is only whitespace', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('test_function', { arg1: 'value1' }); + const functionCall = new AIFunctionCall(AIFunction.SearchVaultFiles, { arg1: 'value1' }); const result = sanitize(' \n\t ', functionCall); @@ -408,7 +409,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should remove exact function call JSON from content', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('write_file', { + const functionCall = new AIFunctionCall(AIFunction.WriteVaultFile, { file_path: 'test.md', content: 'Hello world' }); @@ -422,7 +423,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should remove pretty-printed function call JSON (2 spaces)', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('search_files', { + const functionCall = new AIFunctionCall(AIFunction.SearchVaultFiles, { query: 'test query', limit: 10 }); @@ -437,7 +438,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should remove pretty-printed function call JSON (4 spaces)', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('read_file', { + const functionCall = new AIFunctionCall(AIFunction.ReadVaultFiles, { file_path: 'example.ts' }); @@ -464,7 +465,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should only remove function call JSON, preserving other content', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('delete_file', { + const functionCall = new AIFunctionCall(AIFunction.DeleteVaultFiles, { file_path: 'old.txt' }); const functionCallString = functionCall.toConversationString(); @@ -477,7 +478,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should handle function calls with special characters in arguments', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('write_file', { + const functionCall = new AIFunctionCall(AIFunction.WriteVaultFile, { file_path: 'path/to/file.ts', content: 'const regex = /test.*pattern/g;\nfunction() { return "value"; }' }); @@ -491,7 +492,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should handle function calls with nested objects', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('complex_function', { + const functionCall = new AIFunctionCall(AIFunction.SearchVaultFiles, { config: { nested: { deeply: { @@ -511,7 +512,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should handle function calls with toolId', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('test_function', { arg: 'value' }, 'tool-123'); + const functionCall = new AIFunctionCall(AIFunction.SearchVaultFiles, { arg: 'value' }, 'tool-123'); const functionCallString = functionCall.toConversationString(); const content = functionCallString; @@ -522,7 +523,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should not remove similar but different JSON structures', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('write_file', { + const functionCall = new AIFunctionCall(AIFunction.WriteVaultFile, { file_path: 'test.md' }); @@ -543,12 +544,12 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should handle multiple whitespace variations in pretty-printed JSON', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('test', { key: 'value' }); + const functionCall = new AIFunctionCall(AIFunction.SearchVaultFiles, { key: 'value' }); // Content with irregular whitespace const content = `{ "functionCall": { - "name": "test", + "name": "search_vault_files", "args": { "key": "value" } @@ -562,7 +563,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should preserve natural language content when function call exists', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('write_vault_file', { + const functionCall = new AIFunctionCall(AIFunction.WriteVaultFile, { content: 'Character sheet content', file_name: 'Walker.md' }); @@ -576,7 +577,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should preserve content that does not start with {', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('search_vault_files', { + const functionCall = new AIFunctionCall(AIFunction.SearchVaultFiles, { search_terms: ['Walker'] }); const content = 'Searching for details about Walker in Campaign 2 folder.'; @@ -589,7 +590,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should preserve content that does not end with }', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('test', { key: 'value' }); + const functionCall = new AIFunctionCall(AIFunction.SearchVaultFiles, { key: 'value' }); const content = '{ "incomplete": "json"'; const result = sanitize(content, functionCall); @@ -600,7 +601,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should preserve content with mixed text and JSON-like structures', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('test', { key: 'value' }); + const functionCall = new AIFunctionCall(AIFunction.SearchVaultFiles, { key: 'value' }); const content = 'Here is some info: { "data": "example" } and more text'; const result = sanitize(content, functionCall); @@ -611,7 +612,7 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { it('should only sanitize when content is EXACTLY matching JSON structure', () => { const sanitize = getSanitizeMethod(service); - const functionCall = new AIFunctionCall('write_file', { + const functionCall = new AIFunctionCall(AIFunction.WriteVaultFile, { file_path: 'test.md', content: 'Hello' }); @@ -627,4 +628,22 @@ describe('ChatService - Integration Tests (Sync Methods Only)', () => { expect(sanitize(withSuffix, functionCall)).toBe('- done'); }); }); + + describe('fromString', () => { + it('should convert valid function name strings to AIFunction enum', () => { + expect(fromString('search_vault_files')).toBe(AIFunction.SearchVaultFiles); + expect(fromString('read_vault_files')).toBe(AIFunction.ReadVaultFiles); + expect(fromString('write_vault_file')).toBe(AIFunction.WriteVaultFile); + expect(fromString('delete_vault_files')).toBe(AIFunction.DeleteVaultFiles); + expect(fromString('move_vault_files')).toBe(AIFunction.MoveVaultFiles); + expect(fromString('list_vault_files')).toBe(AIFunction.ListVaultFiles); + expect(fromString('request_web_search')).toBe(AIFunction.RequestWebSearch); + }); + + it('should throw an error for unknown function names', () => { + expect(() => fromString('unknown_function')).toThrow('Unknown function name: unknown_function'); + expect(() => fromString('test_function')).toThrow('Unknown function name: test_function'); + expect(() => fromString('')).toThrow('Unknown function name: '); + }); + }); });