mirror of
https://github.com/zadam/trilium.git
synced 2025-11-05 04:45:47 +01:00
625 lines
21 KiB
TypeScript
625 lines
21 KiB
TypeScript
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
|
import { ChatStorageService } from './chat_storage_service.js';
|
|
import type { Message } from './ai_interface.js';
|
|
|
|
// Mock dependencies
|
|
vi.mock('../notes.js', () => ({
|
|
default: {
|
|
createNewNote: vi.fn()
|
|
}
|
|
}));
|
|
|
|
vi.mock('../sql.js', () => ({
|
|
default: {
|
|
getRow: vi.fn(),
|
|
getRows: vi.fn(),
|
|
execute: vi.fn()
|
|
}
|
|
}));
|
|
|
|
vi.mock('../attributes.js', () => ({
|
|
default: {
|
|
createLabel: vi.fn()
|
|
}
|
|
}));
|
|
|
|
vi.mock('../log.js', () => ({
|
|
default: {
|
|
error: vi.fn(),
|
|
info: vi.fn(),
|
|
warn: vi.fn()
|
|
}
|
|
}));
|
|
|
|
vi.mock('i18next', () => ({
|
|
t: vi.fn((key: string) => {
|
|
switch (key) {
|
|
case 'ai.chat.root_note_title':
|
|
return 'AI Chats';
|
|
case 'ai.chat.root_note_content':
|
|
return 'This note contains all AI chat conversations.';
|
|
case 'ai.chat.new_chat_title':
|
|
return 'New Chat';
|
|
default:
|
|
return key;
|
|
}
|
|
})
|
|
}));
|
|
|
|
describe('ChatStorageService', () => {
|
|
let chatStorageService: ChatStorageService;
|
|
let mockNotes: any;
|
|
let mockSql: any;
|
|
let mockAttributes: any;
|
|
|
|
beforeEach(async () => {
|
|
vi.clearAllMocks();
|
|
chatStorageService = new ChatStorageService();
|
|
|
|
// Get mocked modules
|
|
mockNotes = (await import('../notes.js')).default;
|
|
mockSql = (await import('../sql.js')).default;
|
|
mockAttributes = (await import('../attributes.js')).default;
|
|
});
|
|
|
|
afterEach(() => {
|
|
vi.restoreAllMocks();
|
|
});
|
|
|
|
describe('getOrCreateChatRoot', () => {
|
|
it('should return existing chat root if it exists', async () => {
|
|
mockSql.getRow.mockResolvedValueOnce({ noteId: 'existing-root-123' });
|
|
|
|
const rootId = await chatStorageService.getOrCreateChatRoot();
|
|
|
|
expect(rootId).toBe('existing-root-123');
|
|
expect(mockSql.getRow).toHaveBeenCalledWith(
|
|
'SELECT noteId FROM attributes WHERE name = ? AND value = ?',
|
|
['label', 'triliumChatRoot']
|
|
);
|
|
});
|
|
|
|
it('should create new chat root if it does not exist', async () => {
|
|
mockSql.getRow.mockResolvedValueOnce(null);
|
|
mockNotes.createNewNote.mockReturnValueOnce({
|
|
note: { noteId: 'new-root-123' }
|
|
});
|
|
|
|
const rootId = await chatStorageService.getOrCreateChatRoot();
|
|
|
|
expect(rootId).toBe('new-root-123');
|
|
expect(mockNotes.createNewNote).toHaveBeenCalledWith({
|
|
parentNoteId: 'root',
|
|
title: 'AI Chats',
|
|
type: 'text',
|
|
content: 'This note contains all AI chat conversations.'
|
|
});
|
|
expect(mockAttributes.createLabel).toHaveBeenCalledWith(
|
|
'new-root-123',
|
|
'triliumChatRoot',
|
|
''
|
|
);
|
|
});
|
|
});
|
|
|
|
describe('createChat', () => {
|
|
it('should create a new chat with default title', async () => {
|
|
const mockDate = new Date('2024-01-01T00:00:00Z');
|
|
vi.useFakeTimers();
|
|
vi.setSystemTime(mockDate);
|
|
|
|
mockSql.getRow.mockResolvedValueOnce({ noteId: 'root-123' });
|
|
mockNotes.createNewNote.mockReturnValueOnce({
|
|
note: { noteId: 'chat-123' }
|
|
});
|
|
|
|
const messages: Message[] = [
|
|
{ role: 'user', content: 'Hello' }
|
|
];
|
|
|
|
const result = await chatStorageService.createChat('Test Chat', messages);
|
|
|
|
expect(result).toEqual({
|
|
id: 'chat-123',
|
|
title: 'Test Chat',
|
|
messages,
|
|
noteId: 'chat-123',
|
|
createdAt: mockDate,
|
|
updatedAt: mockDate,
|
|
metadata: {}
|
|
});
|
|
|
|
expect(mockNotes.createNewNote).toHaveBeenCalledWith({
|
|
parentNoteId: 'root-123',
|
|
title: 'Test Chat',
|
|
type: 'code',
|
|
mime: 'application/json',
|
|
content: JSON.stringify({
|
|
messages,
|
|
metadata: {},
|
|
createdAt: mockDate,
|
|
updatedAt: mockDate
|
|
}, null, 2)
|
|
});
|
|
|
|
expect(mockAttributes.createLabel).toHaveBeenCalledWith(
|
|
'chat-123',
|
|
'triliumChat',
|
|
''
|
|
);
|
|
|
|
vi.useRealTimers();
|
|
});
|
|
|
|
it('should create chat with custom metadata', async () => {
|
|
mockSql.getRow.mockResolvedValueOnce({ noteId: 'root-123' });
|
|
mockNotes.createNewNote.mockReturnValueOnce({
|
|
note: { noteId: 'chat-123' }
|
|
});
|
|
|
|
const metadata = {
|
|
model: 'gpt-4',
|
|
provider: 'openai',
|
|
temperature: 0.7
|
|
};
|
|
|
|
const result = await chatStorageService.createChat('Test Chat', [], metadata);
|
|
|
|
expect(result.metadata).toEqual(metadata);
|
|
});
|
|
|
|
it('should generate default title if none provided', async () => {
|
|
mockSql.getRow.mockResolvedValueOnce({ noteId: 'root-123' });
|
|
mockNotes.createNewNote.mockReturnValueOnce({
|
|
note: { noteId: 'chat-123' }
|
|
});
|
|
|
|
const result = await chatStorageService.createChat('');
|
|
|
|
expect(result.title).toContain('New Chat');
|
|
expect(result.title).toMatch(/\d{1,2}\/\d{1,2}\/\d{4}/); // Date pattern
|
|
});
|
|
});
|
|
|
|
describe('getAllChats', () => {
|
|
it('should return all chats with parsed content', async () => {
|
|
const mockChats = [
|
|
{
|
|
noteId: 'chat-1',
|
|
title: 'Chat 1',
|
|
dateCreated: '2024-01-01T00:00:00Z',
|
|
dateModified: '2024-01-01T01:00:00Z',
|
|
content: JSON.stringify({
|
|
messages: [{ role: 'user', content: 'Hello' }],
|
|
metadata: { model: 'gpt-4' },
|
|
createdAt: '2024-01-01T00:00:00Z',
|
|
updatedAt: '2024-01-01T01:00:00Z'
|
|
})
|
|
},
|
|
{
|
|
noteId: 'chat-2',
|
|
title: 'Chat 2',
|
|
dateCreated: '2024-01-02T00:00:00Z',
|
|
dateModified: '2024-01-02T01:00:00Z',
|
|
content: JSON.stringify({
|
|
messages: [{ role: 'user', content: 'Hi' }],
|
|
metadata: { provider: 'anthropic' }
|
|
})
|
|
}
|
|
];
|
|
|
|
mockSql.getRows.mockResolvedValueOnce(mockChats);
|
|
|
|
const result = await chatStorageService.getAllChats();
|
|
|
|
expect(result).toHaveLength(2);
|
|
expect(result[0]).toEqual({
|
|
id: 'chat-1',
|
|
title: 'Chat 1',
|
|
messages: [{ role: 'user', content: 'Hello' }],
|
|
noteId: 'chat-1',
|
|
createdAt: new Date('2024-01-01T00:00:00Z'),
|
|
updatedAt: new Date('2024-01-01T01:00:00Z'),
|
|
metadata: { model: 'gpt-4' }
|
|
});
|
|
|
|
expect(mockSql.getRows).toHaveBeenCalledWith(
|
|
expect.stringContaining('SELECT notes.noteId, notes.title'),
|
|
['label', 'triliumChat']
|
|
);
|
|
});
|
|
|
|
it('should handle chats with invalid JSON content', async () => {
|
|
const mockChats = [
|
|
{
|
|
noteId: 'chat-1',
|
|
title: 'Chat 1',
|
|
dateCreated: '2024-01-01T00:00:00Z',
|
|
dateModified: '2024-01-01T01:00:00Z',
|
|
content: 'invalid json'
|
|
}
|
|
];
|
|
|
|
mockSql.getRows.mockResolvedValueOnce(mockChats);
|
|
const consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
|
|
|
|
const result = await chatStorageService.getAllChats();
|
|
|
|
expect(result).toHaveLength(1);
|
|
expect(result[0]).toEqual({
|
|
id: 'chat-1',
|
|
title: 'Chat 1',
|
|
messages: [],
|
|
noteId: 'chat-1',
|
|
createdAt: new Date('2024-01-01T00:00:00Z'),
|
|
updatedAt: new Date('2024-01-01T01:00:00Z'),
|
|
metadata: {}
|
|
});
|
|
|
|
expect(consoleErrorSpy).toHaveBeenCalledWith('Failed to parse chat content:', expect.any(Error));
|
|
consoleErrorSpy.mockRestore();
|
|
});
|
|
});
|
|
|
|
describe('getChat', () => {
|
|
it('should return specific chat by ID', async () => {
|
|
const mockChat = {
|
|
noteId: 'chat-123',
|
|
title: 'Test Chat',
|
|
dateCreated: '2024-01-01T00:00:00Z',
|
|
dateModified: '2024-01-01T01:00:00Z',
|
|
content: JSON.stringify({
|
|
messages: [{ role: 'user', content: 'Hello' }],
|
|
metadata: { model: 'gpt-4' },
|
|
createdAt: '2024-01-01T00:00:00Z',
|
|
updatedAt: '2024-01-01T01:00:00Z'
|
|
})
|
|
};
|
|
|
|
mockSql.getRow.mockResolvedValueOnce(mockChat);
|
|
|
|
const result = await chatStorageService.getChat('chat-123');
|
|
|
|
expect(result).toEqual({
|
|
id: 'chat-123',
|
|
title: 'Test Chat',
|
|
messages: [{ role: 'user', content: 'Hello' }],
|
|
noteId: 'chat-123',
|
|
createdAt: new Date('2024-01-01T00:00:00Z'),
|
|
updatedAt: new Date('2024-01-01T01:00:00Z'),
|
|
metadata: { model: 'gpt-4' }
|
|
});
|
|
|
|
expect(mockSql.getRow).toHaveBeenCalledWith(
|
|
expect.stringContaining('SELECT notes.noteId, notes.title'),
|
|
['chat-123']
|
|
);
|
|
});
|
|
|
|
it('should return null if chat not found', async () => {
|
|
mockSql.getRow.mockResolvedValueOnce(null);
|
|
|
|
const result = await chatStorageService.getChat('nonexistent');
|
|
|
|
expect(result).toBeNull();
|
|
});
|
|
});
|
|
|
|
describe('updateChat', () => {
|
|
it('should update chat messages and metadata', async () => {
|
|
const existingChat = {
|
|
id: 'chat-123',
|
|
title: 'Test Chat',
|
|
messages: [{ role: 'user', content: 'Hello' }],
|
|
noteId: 'chat-123',
|
|
createdAt: new Date('2024-01-01T00:00:00Z'),
|
|
updatedAt: new Date('2024-01-01T01:00:00Z'),
|
|
metadata: { model: 'gpt-4' }
|
|
};
|
|
|
|
const newMessages: Message[] = [
|
|
{ role: 'user', content: 'Hello' },
|
|
{ role: 'assistant', content: 'Hi there!' }
|
|
];
|
|
|
|
const newMetadata = { provider: 'openai', temperature: 0.7 };
|
|
|
|
// Mock getChat to return existing chat
|
|
vi.spyOn(chatStorageService, 'getChat').mockResolvedValueOnce(existingChat);
|
|
|
|
const mockDate = new Date('2024-01-01T02:00:00Z');
|
|
vi.useFakeTimers();
|
|
vi.setSystemTime(mockDate);
|
|
|
|
const result = await chatStorageService.updateChat(
|
|
'chat-123',
|
|
newMessages,
|
|
'Updated Title',
|
|
newMetadata
|
|
);
|
|
|
|
expect(result).toEqual({
|
|
...existingChat,
|
|
title: 'Updated Title',
|
|
messages: newMessages,
|
|
updatedAt: mockDate,
|
|
metadata: { model: 'gpt-4', provider: 'openai', temperature: 0.7 }
|
|
});
|
|
|
|
expect(mockSql.execute).toHaveBeenCalledWith(
|
|
'UPDATE blobs SET content = ? WHERE blobId = (SELECT blobId FROM notes WHERE noteId = ?)',
|
|
[
|
|
JSON.stringify({
|
|
messages: newMessages,
|
|
metadata: { model: 'gpt-4', provider: 'openai', temperature: 0.7 },
|
|
createdAt: existingChat.createdAt,
|
|
updatedAt: mockDate
|
|
}, null, 2),
|
|
'chat-123'
|
|
]
|
|
);
|
|
|
|
expect(mockSql.execute).toHaveBeenCalledWith(
|
|
'UPDATE notes SET title = ? WHERE noteId = ?',
|
|
['Updated Title', 'chat-123']
|
|
);
|
|
|
|
vi.useRealTimers();
|
|
});
|
|
|
|
it('should return null if chat not found', async () => {
|
|
vi.spyOn(chatStorageService, 'getChat').mockResolvedValueOnce(null);
|
|
|
|
const result = await chatStorageService.updateChat(
|
|
'nonexistent',
|
|
[],
|
|
'Title'
|
|
);
|
|
|
|
expect(result).toBeNull();
|
|
});
|
|
});
|
|
|
|
describe('deleteChat', () => {
|
|
it('should mark chat as deleted', async () => {
|
|
const result = await chatStorageService.deleteChat('chat-123');
|
|
|
|
expect(result).toBe(true);
|
|
expect(mockSql.execute).toHaveBeenCalledWith(
|
|
'UPDATE notes SET isDeleted = 1 WHERE noteId = ?',
|
|
['chat-123']
|
|
);
|
|
});
|
|
|
|
it('should return false on SQL error', async () => {
|
|
mockSql.execute.mockRejectedValueOnce(new Error('SQL error'));
|
|
const consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
|
|
|
|
const result = await chatStorageService.deleteChat('chat-123');
|
|
|
|
expect(result).toBe(false);
|
|
expect(consoleErrorSpy).toHaveBeenCalledWith('Failed to delete chat:', expect.any(Error));
|
|
consoleErrorSpy.mockRestore();
|
|
});
|
|
});
|
|
|
|
describe('recordToolExecution', () => {
|
|
it('should record tool execution in chat metadata', async () => {
|
|
const existingChat = {
|
|
id: 'chat-123',
|
|
title: 'Test Chat',
|
|
messages: [],
|
|
noteId: 'chat-123',
|
|
createdAt: new Date(),
|
|
updatedAt: new Date(),
|
|
metadata: {}
|
|
};
|
|
|
|
vi.spyOn(chatStorageService, 'getChat').mockResolvedValueOnce(existingChat);
|
|
vi.spyOn(chatStorageService, 'updateChat').mockResolvedValueOnce(existingChat);
|
|
|
|
const result = await chatStorageService.recordToolExecution(
|
|
'chat-123',
|
|
'searchNotes',
|
|
'tool-123',
|
|
{ query: 'test' },
|
|
'Found 3 notes'
|
|
);
|
|
|
|
expect(result).toBe(true);
|
|
expect(chatStorageService.updateChat).toHaveBeenCalledWith(
|
|
'chat-123',
|
|
[],
|
|
undefined,
|
|
expect.objectContaining({
|
|
toolExecutions: expect.arrayContaining([
|
|
expect.objectContaining({
|
|
id: 'tool-123',
|
|
name: 'searchNotes',
|
|
arguments: { query: 'test' },
|
|
result: 'Found 3 notes'
|
|
})
|
|
])
|
|
})
|
|
);
|
|
});
|
|
|
|
it('should return false if chat not found', async () => {
|
|
vi.spyOn(chatStorageService, 'getChat').mockResolvedValueOnce(null);
|
|
|
|
const result = await chatStorageService.recordToolExecution(
|
|
'nonexistent',
|
|
'searchNotes',
|
|
'tool-123',
|
|
{ query: 'test' },
|
|
'Result'
|
|
);
|
|
|
|
expect(result).toBe(false);
|
|
});
|
|
});
|
|
|
|
describe('recordSources', () => {
|
|
it('should record sources in chat metadata', async () => {
|
|
const existingChat = {
|
|
id: 'chat-123',
|
|
title: 'Test Chat',
|
|
messages: [],
|
|
noteId: 'chat-123',
|
|
createdAt: new Date(),
|
|
updatedAt: new Date(),
|
|
metadata: {}
|
|
};
|
|
|
|
const sources = [
|
|
{
|
|
noteId: 'note-1',
|
|
title: 'Source Note 1',
|
|
similarity: 0.95
|
|
},
|
|
{
|
|
noteId: 'note-2',
|
|
title: 'Source Note 2',
|
|
similarity: 0.87
|
|
}
|
|
];
|
|
|
|
vi.spyOn(chatStorageService, 'getChat').mockResolvedValueOnce(existingChat);
|
|
vi.spyOn(chatStorageService, 'updateChat').mockResolvedValueOnce(existingChat);
|
|
|
|
const result = await chatStorageService.recordSources('chat-123', sources);
|
|
|
|
expect(result).toBe(true);
|
|
expect(chatStorageService.updateChat).toHaveBeenCalledWith(
|
|
'chat-123',
|
|
[],
|
|
undefined,
|
|
expect.objectContaining({
|
|
sources
|
|
})
|
|
);
|
|
});
|
|
});
|
|
|
|
describe('extractToolExecutionsFromMessages', () => {
|
|
it('should extract tool executions from assistant messages with tool calls', async () => {
|
|
const messages: Message[] = [
|
|
{
|
|
role: 'assistant',
|
|
content: 'I need to search for notes.',
|
|
tool_calls: [
|
|
{
|
|
id: 'call_123',
|
|
type: 'function',
|
|
function: {
|
|
name: 'searchNotes',
|
|
arguments: '{"query": "test"}'
|
|
}
|
|
}
|
|
]
|
|
},
|
|
{
|
|
role: 'tool',
|
|
content: 'Found 2 notes',
|
|
tool_call_id: 'call_123'
|
|
},
|
|
{
|
|
role: 'assistant',
|
|
content: 'Based on the search results...'
|
|
}
|
|
];
|
|
|
|
// Access private method through any cast for testing
|
|
const extractToolExecutions = (chatStorageService as any).extractToolExecutionsFromMessages.bind(chatStorageService);
|
|
const toolExecutions = extractToolExecutions(messages, []);
|
|
|
|
expect(toolExecutions).toHaveLength(1);
|
|
expect(toolExecutions[0]).toEqual(
|
|
expect.objectContaining({
|
|
id: 'call_123',
|
|
name: 'searchNotes',
|
|
arguments: { query: 'test' },
|
|
result: 'Found 2 notes',
|
|
timestamp: expect.any(Date)
|
|
})
|
|
);
|
|
});
|
|
|
|
it('should handle error responses from tools', async () => {
|
|
const messages: Message[] = [
|
|
{
|
|
role: 'assistant',
|
|
content: 'I need to search for notes.',
|
|
tool_calls: [
|
|
{
|
|
id: 'call_123',
|
|
type: 'function',
|
|
function: {
|
|
name: 'searchNotes',
|
|
arguments: '{"query": "test"}'
|
|
}
|
|
}
|
|
]
|
|
},
|
|
{
|
|
role: 'tool',
|
|
content: 'Error: Search service unavailable',
|
|
tool_call_id: 'call_123'
|
|
}
|
|
];
|
|
|
|
const extractToolExecutions = (chatStorageService as any).extractToolExecutionsFromMessages.bind(chatStorageService);
|
|
const toolExecutions = extractToolExecutions(messages, []);
|
|
|
|
expect(toolExecutions).toHaveLength(1);
|
|
expect(toolExecutions[0]).toEqual(
|
|
expect.objectContaining({
|
|
id: 'call_123',
|
|
name: 'searchNotes',
|
|
error: 'Search service unavailable',
|
|
result: 'Error: Search service unavailable'
|
|
})
|
|
);
|
|
});
|
|
|
|
it('should not duplicate existing tool executions', async () => {
|
|
const existingToolExecutions = [
|
|
{
|
|
id: 'call_123',
|
|
name: 'searchNotes',
|
|
arguments: { query: 'existing' },
|
|
result: 'Previous result',
|
|
timestamp: new Date()
|
|
}
|
|
];
|
|
|
|
const messages: Message[] = [
|
|
{
|
|
role: 'assistant',
|
|
content: 'I need to search for notes.',
|
|
tool_calls: [
|
|
{
|
|
id: 'call_123', // Same ID as existing
|
|
type: 'function',
|
|
function: {
|
|
name: 'searchNotes',
|
|
arguments: '{"query": "test"}'
|
|
}
|
|
}
|
|
]
|
|
},
|
|
{
|
|
role: 'tool',
|
|
content: 'Found 2 notes',
|
|
tool_call_id: 'call_123'
|
|
}
|
|
];
|
|
|
|
const extractToolExecutions = (chatStorageService as any).extractToolExecutionsFromMessages.bind(chatStorageService);
|
|
const toolExecutions = extractToolExecutions(messages, existingToolExecutions);
|
|
|
|
expect(toolExecutions).toHaveLength(1);
|
|
expect(toolExecutions[0].arguments).toEqual({ query: 'existing' });
|
|
});
|
|
});
|
|
}); |