mirror of
https://github.com/andy-stack/vaultkeeper-ai.git
synced 2026-07-22 16:30:27 +00:00
- Wrap user messages in <message_to_title> tags to prevent prompt injection - Expand NamePrompt with security guidelines and edge case handling - Add 10-word limit truncation for generated titles - Move conversation save before naming to ensure persistence - Add comprehensive task complexity assessment framework to SystemPrompt - Change stack limit error to warning log for better resilience
139 lines
4.9 KiB
TypeScript
139 lines
4.9 KiB
TypeScript
import { Semaphore } from "Helpers/Semaphore";
|
|
import { Resolve } from "./DependencyService";
|
|
import { Services } from "./Services";
|
|
import type { ConversationFileSystemService } from "./ConversationFileSystemService";
|
|
import type { ConversationNamingService } from "./ConversationNamingService";
|
|
import { Conversation } from "Conversations/Conversation";
|
|
import { ConversationContent } from "Conversations/ConversationContent";
|
|
import { Role } from "Enums/Role";
|
|
import { Notice } from "obsidian";
|
|
import type { EventService } from "./EventService";
|
|
import { Event } from "Enums/Event";
|
|
import { AbortService } from "./AbortService";
|
|
import { Exception } from "Helpers/Exception";
|
|
import type { Attachment } from "Conversations/Attachment";
|
|
import { Reference } from "Conversations/Reference";
|
|
import type { WorkSpaceService } from "./WorkSpaceService";
|
|
import type { AIControllerService } from "./AIControllerService";
|
|
|
|
export interface IChatServiceCallbacks {
|
|
onSubmit: () => void;
|
|
onStreamingUpdate: (streamingMessageId: string | null) => void;
|
|
onThoughtUpdate: (thought: string | null) => void;
|
|
onComplete: () => void;
|
|
onCancel: () => void;
|
|
}
|
|
|
|
export class ChatService {
|
|
|
|
private aiControllerService: AIControllerService;
|
|
private conversationService: ConversationFileSystemService;
|
|
private namingService: ConversationNamingService;
|
|
private workSpaceService: WorkSpaceService;
|
|
private eventService: EventService;
|
|
private abortService: AbortService;
|
|
|
|
private semaphore: Semaphore;
|
|
private semaphoreHeld: boolean = false;
|
|
|
|
constructor() {
|
|
this.aiControllerService = Resolve<AIControllerService>(Services.AIControllerService);
|
|
this.conversationService = Resolve<ConversationFileSystemService>(Services.ConversationFileSystemService);
|
|
this.namingService = Resolve<ConversationNamingService>(Services.ConversationNamingService);
|
|
this.workSpaceService = Resolve<WorkSpaceService>(Services.WorkSpaceService);
|
|
this.eventService = Resolve<EventService>(Services.EventService);
|
|
this.abortService = Resolve<AbortService>(Services.AbortService);
|
|
this.semaphore = new Semaphore(1, false);
|
|
|
|
this.aiControllerService.setSaveCallback(async (conversation) => {
|
|
await this.saveConversation(conversation);
|
|
});
|
|
}
|
|
|
|
public onNameChanged: ((name: string) => void) | undefined = undefined;
|
|
|
|
public async submit(conversation: Conversation, allowDestructiveActions: boolean, userRequest: string, formattedRequest: string, attachments: Attachment[], callbacks: IChatServiceCallbacks) {
|
|
if (!await this.semaphore.wait()) {
|
|
return;
|
|
}
|
|
|
|
this.semaphoreHeld = true;
|
|
|
|
try {
|
|
if (userRequest.trim() === "") {
|
|
return;
|
|
}
|
|
|
|
this.abortService.initialiseAbortController();
|
|
|
|
await this.abortService.abortableOperation(async () => {
|
|
const firstMessage = conversation.contents.length === 0;
|
|
|
|
const conversationContent = new ConversationContent({
|
|
role: Role.User,
|
|
content: this.requestWithContext(formattedRequest),
|
|
displayContent: userRequest
|
|
});
|
|
conversation.contents.push(conversationContent);
|
|
|
|
await this.saveConversation(conversation);
|
|
|
|
if (firstMessage) {
|
|
this.onNameChanged?.(conversation.title); // on change for initial conversation name
|
|
void this.namingService.requestName(conversation, formattedRequest, this.onNameChanged)
|
|
}
|
|
|
|
if (attachments.length > 0) {
|
|
// Add any attachments that came from paste / drop
|
|
conversation.contents.push(new ConversationContent({
|
|
role: Role.User,
|
|
attachments: attachments,
|
|
shouldDisplayContent: false
|
|
}));
|
|
|
|
conversationContent.references = attachments.map(attachment =>
|
|
new Reference(attachment.fileName, attachment.approximateFileSizeMB()));
|
|
}
|
|
await this.saveConversation(conversation);
|
|
|
|
callbacks.onSubmit();
|
|
callbacks.onStreamingUpdate(null);
|
|
|
|
await this.aiControllerService.runMainAgent(conversation, allowDestructiveActions, callbacks);
|
|
});
|
|
} catch (error) {
|
|
if (AbortService.isAbortError(error)) {
|
|
callbacks.onCancel();
|
|
} else {
|
|
Exception.log(error);
|
|
new Notice("Vaultkeeper AI encountered an error");
|
|
}
|
|
} finally {
|
|
this.eventService.trigger(Event.DiffClosed);
|
|
await this.saveConversation(conversation);
|
|
if (this.semaphoreHeld) {
|
|
this.semaphoreHeld = false;
|
|
this.semaphore.release();
|
|
}
|
|
callbacks.onThoughtUpdate(null);
|
|
callbacks.onComplete();
|
|
}
|
|
}
|
|
|
|
public stop() {
|
|
this.abortService.abort("User requested cancellation");
|
|
this.eventService.trigger(Event.DiffClosed);
|
|
}
|
|
|
|
private requestWithContext(request: string) {
|
|
const activeFile = this.workSpaceService.getActiveFile();
|
|
return activeFile ? `${request}\nUser current active file: "${activeFile.path}"` : request;
|
|
}
|
|
|
|
private async saveConversation(conversation: Conversation) {
|
|
const result = await this.conversationService.saveConversation(conversation);
|
|
if (result instanceof Error) {
|
|
new Notice(`Failed to save conversation data for '${conversation.title}'`);
|
|
}
|
|
}
|
|
}
|