refactor: convert AIFunctionCall.name from string to AIFunction enum

- Add AIFunction enum type to AIFunctionCall name property
- Implement fromString helper to convert string to AIFunction
- Update Claude, Gemini, and OpenAI classes to use aiFunctionFromString
- Add type safety to AIFunctionService switch statement
- Update all test files to use AIFunction enum values
This commit is contained in:
Andrew Beal 2025-11-10 20:08:32 +00:00
parent 37db1a8908
commit a30693bfe5
10 changed files with 92 additions and 59 deletions

View file

@ -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<string, unknown>;
public readonly toolId?: string;
constructor(name: string, args: Record<string, unknown>, toolId?: string) {
constructor(name: AIFunction, args: Record<string, unknown>, toolId?: string) {
this.name = name;
this.arguments = args;
this.toolId = toolId;

View file

@ -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<string, unknown>;
functionCall = new AIFunctionCall(
this.accumulatedFunctionName,
aiFunctionFromString(this.accumulatedFunctionName),
args as Record<string, object>,
this.accumulatedFunctionId || undefined
);

View file

@ -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<string, object>
);
}

View file

@ -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<string, unknown>;
functionCall = new AIFunctionCall(
firstToolCall.name,
aiFunctionFromString(firstToolCall.name),
args as Record<string, object>,
firstToolCall.id || undefined
);

View file

@ -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}`);
}

View file

@ -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<object> {
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 };
}

View file

@ -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
});

View file

@ -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
});

View file

@ -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
}

View file

@ -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: ');
});
});
});