diff --git a/.github/workflows/node.js.yml b/.github/workflows/node.js.yml index 272aa8d6..40f0f91b 100644 --- a/.github/workflows/node.js.yml +++ b/.github/workflows/node.js.yml @@ -7,7 +7,7 @@ on: push: branches: ["master", "main"] pull_request: - branches: ["master", "main"] + # Run on all pull requests regardless of target branch jobs: build: diff --git a/CLAUDE.md b/CLAUDE.md index c517d162..3197583f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -233,6 +233,7 @@ This helps ensure thorough testing and provides documentation for QA. - Settings are versioned - migrations may be needed - Local model support available via Ollama/LM Studio - Rate limiting is implemented for all API calls +- For technical debt and known issues, see [`TODO.md`](./TODO.md) ### Architecture Migration Notes diff --git a/MESSAGE_ARCHITECTURE.md b/MESSAGE_ARCHITECTURE.md index 5e151091..08ce5b29 100644 --- a/MESSAGE_ARCHITECTURE.md +++ b/MESSAGE_ARCHITECTURE.md @@ -326,6 +326,81 @@ Return project-specific repository ## Message Lifecycle +### Example: User Message with Context Note + +When a user types "Summarize this note" and attaches "meeting-notes.md": + +1. **Input**: User text + attached file → Chat component +2. **Storage**: MessageRepository stores displayText: "Summarize this note" +3. **Processing**: ContextManager reads the note and creates processedText with proper XML structure +4. **Memory Sync**: Chain memory receives the processed version for LLM +5. **UI Update**: Shows message with context badge, displays only "Summarize this note" +6. **LLM Processing**: AI receives full context and generates response + +### Context XML Format + +All context is wrapped in semantic XML tags for clear structure: + +#### Note Context + +```xml + +meeting-notes +docs/meeting-notes.md +2024-01-15T10:00:00.000Z +2024-01-15T14:30:00.000Z + +[actual note content here] + + +``` + +#### URL Context + +```xml + +https://example.com/article + +[fetched content from URL] + + +``` + +#### Selected Text Context + +```xml + +Source Note Title +path/to/source.md +45 +52 + +[selected text content] + + +``` + +#### Error Cases + +```xml + +filename +path/to/file.ext +[Error: Could not process file] + +``` + +This separation ensures: + +- Clean UI (shows what user typed) +- Rich context for AI (includes note content) +- Reprocessable context on message edits + +For detailed examples, see: + +- `src/core/MessageLifecycle.test.ts` - Complete lifecycle demonstration with context notes +- `src/core/MessageLifecycle.xmltags.test.ts` - XML tag formatting tests and examples + ### 1. Sending a New Message ``` @@ -619,4 +694,6 @@ console.log({ - `src/core/MessageRepository.test.ts` - Repository tests - `src/core/ChatManager.test.ts` - Manager tests +- `src/core/MessageLifecycle.test.ts` - Complete lifecycle examples with context notes +- `src/core/MessageLifecycle.xmltags.test.ts` - XML tag formatting tests and examples - `src/components/chat-components/MessageContext.test.tsx` - Context display tests diff --git a/TODO.md b/TODO.md new file mode 100644 index 00000000..015f8673 --- /dev/null +++ b/TODO.md @@ -0,0 +1,51 @@ +# TODO - Technical Debt & Future Improvements + +This document tracks technical debt items and improvements that need to be addressed in the future. + +## 1. Docs4LLM SSL Error in Projects Mode + +### Issue Description + +Document parsing for projects mode is failing with an SSL error (`net::ERR_SSL_BAD_RECORD_MAC_ALERT`) when trying to upload files to the docs4llm API endpoint. + +### Technical Details + +- **Error Location**: `src/LLMProviders/brevilabsClient.ts:185` in `makeFormDataRequest` method +- **Root Cause**: The method uses native `fetch` API instead of Obsidian's `requestUrl` API +- **Context**: While regular JSON requests were migrated to use `safeFetch` (which uses `requestUrl`) in commit e49aafa to fix CORS issues, the `makeFormDataRequest` method was not updated + +### Why Current Approaches Won't Work + +1. **Backend Constraint**: The `/docs4llm` endpoint only accepts multipart/form-data format with `files: List[UploadFile]` +2. **Obsidian Limitation**: The existing `safeFetch` function is hardcoded for `application/json` content type +3. **No JSON Alternative**: Unlike `pdf4llm` which accepts base64 JSON, there's no JSON endpoint for `docs4llm` + +### Recommended Solution + +Create a new `safeFetchFormData` function that: + +1. Uses Obsidian's `requestUrl` API with proper multipart/form-data configuration +2. Handles FormData objects correctly +3. Bypasses CORS and SSL restrictions like `safeFetch` does for JSON + +### Alternative Solutions + +1. **Backend Modification**: Add a new `/docs4llm-base64` endpoint that accepts base64-encoded JSON payloads +2. **Research Obsidian API**: Investigate if newer versions of Obsidian's `requestUrl` support multipart/form-data +3. **SSL Certificate Fix**: Address the underlying SSL certificate issue (temporary workaround) + +### Impact + +- Users cannot parse non-markdown files (PDFs, Word docs, etc.) in projects mode +- This affects the core functionality of project context loading +- Workaround: Users must ensure their projects only contain markdown files + +### References + +- Related commit: e49aafa (Brevilabs CORS issue #918) +- Forum discussion: https://forum.obsidian.md/t/holo-how-to-add-a-png-image-or-file-to-formdata-in-obsidian-like-below-this-help/73420 +- Backend implementation: `/Users/chaoyang/webapps/brevilabs-api/app/main.py:1039` + +--- + +_Last updated: 2025-07-18_ diff --git a/__mocks__/obsidian.js b/__mocks__/obsidian.js index 7075beb3..8f382ddc 100644 --- a/__mocks__/obsidian.js +++ b/__mocks__/obsidian.js @@ -46,4 +46,47 @@ module.exports = { read: jest.fn(), }, })), + ItemView: jest.fn().mockImplementation(function () { + this.containerEl = document.createElement("div"); + this.onOpen = jest.fn(); + this.onClose = jest.fn(); + this.getDisplayText = jest.fn().mockReturnValue("Mock View"); + this.getViewType = jest.fn().mockReturnValue("mock-view"); + this.getIcon = jest.fn().mockReturnValue("document"); + }), + Notice: jest.fn().mockImplementation(function (message) { + this.message = message; + this.noticeEl = document.createElement("div"); + this.hide = jest.fn(); + }), + TFile: jest.fn().mockImplementation(function (path) { + this.path = path; + this.name = path.split("/").pop(); + this.basename = this.name.replace(/\.[^/.]+$/, ""); + this.extension = path.split(".").pop(); + }), + WorkspaceLeaf: jest.fn().mockImplementation(function () { + this.view = null; + this.setViewState = jest.fn(); + this.detach = jest.fn(); + this.getViewState = jest.fn().mockReturnValue({}); + }), +}; + +// Mock the global app object +global.app = { + vault: { + getAbstractFileByPath: jest.fn().mockReturnValue({ + name: "test-file.md", + path: "test-file.md", + }), + read: jest.fn().mockResolvedValue("test content"), + modify: jest.fn().mockResolvedValue(undefined), + }, + workspace: { + getActiveFile: jest.fn().mockReturnValue(null), + getLeaf: jest.fn().mockReturnValue({ + openFile: jest.fn().mockResolvedValue(undefined), + }), + }, }; diff --git a/manifest.json b/manifest.json index 8af2d3f3..b6f0d557 100644 --- a/manifest.json +++ b/manifest.json @@ -1,7 +1,7 @@ { "id": "copilot", "name": "Copilot", - "version": "2.9.4", + "version": "3.0.0-preview-250720", "minAppVersion": "0.15.0", "description": "Your AI Copilot: Chat with Your Second Brain, Learn Faster, Work Smarter.", "author": "Logan Yang", diff --git a/src/LLMProviders/chainManager.ts b/src/LLMProviders/chainManager.ts index faeb6aa9..72969c3e 100644 --- a/src/LLMProviders/chainManager.ts +++ b/src/LLMProviders/chainManager.ts @@ -13,7 +13,8 @@ import { LLMChainRunner, ProjectChainRunner, VaultQAChainRunner, -} from "@/LLMProviders/chainRunner"; + AutonomousAgentChainRunner, +} from "@/LLMProviders/chainRunner/index"; import { logError, logInfo } from "@/logger"; import { HybridRetriever } from "@/search/hybridRetriever"; import VectorStoreManager from "@/search/vectorStoreManager"; @@ -268,12 +269,18 @@ export default class ChainManager { private getChainRunner(): ChainRunner { const chainType = getChainType(); + const settings = getSettings(); + switch (chainType) { case ChainType.LLM_CHAIN: return new LLMChainRunner(this); case ChainType.VAULT_QA_CHAIN: return new VaultQAChainRunner(this); case ChainType.COPILOT_PLUS_CHAIN: + // Use AutonomousAgentChainRunner if the setting is enabled + if (settings.enableAutonomousAgent) { + return new AutonomousAgentChainRunner(this); + } return new CopilotPlusChainRunner(this); case ChainType.PROJECT_CHAIN: return new ProjectChainRunner(this); diff --git a/src/LLMProviders/chainRunner.ts b/src/LLMProviders/chainRunner.ts deleted file mode 100644 index ac3da297..00000000 --- a/src/LLMProviders/chainRunner.ts +++ /dev/null @@ -1,983 +0,0 @@ -import { getCurrentProject } from "@/aiParams"; -import { getStandaloneQuestion } from "@/chainUtils"; -import { - ABORT_REASON, - AI_SENDER, - EMPTY_INDEX_ERROR_MESSAGE, - LOADING_MESSAGES, - MAX_CHARS_FOR_LOCAL_SEARCH_CONTEXT, - ModelCapability, -} from "@/constants"; -import { - ImageBatchProcessor, - ImageContent, - ImageProcessingResult, - MessageContent, -} from "@/imageProcessing/imageProcessor"; -import { BrevilabsClient } from "@/LLMProviders/brevilabsClient"; -import { logError, logInfo, logWarn } from "@/logger"; -import { HybridRetriever } from "@/search/hybridRetriever"; -import { getSettings, getSystemPrompt } from "@/settings/model"; -import { ChatMessage } from "@/types/message"; -import { ToolManager } from "@/tools/toolManager"; -import { - err2String, - extractChatHistory, - extractUniqueTitlesFromDocs, - extractYoutubeUrl, - formatDateTime, - getApiErrorMessage, - getMessageRole, - withSuppressedTokenWarnings, -} from "@/utils"; -import { BaseChatModel } from "@langchain/core/language_models/chat_models"; -import { Notice } from "obsidian"; -import ChainManager from "./chainManager"; -import { COPILOT_TOOL_NAMES, IntentAnalyzer } from "./intentAnalyzer"; -import ProjectManager from "./projectManager"; - -class ThinkBlockStreamer { - private hasOpenThinkBlock = false; - private fullResponse = ""; - - constructor(private updateCurrentAiMessage: (message: string) => void) {} - - private handleClaude37Chunk(content: any[]) { - let textContent = ""; - for (const item of content) { - switch (item.type) { - case "text": - textContent += item.text; - break; - case "thinking": - if (!this.hasOpenThinkBlock) { - this.fullResponse += "\n"; - this.hasOpenThinkBlock = true; - } - this.fullResponse += item.thinking; - this.updateCurrentAiMessage(this.fullResponse); - return true; // Indicate we handled a thinking chunk - } - } - if (textContent) { - this.fullResponse += textContent; - } - return false; // No thinking chunk handled - } - - private handleDeepseekChunk(chunk: any) { - // Handle standard string content - if (typeof chunk.content === "string") { - this.fullResponse += chunk.content; - } - - // Handle deepseek reasoning/thinking content - if (chunk.additional_kwargs?.reasoning_content) { - if (!this.hasOpenThinkBlock) { - this.fullResponse += "\n"; - this.hasOpenThinkBlock = true; - } - this.fullResponse += chunk.additional_kwargs.reasoning_content; - return true; // Indicate we handled a thinking chunk - } - return false; // No thinking chunk handled - } - - processChunk(chunk: any) { - let handledThinking = false; - - // Handle Claude 3.7 array-based content - if (Array.isArray(chunk.content)) { - handledThinking = this.handleClaude37Chunk(chunk.content); - } else { - // Handle deepseek format - handledThinking = this.handleDeepseekChunk(chunk); - } - - // Close think block if we have one open and didn't handle thinking content - if (this.hasOpenThinkBlock && !handledThinking) { - this.fullResponse += ""; - this.hasOpenThinkBlock = false; - } - - this.updateCurrentAiMessage(this.fullResponse); - } - - close() { - // Make sure to close any open think block at the end - if (this.hasOpenThinkBlock) { - this.fullResponse += ""; - this.updateCurrentAiMessage(this.fullResponse); - } - return this.fullResponse; - } -} - -export interface ChainRunner { - run( - userMessage: ChatMessage, - abortController: AbortController, - updateCurrentAiMessage: (message: string) => void, - addMessage: (message: ChatMessage) => void, - options: { - debug?: boolean; - ignoreSystemMessage?: boolean; - updateLoading?: (loading: boolean) => void; - } - ): Promise; -} - -abstract class BaseChainRunner implements ChainRunner { - protected chainManager: ChainManager; - - constructor(chainManager: ChainManager) { - this.chainManager = chainManager; - } - - abstract run( - userMessage: ChatMessage, - abortController: AbortController, - updateCurrentAiMessage: (message: string) => void, - addMessage: (message: ChatMessage) => void, - options: { - debug?: boolean; - ignoreSystemMessage?: boolean; - updateLoading?: (loading: boolean) => void; - } - ): Promise; - - protected async handleResponse( - fullAIResponse: string, - userMessage: ChatMessage, - abortController: AbortController, - addMessage: (message: ChatMessage) => void, - updateCurrentAiMessage: (message: string) => void, - sources?: { title: string; score: number }[] - ) { - // Save to memory and add message if we have a response - // Skip only if it's a NEW_CHAT abort (clearing everything) - if ( - fullAIResponse && - !(abortController.signal.aborted && abortController.signal.reason === ABORT_REASON.NEW_CHAT) - ) { - await this.chainManager.memoryManager - .getMemory() - .saveContext({ input: userMessage.message }, { output: fullAIResponse }); - - addMessage({ - message: fullAIResponse, - sender: AI_SENDER, - isVisible: true, - timestamp: formatDateTime(new Date()), - sources: sources, - }); - - // Clear the streaming message since it's now in chat history - updateCurrentAiMessage(""); - } else if (abortController.signal.reason === ABORT_REASON.NEW_CHAT) { - // Also clear if it's a new chat - updateCurrentAiMessage(""); - } - logInfo( - "==== Chat Memory ====\n", - (this.chainManager.memoryManager.getMemory().chatHistory as any).messages.map( - (m: any) => m.content - ) - ); - logInfo("==== Final AI Response ====\n", fullAIResponse); - return fullAIResponse; - } - - protected async handleError( - error: any, - addMessage?: (message: ChatMessage) => void, - updateCurrentAiMessage?: (message: string) => void - ) { - const msg = err2String(error); - logError("Error during LLM invocation:", msg); - const errorData = error?.response?.data?.error || msg; - const errorCode = errorData?.code || msg; - let errorMessage = ""; - - // Check for specific error messages - if (error?.message?.includes("Invalid license key")) { - errorMessage = "Invalid Copilot Plus license key. Please check your license key in settings."; - } else if (errorCode === "model_not_found") { - errorMessage = - "You do not have access to this model or the model does not exist, please check with your API provider."; - } else { - errorMessage = `${errorCode}`; - } - - logError(errorData); - - if (addMessage && updateCurrentAiMessage) { - updateCurrentAiMessage(""); - - // remove langchain troubleshooting URL from error message - const ignoreEndIndex = errorMessage.search("Troubleshooting URL"); - errorMessage = ignoreEndIndex !== -1 ? errorMessage.slice(0, ignoreEndIndex) : errorMessage; - - // add more user guide for invalid API key - if (msg.search(/401|invalid|not valid/gi) !== -1) { - errorMessage = - "Something went wrong. Please check if you have set your API key." + - "\nPath: Settings > copilot plugin > Basic Tab > Set Keys." + - "\nOr check model config" + - "\nError Details: " + - errorMessage; - } - - addMessage({ - message: errorMessage, - isErrorMessage: true, - sender: AI_SENDER, - isVisible: true, - timestamp: formatDateTime(new Date()), - }); - } else { - // Fallback to Notice if message handlers aren't provided - new Notice(errorMessage); - logError(errorData); - } - } -} - -class LLMChainRunner extends BaseChainRunner { - async run( - userMessage: ChatMessage, - abortController: AbortController, - updateCurrentAiMessage: (message: string) => void, - addMessage: (message: ChatMessage) => void, - options: { - debug?: boolean; - ignoreSystemMessage?: boolean; - updateLoading?: (loading: boolean) => void; - } - ): Promise { - const streamer = new ThinkBlockStreamer(updateCurrentAiMessage); - - try { - // Get chat history from memory - const memory = this.chainManager.memoryManager.getMemory(); - const memoryVariables = await memory.loadMemoryVariables({}); - const chatHistory = extractChatHistory(memoryVariables); - - // Create messages array starting with system message - const messages: any[] = []; - - // Add system message if available - const systemPrompt = getSystemPrompt(); - const chatModel = this.chainManager.chatModelManager.getChatModel(); - - if (systemPrompt) { - messages.push({ - role: getMessageRole(chatModel), - content: systemPrompt, - }); - } - - // Add chat history - for (const entry of chatHistory) { - messages.push({ role: entry.role, content: entry.content }); - } - - // Add current user message - messages.push({ - role: "user", - content: userMessage.message, - }); - - logInfo("==== Final Request to AI ====\n", messages); - - // Stream with abort signal - const chatStream = await withSuppressedTokenWarnings(() => - this.chainManager.chatModelManager.getChatModel().stream(messages, { - signal: abortController.signal, - }) - ); - - for await (const chunk of chatStream) { - if (abortController.signal.aborted) { - logInfo("Stream iteration aborted", { reason: abortController.signal.reason }); - break; - } - streamer.processChunk(chunk); - } - } catch (error: any) { - // Check if the error is due to abort signal - if (error.name === "AbortError" || abortController.signal.aborted) { - logInfo("Stream aborted by user", { reason: abortController.signal.reason }); - // Don't show error message for user-initiated aborts - } else { - await this.handleError(error, addMessage, updateCurrentAiMessage); - } - } - - // Always return the response, even if partial - const response = streamer.close(); - - // Only skip saving if it's a new chat (clearing everything) - if (abortController.signal.aborted && abortController.signal.reason === ABORT_REASON.NEW_CHAT) { - updateCurrentAiMessage(""); - return ""; - } - - return this.handleResponse( - response, - userMessage, - abortController, - addMessage, - updateCurrentAiMessage - ); - } -} - -class VaultQAChainRunner extends BaseChainRunner { - async run( - userMessage: ChatMessage, - abortController: AbortController, - updateCurrentAiMessage: (message: string) => void, - addMessage: (message: ChatMessage) => void, - options: { - debug?: boolean; - ignoreSystemMessage?: boolean; - updateLoading?: (loading: boolean) => void; - } - ): Promise { - const streamer = new ThinkBlockStreamer(updateCurrentAiMessage); - - try { - // Add check for empty index - const indexEmpty = await this.chainManager.vectorStoreManager.isIndexEmpty(); - if (indexEmpty) { - return this.handleResponse( - EMPTY_INDEX_ERROR_MESSAGE, - userMessage, - abortController, - addMessage, - updateCurrentAiMessage - ); - } - - // Get chat history from memory - const memory = this.chainManager.memoryManager.getMemory(); - const memoryVariables = await memory.loadMemoryVariables({}); - const chatHistory = extractChatHistory(memoryVariables); - - // Generate standalone question from user message + chat history - // This is similar to what the conversational retrieval chain does - let standaloneQuestion = userMessage.message; - if (chatHistory.length > 0) { - // For simplicity, we'll use the original question directly - // The original chain would rephrase it, but this approach should work for most cases - standaloneQuestion = userMessage.message; - } - - // Create retriever (similar to how it's done in chainManager) - const retriever = new HybridRetriever({ - minSimilarityScore: 0.01, - maxK: getSettings().maxSourceChunks, - salientTerms: [], - }); - - // Retrieve relevant documents - const retrievedDocs = await retriever.getRelevantDocuments(standaloneQuestion); - - // Store retrieved documents for sources - this.chainManager.storeRetrieverDocuments(retrievedDocs); - - // Format documents as context - const context = retrievedDocs.map((doc: any) => doc.pageContent).join("\n\n"); - - // Create messages array - const messages: any[] = []; - - // Add system message with QA instruction - const systemPrompt = getSystemPrompt(); - const qaInstructions = - "\n\nAnswer the question with as detailed as possible based only on the following context:\n" + - context; - const fullSystemMessage = systemPrompt + qaInstructions; - - const chatModel = this.chainManager.chatModelManager.getChatModel(); - if (fullSystemMessage) { - messages.push({ - role: getMessageRole(chatModel), - content: fullSystemMessage, - }); - } - - // Add chat history - for (const entry of chatHistory) { - messages.push({ role: entry.role, content: entry.content }); - } - - // Add current user question - messages.push({ - role: "user", - content: userMessage.message, - }); - - logInfo("==== Final Request to AI ====\n", messages); - - // Stream with abort signal - const chatStream = await withSuppressedTokenWarnings(() => - this.chainManager.chatModelManager.getChatModel().stream(messages, { - signal: abortController.signal, - }) - ); - - for await (const chunk of chatStream) { - if (abortController.signal.aborted) { - logInfo("VaultQA stream iteration aborted", { reason: abortController.signal.reason }); - break; - } - streamer.processChunk(chunk); - } - } catch (error: any) { - // Check if the error is due to abort signal - if (error.name === "AbortError" || abortController.signal.aborted) { - logInfo("VaultQA stream aborted by user", { reason: abortController.signal.reason }); - // Don't show error message for user-initiated aborts - } else { - await this.handleError(error, addMessage, updateCurrentAiMessage); - } - } - - // Always get the response, even if partial - let fullAIResponse = streamer.close(); - - // Only skip saving if it's a new chat (clearing everything) - if (abortController.signal.aborted && abortController.signal.reason === ABORT_REASON.NEW_CHAT) { - updateCurrentAiMessage(""); - return ""; - } - - // Add sources to the response - fullAIResponse = this.addSourcestoResponse(fullAIResponse); - - return this.handleResponse( - fullAIResponse, - userMessage, - abortController, - addMessage, - updateCurrentAiMessage - ); - } - - private addSourcestoResponse(response: string): string { - const docTitles = extractUniqueTitlesFromDocs(this.chainManager.getRetrievedDocuments()); - if (docTitles.length > 0) { - const links = docTitles.map((title) => `- [[${title}]]`).join("\n"); - response += "\n\n#### Sources:\n\n" + links; - } - return response; - } -} - -class CopilotPlusChainRunner extends BaseChainRunner { - private isYoutubeOnlyMessage(message: string): boolean { - const trimmedMessage = message.trim(); - const hasYoutubeCommand = trimmedMessage.includes("@youtube"); - const youtubeUrl = extractYoutubeUrl(trimmedMessage); - - // Check if message only contains @youtube command and a valid URL - const words = trimmedMessage - .split(/\s+/) - .filter((word) => word !== "@youtube" && word.length > 0); - - return hasYoutubeCommand && youtubeUrl !== null && words.length === 1; - } - - private async processImageUrls(urls: string[]): Promise { - const failedImages: string[] = []; - const processedImages = await ImageBatchProcessor.processUrlBatch( - urls, - failedImages, - this.chainManager.app.vault - ); - ImageBatchProcessor.showFailedImagesNotice(failedImages); - return processedImages; - } - - private async processChatInputImages(content: MessageContent[]): Promise { - const failedImages: string[] = []; - const processedImages = await ImageBatchProcessor.processChatImageBatch( - content, - failedImages, - this.chainManager.app.vault - ); - ImageBatchProcessor.showFailedImagesNotice(failedImages); - return processedImages; - } - - private async extractEmbeddedImages(content: string): Promise { - const imageRegex = /!\[\[(.*?\.(png|jpg|jpeg|gif|webp|bmp|svg))\]\]/g; - const matches = [...content.matchAll(imageRegex)]; - const images = matches.map((match) => match[1]); - return images; - } - - private async buildMessageContent( - textContent: string, - userMessage: ChatMessage - ): Promise { - const failureMessages: string[] = []; - const successfulImages: ImageContent[] = []; - const settings = getSettings(); - - // Collect all image sources - const imageSources: { urls: string[]; type: string }[] = []; - - // Safely check and add context URLs - const contextUrls = userMessage.context?.urls; - if (contextUrls && contextUrls.length > 0) { - imageSources.push({ urls: contextUrls, type: "context" }); - } - - // Process embedded images only if setting is enabled - if (settings.passMarkdownImages) { - const embeddedImages = await this.extractEmbeddedImages(textContent); - if (embeddedImages.length > 0) { - imageSources.push({ urls: embeddedImages, type: "embedded" }); - } - } - - // Process all image sources - for (const source of imageSources) { - const result = await this.processImageUrls(source.urls); - successfulImages.push(...result.successfulImages); - failureMessages.push(...result.failureDescriptions); - } - - // Process existing chat content images if present - const existingContent = userMessage.content; - if (existingContent && existingContent.length > 0) { - const result = await this.processChatInputImages(existingContent); - successfulImages.push(...result.successfulImages); - failureMessages.push(...result.failureDescriptions); - } - - // Let the LLM know about the image processing failures - let finalText = textContent; - if (failureMessages.length > 0) { - finalText = `${textContent}\n\nNote: \n${failureMessages.join("\n")}\n`; - } - - const messageContent: MessageContent[] = [ - { - type: "text", - text: finalText, - }, - ]; - - // Add successful images after the text content - if (successfulImages.length > 0) { - messageContent.push(...successfulImages); - } - - return messageContent; - } - - private hasCapability(model: BaseChatModel, capability: ModelCapability): boolean { - const modelName = (model as any).modelName || (model as any).model || ""; - const customModel = this.chainManager.chatModelManager.findModelByName(modelName); - return customModel?.capabilities?.includes(capability) ?? false; - } - - private isMultimodalModel(model: BaseChatModel): boolean { - return this.hasCapability(model, ModelCapability.VISION); - } - - private async streamMultimodalResponse( - textContent: string, - userMessage: ChatMessage, - abortController: AbortController, - updateCurrentAiMessage: (message: string) => void - ): Promise { - // Get chat history - const memory = this.chainManager.memoryManager.getMemory(); - const memoryVariables = await memory.loadMemoryVariables({}); - const chatHistory = extractChatHistory(memoryVariables); - - // Create messages array starting with system message - const messages: any[] = []; - - // Add system message if available - let fullSystemMessage = await this.getSystemPrompt(); - - // Add chat history context to system message if exists - if (chatHistory.length > 0) { - fullSystemMessage += - "\n\nThe following is the relevant conversation history. Use this context to maintain consistency in your responses:"; - } - - // Get chat model for role determination for O-series models - const chatModel = this.chainManager.chatModelManager.getChatModel(); - - // Add the combined system message with appropriate role - if (fullSystemMessage) { - messages.push({ - role: getMessageRole(chatModel), - content: `${fullSystemMessage}\nIMPORTANT: Maintain consistency with previous responses in the conversation. If you've provided information about a person or topic before, use that same information in follow-up questions.`, - }); - } - - // Add chat history - for (const entry of chatHistory) { - messages.push({ role: entry.role, content: entry.content }); - } - - // Get the current chat model - const chatModelCurrent = this.chainManager.chatModelManager.getChatModel(); - const isMultimodalCurrent = this.isMultimodalModel(chatModelCurrent); - - // Build message content with text and images for multimodal models, or just text for text-only models - const content = isMultimodalCurrent - ? await this.buildMessageContent(textContent, userMessage) - : textContent; - - // Add current user message - messages.push({ - role: "user", - content, - }); - - const enhancedUserMessage = content instanceof Array ? (content[0] as any).text : content; - logInfo("Enhanced user message: ", enhancedUserMessage); - logInfo("==== Final Request to AI ====\n", messages); - const streamer = new ThinkBlockStreamer(updateCurrentAiMessage); - - // Wrap the stream call with warning suppression - const chatStream = await withSuppressedTokenWarnings(() => - this.chainManager.chatModelManager.getChatModel().stream(messages, { - signal: abortController.signal, - }) - ); - - for await (const chunk of chatStream) { - if (abortController.signal.aborted) { - logInfo("CopilotPlus multimodal stream iteration aborted", { - reason: abortController.signal.reason, - }); - break; - } - streamer.processChunk(chunk); - } - - return streamer.close(); - } - - async run( - userMessage: ChatMessage, - abortController: AbortController, - updateCurrentAiMessage: (message: string) => void, - addMessage: (message: ChatMessage) => void, - options: { - debug?: boolean; - ignoreSystemMessage?: boolean; - updateLoading?: (loading: boolean) => void; - updateLoadingMessage?: (message: string) => void; - } - ): Promise { - const { updateLoadingMessage } = options; - let fullAIResponse = ""; - let sources: { title: string; score: number }[] = []; - let currentPartialResponse = ""; - - // Wrapper to track partial response - const trackAndUpdateAiMessage = (message: string) => { - currentPartialResponse = message; - updateCurrentAiMessage(message); - }; - - try { - // Check if this is a YouTube-only message - if (this.isYoutubeOnlyMessage(userMessage.message)) { - const url = extractYoutubeUrl(userMessage.message); - const failMessage = - "Transcript not available. Only videos with the auto transcript option turned on are supported at the moment."; - if (url) { - try { - const response = await BrevilabsClient.getInstance().youtube4llm(url); - if (response.response.transcript) { - return this.handleResponse( - response.response.transcript, - userMessage, - abortController, - addMessage, - updateCurrentAiMessage - ); - } - return this.handleResponse( - failMessage, - userMessage, - abortController, - addMessage, - updateCurrentAiMessage - ); - } catch (error) { - logError("Error processing YouTube video:", error); - return this.handleResponse( - failMessage, - userMessage, - abortController, - addMessage, - updateCurrentAiMessage - ); - } - } - } - - logInfo("==== Step 1: Analyzing intent ===="); - let toolCalls; - // Use the original message for intent analysis - const messageForAnalysis = userMessage.originalMessage || userMessage.message; - try { - toolCalls = await IntentAnalyzer.analyzeIntent(messageForAnalysis); - } catch (error: any) { - return this.handleResponse( - getApiErrorMessage(error), - userMessage, - abortController, - addMessage, - updateCurrentAiMessage - ); - } - - // Use the same removeAtCommands logic as IntentAnalyzer - const cleanedUserMessage = userMessage.message - .split(" ") - .filter((word) => !COPILOT_TOOL_NAMES.includes(word.toLowerCase())) - .join(" ") - .trim(); - - const toolOutputs = await this.executeToolCalls(toolCalls, updateLoadingMessage); - const localSearchResult = toolOutputs.find( - (output) => output.tool === "localSearch" && output.output && output.output.length > 0 - ); - - // Format chat history from memory - const memory = this.chainManager.memoryManager.getMemory(); - const memoryVariables = await memory.loadMemoryVariables({}); - const chatHistory = extractChatHistory(memoryVariables); - - if (localSearchResult) { - logInfo("==== Step 2: Processing local search results ===="); - const documents = JSON.parse(localSearchResult.output); - - logInfo("==== Step 3: Condensing Question ===="); - const standaloneQuestion = await getStandaloneQuestion(cleanedUserMessage, chatHistory); - logInfo("Condensed standalone question: ", standaloneQuestion); - - logInfo("==== Step 4: Preparing context ===="); - const timeExpression = this.getTimeExpression(toolCalls); - const context = this.prepareLocalSearchResult(documents, timeExpression); - - const currentTimeOutputs = toolOutputs.filter((output) => output.tool === "getCurrentTime"); - const enhancedQuestion = this.prepareEnhancedUserMessage( - standaloneQuestion, - currentTimeOutputs - ); - - logInfo(context); - logInfo("==== Step 5: Invoking QA Chain ===="); - const qaPrompt = await this.chainManager.promptManager.getQAPrompt({ - question: enhancedQuestion, - context, - systemMessage: "", // System prompt is added separately in streamMultimodalResponse - }); - - fullAIResponse = await this.streamMultimodalResponse( - qaPrompt, - userMessage, - abortController, - trackAndUpdateAiMessage - ); - - // Append sources to the response - sources = this.getSources(documents); - } else { - // Enhance with tool outputs. - const enhancedUserMessage = this.prepareEnhancedUserMessage( - cleanedUserMessage, - toolOutputs - ); - // If no results, default to LLM Chain - logInfo("No local search results. Using standard LLM Chain."); - - fullAIResponse = await this.streamMultimodalResponse( - enhancedUserMessage, - userMessage, - abortController, - trackAndUpdateAiMessage - ); - } - } catch (error: any) { - // Reset loading message to default - updateLoadingMessage?.(LOADING_MESSAGES.DEFAULT); - - // Check if the error is due to abort signal - if (error.name === "AbortError" || abortController.signal.aborted) { - logInfo("CopilotPlus stream aborted by user", { reason: abortController.signal.reason }); - // Don't show error message for user-initiated aborts - } else { - await this.handleError(error, addMessage, updateCurrentAiMessage); - } - } - - // Only skip saving if it's a new chat (clearing everything) - if (abortController.signal.aborted && abortController.signal.reason === ABORT_REASON.NEW_CHAT) { - updateCurrentAiMessage(""); - return ""; - } - - // If aborted but not a new chat, use the partial response - if (abortController.signal.aborted && currentPartialResponse) { - fullAIResponse = currentPartialResponse; - } - - return this.handleResponse( - fullAIResponse, - userMessage, - abortController, - addMessage, - updateCurrentAiMessage, - sources - ); - } - - private getSources(documents: any): { title: string; score: number }[] { - if (!documents || !Array.isArray(documents)) { - logWarn("No valid documents provided to getSources"); - return []; - } - return this.sortUniqueDocsByScore(documents); - } - - private sortUniqueDocsByScore(documents: any[]): any[] { - const uniqueDocs = new Map(); - - // Iterate through all documents - for (const doc of documents) { - if (!doc.title || (!doc?.score && !doc?.rerank_score)) { - logWarn("Invalid document structure:", doc); - continue; - } - - const currentDoc = uniqueDocs.get(doc.title); - const isReranked = doc && "rerank_score" in doc; - const docScore = isReranked ? doc.rerank_score : doc.score; - - // If the title doesn't exist in the map, or if the new doc has a higher score, update the map - if (!currentDoc || docScore > (currentDoc.score ?? 0)) { - uniqueDocs.set(doc.title, { - title: doc.title, - score: docScore, - isReranked: isReranked, - }); - } - } - - // Convert the map values back to an array and sort by score in descending order - return Array.from(uniqueDocs.values()).sort((a, b) => (b.score ?? 0) - (a.score ?? 0)); - } - - private async executeToolCalls( - toolCalls: any[], - updateLoadingMessage?: (message: string) => void - ) { - const toolOutputs = []; - for (const toolCall of toolCalls) { - logInfo(`==== Step 2: Calling tool: ${toolCall.tool.name} ====`); - if (toolCall.tool.name === "localSearch") { - updateLoadingMessage?.(LOADING_MESSAGES.READING_FILES); - } else if (toolCall.tool.name === "webSearch") { - updateLoadingMessage?.(LOADING_MESSAGES.SEARCHING_WEB); - } else if (toolCall.tool.name === "getFileTree") { - updateLoadingMessage?.(LOADING_MESSAGES.READING_FILE_TREE); - } - const output = await ToolManager.callTool(toolCall.tool, toolCall.args); - toolOutputs.push({ tool: toolCall.tool.name, output }); - } - return toolOutputs; - } - - private prepareEnhancedUserMessage(userMessage: string, toolOutputs: any[]) { - let context = ""; - if (toolOutputs.length > 0) { - const validOutputs = toolOutputs.filter((output) => output.output != null); - if (validOutputs.length > 0) { - context = - "\n\n# Additional context:\n\n" + - validOutputs - .map( - (output) => - `<${output.tool}>\n${typeof output.output !== "string" ? JSON.stringify(output.output) : output.output}\n` - ) - .join("\n\n"); - } - } - return `${userMessage}${context}`; - } - - private getTimeExpression(toolCalls: any[]): string { - const timeRangeCall = toolCalls.find((call) => call.tool.name === "getTimeRangeMs"); - return timeRangeCall ? timeRangeCall.args.timeExpression : ""; - } - - private prepareLocalSearchResult(documents: any[], timeExpression: string): string { - // First filter documents with includeInContext - const includedDocs = documents.filter((doc) => doc.includeInContext); - - // Calculate total content length - const totalLength = includedDocs.reduce((sum, doc) => sum + doc.content.length, 0); - - // If total length exceeds threshold, calculate truncation ratio - let truncatedDocs = includedDocs; - if (totalLength > MAX_CHARS_FOR_LOCAL_SEARCH_CONTEXT) { - const truncationRatio = MAX_CHARS_FOR_LOCAL_SEARCH_CONTEXT / totalLength; - logInfo("Truncating documents to fit context length. Truncation ratio:", truncationRatio); - truncatedDocs = includedDocs.map((doc) => ({ - ...doc, - content: doc.content.slice(0, Math.floor(doc.content.length * truncationRatio)), - })); - } - - const formattedDocs = truncatedDocs - .map((doc: any) => `Note in Vault: ${doc.content}`) - .join("\n\n"); - - return timeExpression - ? `Local Search Result for ${timeExpression}:\n${formattedDocs}` - : `Local Search Result:\n${formattedDocs}`; - } - - protected async getSystemPrompt(): Promise { - return getSystemPrompt(); - } -} - -class ProjectChainRunner extends CopilotPlusChainRunner { - protected async getSystemPrompt(): Promise { - let finalPrompt = getSystemPrompt(); - const projectConfig = getCurrentProject(); - if (!projectConfig) { - return finalPrompt; - } - - // Get context asynchronously - const context = await ProjectManager.instance.getProjectContext(projectConfig.id); - finalPrompt = `${finalPrompt}\n\n\n${projectConfig.systemPrompt}\n`; - - // TODO: Move project context out of the system prompt and into the user prompt. - if (context) { - finalPrompt = `${finalPrompt}\n\n \n${context}\n`; - } - - return finalPrompt; - } -} - -export { CopilotPlusChainRunner, LLMChainRunner, ProjectChainRunner, VaultQAChainRunner }; diff --git a/src/LLMProviders/chainRunner/AutonomousAgentChainRunner.ts b/src/LLMProviders/chainRunner/AutonomousAgentChainRunner.ts new file mode 100644 index 00000000..e21dd039 --- /dev/null +++ b/src/LLMProviders/chainRunner/AutonomousAgentChainRunner.ts @@ -0,0 +1,519 @@ +import { MessageContent } from "@/imageProcessing/imageProcessor"; +import { logError, logInfo, logWarn } from "@/logger"; +import { checkIsPlusUser } from "@/plusUtils"; +import { getSettings, getSystemPrompt } from "@/settings/model"; +import { extractParametersFromZod, SimpleTool } from "@/tools/SimpleTool"; +import { ToolRegistry } from "@/tools/ToolRegistry"; +import { initializeBuiltinTools } from "@/tools/builtinTools"; +import { ChatMessage } from "@/types/message"; +import { getMessageRole, withSuppressedTokenWarnings } from "@/utils"; +import { processToolResults } from "@/utils/toolResultUtils"; +import { CopilotPlusChainRunner } from "./CopilotPlusChainRunner"; +import { addChatHistoryToMessages } from "./utils/chatHistoryUtils"; +import { messageRequiresTools, ModelAdapter, ModelAdapterFactory } from "./utils/modelAdapter"; +import { ThinkBlockStreamer } from "./utils/ThinkBlockStreamer"; +import { createToolCallMarker, updateToolCallMarker } from "./utils/toolCallParser"; +import { + deduplicateSources, + executeSequentialToolCall, + getToolConfirmtionMessage, + getToolDisplayName, + getToolEmoji, + logToolCall, + logToolResult, + ToolExecutionResult, +} from "./utils/toolExecution"; + +import { parseXMLToolCalls, stripToolCallXML } from "./utils/xmlParsing"; + +export class AutonomousAgentChainRunner extends CopilotPlusChainRunner { + private llmFormattedMessages: string[] = []; // Track LLM-formatted messages for memory + + private getAvailableTools(): SimpleTool[] { + const settings = getSettings(); + const registry = ToolRegistry.getInstance(); + + // Initialize tools if not already done + if (registry.getAllTools().length === 0) { + initializeBuiltinTools(this.chainManager.app?.vault); + } + + // Get enabled tool IDs from settings + const enabledToolIds = new Set(settings.autonomousAgentEnabledToolIds || []); + + // Get all enabled tools from registry + return registry.getEnabledTools(enabledToolIds, !!this.chainManager.app?.vault); + } + + private generateToolDescriptions(): string { + const tools = this.getAvailableTools(); + return tools + .map((tool) => { + let params = ""; + + // All tools now have Zod schema + const parameters = extractParametersFromZod(tool.schema); + if (Object.keys(parameters).length > 0) { + params = Object.entries(parameters) + .map(([key, description]) => `<${key}>${description}`) + .join("\n"); + } + + return `<${tool.name}> +${tool.description} + +${params} + +`; + }) + .join("\n\n"); + } + + private buildIterationDisplay( + iterationHistory: string[], + currentIteration: number, + currentMessage: string + ): string { + // Simply join all history without headers or separators + const allParts = [...iterationHistory]; + + // Add current message if present + if (currentMessage) { + allParts.push(currentMessage); + } + + // Join with simple spacing + return allParts.join("\n\n"); + } + + private generateSystemPrompt(): string { + const basePrompt = getSystemPrompt(); + const toolDescriptions = this.generateToolDescriptions(); + const availableTools = this.getAvailableTools(); + const toolNames = availableTools.map((tool) => tool.name); + + // Get tool metadata for custom instructions + const registry = ToolRegistry.getInstance(); + const toolMetadata = availableTools + .map((tool) => registry.getToolMetadata(tool.name)) + .filter((meta): meta is NonNullable => meta !== undefined); + + // Use model adapter for clean model-specific handling + const chatModel = this.chainManager.chatModelManager.getChatModel(); + const adapter = ModelAdapterFactory.createAdapter(chatModel); + + return adapter.enhanceSystemPrompt(basePrompt, toolDescriptions, toolNames, toolMetadata); + } + + async run( + userMessage: ChatMessage, + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void, + addMessage: (message: ChatMessage) => void, + options: { + debug?: boolean; + ignoreSystemMessage?: boolean; + updateLoading?: (loading: boolean) => void; + updateLoadingMessage?: (message: string) => void; + } + ): Promise { + let fullAIResponse = ""; + const conversationMessages: any[] = []; + const iterationHistory: string[] = []; // Track all iterations for display + const collectedSources: { title: string; path: string; score: number }[] = []; // Collect sources from localSearch + this.llmFormattedMessages = []; // Reset LLM messages for this run + const isPlusUser = await checkIsPlusUser(); + if (!isPlusUser) { + await this.handleError(new Error("Invalid license key"), addMessage, updateCurrentAiMessage); + return ""; + } + + try { + // Get chat history from memory + const memory = this.chainManager.memoryManager.getMemory(); + const memoryVariables = await memory.loadMemoryVariables({}); + // Use raw history to preserve multimodal content + const rawHistory = memoryVariables.history || []; + + // Build initial conversation messages + const customSystemPrompt = this.generateSystemPrompt(); + const chatModel = this.chainManager.chatModelManager.getChatModel(); + const adapter = ModelAdapterFactory.createAdapter(chatModel); + + if (customSystemPrompt) { + conversationMessages.push({ + role: getMessageRole(chatModel), + content: customSystemPrompt, + }); + } + + // Add chat history - safely handle different message formats + addChatHistoryToMessages(rawHistory, conversationMessages); + + // Check if the model supports multimodal (vision) capability + const isMultimodal = this.isMultimodalModel(chatModel); + + // Add current user message with model-specific enhancements + const requiresTools = messageRequiresTools(userMessage.message); + const enhancedUserMessage = adapter.enhanceUserMessage(userMessage.message, requiresTools); + + // Build message content with images if multimodal, otherwise just use text + const content: string | MessageContent[] = isMultimodal + ? await this.buildMessageContent(enhancedUserMessage, userMessage) + : enhancedUserMessage; + + conversationMessages.push({ + role: "user", + content, + }); + + // Store original user prompt for tools that need it + const originalUserPrompt = userMessage.originalMessage || userMessage.message; + + // Autonomous agent loop + const maxIterations = getSettings().autonomousAgentMaxIterations; // Get from settings + let iteration = 0; + + while (iteration < maxIterations) { + if (abortController.signal.aborted) { + break; + } + + iteration++; + logInfo(`=== Autonomous Agent Iteration ${iteration} ===`); + + // Store tool call messages for this iteration (declared here so it's accessible in the streaming callback) + const currentIterationToolCallMessages: string[] = []; + + // Get AI response + const response = await this.streamResponse( + conversationMessages, + abortController, + (message) => { + // Show tool calls as indicators during streaming for clarity, preserve think blocks + const cleanedMessage = stripToolCallXML(message); + // Build display with ALL content including tool calls from history + const displayParts = []; + + // Add all iteration history (which includes tool call markers) + displayParts.push(...iterationHistory); + + // Add current iteration's tool calls if any + if (currentIterationToolCallMessages.length > 0) { + displayParts.push(currentIterationToolCallMessages.join("\n")); + } + + // Add the current streaming message + if (cleanedMessage.trim()) { + displayParts.push(cleanedMessage); + } + + const currentDisplay = displayParts.join("\n\n"); + updateCurrentAiMessage(currentDisplay); + }, + adapter + ); + + if (!response) break; + + // Parse tool calls from the response + const toolCalls = parseXMLToolCalls(response); + + // Use model adapter to detect and handle premature responses + const prematureResponseResult = adapter.detectPrematureResponse?.(response); + if (prematureResponseResult?.hasPremature && iteration === 1) { + if (prematureResponseResult.type === "before") { + logWarn("⚠️ Model provided premature response BEFORE tool calls!"); + logWarn("Sanitizing response to keep only tool calls for first iteration"); + } else if (prematureResponseResult.type === "after") { + logWarn("⚠️ Model provided hallucinated response AFTER tool calls!"); + logWarn("Truncating response at last tool call for first iteration"); + } + } + + if (toolCalls.length === 0) { + // No tool calls, this is the final response + // Strip any tool call XML from final response but preserve think blocks + const cleanedResponse = stripToolCallXML(response); + + // Build full response from history (which includes tool call markers) and final response + const allParts = [...iterationHistory]; + if (cleanedResponse.trim()) { + allParts.push(cleanedResponse); + } + fullAIResponse = allParts.join("\n\n"); + + // Add final response to LLM messages + this.llmFormattedMessages.push(response); + break; + } + + // Use model adapter to sanitize response if needed + let sanitizedResponse = response; + if (adapter.sanitizeResponse && prematureResponseResult?.hasPremature) { + sanitizedResponse = adapter.sanitizeResponse(response, iteration); + } + + // Store this iteration's response (AI reasoning) with tool indicators and think blocks preserved + const responseForHistory: string = stripToolCallXML(sanitizedResponse); + + // Only add to history if there's meaningful content + if (responseForHistory.trim()) { + iterationHistory.push(responseForHistory); + } + + // Execute tool calls and show progress + const toolResults: ToolExecutionResult[] = []; + const toolCallIdMap = new Map(); // Map index to tool call ID + + for (let i = 0; i < toolCalls.length; i++) { + const toolCall = toolCalls[i]; + if (abortController.signal.aborted) break; + + // Log tool call details for debugging + logToolCall(toolCall, iteration); + + // Find the tool to check if it's a background tool + const availableTools = this.getAvailableTools(); + const tool = availableTools.find((t) => t.name === toolCall.name); + const isBackgroundTool = tool?.isBackground || false; + + let toolCallId: string | undefined; + + // Only show tool calling message for non-background tools + if (!isBackgroundTool) { + // Create tool calling message with structured marker + const toolEmoji = getToolEmoji(toolCall.name); + const toolDisplayName = getToolDisplayName(toolCall.name); + const confirmationMessage = getToolConfirmtionMessage(toolCall.name); + + // Generate unique ID for this tool call + toolCallId = `${toolCall.name}-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`; + toolCallIdMap.set(i, toolCallId); + + // Create structured tool call marker + const toolCallMarker = createToolCallMarker( + toolCallId, + toolCall.name, + toolDisplayName, + toolEmoji, + confirmationMessage || "", + true, // isExecuting + "", // content (empty for now) + "" // result (empty until execution completes) + ); + + currentIterationToolCallMessages.push(toolCallMarker); + + // Show all history plus all tool call messages + const currentDisplay = [...iterationHistory, ...currentIterationToolCallMessages].join( + "\n\n" + ); + updateCurrentAiMessage(currentDisplay); + } + + const result = await executeSequentialToolCall( + toolCall, + availableTools, + originalUserPrompt + ); + toolResults.push(result); + + // Update the tool call marker with the result if we have an ID + if (toolCallId && !isBackgroundTool) { + // Update the specific tool call message + const messageIndex = currentIterationToolCallMessages.findIndex((msg) => + msg.includes(toolCallId) + ); + if (messageIndex !== -1) { + currentIterationToolCallMessages[messageIndex] = updateToolCallMarker( + currentIterationToolCallMessages[messageIndex], + toolCallId, + result.result + ); + } + + // Update the display with the result + const currentDisplay = [...iterationHistory, ...currentIterationToolCallMessages].join( + "\n\n" + ); + updateCurrentAiMessage(currentDisplay); + } + + // Log tool result + logToolResult(toolCall.name, result); + + // Collect sources from localSearch results + if (toolCall.name === "localSearch" && result.success) { + try { + const searchResults = JSON.parse(result.result); + if (Array.isArray(searchResults)) { + const sources = searchResults.map((doc: any) => ({ + title: doc.title || doc.path, + path: doc.path || doc.title || "", + score: doc.rerank_score || doc.score || 0, + })); + collectedSources.push(...sources); + } + } catch (e) { + logWarn("Failed to parse localSearch results for sources:", e); + } + } + } + + // Add all tool call messages to history so they persist + if (currentIterationToolCallMessages.length > 0) { + const toolCallsString = currentIterationToolCallMessages.join("\n"); + iterationHistory.push(toolCallsString); + } + + // Don't add tool results to display - they're internal only + + // Track LLM-formatted messages for memory + // Add the assistant's response with tool calls + this.llmFormattedMessages.push(response); + + // Add tool results in LLM format (truncated for memory) + if (toolResults.length > 0) { + const toolResultsForLLM = processToolResults(toolResults, true); // truncated for memory + if (toolResultsForLLM) { + this.llmFormattedMessages.push(toolResultsForLLM); + } + } + + // Add AI response to conversation for next iteration + conversationMessages.push({ + role: "assistant", + content: response, + }); + + // Add tool results as user messages for next iteration (full results for current turn) + const toolResultsForConversation = processToolResults(toolResults, false); // full results + + conversationMessages.push({ + role: "user", + content: toolResultsForConversation, + }); + + logInfo("Tool results added to conversation:", toolResultsForConversation); + } + + // If we hit max iterations, add a message explaining the limit was reached + if (iteration >= maxIterations && !fullAIResponse) { + logWarn( + `Autonomous agent reached maximum iterations (${maxIterations}) without completing the task` + ); + + const limitMessage = + `\n\nI've reached the maximum number of iterations (${maxIterations}) for this task. ` + + "I attempted to gather information using various tools but couldn't complete the analysis within the iteration limit. " + + "You may want to try a more specific question or break down your request into smaller parts."; + + fullAIResponse = iterationHistory.join("\n\n") + limitMessage; + } + } catch (error: any) { + if (error.name === "AbortError" || abortController.signal.aborted) { + logInfo("Autonomous agent stream aborted by user", { + reason: abortController.signal.reason, + }); + } else { + logError("Autonomous agent failed, falling back to regular Plus mode:", error); + + // Fallback to regular CopilotPlusChainRunner + try { + const fallbackRunner = new CopilotPlusChainRunner(this.chainManager); + return await fallbackRunner.run( + userMessage, + abortController, + updateCurrentAiMessage, + addMessage, + options + ); + } catch (fallbackError) { + logError("Fallback to regular Plus mode also failed:", fallbackError); + await this.handleError(fallbackError, addMessage, updateCurrentAiMessage); + return ""; + } + } + } + + // Handle response like the parent class, with sources if we found any + const uniqueSources = deduplicateSources(collectedSources); + + // Create LLM-formatted output for memory + const llmFormattedOutput = this.llmFormattedMessages.join("\n\n"); + + // If we somehow don't have a fullAIResponse but have iteration history, use that + if (!fullAIResponse && iterationHistory.length > 0) { + logWarn("fullAIResponse was empty, using iteration history"); + fullAIResponse = iterationHistory.join("\n\n"); + } + + return this.handleResponse( + fullAIResponse, + userMessage, + abortController, + addMessage, + updateCurrentAiMessage, + uniqueSources.length > 0 ? uniqueSources : undefined, + llmFormattedOutput + ); + } + + private async streamResponse( + messages: any[], + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void, + adapter: ModelAdapter + ): Promise { + const streamer = new ThinkBlockStreamer(updateCurrentAiMessage, adapter); + + const maxRetries = 2; + let retryCount = 0; + + while (retryCount <= maxRetries) { + try { + const chatStream = await withSuppressedTokenWarnings(() => + this.chainManager.chatModelManager.getChatModel().stream(messages, { + signal: abortController.signal, + }) + ); + + for await (const chunk of chatStream) { + if (abortController.signal.aborted) { + break; + } + streamer.processChunk(chunk); + } + + return streamer.close(); + } catch (error) { + if (error.name === "AbortError" || abortController.signal.aborted) { + return streamer.close(); + } + + // Check if this is an overloaded error that we should retry + const isOverloadedError = + error?.message?.includes("overloaded") || + error?.message?.includes("Overloaded") || + error?.error?.type === "overloaded_error"; + + if (isOverloadedError && retryCount < maxRetries) { + retryCount++; + logInfo( + `Retrying autonomous agent request (attempt ${retryCount}/${maxRetries + 1}) due to overloaded error` + ); + + // Wait before retrying (exponential backoff) + await new Promise((resolve) => setTimeout(resolve, 1000 * retryCount)); + continue; + } + + throw error; + } + } + + // This should never be reached, but just in case + return streamer.close(); + } +} diff --git a/src/LLMProviders/chainRunner/BaseChainRunner.ts b/src/LLMProviders/chainRunner/BaseChainRunner.ts new file mode 100644 index 00000000..2728ca9f --- /dev/null +++ b/src/LLMProviders/chainRunner/BaseChainRunner.ts @@ -0,0 +1,142 @@ +import { ABORT_REASON, AI_SENDER } from "@/constants"; +import { logError, logInfo } from "@/logger"; +import { ChatMessage } from "@/types/message"; +import { err2String, formatDateTime } from "@/utils"; +import { Notice } from "obsidian"; +import ChainManager from "../chainManager"; + +export interface ChainRunner { + run( + userMessage: ChatMessage, + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void, + addMessage: (message: ChatMessage) => void, + options: { + debug?: boolean; + ignoreSystemMessage?: boolean; + updateLoading?: (loading: boolean) => void; + } + ): Promise; +} + +export abstract class BaseChainRunner implements ChainRunner { + protected chainManager: ChainManager; + + constructor(chainManager: ChainManager) { + this.chainManager = chainManager; + } + + abstract run( + userMessage: ChatMessage, + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void, + addMessage: (message: ChatMessage) => void, + options: { + debug?: boolean; + ignoreSystemMessage?: boolean; + updateLoading?: (loading: boolean) => void; + } + ): Promise; + + protected async handleResponse( + fullAIResponse: string, + userMessage: ChatMessage, + abortController: AbortController, + addMessage: (message: ChatMessage) => void, + updateCurrentAiMessage: (message: string) => void, + sources?: { title: string; path: string; score: number }[], + llmFormattedOutput?: string + ) { + // Save to memory and add message if we have a response + // Skip only if it's a NEW_CHAT abort (clearing everything) + if ( + fullAIResponse && + !(abortController.signal.aborted && abortController.signal.reason === ABORT_REASON.NEW_CHAT) + ) { + // Use saveContext for atomic operation and proper memory management + // Note: LangChain's memory expects text content, not multimodal arrays, so multimodal content is not saved + await this.chainManager.memoryManager + .getMemory() + .saveContext( + { input: userMessage.message }, + { output: llmFormattedOutput || fullAIResponse } + ); + + addMessage({ + message: fullAIResponse, + sender: AI_SENDER, + isVisible: true, + timestamp: formatDateTime(new Date()), + sources: sources, + }); + + // Clear the streaming message since it's now in chat history + updateCurrentAiMessage(""); + } else if (abortController.signal.reason === ABORT_REASON.NEW_CHAT) { + // Also clear if it's a new chat + updateCurrentAiMessage(""); + } + logInfo( + "==== Chat Memory ====\n", + (this.chainManager.memoryManager.getMemory().chatHistory as any).messages.map( + (m: any) => m.content + ) + ); + logInfo("==== Final AI Response ====\n", fullAIResponse); + return fullAIResponse; + } + + protected async handleError( + error: any, + addMessage?: (message: ChatMessage) => void, + updateCurrentAiMessage?: (message: string) => void + ) { + const msg = err2String(error); + logError("Error during LLM invocation:", msg); + const errorData = error?.response?.data?.error || msg; + const errorCode = errorData?.code || msg; + let errorMessage = ""; + + // Check for specific error messages + if (error?.message?.includes("Invalid license key")) { + errorMessage = "Invalid Copilot Plus license key. Please check your license key in settings."; + } else if (errorCode === "model_not_found") { + errorMessage = + "You do not have access to this model or the model does not exist, please check with your API provider."; + } else { + errorMessage = `${errorCode}`; + } + + logError(errorData); + + if (addMessage && updateCurrentAiMessage) { + updateCurrentAiMessage(""); + + // remove langchain troubleshooting URL from error message + const ignoreEndIndex = errorMessage.search("Troubleshooting URL"); + errorMessage = ignoreEndIndex !== -1 ? errorMessage.slice(0, ignoreEndIndex) : errorMessage; + + // add more user guide for invalid API key + if (msg.search(/401|invalid|not valid/gi) !== -1) { + errorMessage = + "Something went wrong. Please check if you have set your API key." + + "\nPath: Settings > copilot plugin > Basic Tab > Set Keys." + + "\nOr check model config" + + "\nError Details: " + + errorMessage; + } + + addMessage({ + message: errorMessage, + isErrorMessage: true, + sender: AI_SENDER, + isVisible: true, + timestamp: formatDateTime(new Date()), + }); + } else { + // Fallback to Notice if message handlers aren't provided + new Notice(errorMessage); + logError(errorData); + } + } +} diff --git a/src/LLMProviders/chainRunner/CopilotPlusChainRunner.ts b/src/LLMProviders/chainRunner/CopilotPlusChainRunner.ts new file mode 100644 index 00000000..bd443981 --- /dev/null +++ b/src/LLMProviders/chainRunner/CopilotPlusChainRunner.ts @@ -0,0 +1,683 @@ +import { getStandaloneQuestion } from "@/chainUtils"; +import { + ABORT_REASON, + COMPOSER_OUTPUT_INSTRUCTIONS, + LOADING_MESSAGES, + MAX_CHARS_FOR_LOCAL_SEARCH_CONTEXT, + ModelCapability, +} from "@/constants"; +import { + ImageBatchProcessor, + ImageContent, + ImageProcessingResult, + MessageContent, +} from "@/imageProcessing/imageProcessor"; +import { BrevilabsClient } from "@/LLMProviders/brevilabsClient"; +import { logError, logInfo, logWarn } from "@/logger"; +import { getSettings, getSystemPrompt } from "@/settings/model"; +import { ToolManager } from "@/tools/toolManager"; +import { writeToFileTool } from "@/tools/ComposerTools"; +import { ChatMessage } from "@/types/message"; +import { + extractYoutubeUrl, + getApiErrorMessage, + getMessageRole, + withSuppressedTokenWarnings, +} from "@/utils"; +import { BaseChatModel } from "@langchain/core/language_models/chat_models"; +import { COPILOT_TOOL_NAMES, IntentAnalyzer } from "../intentAnalyzer"; +import { BaseChainRunner } from "./BaseChainRunner"; +import { ActionBlockStreamer } from "./utils/ActionBlockStreamer"; +import { ThinkBlockStreamer } from "./utils/ThinkBlockStreamer"; +import { + addChatHistoryToMessages, + processRawChatHistory, + processedMessagesToTextOnly, +} from "./utils/chatHistoryUtils"; + +export class CopilotPlusChainRunner extends BaseChainRunner { + private isYoutubeOnlyMessage(message: string): boolean { + const trimmedMessage = message.trim(); + const hasYoutubeCommand = trimmedMessage.includes("@youtube"); + const youtubeUrl = extractYoutubeUrl(trimmedMessage); + + // Check if message only contains @youtube command and a valid URL + const words = trimmedMessage + .split(/\s+/) + .filter((word) => word !== "@youtube" && word.length > 0); + + return hasYoutubeCommand && youtubeUrl !== null && words.length === 1; + } + + private async processImageUrls(urls: string[]): Promise { + const failedImages: string[] = []; + const processedImages = await ImageBatchProcessor.processUrlBatch( + urls, + failedImages, + this.chainManager.app.vault + ); + ImageBatchProcessor.showFailedImagesNotice(failedImages); + return processedImages; + } + + private async processChatInputImages(content: MessageContent[]): Promise { + const failedImages: string[] = []; + const processedImages = await ImageBatchProcessor.processChatImageBatch( + content, + failedImages, + this.chainManager.app.vault + ); + ImageBatchProcessor.showFailedImagesNotice(failedImages); + return processedImages; + } + + private async extractEmbeddedImages(content: string, sourcePath?: string): Promise { + // Match both wiki-style ![[image.ext]] and standard markdown ![alt](image.ext) + const wikiImageRegex = /!\[\[(.*?\.(png|jpg|jpeg|gif|webp|bmp|svg))\]\]/g; + // Updated regex to handle URLs with or without file extensions + const markdownImageRegex = /!\[.*?\]\(([^)]+)\)/g; + + const resolvedImages: string[] = []; + + // Process wiki-style images + const wikiMatches = [...content.matchAll(wikiImageRegex)]; + for (const match of wikiMatches) { + const imageName = match[1]; + + // If we have a source path and access to the app, resolve the wikilink + if (sourcePath) { + const resolvedFile = app.metadataCache.getFirstLinkpathDest(imageName, sourcePath); + + if (resolvedFile) { + // Use the resolved path + resolvedImages.push(resolvedFile.path); + } else { + // If file not found, log a warning but still include the raw filename + logWarn(`Could not resolve embedded image: ${imageName} from source: ${sourcePath}`); + resolvedImages.push(imageName); + } + } else { + // Fallback to raw filename if no source path available + resolvedImages.push(imageName); + } + } + + // Process standard markdown images + const mdMatches = [...content.matchAll(markdownImageRegex)]; + for (const match of mdMatches) { + const imagePath = match[1].trim(); + + // Skip empty paths + if (!imagePath) continue; + + // Handle external URLs (http://, https://, etc.) + if (imagePath.match(/^https?:\/\//)) { + // Include external URLs - they will be processed by processImageUrls + // The ImageProcessor will validate if it's actually an image + resolvedImages.push(imagePath); + continue; + } + + // For local paths, resolve them using Obsidian's metadata cache + // Let ImageBatchProcessor handle validation of whether it's actually an image + // Clean up the path (remove any leading ./ or /) + const cleanPath = imagePath.replace(/^\.\//, "").replace(/^\//, ""); + + // If we have a source path and access to the app, resolve the path + if (sourcePath) { + const resolvedFile = app.metadataCache.getFirstLinkpathDest(cleanPath, sourcePath); + + if (resolvedFile) { + // Use the resolved path + resolvedImages.push(resolvedFile.path); + } else { + // If file not found, still include the raw path + // Let ImageBatchProcessor handle validation + resolvedImages.push(cleanPath); + } + } else { + // Fallback to raw path if no source path available + resolvedImages.push(cleanPath); + } + } + + return resolvedImages; + } + + protected async buildMessageContent( + textContent: string, + userMessage: ChatMessage + ): Promise { + const failureMessages: string[] = []; + const successfulImages: ImageContent[] = []; + const settings = getSettings(); + + // Collect all image sources + const imageSources: { urls: string[]; type: string }[] = []; + + // Safely check and add context URLs + const contextUrls = userMessage.context?.urls; + if (contextUrls && contextUrls.length > 0) { + imageSources.push({ urls: contextUrls, type: "context" }); + } + + // Process embedded images only if setting is enabled + if (settings.passMarkdownImages) { + // Determine source path for resolving wikilinks + let sourcePath: string | undefined; + + // First, check if we have context notes + if (userMessage.context?.notes && userMessage.context.notes.length > 0) { + // Use the first note in context as the source path + sourcePath = userMessage.context.notes[0].path; + } else { + // Fallback to active file if no context notes + const activeFile = this.chainManager.app?.workspace.getActiveFile(); + if (activeFile) { + sourcePath = activeFile.path; + } + } + + const embeddedImages = await this.extractEmbeddedImages(textContent, sourcePath); + if (embeddedImages.length > 0) { + imageSources.push({ urls: embeddedImages, type: "embedded" }); + } + } + + // Process all image sources + for (const source of imageSources) { + const result = await this.processImageUrls(source.urls); + successfulImages.push(...result.successfulImages); + failureMessages.push(...result.failureDescriptions); + } + + // Process existing chat content images if present + const existingContent = userMessage.content; + if (existingContent && existingContent.length > 0) { + const result = await this.processChatInputImages(existingContent); + successfulImages.push(...result.successfulImages); + failureMessages.push(...result.failureDescriptions); + } + + // Let the LLM know about the image processing failures + let finalText = textContent; + if (failureMessages.length > 0) { + finalText = `${textContent}\n\nNote: \n${failureMessages.join("\n")}\n`; + } + + const messageContent: MessageContent[] = [ + { + type: "text", + text: finalText, + }, + ]; + + // Add successful images after the text content + if (successfulImages.length > 0) { + messageContent.push(...successfulImages); + } + + return messageContent; + } + + private hasCapability(model: BaseChatModel, capability: ModelCapability): boolean { + const modelName = (model as any).modelName || (model as any).model || ""; + const customModel = this.chainManager.chatModelManager.findModelByName(modelName); + return customModel?.capabilities?.includes(capability) ?? false; + } + + protected isMultimodalModel(model: BaseChatModel): boolean { + return this.hasCapability(model, ModelCapability.VISION); + } + + /** + * If userMessage.message contains '@composer', append COMPOSER_OUTPUT_INSTRUCTIONS to the text content. + * Handles both string and MessageContent[] types. + */ + private appendComposerInstructionsIfNeeded(content: string, userMessage: ChatMessage): string { + if (!userMessage.message || !userMessage.message.includes("@composer")) { + return content; + } + const composerPrompt = `\n${COMPOSER_OUTPUT_INSTRUCTIONS}\n`; + return `${content}\n\n${composerPrompt}`; + } + + private async streamMultimodalResponse( + textContent: string, + userMessage: ChatMessage, + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void + ): Promise { + // Get chat history + const memory = this.chainManager.memoryManager.getMemory(); + const memoryVariables = await memory.loadMemoryVariables({}); + // Use the raw history array which contains BaseMessage objects + const rawHistory = memoryVariables.history || []; + + // Create messages array starting with system message + const messages: any[] = []; + + // Add system message if available + let fullSystemMessage = await this.getSystemPrompt(); + + // Add chat history context to system message if exists + if (rawHistory.length > 0) { + fullSystemMessage += + "\n\nThe following is the relevant conversation history. Use this context to maintain consistency in your responses:"; + } + + // Get chat model for role determination for O-series models + const chatModel = this.chainManager.chatModelManager.getChatModel(); + + // Add the combined system message with appropriate role + if (fullSystemMessage) { + messages.push({ + role: getMessageRole(chatModel), + content: `${fullSystemMessage}\nIMPORTANT: Maintain consistency with previous responses in the conversation. If you've provided information about a person or topic before, use that same information in follow-up questions.`, + }); + } + + // Add chat history - safely handle different message formats + addChatHistoryToMessages(rawHistory, messages); + + // Get the current chat model + const chatModelCurrent = this.chainManager.chatModelManager.getChatModel(); + const isMultimodalCurrent = this.isMultimodalModel(chatModelCurrent); + + // Build message content with text and images for multimodal models, or just text for text-only models + const content: string | MessageContent[] = isMultimodalCurrent + ? await this.buildMessageContent(textContent, userMessage) + : textContent; + + // Add current user message + messages.push({ + role: "user", + content, + }); + + const enhancedUserMessage = content instanceof Array ? (content[0] as any).text : content; + logInfo("Enhanced user message: ", enhancedUserMessage); + logInfo("==== Final Request to AI ====\n", messages); + const actionStreamer = new ActionBlockStreamer(ToolManager, writeToFileTool); + const thinkStreamer = new ThinkBlockStreamer(updateCurrentAiMessage); + + // Wrap the stream call with warning suppression + const chatStream = await withSuppressedTokenWarnings(() => + this.chainManager.chatModelManager.getChatModel().stream(messages, { + signal: abortController.signal, + }) + ); + + for await (const chunk of chatStream) { + if (abortController.signal.aborted) { + logInfo("CopilotPlus multimodal stream iteration aborted", { + reason: abortController.signal.reason, + }); + break; + } + for await (const processedChunk of actionStreamer.processChunk(chunk)) { + thinkStreamer.processChunk(processedChunk); + } + } + return thinkStreamer.close(); + } + + async run( + userMessage: ChatMessage, + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void, + addMessage: (message: ChatMessage) => void, + options: { + debug?: boolean; + ignoreSystemMessage?: boolean; + updateLoading?: (loading: boolean) => void; + updateLoadingMessage?: (message: string) => void; + } + ): Promise { + const { updateLoadingMessage } = options; + let fullAIResponse = ""; + let sources: { title: string; path: string; score: number }[] = []; + let currentPartialResponse = ""; + + // Wrapper to track partial response + const trackAndUpdateAiMessage = (message: string) => { + currentPartialResponse = message; + updateCurrentAiMessage(message); + }; + + try { + // Check if this is a YouTube-only message + if (this.isYoutubeOnlyMessage(userMessage.message)) { + const url = extractYoutubeUrl(userMessage.message); + const failMessage = + "Transcript not available. Only videos with the auto transcript option turned on are supported at the moment."; + if (url) { + try { + const response = await BrevilabsClient.getInstance().youtube4llm(url); + if (response.response.transcript) { + return this.handleResponse( + response.response.transcript, + userMessage, + abortController, + addMessage, + updateCurrentAiMessage + ); + } + return this.handleResponse( + failMessage, + userMessage, + abortController, + addMessage, + updateCurrentAiMessage + ); + } catch (error) { + logError("Error processing YouTube video:", error); + return this.handleResponse( + failMessage, + userMessage, + abortController, + addMessage, + updateCurrentAiMessage + ); + } + } + } + + logInfo("==== Step 1: Analyzing intent ===="); + let toolCalls; + // Use the original message for intent analysis + const messageForAnalysis = userMessage.originalMessage || userMessage.message; + try { + toolCalls = await IntentAnalyzer.analyzeIntent(messageForAnalysis); + } catch (error: any) { + return this.handleResponse( + getApiErrorMessage(error), + userMessage, + abortController, + addMessage, + updateCurrentAiMessage + ); + } + + // Use the same removeAtCommands logic as IntentAnalyzer + const cleanedUserMessage = userMessage.message + .split(" ") + .filter((word) => !COPILOT_TOOL_NAMES.includes(word.toLowerCase())) + .join(" ") + .trim(); + + const toolOutputs = await this.executeToolCalls(toolCalls, updateLoadingMessage); + + // Extract sources from localSearch if present + const localSearchResult = toolOutputs.find( + (output) => output.tool === "localSearch" && output.output != null + ); + + let hasLocalSearchWithResults = false; + if (localSearchResult) { + try { + const documents = JSON.parse(localSearchResult.output); + if (Array.isArray(documents) && documents.length > 0) { + hasLocalSearchWithResults = true; + sources = this.getSources(documents); + } + } catch (error) { + logWarn("Failed to parse localSearch results for sources:", error); + } + } + + // Format chat history from memory + const memory = this.chainManager.memoryManager.getMemory(); + const memoryVariables = await memory.loadMemoryVariables({}); + const rawHistory = memoryVariables.history || []; + + // Process history consistently - same data used for both LLM and question condensing + const processedHistory = processRawChatHistory(rawHistory); + const chatHistory = processedMessagesToTextOnly(processedHistory); + + // Get standalone question if we have chat history + let questionForEnhancement = cleanedUserMessage; + if (chatHistory.length > 0) { + logInfo("==== Condensing Question ===="); + questionForEnhancement = await getStandaloneQuestion(cleanedUserMessage, chatHistory); + logInfo("Condensed standalone question: ", questionForEnhancement); + } + + // Enhance with ALL tool outputs including localSearch + let enhancedUserMessage = this.prepareEnhancedUserMessage( + questionForEnhancement, + toolOutputs, + toolCalls + ); + + // If localSearch has actual results and no other tools, add QA-style instruction to maintain same behavior + const hasOtherTools = toolOutputs.some( + (output) => output.tool !== "localSearch" && output.output != null + ); + if (hasLocalSearchWithResults && !hasOtherTools) { + // The QA format is already handled in prepareEnhancedUserMessage, just add the instruction + enhancedUserMessage = `Answer the question with as detailed as possible based only on the following context:\n${enhancedUserMessage}`; + } + + // Append composer instruction to the end of text prompt to enhance instruction following. + enhancedUserMessage = this.appendComposerInstructionsIfNeeded( + enhancedUserMessage, + userMessage + ); + + logInfo("==== Invoking LLM with all tool results ===="); + fullAIResponse = await this.streamMultimodalResponse( + enhancedUserMessage, + userMessage, + abortController, + trackAndUpdateAiMessage + ); + } catch (error: any) { + // Reset loading message to default + updateLoadingMessage?.(LOADING_MESSAGES.DEFAULT); + + // Check if the error is due to abort signal + if (error.name === "AbortError" || abortController.signal.aborted) { + logInfo("CopilotPlus stream aborted by user", { reason: abortController.signal.reason }); + // Don't show error message for user-initiated aborts + } else { + await this.handleError(error, addMessage, updateCurrentAiMessage); + } + } + + // Only skip saving if it's a new chat (clearing everything) + if (abortController.signal.aborted && abortController.signal.reason === ABORT_REASON.NEW_CHAT) { + updateCurrentAiMessage(""); + return ""; + } + + // If aborted but not a new chat, use the partial response + if (abortController.signal.aborted && currentPartialResponse) { + fullAIResponse = currentPartialResponse; + } + + return this.handleResponse( + fullAIResponse, + userMessage, + abortController, + addMessage, + updateCurrentAiMessage, + sources + ); + } + + private getSources(documents: any): { title: string; path: string; score: number }[] { + if (!documents || !Array.isArray(documents)) { + logWarn("No valid documents provided to getSources"); + return []; + } + return this.sortUniqueDocsByScore(documents); + } + + private sortUniqueDocsByScore(documents: any[]): any[] { + const uniqueDocs = new Map(); + + // Iterate through all documents + for (const doc of documents) { + if (!doc.title || (!doc?.score && !doc?.rerank_score)) { + logWarn("Invalid document structure:", doc); + continue; + } + + // Use path as the unique key, falling back to title if path is not available + const key = doc.path || doc.title; + const currentDoc = uniqueDocs.get(key); + const isReranked = doc && "rerank_score" in doc; + const docScore = isReranked ? doc.rerank_score : doc.score; + + // If the document doesn't exist in the map, or if the new doc has a higher score, update the map + if (!currentDoc || docScore > (currentDoc.score ?? 0)) { + uniqueDocs.set(key, { + title: doc.title, + path: doc.path || doc.title, // Use path if available, otherwise use title + score: docScore, + isReranked: isReranked, + }); + } + } + + // Convert the map values back to an array and sort by score in descending order + return Array.from(uniqueDocs.values()).sort((a, b) => (b.score ?? 0) - (a.score ?? 0)); + } + + private async executeToolCalls( + toolCalls: any[], + updateLoadingMessage?: (message: string) => void + ) { + const toolOutputs = []; + for (const toolCall of toolCalls) { + logInfo(`==== Step 2: Calling tool: ${toolCall.tool.name} ====`); + if (toolCall.tool.name === "localSearch") { + updateLoadingMessage?.(LOADING_MESSAGES.READING_FILES); + } else if (toolCall.tool.name === "webSearch") { + updateLoadingMessage?.(LOADING_MESSAGES.SEARCHING_WEB); + } else if (toolCall.tool.name === "getFileTree") { + updateLoadingMessage?.(LOADING_MESSAGES.READING_FILE_TREE); + } + const output = await ToolManager.callTool(toolCall.tool, toolCall.args); + toolOutputs.push({ tool: toolCall.tool.name, output }); + } + return toolOutputs; + } + + private prepareEnhancedUserMessage(userMessage: string, toolOutputs: any[], toolCalls?: any[]) { + let context = ""; + let hasLocalSearchWithResults = false; + + // Check if localSearch has actual results (non-empty documents array) + const localSearchOutput = toolOutputs.find( + (output) => output.tool === "localSearch" && output.output != null + ); + + if (localSearchOutput && typeof localSearchOutput.output === "string") { + try { + const documents = JSON.parse(localSearchOutput.output); + if (Array.isArray(documents) && documents.length > 0) { + hasLocalSearchWithResults = true; + } + } catch { + // Invalid JSON or parsing error + } + } + + if (toolOutputs.length > 0) { + const validOutputs = toolOutputs.filter((output) => output.output != null); + if (validOutputs.length > 0) { + // Don't add "Additional context" header if only localSearch with results to maintain QA format + const contextHeader = + hasLocalSearchWithResults && validOutputs.length === 1 + ? "" + : "\n\n# Additional context:\n\n"; + context = + contextHeader + + validOutputs + .map((output) => { + let content = output.output; + + // Special formatting for localSearch results + if (output.tool === "localSearch" && typeof content === "string") { + try { + const documents = JSON.parse(content); + if (Array.isArray(documents) && documents.length > 0) { + // Get time expression from toolCalls if available + const timeExpression = toolCalls ? this.getTimeExpression(toolCalls) : ""; + const formattedContent = this.prepareLocalSearchResult( + documents, + timeExpression + ); + content = formattedContent; + } + } catch (error) { + // If parsing fails, use the raw output + logWarn("Failed to parse localSearch output for formatting:", error); + } + } + + // Ensure content is string + if (typeof content !== "string") { + content = JSON.stringify(content); + } + + // Only wrap in XML tags if there are multiple tools + if (validOutputs.length > 1) { + return `<${output.tool}>\n${content}\n`; + } else if (output.tool === "localSearch" && hasLocalSearchWithResults) { + // For localSearch with results only, don't wrap in XML to maintain QA format + return content; + } else { + return `<${output.tool}>\n${content}\n`; + } + }) + .join("\n\n"); + } + } + + // For QA format when only localSearch with results is present + if (hasLocalSearchWithResults && toolOutputs.filter((o) => o.output != null).length === 1) { + return `${context}\n\nQuestion: ${userMessage}`; + } + + return `${userMessage}${context}`; + } + + private getTimeExpression(toolCalls: any[]): string { + const timeRangeCall = toolCalls.find((call) => call.tool.name === "getTimeRangeMs"); + return timeRangeCall ? timeRangeCall.args.timeExpression : ""; + } + + private prepareLocalSearchResult(documents: any[], timeExpression: string): string { + // First filter documents with includeInContext + const includedDocs = documents.filter((doc) => doc.includeInContext); + + // Calculate total content length + const totalLength = includedDocs.reduce((sum, doc) => sum + doc.content.length, 0); + + // If total length exceeds threshold, calculate truncation ratio + let truncatedDocs = includedDocs; + if (totalLength > MAX_CHARS_FOR_LOCAL_SEARCH_CONTEXT) { + const truncationRatio = MAX_CHARS_FOR_LOCAL_SEARCH_CONTEXT / totalLength; + logInfo("Truncating documents to fit context length. Truncation ratio:", truncationRatio); + truncatedDocs = includedDocs.map((doc) => ({ + ...doc, + content: doc.content.slice(0, Math.floor(doc.content.length * truncationRatio)), + })); + } + + const formattedDocs = truncatedDocs + .map((doc: any) => `Note in Vault: ${doc.content}`) + .join("\n\n"); + + return timeExpression + ? `Local Search Result for ${timeExpression}:\n${formattedDocs}` + : `Local Search Result:\n${formattedDocs}`; + } + + protected async getSystemPrompt(): Promise { + return getSystemPrompt(); + } +} diff --git a/src/LLMProviders/chainRunner/LLMChainRunner.ts b/src/LLMProviders/chainRunner/LLMChainRunner.ts new file mode 100644 index 00000000..13b667f8 --- /dev/null +++ b/src/LLMProviders/chainRunner/LLMChainRunner.ts @@ -0,0 +1,97 @@ +import { ABORT_REASON } from "@/constants"; +import { logInfo } from "@/logger"; +import { getSystemPrompt } from "@/settings/model"; +import { ChatMessage } from "@/types/message"; +import { extractChatHistory, getMessageRole, withSuppressedTokenWarnings } from "@/utils"; +import { BaseChainRunner } from "./BaseChainRunner"; +import { ThinkBlockStreamer } from "./utils/ThinkBlockStreamer"; + +export class LLMChainRunner extends BaseChainRunner { + async run( + userMessage: ChatMessage, + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void, + addMessage: (message: ChatMessage) => void, + options: { + debug?: boolean; + ignoreSystemMessage?: boolean; + updateLoading?: (loading: boolean) => void; + } + ): Promise { + const streamer = new ThinkBlockStreamer(updateCurrentAiMessage); + + try { + // Get chat history from memory + const memory = this.chainManager.memoryManager.getMemory(); + const memoryVariables = await memory.loadMemoryVariables({}); + const chatHistory = extractChatHistory(memoryVariables); + + // Create messages array starting with system message + const messages: any[] = []; + + // Add system message if available + const systemPrompt = getSystemPrompt(); + const chatModel = this.chainManager.chatModelManager.getChatModel(); + + if (systemPrompt) { + messages.push({ + role: getMessageRole(chatModel), + content: systemPrompt, + }); + } + + // Add chat history + for (const entry of chatHistory) { + messages.push({ role: entry.role, content: entry.content }); + } + + // Add current user message + messages.push({ + role: "user", + content: userMessage.message, + }); + + logInfo("==== Final Request to AI ====\n", messages); + + // Stream with abort signal + const chatStream = await withSuppressedTokenWarnings(() => + this.chainManager.chatModelManager.getChatModel().stream(messages, { + signal: abortController.signal, + }) + ); + + for await (const chunk of chatStream) { + if (abortController.signal.aborted) { + logInfo("Stream iteration aborted", { reason: abortController.signal.reason }); + break; + } + streamer.processChunk(chunk); + } + } catch (error: any) { + // Check if the error is due to abort signal + if (error.name === "AbortError" || abortController.signal.aborted) { + logInfo("Stream aborted by user", { reason: abortController.signal.reason }); + // Don't show error message for user-initiated aborts + } else { + await this.handleError(error, addMessage, updateCurrentAiMessage); + } + } + + // Always return the response, even if partial + const response = streamer.close(); + + // Only skip saving if it's a new chat (clearing everything) + if (abortController.signal.aborted && abortController.signal.reason === ABORT_REASON.NEW_CHAT) { + updateCurrentAiMessage(""); + return ""; + } + + return this.handleResponse( + response, + userMessage, + abortController, + addMessage, + updateCurrentAiMessage + ); + } +} diff --git a/src/LLMProviders/chainRunner/ProjectChainRunner.ts b/src/LLMProviders/chainRunner/ProjectChainRunner.ts new file mode 100644 index 00000000..443cd407 --- /dev/null +++ b/src/LLMProviders/chainRunner/ProjectChainRunner.ts @@ -0,0 +1,25 @@ +import { getCurrentProject } from "@/aiParams"; +import { getSystemPrompt } from "@/settings/model"; +import ProjectManager from "../projectManager"; +import { CopilotPlusChainRunner } from "./CopilotPlusChainRunner"; + +export class ProjectChainRunner extends CopilotPlusChainRunner { + protected async getSystemPrompt(): Promise { + let finalPrompt = getSystemPrompt(); + const projectConfig = getCurrentProject(); + if (!projectConfig) { + return finalPrompt; + } + + // Get context asynchronously + const context = await ProjectManager.instance.getProjectContext(projectConfig.id); + finalPrompt = `${finalPrompt}\n\n\n${projectConfig.systemPrompt}\n`; + + // TODO: Move project context out of the system prompt and into the user prompt. + if (context) { + finalPrompt = `${finalPrompt}\n\n \n${context}\n`; + } + + return finalPrompt; + } +} diff --git a/src/LLMProviders/chainRunner/README.md b/src/LLMProviders/chainRunner/README.md new file mode 100644 index 00000000..527b9da8 --- /dev/null +++ b/src/LLMProviders/chainRunner/README.md @@ -0,0 +1,640 @@ +# Chain Runner Architecture & Tool Calling System + +This directory contains the refactored chain runner system for Obsidian Copilot, providing multiple chain execution strategies with different tool calling approaches. + +## Overview + +The chain runner system provides two distinct tool calling approaches: + +1. **Legacy Tool Calling** (CopilotPlusChainRunner) - Uses Brevilabs API for intent analysis +2. **Autonomous Agent** (AutonomousAgentChainRunner) - Uses XML-based tool calling + +## Architecture + +``` +chainRunner/ +├── BaseChainRunner.ts # Abstract base class with shared functionality +├── LLMChainRunner.ts # Basic LLM interaction (no tools) +├── VaultQAChainRunner.ts # Vault-only Q&A with retrieval +├── CopilotPlusChainRunner.ts # Legacy tool calling system +├── ProjectChainRunner.ts # Project-aware extension of Plus +├── AutonomousAgentChainRunner.ts # XML-based autonomous agent tool calling +├── index.ts # Main exports +└── utils/ + ├── ThinkBlockStreamer.ts # Handles thinking content from models + ├── xmlParsing.ts # XML tool call parsing utilities + ├── toolExecution.ts # Tool execution helpers + └── modelAdapter.ts # Model-specific adaptations +``` + +## Tool Calling Systems Comparison + +### 1. Legacy Tool Calling (CopilotPlusChainRunner) + +**How it works:** + +- Uses Brevilabs API (`IntentAnalyzer.analyzeIntent()`) to analyze user intent +- Determines which tools to call based on the analysis +- Executes tools synchronously before sending to LLM +- Enhances user message with tool outputs as context + +**Flow:** + +``` +User Message → Intent Analysis → Tool Execution → Enhanced Prompt → LLM Response +``` + +**Example:** + +```typescript +// 1. Analyze intent +const toolCalls = await IntentAnalyzer.analyzeIntent(message); + +// 2. Execute tools +const toolOutputs = await this.executeToolCalls(toolCalls); + +// 3. Enhance message with context +const enhancedMessage = this.prepareEnhancedUserMessage(message, toolOutputs); + +// 4. Send to LLM +const response = await this.streamMultimodalResponse(enhancedMessage, ...); +``` + +**Tools Available:** + +- `localSearch` - Search vault content +- `webSearch` - Search the web +- `getCurrentTime` - Get current time +- `getFileTree` - Get file structure +- `pomodoroTool` - Pomodoro timer +- `youtubeTranscription` - YouTube video transcription + +### 2. Autonomous Agent (AutonomousAgentChainRunner) + +**How it works:** + +- **No LangChain dependency** - Uses simple tool interface with XML-based tool calling +- AI decides autonomously which tools to use via structured XML format +- Iterative loop where AI can call multiple tools in sequence +- Each tool result informs the next decision + +**Flow:** + +``` +User Message → AI Reasoning → XML Tool Call → Tool Execution → +AI Analysis → More Tools? → Final Response +``` + +**XML Tool Call Format:** + +```xml + +localSearch + +{ + "query": "machine learning notes", + "salientTerms": ["machine", "learning", "AI", "algorithms"] +} + + +``` + +**Sequential Loop:** + +```typescript +while (iteration < maxIterations) { + // 1. Get AI response + const response = await this.streamResponse(messages); + + // 2. Parse XML tool calls + const toolCalls = parseXMLToolCalls(response); + + if (toolCalls.length === 0) { + // No tools needed - final response + break; + } + + // 3. Execute each tool + for (const toolCall of toolCalls) { + const result = await executeSequentialToolCall(toolCall, availableTools); + toolResults.push(result); + } + + // 4. Add results to conversation for next iteration + messages.push({ role: "user", content: toolResultsForConversation }); +} +``` + +## Key Differences + +| Aspect | Legacy (Plus) | Autonomous Agent | +| ------------------ | ----------------------- | ------------------------------------- | +| **Tool Decision** | Brevilabs API analysis | AI decides autonomously | +| **Tool Execution** | Pre-LLM, synchronous | During conversation, iterative | +| **Tool Format** | SimpleTool interface | XML-based structured format | +| **Reasoning** | Intent analysis → tools | AI reasoning → tools → more reasoning | +| **Iterations** | Single pass | Up to 4 iterations | +| **Tool Chaining** | Limited | Full chaining support | + +## SimpleTool Interface + +### Overview + +The SimpleTool interface provides a clean, type-safe way to define tools with Zod validation: + +```typescript +interface SimpleTool { + name: string; + description: string; + schema: TSchema; + call: (args: z.infer) => Promise; + timeoutMs?: number; + isBackground?: boolean; +} +``` + +### Creating Tools + +All tools are created using the unified `createTool` function with Zod schemas: + +#### Tool with No Parameters + +```typescript +const indexTool = createTool({ + name: "indexVault", + description: "Index the vault to the Copilot index", + schema: z.void(), // No parameters + handler: async () => { + // Tool implementation + return "Indexing complete"; + }, + isBackground: true, // Optional: hide from user +}); +``` + +#### Tool with Parameters + +```typescript +// Define schema with validation rules +const searchSchema = z.object({ + query: z.string().min(1).describe("The search query"), + salientTerms: z.array(z.string()).min(1).describe("Key terms extracted from query"), + timeRange: z + .object({ + startTime: z.any(), + endTime: z.any(), + }) + .optional() + .describe("Time range for search"), +}); + +// Create tool with automatic validation +const searchTool = createTool({ + name: "localSearch", + description: "Search for notes based on query and time range", + schema: searchSchema, + handler: async ({ query, salientTerms, timeRange }) => { + // Handler receives fully typed and validated arguments + // TypeScript knows the exact types from the schema + return performSearch(query, salientTerms, timeRange); + }, + timeoutMs: 30000, // Optional: custom timeout +}); +``` + +### Benefits of Unified Zod Approach + +1. **Type Safety**: Full TypeScript type inference from schemas +2. **Runtime Validation**: All inputs validated before reaching handler +3. **Consistent Interface**: One way to create all tools +4. **Better Error Messages**: Zod provides detailed validation errors +5. **No Any Types**: Everything is properly typed +6. **Simpler Codebase**: No need to maintain multiple tool creation methods + +### Advanced Zod Patterns + +#### Complex Validation + +```typescript +const emailToolSchema = z.object({ + to: z.string().email().describe("Recipient email"), + subject: z.string().min(1).max(100).describe("Email subject"), + body: z.string().min(1).describe("Email content"), + cc: z.array(z.string().email()).optional().describe("CC recipients"), +}); +``` + +#### Union Types for Actions + +```typescript +const actionSchema = z.discriminatedUnion("type", [ + z.object({ + type: z.literal("search"), + query: z.string().min(1), + }), + z.object({ + type: z.literal("create"), + content: z.string().min(1), + tags: z.array(z.string()).default([]), + }), + z.object({ + type: z.literal("delete"), + id: z.string().uuid(), + }), +]); + +const actionTool = createTool({ + name: "performAction", + description: "Perform various actions", + schema: actionSchema, + handler: async (action) => { + // TypeScript knows exactly which type based on discriminator + switch (action.type) { + case "search": + return search(action.query); + case "create": + return create(action.content, action.tags); + case "delete": + return deleteItem(action.id); + } + }, +}); +``` + +#### Custom Validation + +```typescript +const filePathSchema = z + .string() + .refine((val) => val.endsWith(".md") || val.endsWith(".canvas"), { + message: "File must be .md or .canvas", + }) + .refine((val) => !val.includes(".."), { message: "Path traversal not allowed" }) + .describe("Path to markdown or canvas file"); +``` + +#### Transformations + +```typescript +const dateToolSchema = z.object({ + date: z + .string() + .describe("Date in ISO format or natural language") + .transform((str) => new Date(str)), + timezone: z.string().default("UTC").describe("Timezone identifier"), +}); +``` + +### Schema Composition + +```typescript +// Base schemas that can be reused +const timeRangeSchema = z + .object({ + startTime: z.date(), + endTime: z.date(), + }) + .refine((data) => data.endTime > data.startTime, { + message: "End time must be after start time", + }); + +const paginationSchema = z.object({ + page: z.number().int().positive().default(1), + pageSize: z.number().int().positive().max(100).default(20), +}); + +// Compose into larger schemas +const searchWithPaginationSchema = z.object({ + query: z.string().min(1).describe("Search query"), + filters: z.record(z.string()).optional().describe("Additional filters"), + timeRange: timeRangeSchema.optional(), + pagination: paginationSchema, +}); +``` + +### Default Values + +```typescript +const configSchema = z.object({ + temperature: z.number().min(0).max(2).default(0.7), + maxTokens: z.number().int().positive().default(1000), + model: z.enum(["gpt-4", "gpt-3.5-turbo"]).default("gpt-4"), +}); + +// Handler receives object with defaults applied +const configTool = createTool({ + name: "updateConfig", + schema: configSchema, + handler: async (config) => { + // config.temperature is always defined (0.7 if not provided) + // config.maxTokens is always defined (1000 if not provided) + // config.model is always defined ("gpt-4" if not provided) + return updateConfiguration(config); + }, +}); +``` + +### Handling Validation Errors with Retry + +When AI-generated parameters fail Zod validation, the tool execution will return a formatted error. The autonomous agent automatically handles this through its iterative loop: + +```typescript +// Example tool with strict validation +const searchToolWithValidation = createTool({ + name: "searchNotes", + description: "Search notes with specific criteria", + schema: z.object({ + query: z.string().min(2, "Query must be at least 2 characters"), + limit: z.number().int().min(1).max(100), + sortBy: z.enum(["relevance", "date", "title"]), + }), + handler: async ({ query, limit, sortBy }) => { + return performSearch(query, limit, sortBy); + }, +}); + +// When the AI provides invalid parameters: +// Input: { query: "a", limit: 200, sortBy: "random" } +// +// The flow: +// 1. Tool execution catches Zod validation error +// 2. Returns: "Tool searchNotes validation failed: query: Query must be at least 2 characters, +// limit: Number must be less than or equal to 100, sortBy: Invalid enum value" +// 3. This error is added to the conversation as a user message +// 4. The AI sees the error in the next iteration and can retry with corrected parameters +// 5. The autonomous agent continues up to 4 iterations, allowing multiple retry attempts + +// Example conversation flow: +// Iteration 1: AI calls tool with invalid params → receives error +// Iteration 2: AI understands error and retries with { query: "search term", limit: 50, sortBy: "date" } → success +``` + +The validation errors are automatically formatted to be clear and actionable, helping the AI self-correct. The autonomous agent's iterative design naturally provides retry capability with the AI learning from each error. + +## XML Tool Calling Details + +### Tool Call Parsing (`xmlParsing.ts`) + +```typescript +// Parse XML tool calls from AI response +function parseXMLToolCalls(text: string): ToolCall[] { + const regex = /([\s\S]*?)<\/use_tool>/g; + // Extracts name and args from XML structure +} + +// Strip tool calls from display +function stripToolCallXML(text: string): string { + // Removes XML tool blocks and code blocks for clean display +} +``` + +### Tool Execution (`toolExecution.ts`) + +```typescript +// Execute individual tool with timeout and error handling +async function executeSequentialToolCall( + toolCall: ToolCall, + availableTools: any[] +): Promise { + // 30-second timeout per tool + // Error handling and validation + // Result formatting +} +``` + +### Available Tools in Sequential Mode + +All tools from the legacy system plus autonomous decision-making: + +- **localSearch** - Vault content search with salient terms +- **webSearch** - Web search with chat history context +- **getFileTree** - File structure exploration +- **getCurrentTime** - Time-based queries +- **pomodoroTool** - Productivity timer +- **indexTool** - Vault indexing operations +- **youtubeTranscription** - Video content analysis + +### System Prompt Engineering + +The Autonomous Agent mode uses a comprehensive system prompt that: + +1. **Explains the XML format** with exact examples +2. **Provides tool descriptions** with parameter details +3. **Sets expectations** for reasoning and tool chaining +4. **Includes critical requirements** (e.g., salientTerms for localSearch) + +Example system prompt section: + +``` +When you need to use a tool, format it EXACTLY like this: + +localSearch + +{ + "query": "piano learning", + "salientTerms": ["piano", "learning", "practice", "music"] +} + + + +CRITICAL: For localSearch, you MUST always provide both "query" (string) and "salientTerms" (array of strings). +``` + +## Benefits of Autonomous Agent + +1. **Autonomous Tool Selection** - AI decides what tools to use without pre-analysis +2. **Tool Chaining** - Can use results from one tool to inform the next +3. **Complex Workflows** - Multi-step reasoning with tool support +4. **Model Agnostic** - Works with any LLM that can follow XML format +5. **No External Dependencies** - No Brevilabs API required +6. **Transparency** - User can see the AI's reasoning process + +## Usage + +### Enable Autonomous Agent + +```typescript +// In settings +settings.enableAutonomousAgent = true; + +// ChainManager automatically selects the appropriate runner +const runner = chainManager.getChainRunner(); // Returns AutonomousAgentChainRunner +``` + +### Example Query Flow + +**User Input:** "Find my notes about machine learning and research current best practices" + +**Autonomous Agent Process:** + +1. **Iteration 1**: AI reasons about the task → calls `localSearch` for ML notes +2. **Iteration 2**: Analyzes vault results → calls `webSearch` for current practices +3. **Iteration 3**: Synthesizes both sources → provides comprehensive response + +**Legacy Process:** + +1. Intent analysis determines both tools needed +2. Executes both tools +3. Single LLM call with all context + +## Error Handling & Fallbacks + +### Autonomous Agent Fallbacks + +```typescript +try { + // Sequential thinking execution +} catch (error) { + // Automatic fallback to CopilotPlusChainRunner + const fallbackRunner = new CopilotPlusChainRunner(this.chainManager); + return await fallbackRunner.run(/* same parameters */); +} +``` + +### Tool Execution Safeguards + +- 30-second timeout per tool +- Graceful error handling with descriptive messages +- Tool availability validation +- Result validation and formatting + +## Model Adapter Pattern + +### Overview + +The Model Adapter pattern handles model-specific quirks and requirements cleanly, keeping the core logic model-agnostic. + +### Architecture + +```typescript +interface ModelAdapter { + enhanceSystemPrompt(basePrompt: string, toolDescriptions: string): string; + enhanceUserMessage(message: string, requiresTools: boolean): string; + parseToolCalls?(response: string): any[]; // Future extension + needsSpecialHandling(): boolean; + sanitizeResponse?(response: string, iteration: number): string; + shouldTruncateStreaming?(partialResponse: string): boolean; + detectPrematureResponse?(response: string): { + hasPremature: boolean; + type: "before" | "after" | null; + }; +} +``` + +### Current Adapters + +1. **BaseModelAdapter** - Default behavior for well-behaved models +2. **GPTModelAdapter** - Aggressive prompting for GPT models that often skip tool calls +3. **ClaudeModelAdapter** - Specialized handling for Claude thinking models (3.7 Sonnet, Claude 4) +4. **GeminiModelAdapter** - Ready for Gemini-specific handling + +### Adding a New Model + +```typescript +class NewModelAdapter extends BaseModelAdapter { + enhanceSystemPrompt(basePrompt: string, toolDescriptions: string): string { + const base = super.enhanceSystemPrompt(basePrompt, toolDescriptions); + return base + "\n\n[Model-specific instructions here]"; + } + + enhanceUserMessage(message: string, requiresTools: boolean): string { + // Add model-specific hints if needed + return requiresTools ? `${message}\n[Model-specific hint]` : message; + } +} + +// Register in ModelAdapterFactory +if (modelName.includes("newmodel")) { + return new NewModelAdapter(modelName); +} +``` + +### Claude Model Adapter Features + +The `ClaudeModelAdapter` includes specialized handling for Claude thinking models: + +#### Thinking Model Support + +- **Claude 3.7 Sonnet** and **Claude 4** - Automatic thinking mode configuration +- **Think Block Preservation** - Maintains valuable reasoning context in responses +- **Temperature Control** - Disables temperature for thinking models (as required by API) + +#### Claude 4 Hallucination Prevention + +Claude 4 has a tendency to write complete responses immediately after tool calls instead of waiting for results. The adapter addresses this with: + +```typescript +// Enhanced prompting with explicit autonomous agent pattern +enhanceSystemPrompt(basePrompt: string, toolDescriptions: string): string { + if (this.isClaudeSonnet4()) { + // Add specific instructions for Claude 4: + // - Brief sentence + tool calls + STOP pattern + // - Explicit warnings about premature responses + // - Clear autonomous agent iteration guidance + } +} + +// Detection of premature responses +detectPrematureResponse(response: string): { + hasPremature: boolean; + type: "before" | "after" | null; +} { + // Allows brief sentences before tool calls (up to 2 sentences, 200 chars) + // Detects substantial content after tool calls (forbidden) + // Uses threshold-based detection for generalizability +} + +// Response sanitization +sanitizeResponse(response: string, iteration: number): string { + // Preserves ALL think blocks + // Removes substantial non-thinking content after tool calls + // Only applies to first iteration when hallucination occurs +} + +// Streaming truncation +shouldTruncateStreaming(partialResponse: string): boolean { + // Prevents streaming of hallucinated content to users + // Truncates at last complete tool call when threshold exceeded +} +``` + +#### Flow Improvement + +The adapter creates a better conversational flow by allowing brief explanatory sentences before tool calls: + +``` +[Think block] +I'll search your vault and web for piano practice information. +🔍 Calling vault search... +[Think block] +Let me gather more specific information about practice routines. +🌐 Calling web search... +[Think block] +[final answer] +``` + +### Benefits + +1. **Separation of Concerns** - Model quirks isolated from core logic +2. **Maintainability** - Easy to find and update model-specific code +3. **Extensibility** - Simple to add support for new models +4. **Testing** - Model adapters can be unit tested independently +5. **Clean Core** - Autonomous agent logic remains model-agnostic +6. **Hallucination Prevention** - Specialized handling for problematic models +7. **Streaming Protection** - Prevents bad content from reaching users +8. **Generalizable Solutions** - Uses threshold-based detection over regex patterns + +## Future Considerations + +1. **Tool Discovery** - Dynamic tool registration +2. **Custom Tools** - User-defined tool capabilities +3. **Parallel Execution** - Multiple tools simultaneously +4. **Tool Result Caching** - Avoid redundant calls +5. **Advanced Reasoning** - More sophisticated decision trees +6. **Tool Permissions** - User control over tool access +7. **Alternative Parsing** - Model adapters could handle non-XML formats +8. **Response Validation** - Adapters could validate model outputs +9. **Model-Specific Optimizations** - Expand adapter capabilities for emerging models +10. **Hallucination Detection** - More sophisticated premature response detection + +The autonomous agent approach represents a significant evolution from traditional tool calling, enabling more sophisticated AI reasoning and autonomous task completion. diff --git a/src/LLMProviders/chainRunner/VaultQAChainRunner.ts b/src/LLMProviders/chainRunner/VaultQAChainRunner.ts new file mode 100644 index 00000000..e3b249d0 --- /dev/null +++ b/src/LLMProviders/chainRunner/VaultQAChainRunner.ts @@ -0,0 +1,161 @@ +import { ABORT_REASON, EMPTY_INDEX_ERROR_MESSAGE, RETRIEVED_DOCUMENT_TAG } from "@/constants"; +import { logInfo } from "@/logger"; +import { HybridRetriever } from "@/search/hybridRetriever"; +import { getSettings, getSystemPrompt } from "@/settings/model"; +import { ChatMessage } from "@/types/message"; +import { + extractChatHistory, + extractUniqueTitlesFromDocs, + getMessageRole, + withSuppressedTokenWarnings, +} from "@/utils"; +import { BaseChainRunner } from "./BaseChainRunner"; +import { ThinkBlockStreamer } from "./utils/ThinkBlockStreamer"; + +export class VaultQAChainRunner extends BaseChainRunner { + async run( + userMessage: ChatMessage, + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void, + addMessage: (message: ChatMessage) => void, + options: { + debug?: boolean; + ignoreSystemMessage?: boolean; + updateLoading?: (loading: boolean) => void; + } + ): Promise { + const streamer = new ThinkBlockStreamer(updateCurrentAiMessage); + + try { + // Add check for empty index + const indexEmpty = await this.chainManager.vectorStoreManager.isIndexEmpty(); + if (indexEmpty) { + return this.handleResponse( + EMPTY_INDEX_ERROR_MESSAGE, + userMessage, + abortController, + addMessage, + updateCurrentAiMessage + ); + } + + // Get chat history from memory + const memory = this.chainManager.memoryManager.getMemory(); + const memoryVariables = await memory.loadMemoryVariables({}); + const chatHistory = extractChatHistory(memoryVariables); + + // Generate standalone question from user message + chat history + // This is similar to what the conversational retrieval chain does + let standaloneQuestion = userMessage.message; + if (chatHistory.length > 0) { + // For simplicity, we'll use the original question directly + // The original chain would rephrase it, but this approach should work for most cases + standaloneQuestion = userMessage.message; + } + + // Create retriever (similar to how it's done in chainManager) + const retriever = new HybridRetriever({ + minSimilarityScore: 0.01, + maxK: getSettings().maxSourceChunks, + salientTerms: [], + }); + + // Retrieve relevant documents + const retrievedDocs = await retriever.getRelevantDocuments(standaloneQuestion); + + // Store retrieved documents for sources + this.chainManager.storeRetrieverDocuments(retrievedDocs); + + // Format documents as context with XML tags + const context = retrievedDocs + .map( + (doc: any) => + `<${RETRIEVED_DOCUMENT_TAG}>\n${doc.pageContent}\n` + ) + .join("\n\n"); + + // Create messages array + const messages: any[] = []; + + // Add system message with QA instruction + const systemPrompt = getSystemPrompt(); + const qaInstructions = + "\n\nAnswer the question with as detailed as possible based only on the following context:\n" + + context; + const fullSystemMessage = systemPrompt + qaInstructions; + + const chatModel = this.chainManager.chatModelManager.getChatModel(); + if (fullSystemMessage) { + messages.push({ + role: getMessageRole(chatModel), + content: fullSystemMessage, + }); + } + + // Add chat history + for (const entry of chatHistory) { + messages.push({ role: entry.role, content: entry.content }); + } + + // Add current user question + messages.push({ + role: "user", + content: userMessage.message, + }); + + logInfo("==== Final Request to AI ====\n", messages); + + // Stream with abort signal + const chatStream = await withSuppressedTokenWarnings(() => + this.chainManager.chatModelManager.getChatModel().stream(messages, { + signal: abortController.signal, + }) + ); + + for await (const chunk of chatStream) { + if (abortController.signal.aborted) { + logInfo("VaultQA stream iteration aborted", { reason: abortController.signal.reason }); + break; + } + streamer.processChunk(chunk); + } + } catch (error: any) { + // Check if the error is due to abort signal + if (error.name === "AbortError" || abortController.signal.aborted) { + logInfo("VaultQA stream aborted by user", { reason: abortController.signal.reason }); + // Don't show error message for user-initiated aborts + } else { + await this.handleError(error, addMessage, updateCurrentAiMessage); + } + } + + // Always get the response, even if partial + let fullAIResponse = streamer.close(); + + // Only skip saving if it's a new chat (clearing everything) + if (abortController.signal.aborted && abortController.signal.reason === ABORT_REASON.NEW_CHAT) { + updateCurrentAiMessage(""); + return ""; + } + + // Add sources to the response + fullAIResponse = this.addSourcestoResponse(fullAIResponse); + + return this.handleResponse( + fullAIResponse, + userMessage, + abortController, + addMessage, + updateCurrentAiMessage + ); + } + + private addSourcestoResponse(response: string): string { + const docTitles = extractUniqueTitlesFromDocs(this.chainManager.getRetrievedDocuments()); + if (docTitles.length > 0) { + const links = docTitles.map((title) => `- [[${title}]]`).join("\n"); + response += "\n\n#### Sources:\n\n" + links; + } + return response; + } +} diff --git a/src/LLMProviders/chainRunner/index.ts b/src/LLMProviders/chainRunner/index.ts new file mode 100644 index 00000000..1103cb1b --- /dev/null +++ b/src/LLMProviders/chainRunner/index.ts @@ -0,0 +1,22 @@ +// Main exports for chain runners +export type { ChainRunner } from "./BaseChainRunner"; +export { BaseChainRunner } from "./BaseChainRunner"; +export { LLMChainRunner } from "./LLMChainRunner"; +export { VaultQAChainRunner } from "./VaultQAChainRunner"; +export { CopilotPlusChainRunner } from "./CopilotPlusChainRunner"; +export { ProjectChainRunner } from "./ProjectChainRunner"; +export { AutonomousAgentChainRunner } from "./AutonomousAgentChainRunner"; + +// Utility exports (for internal use or testing) +export { ThinkBlockStreamer } from "./utils/ThinkBlockStreamer"; +export { parseXMLToolCalls, stripToolCallXML } from "./utils/xmlParsing"; +export type { ToolCall } from "./utils/xmlParsing"; +export { + executeSequentialToolCall, + getToolDisplayName, + getToolEmoji, + logToolCall, + logToolResult, + deduplicateSources, +} from "./utils/toolExecution"; +export type { ToolExecutionResult } from "./utils/toolExecution"; diff --git a/src/LLMProviders/chainRunner/utils/ActionBlockStreamer.test.ts b/src/LLMProviders/chainRunner/utils/ActionBlockStreamer.test.ts new file mode 100644 index 00000000..aaaeb6fa --- /dev/null +++ b/src/LLMProviders/chainRunner/utils/ActionBlockStreamer.test.ts @@ -0,0 +1,232 @@ +import { ActionBlockStreamer } from "./ActionBlockStreamer"; +import { ToolManager } from "@/tools/toolManager"; +import { ToolResultFormatter } from "@/tools/ToolResultFormatter"; + +// Mock the ToolManager and ToolResultFormatter +jest.mock("@/tools/toolManager"); +jest.mock("@/tools/ToolResultFormatter"); + +const MockedToolManager = ToolManager as jest.Mocked; +const MockedToolResultFormatter = ToolResultFormatter as jest.Mocked; + +describe("ActionBlockStreamer", () => { + let writeToFileTool: any; + let streamer: ActionBlockStreamer; + + beforeEach(() => { + writeToFileTool = { name: "writeToFile" }; + MockedToolManager.callTool.mockClear(); + + // Mock ToolResultFormatter to return the raw result without "File change result: " prefix + MockedToolResultFormatter.format = jest.fn((_toolName, result) => result); + + streamer = new ActionBlockStreamer(MockedToolManager, writeToFileTool); + }); + + // Helper function to process chunks and collect results + async function processChunks(chunks: { content: string | null }[]) { + const outputContents: any[] = []; + for (const chunk of chunks) { + for await (const result of streamer.processChunk(chunk)) { + // Always push the content, even if it's null, undefined, or empty string + outputContents.push(result.content); + } + } + return outputContents; + } + + it("should pass through chunks without writeToFile tags unchanged", async () => { + const chunks = [{ content: "Hello " }, { content: "world, this is " }, { content: "a test." }]; + const output = await processChunks(chunks); + + // All chunks should be yielded as-is + expect(output).toEqual(["Hello ", "world, this is ", "a test."]); + + // No tool calls should be made + expect(MockedToolManager.callTool).not.toHaveBeenCalled(); + }); + + it("should handle a complete writeToFile block in a single chunk", async () => { + MockedToolManager.callTool.mockResolvedValue("File written successfully."); + const chunks = [ + { + content: + "Some text before file.txtcontent and some text after.", + }, + ]; + const output = await processChunks(chunks); + + // Should yield original chunk plus tool result + expect(output).toEqual([ + "Some text before file.txtcontent and some text after.", + "\nFile written successfully.\n", + ]); + + expect(MockedToolManager.callTool).toHaveBeenCalledWith(writeToFileTool, { + path: "file.txt", + content: "content", + }); + }); + + it("should handle a complete xml-wrapped writeToFile block", async () => { + MockedToolManager.callTool.mockResolvedValue("XML file written."); + const chunks = [ + { + content: + "```xml\nfile.xmlxml content\n```", + }, + ]; + const output = await processChunks(chunks); + + expect(output).toEqual([ + "```xml\nfile.xmlxml content\n```", + "\nXML file written.\n", + ]); + + expect(MockedToolManager.callTool).toHaveBeenCalledWith(writeToFileTool, { + path: "file.xml", + content: "xml content", + }); + }); + + it("should handle a writeToFile block split across multiple chunks", async () => { + MockedToolManager.callTool.mockResolvedValue("Split file written."); + const chunks = [ + { content: "Here is a file split.txt" }, + { content: "split content" }, + { content: " That was it." }, + ]; + const output = await processChunks(chunks); + + // All chunks should be yielded as-is, plus tool result when complete block is detected + expect(output).toEqual([ + "Here is a file split.txt", + "split content", + " That was it.", + "\nSplit file written.\n", + ]); + + expect(MockedToolManager.callTool).toHaveBeenCalledWith(writeToFileTool, { + path: "split.txt", + content: "split content", + }); + }); + + it("should handle multiple writeToFile blocks in the stream", async () => { + MockedToolManager.callTool + .mockResolvedValueOnce("File 1 written.") + .mockResolvedValueOnce("File 2 written."); + const chunks = [ + { + content: + "f1.txtc1Some textf2.txtc2", + }, + ]; + const output = await processChunks(chunks); + + // Should yield original chunk plus both tool results + expect(output).toEqual([ + "f1.txtc1Some textf2.txtc2", + "\nFile 1 written.\n", + "\nFile 2 written.\n", + ]); + + expect(MockedToolManager.callTool).toHaveBeenCalledTimes(2); + expect(MockedToolManager.callTool).toHaveBeenCalledWith(writeToFileTool, { + path: "f1.txt", + content: "c1", + }); + expect(MockedToolManager.callTool).toHaveBeenCalledWith(writeToFileTool, { + path: "f2.txt", + content: "c2", + }); + }); + + it("should handle unclosed tags without calling tools", async () => { + const chunks = [ + { content: "Starting... unclosed.txt" }, + { content: "this will not be closed" }, + ]; + const output = await processChunks(chunks); + + // Should yield all chunks as-is + expect(output).toEqual([ + "Starting... unclosed.txt", + "this will not be closed", + ]); + + // No tool calls should be made for incomplete blocks + expect(MockedToolManager.callTool).not.toHaveBeenCalled(); + }); + + it("should handle tool call errors gracefully", async () => { + MockedToolManager.callTool.mockRejectedValue(new Error("Tool error")); + const chunks = [ + { + content: "error.txtcontent", + }, + ]; + const output = await processChunks(chunks); + + expect(output).toEqual([ + "error.txtcontent", + "\nError: Tool error\n", + ]); + + expect(MockedToolManager.callTool).toHaveBeenCalledWith(writeToFileTool, { + path: "error.txt", + content: "content", + }); + }); + + it("should handle chunks with different content types", async () => { + const chunks: any[] = [ + { content: "Hello" }, + { content: null }, + { content: "" }, + { content: " World" }, + ]; + const output = await processChunks(chunks); + + // Should yield all chunks as-is, null content is yielded but not added to buffer + expect(output).toEqual(["Hello", null, "", " World"]); + }); + + it("should handle whitespace in path and content", async () => { + MockedToolManager.callTool.mockResolvedValue("Whitespace handled."); + const chunks = [ + { + content: + " spaced.txt content with spaces ", + }, + ]; + await processChunks(chunks); + + expect(MockedToolManager.callTool).toHaveBeenCalledWith(writeToFileTool, { + path: "spaced.txt", + content: "content with spaces", + }); + }); + + it("should handle malformed blocks gracefully", async () => { + MockedToolManager.callTool.mockResolvedValue("Malformed handled."); + const chunks = [ + { + content: "missing-content.txt", + }, + ]; + const output = await processChunks(chunks); + + // Should yield chunk as-is plus tool result + expect(output).toEqual([ + "missing-content.txt", + "\nMalformed handled.\n", + ]); + + // Tool should be called with undefined content + expect(MockedToolManager.callTool).toHaveBeenCalledWith(writeToFileTool, { + path: "missing-content.txt", + content: undefined, + }); + }); +}); diff --git a/src/LLMProviders/chainRunner/utils/ActionBlockStreamer.ts b/src/LLMProviders/chainRunner/utils/ActionBlockStreamer.ts new file mode 100644 index 00000000..7ac41888 --- /dev/null +++ b/src/LLMProviders/chainRunner/utils/ActionBlockStreamer.ts @@ -0,0 +1,93 @@ +import { ToolManager } from "@/tools/toolManager"; +import { ToolResultFormatter } from "@/tools/ToolResultFormatter"; + +/** + * ActionBlockStreamer processes streaming chunks to detect and handle writeToFile blocks. + * + * 1. Accumulates chunks in a buffer + * 2. Detects complete writeToFile blocks + * 3. Calls the writeToFile tool when a complete block is found + * 4. Returns chunks as-is otherwise + */ +export class ActionBlockStreamer { + private buffer = ""; + + constructor( + private toolManager: typeof ToolManager, + private writeToFileTool: any + ) {} + + private findCompleteBlock(str: string) { + // Regex for both formats + const regex = /[\s\S]*?<\/writeToFile>/; + const match = str.match(regex); + + if (!match || match.index === undefined) { + return null; + } + + return { + block: match[0], + endIdx: match.index + match[0].length, + }; + } + + async *processChunk(chunk: any): AsyncGenerator { + // Handle different chunk formats + let chunkContent = ""; + + // Handle Claude thinking model array-based content + if (Array.isArray(chunk.content)) { + for (const item of chunk.content) { + if (item.type === "text" && item.text != null) { + chunkContent += item.text; + } + } + } + // Handle standard string content + else if (chunk.content != null) { + chunkContent = chunk.content; + } + + // Add to buffer + if (chunkContent) { + this.buffer += chunkContent; + } + + // Yield the original chunk as-is + yield chunk; + + // Process all complete blocks in the buffer + let blockInfo = this.findCompleteBlock(this.buffer); + + while (blockInfo) { + const { block, endIdx } = blockInfo; + + // Extract content from the block + const pathMatch = block.match(/([\s\S]*?)<\/path>/); + const contentMatch = block.match(/([\s\S]*?)<\/content>/); + const filePath = pathMatch ? pathMatch[1].trim() : undefined; + const fileContent = contentMatch ? contentMatch[1].trim() : undefined; + + // Call the tool + try { + const result = await this.toolManager.callTool(this.writeToFileTool, { + path: filePath, + content: fileContent, + }); + + // Format tool result using ToolResultFormatter for consistency with agent mode + const formattedResult = ToolResultFormatter.format("writeToFile", result); + yield { ...chunk, content: `\n${formattedResult}\n` }; + } catch (err: any) { + yield { ...chunk, content: `\nError: ${err?.message || err}\n` }; + } + + // Remove processed block from buffer + this.buffer = this.buffer.substring(endIdx); + + // Check for another complete block in the remaining buffer + blockInfo = this.findCompleteBlock(this.buffer); + } + } +} diff --git a/src/LLMProviders/chainRunner/utils/ThinkBlockStreamer.ts b/src/LLMProviders/chainRunner/utils/ThinkBlockStreamer.ts new file mode 100644 index 00000000..5aaee1df --- /dev/null +++ b/src/LLMProviders/chainRunner/utils/ThinkBlockStreamer.ts @@ -0,0 +1,124 @@ +import { ModelAdapter } from "./modelAdapter"; + +/** + * ThinkBlockStreamer handles streaming content from various LLM providers + * that support thinking/reasoning modes (like Claude and Deepseek). + */ +export class ThinkBlockStreamer { + private hasOpenThinkBlock = false; + private fullResponse = ""; + private shouldTruncate = false; + + constructor( + private updateCurrentAiMessage: (message: string) => void, + private modelAdapter?: ModelAdapter + ) {} + + private handleClaude37Chunk(content: any[]) { + let textContent = ""; + for (const item of content) { + switch (item.type) { + case "text": + textContent += item.text; + break; + case "thinking": + if (!this.hasOpenThinkBlock) { + this.fullResponse += "\n"; + this.hasOpenThinkBlock = true; + } + // Guard against undefined thinking content + if (item.thinking !== undefined) { + this.fullResponse += item.thinking; + } + this.updateCurrentAiMessage(this.fullResponse); + return true; // Indicate we handled a thinking chunk + } + } + if (textContent) { + this.fullResponse += textContent; + } + return false; // No thinking chunk handled + } + + private handleDeepseekChunk(chunk: any) { + // Handle standard string content + if (typeof chunk.content === "string") { + this.fullResponse += chunk.content; + } + + // Handle deepseek reasoning/thinking content + if (chunk.additional_kwargs?.reasoning_content) { + if (!this.hasOpenThinkBlock) { + this.fullResponse += "\n"; + this.hasOpenThinkBlock = true; + } + // Guard against undefined reasoning content + if (chunk.additional_kwargs.reasoning_content !== undefined) { + this.fullResponse += chunk.additional_kwargs.reasoning_content; + } + return true; // Indicate we handled a thinking chunk + } + return false; // No thinking chunk handled + } + + processChunk(chunk: any) { + // If we've already decided to truncate, don't process more chunks + if (this.shouldTruncate) { + return; + } + + let handledThinking = false; + + // Handle Claude 3.7 array-based content + if (Array.isArray(chunk.content)) { + handledThinking = this.handleClaude37Chunk(chunk.content); + } else { + // Handle deepseek format + handledThinking = this.handleDeepseekChunk(chunk); + } + + // Close think block if we have one open and didn't handle thinking content + if (this.hasOpenThinkBlock && !handledThinking) { + this.fullResponse += ""; + this.hasOpenThinkBlock = false; + } + + // Check if we should truncate streaming based on model adapter + if (this.modelAdapter?.shouldTruncateStreaming?.(this.fullResponse)) { + this.shouldTruncate = true; + // Find the last complete tool call to truncate cleanly + this.fullResponse = this.truncateToLastCompleteToolCall(this.fullResponse); + } + + this.updateCurrentAiMessage(this.fullResponse); + } + + private truncateToLastCompleteToolCall(response: string): string { + // Find the last complete tag + const lastCompleteToolEnd = response.lastIndexOf(""); + + if (lastCompleteToolEnd === -1) { + // No complete tool calls found, return original response + return response; + } + + // Truncate to after the last complete tool call + const truncated = response.substring(0, lastCompleteToolEnd + "".length); + + // Use model adapter to sanitize if available + if (this.modelAdapter?.sanitizeResponse) { + return this.modelAdapter.sanitizeResponse(truncated, 1); + } + + return truncated; + } + + close() { + // Make sure to close any open think block at the end + if (this.hasOpenThinkBlock) { + this.fullResponse += ""; + this.updateCurrentAiMessage(this.fullResponse); + } + return this.fullResponse; + } +} diff --git a/src/LLMProviders/chainRunner/utils/chatHistoryUtils.test.ts b/src/LLMProviders/chainRunner/utils/chatHistoryUtils.test.ts new file mode 100644 index 00000000..1e028f47 --- /dev/null +++ b/src/LLMProviders/chainRunner/utils/chatHistoryUtils.test.ts @@ -0,0 +1,245 @@ +import { + processRawChatHistory, + addChatHistoryToMessages, + processedMessagesToTextOnly, +} from "./chatHistoryUtils"; + +describe("chatHistoryUtils", () => { + describe("processRawChatHistory", () => { + it("should process BaseMessage objects correctly", () => { + const rawHistory = [ + { + _getType: () => "human", + content: "Hello", + }, + { + _getType: () => "ai", + content: "Hi there!", + }, + ]; + + const result = processRawChatHistory(rawHistory); + + expect(result).toEqual([ + { role: "user", content: "Hello" }, + { role: "assistant", content: "Hi there!" }, + ]); + }); + + it("should handle multimodal content in BaseMessage objects", () => { + const multimodalContent = [ + { type: "text", text: "What is this?" }, + { type: "image_url", image_url: { url: "data:image/jpeg;base64,..." } }, + ]; + + const rawHistory = [ + { + _getType: () => "human", + content: multimodalContent, + }, + { + _getType: () => "ai", + content: "This is an image of a cat.", + }, + ]; + + const result = processRawChatHistory(rawHistory); + + expect(result).toEqual([ + { role: "user", content: multimodalContent }, + { role: "assistant", content: "This is an image of a cat." }, + ]); + }); + + it("should skip system messages", () => { + const rawHistory = [ + { + _getType: () => "system", + content: "You are a helpful assistant", + }, + { + _getType: () => "human", + content: "Hello", + }, + ]; + + const result = processRawChatHistory(rawHistory); + + expect(result).toEqual([{ role: "user", content: "Hello" }]); + }); + + it("should handle legacy message formats", () => { + const rawHistory = [ + { + role: "human", + content: "Hello", + }, + { + role: "ai", + content: "Hi!", + }, + { + sender: "user", + content: "How are you?", + }, + { + sender: "AI", + content: "I am doing well!", + }, + ]; + + const result = processRawChatHistory(rawHistory); + + expect(result).toEqual([ + { role: "user", content: "Hello" }, + { role: "assistant", content: "Hi!" }, + { role: "user", content: "How are you?" }, + { role: "assistant", content: "I am doing well!" }, + ]); + }); + + it("should handle null and undefined messages", () => { + const rawHistory = [ + null, + undefined, + { + _getType: () => "human", + content: "Hello", + }, + {}, + { content: undefined }, + ]; + + const result = processRawChatHistory(rawHistory); + + expect(result).toEqual([{ role: "user", content: "Hello" }]); + }); + + it("should handle messages with unknown types", () => { + const rawHistory = [ + { + _getType: () => "unknown", + content: "Some content", + }, + { + role: "unknown", + content: "Other content", + }, + ]; + + const result = processRawChatHistory(rawHistory); + + expect(result).toEqual([]); + }); + }); + + describe("addChatHistoryToMessages", () => { + it("should add processed messages to target array", () => { + const rawHistory = [ + { + _getType: () => "human", + content: "Hello", + }, + { + _getType: () => "ai", + content: "Hi there!", + }, + ]; + + const messages: any[] = []; + addChatHistoryToMessages(rawHistory, messages); + + expect(messages).toEqual([ + { role: "user", content: "Hello" }, + { role: "assistant", content: "Hi there!" }, + ]); + }); + + it("should append to existing messages", () => { + const rawHistory = [ + { + _getType: () => "human", + content: "Hello", + }, + ]; + + const messages = [{ role: "system", content: "You are helpful" }]; + + addChatHistoryToMessages(rawHistory, messages); + + expect(messages).toEqual([ + { role: "system", content: "You are helpful" }, + { role: "user", content: "Hello" }, + ]); + }); + }); + + describe("processedMessagesToTextOnly", () => { + it("should handle string content", () => { + const processedMessages = [ + { role: "user" as const, content: "Hello" }, + { role: "assistant" as const, content: "Hi there!" }, + ]; + + const result = processedMessagesToTextOnly(processedMessages); + + expect(result).toEqual([ + { role: "user", content: "Hello" }, + { role: "assistant", content: "Hi there!" }, + ]); + }); + + it("should extract text from multimodal content", () => { + const processedMessages = [ + { + role: "user" as const, + content: [ + { type: "text", text: "What is this?" }, + { type: "image_url", image_url: { url: "data:image/jpeg;base64,..." } }, + ], + }, + { + role: "assistant" as const, + content: "This is a cat.", + }, + ]; + + const result = processedMessagesToTextOnly(processedMessages); + + expect(result).toEqual([ + { role: "user", content: "What is this?" }, + { role: "assistant", content: "This is a cat." }, + ]); + }); + + it("should handle multiple text parts in multimodal content", () => { + const processedMessages = [ + { + role: "user" as const, + content: [ + { type: "text", text: "First part." }, + { type: "image_url", image_url: { url: "data:image/jpeg;base64,..." } }, + { type: "text", text: "Second part." }, + ], + }, + ]; + + const result = processedMessagesToTextOnly(processedMessages); + + expect(result).toEqual([{ role: "user", content: "First part. Second part." }]); + }); + + it("should handle image-only content", () => { + const processedMessages = [ + { + role: "user" as const, + content: [{ type: "image_url", image_url: { url: "data:image/jpeg;base64,..." } }], + }, + ]; + + const result = processedMessagesToTextOnly(processedMessages); + + expect(result).toEqual([{ role: "user", content: "[Image content]" }]); + }); + }); +}); diff --git a/src/LLMProviders/chainRunner/utils/chatHistoryUtils.ts b/src/LLMProviders/chainRunner/utils/chatHistoryUtils.ts new file mode 100644 index 00000000..f868ccfb --- /dev/null +++ b/src/LLMProviders/chainRunner/utils/chatHistoryUtils.ts @@ -0,0 +1,112 @@ +/** + * Utility functions for safely processing chat history from LangChain memory + */ + +export interface ProcessedMessage { + role: "user" | "assistant"; + content: any; // string or MessageContent[] +} + +/** + * Safely process raw history from LangChain memory, handling both BaseMessage + * objects and legacy formats while preserving multimodal content + * + * @param rawHistory Array of messages from memory.loadMemoryVariables() + * @returns Array of processed messages safe for LLM consumption + */ +export function processRawChatHistory(rawHistory: any[]): ProcessedMessage[] { + const messages: ProcessedMessage[] = []; + + for (const message of rawHistory) { + if (!message) continue; + + // Check if this is a BaseMessage with _getType method + if (typeof message._getType === "function") { + const messageType = message._getType(); + + // Only process human and AI messages + if (messageType === "human") { + messages.push({ role: "user", content: message.content }); + } else if (messageType === "ai") { + messages.push({ role: "assistant", content: message.content }); + } + // Skip system messages and unknown types + } else if (message.content !== undefined) { + // Fallback for other message formats - try to infer role + const role = inferMessageRole(message); + if (role) { + messages.push({ role, content: message.content }); + } + } + } + + return messages; +} + +/** + * Try to infer the role from various message format properties + * @returns 'user' | 'assistant' | null + */ +function inferMessageRole(message: any): "user" | "assistant" | null { + // Check various properties that might indicate the role + if (message.role === "human" || message.role === "user" || message.sender === "user") { + return "user"; + } else if (message.role === "ai" || message.role === "assistant" || message.sender === "AI") { + return "assistant"; + } + + // Can't determine role + return null; +} + +/** + * Add processed chat history to messages array for LLM consumption + * This is a convenience function that combines processing and adding + * + * @param rawHistory Raw history from memory + * @param messages Target messages array to add to + */ +export function addChatHistoryToMessages(rawHistory: any[], messages: any[]): void { + const processedHistory = processRawChatHistory(rawHistory); + for (const msg of processedHistory) { + messages.push({ role: msg.role, content: msg.content }); + } +} + +export interface ChatHistoryEntry { + role: "user" | "assistant"; + content: string; +} + +/** + * Convert processed messages to text-only format for question condensing + * This extracts just the text content from potentially multimodal messages + * + * @param processedMessages Messages processed by processRawChatHistory + * @returns Array of text-only chat history entries + */ +export function processedMessagesToTextOnly( + processedMessages: ProcessedMessage[] +): ChatHistoryEntry[] { + return processedMessages.map((msg) => { + let textContent: string; + + if (typeof msg.content === "string") { + textContent = msg.content; + } else if (Array.isArray(msg.content)) { + // Extract text from multimodal content + const textParts = msg.content + .filter((item: any) => item.type === "text") + .map((item: any) => item.text || "") + .join(" "); + textContent = textParts || "[Image content]"; + } else { + textContent = String(msg.content || ""); + } + + return { + role: msg.role, + content: textContent, + }; + }); +} diff --git a/src/LLMProviders/chainRunner/utils/imageExtraction.test.ts b/src/LLMProviders/chainRunner/utils/imageExtraction.test.ts new file mode 100644 index 00000000..3fb25ab1 --- /dev/null +++ b/src/LLMProviders/chainRunner/utils/imageExtraction.test.ts @@ -0,0 +1,342 @@ +// Test for the image extraction logic +describe("Image extraction from content", () => { + // Mock the global app object + const mockApp = { + metadataCache: { + getFirstLinkpathDest: jest.fn(), + }, + }; + + (global as any).app = mockApp; + + // Mock logger (removed since we're no longer logging warnings) + + beforeEach(() => { + jest.clearAllMocks(); + }); + + // Helper function that replicates the extractEmbeddedImages logic + async function extractEmbeddedImages(content: string, sourcePath?: string): Promise { + // Match both wiki-style ![[image.ext]] and standard markdown ![alt](image.ext) + const wikiImageRegex = /!\[\[(.*?\.(png|jpg|jpeg|gif|webp|bmp|svg))\]\]/g; + // Updated regex to handle URLs with or without file extensions + const markdownImageRegex = /!\[.*?\]\(([^)]+)\)/g; + + const resolvedImages: string[] = []; + + // Process wiki-style images + const wikiMatches = [...content.matchAll(wikiImageRegex)]; + for (const match of wikiMatches) { + const imageName = match[1]; + + // If we have a source path and access to the app, resolve the wikilink + if (sourcePath) { + const resolvedFile = mockApp.metadataCache.getFirstLinkpathDest(imageName, sourcePath); + + if (resolvedFile) { + // Use the resolved path + resolvedImages.push(resolvedFile.path); + } else { + // If file not found, still include the raw filename + resolvedImages.push(imageName); + } + } else { + // Fallback to raw filename if no source path available + resolvedImages.push(imageName); + } + } + + // Process standard markdown images + const mdMatches = [...content.matchAll(markdownImageRegex)]; + for (const match of mdMatches) { + const imagePath = match[1].trim(); + + // Skip empty paths + if (!imagePath) continue; + + // Handle external URLs (http://, https://, etc.) + if (imagePath.match(/^https?:\/\//)) { + // Include external URLs - they will be processed by processImageUrls + // The ImageProcessor will validate if it's actually an image + resolvedImages.push(imagePath); + continue; + } + + // For local paths, resolve them using Obsidian's metadata cache + // Let ImageBatchProcessor handle validation of whether it's actually an image + // Clean up the path (remove any leading ./ or /) + const cleanPath = imagePath.replace(/^\.\//, "").replace(/^\//, ""); + + // If we have a source path and access to the app, resolve the path + if (sourcePath) { + const resolvedFile = mockApp.metadataCache.getFirstLinkpathDest(cleanPath, sourcePath); + + if (resolvedFile) { + // Use the resolved path + resolvedImages.push(resolvedFile.path); + } else { + // If file not found, still include the raw path + // Let ImageBatchProcessor handle validation + resolvedImages.push(cleanPath); + } + } else { + // Fallback to raw path if no source path available + resolvedImages.push(cleanPath); + } + } + + return resolvedImages; + } + + describe("Wiki-style image syntax ![[image]]", () => { + it("should extract wiki-style images with extensions", async () => { + const content = "Here is an image ![[screenshot.png]] and another ![[diagram.jpg]]"; + const sourcePath = "notes/test.md"; + + mockApp.metadataCache.getFirstLinkpathDest + .mockReturnValueOnce({ path: "attachments/screenshot.png" }) + .mockReturnValueOnce({ path: "images/diagram.jpg" }); + + const result = await extractEmbeddedImages(content, sourcePath); + + expect(result).toEqual(["attachments/screenshot.png", "images/diagram.jpg"]); + expect(mockApp.metadataCache.getFirstLinkpathDest).toHaveBeenCalledWith( + "screenshot.png", + sourcePath + ); + expect(mockApp.metadataCache.getFirstLinkpathDest).toHaveBeenCalledWith( + "diagram.jpg", + sourcePath + ); + }); + + it("should handle wiki-style images with paths", async () => { + const content = "Image with path ![[folder/image.png]]"; + const sourcePath = "notes/test.md"; + + mockApp.metadataCache.getFirstLinkpathDest.mockReturnValueOnce({ + path: "resolved/folder/image.png", + }); + + const result = await extractEmbeddedImages(content, sourcePath); + + expect(result).toEqual(["resolved/folder/image.png"]); + expect(mockApp.metadataCache.getFirstLinkpathDest).toHaveBeenCalledWith( + "folder/image.png", + sourcePath + ); + }); + + it("should fallback to raw filename when file not found", async () => { + const content = "Missing image ![[missing.png]]"; + const sourcePath = "notes/test.md"; + + mockApp.metadataCache.getFirstLinkpathDest.mockReturnValueOnce(null); + + const result = await extractEmbeddedImages(content, sourcePath); + + expect(result).toEqual(["missing.png"]); + }); + }); + + describe("Markdown image syntax ![](image)", () => { + it("should extract markdown images with simple paths", async () => { + const content = "Here is ![alt text](image.png) and ![](photo.jpg)"; + const sourcePath = "notes/test.md"; + + mockApp.metadataCache.getFirstLinkpathDest + .mockReturnValueOnce({ path: "attachments/image.png" }) + .mockReturnValueOnce({ path: "attachments/photo.jpg" }); + + const result = await extractEmbeddedImages(content, sourcePath); + + expect(result).toEqual(["attachments/image.png", "attachments/photo.jpg"]); + expect(mockApp.metadataCache.getFirstLinkpathDest).toHaveBeenCalledWith( + "image.png", + sourcePath + ); + expect(mockApp.metadataCache.getFirstLinkpathDest).toHaveBeenCalledWith( + "photo.jpg", + sourcePath + ); + }); + + it("should handle markdown images with relative paths", async () => { + const content = "Relative paths ![](./images/test.png) and ![](../assets/icon.svg)"; + const sourcePath = "notes/test.md"; + + mockApp.metadataCache.getFirstLinkpathDest + .mockReturnValueOnce({ path: "images/test.png" }) + .mockReturnValueOnce({ path: "assets/icon.svg" }); + + const result = await extractEmbeddedImages(content, sourcePath); + + expect(result).toEqual(["images/test.png", "assets/icon.svg"]); + expect(mockApp.metadataCache.getFirstLinkpathDest).toHaveBeenCalledWith( + "images/test.png", + sourcePath + ); + expect(mockApp.metadataCache.getFirstLinkpathDest).toHaveBeenCalledWith( + "../assets/icon.svg", + sourcePath + ); + }); + + it("should include external URLs", async () => { + const content = + "External images ![](https://example.com/image.png) and ![](http://site.com/pic.jpg)"; + const sourcePath = "notes/test.md"; + + const result = await extractEmbeddedImages(content, sourcePath); + + expect(result).toEqual(["https://example.com/image.png", "http://site.com/pic.jpg"]); + expect(mockApp.metadataCache.getFirstLinkpathDest).not.toHaveBeenCalled(); + }); + + it("should include external URLs with query parameters", async () => { + const content = + "Unsplash image ![](https://images.unsplash.com/photo-1746555697990-3a405a5152b9?q=80&w=1587&auto=format&fit=crop&ixlib=rb-4.1.0&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D)"; + const sourcePath = "notes/test.md"; + + const result = await extractEmbeddedImages(content, sourcePath); + + expect(result).toEqual([ + "https://images.unsplash.com/photo-1746555697990-3a405a5152b9?q=80&w=1587&auto=format&fit=crop&ixlib=rb-4.1.0&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D", + ]); + expect(mockApp.metadataCache.getFirstLinkpathDest).not.toHaveBeenCalled(); + }); + + it("should clean leading slashes from paths", async () => { + const content = "Absolute path ![](/images/absolute.png)"; + const sourcePath = "notes/test.md"; + + mockApp.metadataCache.getFirstLinkpathDest.mockReturnValueOnce({ + path: "images/absolute.png", + }); + + const result = await extractEmbeddedImages(content, sourcePath); + + expect(result).toEqual(["images/absolute.png"]); + expect(mockApp.metadataCache.getFirstLinkpathDest).toHaveBeenCalledWith( + "images/absolute.png", + sourcePath + ); + }); + }); + + describe("Mixed syntaxes", () => { + it("should extract both wiki and markdown images", async () => { + const content = "Wiki ![[wiki.png]] and markdown ![](markdown.jpg) together"; + const sourcePath = "notes/test.md"; + + mockApp.metadataCache.getFirstLinkpathDest + .mockReturnValueOnce({ path: "attachments/wiki.png" }) + .mockReturnValueOnce({ path: "attachments/markdown.jpg" }); + + const result = await extractEmbeddedImages(content, sourcePath); + + expect(result).toEqual(["attachments/wiki.png", "attachments/markdown.jpg"]); + expect(mockApp.metadataCache.getFirstLinkpathDest).toHaveBeenCalledTimes(2); + }); + + it("should handle multiple images of both types", async () => { + const content = ` + ![[first.png]] some text ![](second.jpg) + More text ![[third.gif]] and ![alt](fourth.svg) + `; + const sourcePath = "notes/test.md"; + + mockApp.metadataCache.getFirstLinkpathDest + .mockReturnValueOnce({ path: "imgs/first.png" }) + .mockReturnValueOnce({ path: "imgs/third.gif" }) + .mockReturnValueOnce({ path: "imgs/second.jpg" }) + .mockReturnValueOnce({ path: "imgs/fourth.svg" }); + + const result = await extractEmbeddedImages(content, sourcePath); + + expect(result).toEqual([ + "imgs/first.png", + "imgs/third.gif", + "imgs/second.jpg", + "imgs/fourth.svg", + ]); + }); + }); + + describe("Edge cases", () => { + it("should extract all markdown links regardless of extension", async () => { + const content = "Valid ![[image.png]] but not ![[document.pdf]] or ![](script.js)"; + const sourcePath = "notes/test.md"; + + mockApp.metadataCache.getFirstLinkpathDest + .mockReturnValueOnce({ path: "attachments/image.png" }) + .mockReturnValueOnce(null); // script.js not found + + const result = await extractEmbeddedImages(content, sourcePath); + + // Wiki-style ![[document.pdf]] is skipped because it doesn't have image extension + // But markdown ![](script.js) is included - ImageBatchProcessor will validate + expect(result).toEqual(["attachments/image.png", "script.js"]); + expect(mockApp.metadataCache.getFirstLinkpathDest).toHaveBeenCalledTimes(2); + }); + + it("should handle content with no images", async () => { + const content = "Just text with no images"; + const sourcePath = "notes/test.md"; + + const result = await extractEmbeddedImages(content, sourcePath); + + expect(result).toEqual([]); + expect(mockApp.metadataCache.getFirstLinkpathDest).not.toHaveBeenCalled(); + }); + + it("should handle no source path", async () => { + const content = "Image ![[test.png]] and ![](other.jpg)"; + + const result = await extractEmbeddedImages(content); + + expect(result).toEqual(["test.png", "other.jpg"]); + expect(mockApp.metadataCache.getFirstLinkpathDest).not.toHaveBeenCalled(); + }); + + it("should handle malformed markdown syntax", async () => { + const content = "Malformed ![](no-extension) and ![]() empty"; + const sourcePath = "notes/test.md"; + + mockApp.metadataCache.getFirstLinkpathDest.mockReturnValueOnce(null); // no-extension not found + + const result = await extractEmbeddedImages(content, sourcePath); + + // no-extension is included, empty path is skipped + expect(result).toEqual(["no-extension"]); + expect(mockApp.metadataCache.getFirstLinkpathDest).toHaveBeenCalledTimes(1); + }); + + it("should include all markdown links for local paths (validation happens in ImageBatchProcessor)", async () => { + const content = "Document ![](document.pdf) and text ![](notes.md)"; + const sourcePath = "notes/test.md"; + + mockApp.metadataCache.getFirstLinkpathDest + .mockReturnValueOnce(null) // document.pdf not found + .mockReturnValueOnce(null); // notes.md not found + + const result = await extractEmbeddedImages(content, sourcePath); + + expect(result).toEqual(["document.pdf", "notes.md"]); + expect(mockApp.metadataCache.getFirstLinkpathDest).toHaveBeenCalledTimes(2); + }); + + it("should handle special characters in filenames", async () => { + const content = "Special chars ![[image (1).png]] and ![](file-name_2.jpg)"; + const sourcePath = "notes/test.md"; + + mockApp.metadataCache.getFirstLinkpathDest + .mockReturnValueOnce({ path: "attachments/image (1).png" }) + .mockReturnValueOnce({ path: "attachments/file-name_2.jpg" }); + + const result = await extractEmbeddedImages(content, sourcePath); + + expect(result).toEqual(["attachments/image (1).png", "attachments/file-name_2.jpg"]); + }); + }); +}); diff --git a/src/LLMProviders/chainRunner/utils/modelAdapter.test.ts b/src/LLMProviders/chainRunner/utils/modelAdapter.test.ts new file mode 100644 index 00000000..24286e21 --- /dev/null +++ b/src/LLMProviders/chainRunner/utils/modelAdapter.test.ts @@ -0,0 +1,155 @@ +import { ModelAdapterFactory } from "./modelAdapter"; +import { ToolMetadata } from "@/tools/ToolRegistry"; + +describe("ModelAdapter", () => { + describe("enhanceSystemPrompt", () => { + const basePrompt = "You are a helpful assistant."; + const toolDescriptions = "Tool descriptions here"; + + const createToolMetadata = (id: string, instructions: string): ToolMetadata => ({ + id, + displayName: `${id} Display`, + description: `${id} description`, + category: "custom", + customPromptInstructions: instructions, + }); + + it("should include tool instructions when tools are enabled", () => { + const mockModel = { modelName: "gpt-4" } as any; + const adapter = ModelAdapterFactory.createAdapter(mockModel); + + const toolMetadata: ToolMetadata[] = [ + createToolMetadata("localSearch", "LocalSearch specific instructions"), + createToolMetadata("webSearch", "WebSearch specific instructions"), + createToolMetadata("writeToFile", "WriteToFile specific instructions"), + ]; + + const enhancedPrompt = adapter.enhanceSystemPrompt( + basePrompt, + toolDescriptions, + ["localSearch", "webSearch", "writeToFile"], + toolMetadata + ); + + // Check that tool instructions are included + expect(enhancedPrompt).toContain("LocalSearch specific instructions"); + expect(enhancedPrompt).toContain("WebSearch specific instructions"); + expect(enhancedPrompt).toContain("WriteToFile specific instructions"); + }); + + it("should only include instructions for enabled tools", () => { + const mockModel = { modelName: "gpt-4" } as any; + const adapter = ModelAdapterFactory.createAdapter(mockModel); + + const toolMetadata: ToolMetadata[] = [ + createToolMetadata("localSearch", "LocalSearch specific instructions"), + createToolMetadata("webSearch", "WebSearch specific instructions"), + createToolMetadata("writeToFile", "WriteToFile specific instructions"), + ]; + + // Only pass localSearch as enabled + const enhancedPrompt = adapter.enhanceSystemPrompt( + basePrompt, + toolDescriptions, + ["localSearch"], + [toolMetadata[0]] // Only localSearch metadata + ); + + // Should include localSearch instructions + expect(enhancedPrompt).toContain("LocalSearch specific instructions"); + + // Should NOT include other tool instructions + expect(enhancedPrompt).not.toContain("WebSearch specific instructions"); + expect(enhancedPrompt).not.toContain("WriteToFile specific instructions"); + }); + + it("should include base structure elements", () => { + const mockModel = { modelName: "gpt-4" } as any; + const adapter = ModelAdapterFactory.createAdapter(mockModel); + + const enhancedPrompt = adapter.enhanceSystemPrompt(basePrompt, toolDescriptions, [], []); + + // Check base sections exist + expect(enhancedPrompt).toContain("# Autonomous Agent Mode"); + expect(enhancedPrompt).toContain("## Time-based Queries"); + expect(enhancedPrompt).toContain("## General Guidelines"); + }); + + it("should handle GPT-specific enhancements", () => { + const mockModel = { modelName: "gpt-4" } as any; + const adapter = ModelAdapterFactory.createAdapter(mockModel); + + const enhancedPrompt = adapter.enhanceSystemPrompt(basePrompt, toolDescriptions, [], []); + + // Check GPT-specific sections + expect(enhancedPrompt).toContain("CRITICAL FOR GPT MODELS"); + expect(enhancedPrompt).toContain("FINAL REMINDER FOR GPT MODELS"); + }); + + it("should handle Claude-specific enhancements", () => { + const mockModel = { modelName: "claude-3-7-sonnet" } as any; + const adapter = ModelAdapterFactory.createAdapter(mockModel); + + const enhancedPrompt = adapter.enhanceSystemPrompt(basePrompt, toolDescriptions, [], []); + + // Check Claude-specific sections + expect(enhancedPrompt).toContain("IMPORTANT FOR CLAUDE THINKING MODELS"); + }); + + it("should handle Gemini-specific enhancements", () => { + const mockModel = { modelName: "gemini-pro" } as any; + const adapter = ModelAdapterFactory.createAdapter(mockModel); + + const enhancedPrompt = adapter.enhanceSystemPrompt(basePrompt, toolDescriptions, [], []); + + // Check Gemini-specific sections + expect(enhancedPrompt).toContain("CRITICAL INSTRUCTIONS FOR GEMINI"); + }); + + it("should exclude instructions when no metadata provided", () => { + const mockModel = { modelName: "gpt-4" } as any; + const adapter = ModelAdapterFactory.createAdapter(mockModel); + + const enhancedPrompt = adapter.enhanceSystemPrompt( + basePrompt, + toolDescriptions, + ["localSearch", "webSearch"], + [] // No metadata + ); + + // Should not include any tool-specific instructions + expect(enhancedPrompt).not.toContain("LocalSearch specific instructions"); + expect(enhancedPrompt).not.toContain("WebSearch specific instructions"); + }); + + it("should include composer-specific examples for GPT when file tools are enabled", () => { + const mockModel = { modelName: "gpt-4" } as any; + const adapter = ModelAdapterFactory.createAdapter(mockModel); + + const enhancedPrompt = adapter.enhanceSystemPrompt( + basePrompt, + toolDescriptions, + ["replaceInFile", "writeToFile"], + [] + ); + + // Check for composer-specific GPT instructions + expect(enhancedPrompt).toContain("FILE EDITING WITH COMPOSER TOOLS"); + expect(enhancedPrompt).toContain("fix the typo"); + expect(enhancedPrompt).toContain("add item 4 to the list"); + expect(enhancedPrompt).toContain("diff parameter MUST contain"); + }); + + it("should enhance file editing messages for GPT", () => { + const mockModel = { modelName: "gpt-4" } as any; + const adapter = ModelAdapterFactory.createAdapter(mockModel); + + const editMessage = "fix the typo in my note"; + const enhanced = adapter.enhanceUserMessage(editMessage, true); + + expect(enhanced).toContain("GPT REMINDER"); + expect(enhanced).toContain("replaceInFile"); + expect(enhanced).toContain("SEARCH/REPLACE blocks"); + }); + }); +}); diff --git a/src/LLMProviders/chainRunner/utils/modelAdapter.ts b/src/LLMProviders/chainRunner/utils/modelAdapter.ts new file mode 100644 index 00000000..ca84112e --- /dev/null +++ b/src/LLMProviders/chainRunner/utils/modelAdapter.ts @@ -0,0 +1,697 @@ +import { BaseChatModel } from "@langchain/core/language_models/chat_models"; +import { logInfo } from "@/logger"; +import { ToolMetadata } from "@/tools/ToolRegistry"; + +/** + * Model-specific adaptations for autonomous agent + * Handles quirks and requirements of different LLM providers + */ + +export interface ModelAdapter { + /** + * Enhance system prompt with model-specific instructions + * @param basePrompt - The base system prompt to enhance + * @param toolDescriptions - Available tool descriptions to include + * @param availableToolNames - List of enabled tool names (for backward compatibility) + * @param toolMetadata - Tool metadata including custom instructions + * @returns The enhanced system prompt + */ + enhanceSystemPrompt( + basePrompt: string, + toolDescriptions: string, + availableToolNames?: string[], + toolMetadata?: ToolMetadata[] + ): string; + + /** + * Enhance user message if needed for specific models + * @param message - The user's message to enhance + * @param requiresTools - Whether the message likely requires tool usage + * @returns The enhanced user message + */ + enhanceUserMessage(message: string, requiresTools: boolean): string; + + /** + * Parse tool calls from model response (future: handle different formats) + * @param response - The model's response text containing tool calls + * @returns Array of parsed tool calls + */ + parseToolCalls?(response: string): any[]; + + /** + * Check if model needs special handling + * @returns True if the model requires special handling beyond base behavior + */ + needsSpecialHandling(): boolean; + + /** + * Sanitize response for autonomous agent mode (e.g., remove premature content) + * @param response - The model's response text to sanitize + * @param iteration - The current iteration number in the autonomous agent loop + * @returns The sanitized response text + */ + sanitizeResponse?(response: string, iteration: number): string; + + /** + * Check if streaming should be truncated early for this model + * @param partialResponse - The partial response text received during streaming + * @returns True if streaming should be truncated at this point + */ + shouldTruncateStreaming?(partialResponse: string): boolean; + + /** + * Detect premature responses for this model + * @param response - The model's response text to analyze + * @returns Object indicating if premature content was detected and its type + */ + detectPrematureResponse?(response: string): { + hasPremature: boolean; + type: "before" | "after" | null; + }; +} + +/** + * Base adapter with default behavior (no modifications) + */ +class BaseModelAdapter implements ModelAdapter { + constructor(protected modelName: string) {} + + private buildToolSpecificInstructions(toolMetadata: ToolMetadata[]): string { + const instructions: string[] = []; + + // Collect all custom instructions from tool metadata + for (const meta of toolMetadata) { + if (meta.customPromptInstructions) { + instructions.push(meta.customPromptInstructions); + } + } + + return instructions.length > 0 ? instructions.join("\n\n") : ""; + } + + enhanceSystemPrompt( + basePrompt: string, + toolDescriptions: string, + availableToolNames?: string[], + toolMetadata?: ToolMetadata[] + ): string { + const metadata = toolMetadata || []; + + // Build tool-specific instructions from metadata + const toolSpecificInstructions = this.buildToolSpecificInstructions(metadata); + + return `${basePrompt} + +# Autonomous Agent Mode + +You are now in autonomous agent mode. You can use tools to gather information and complete tasks step by step. + +When you need to use a tool, format it EXACTLY like this: + +tool_name_here +value +["array", "values"] + + +IMPORTANT: Use the EXACT parameter names as shown in the tool descriptions below. Do NOT use generic names like "param1" or "param". + +Available tools: +${toolDescriptions} + +# Tool Usage Guidelines + +## Time-based Queries +When users ask about temporal periods (e.g., "what did I do last month", "show me notes from last week"), you MUST: +1. First call getTimeRangeMs to convert the time expression to a proper time range +2. Then use localSearch with the timeRange parameter from step 1 +3. For salientTerms, ONLY use words that exist in the user's original query (excluding time expressions) + +Example for "what did I do last month": +1. Call getTimeRangeMs with timeExpression: "last month" +2. Use localSearch with query matching the user's question +3. salientTerms: [] - empty because "what", "I", "do" are not meaningful search terms + +Example for "meetings about project X last week": +1. Call getTimeRangeMs with timeExpression: "last week" +2. Use localSearch with query "meetings about project X" +3. salientTerms: ["meetings", "project", "X"] - these words exist in the original query + +${toolSpecificInstructions ? toolSpecificInstructions + "\n\n" : ""}## General Guidelines +- NEVER mention tool names like "localSearch", "webSearch", etc. in your responses. Use natural language like "searching your vault", "searching the web", etc. + +You can use multiple tools in sequence. After each tool execution, you'll receive the results and can decide whether to use more tools or provide your final response. + +Always explain your reasoning before using tools. Be conversational and clear about what you're doing. +When you've gathered enough information, provide your final response without any tool calls. + +IMPORTANT: Do not include any code blocks (\`\`\`) or tool_code blocks in your responses. Only use the format for tool calls. + +NOTE: Use individual XML parameter tags. For arrays, use JSON format like ["item1", "item2"].`; + } + + enhanceUserMessage(message: string, requiresTools: boolean): string { + return message; + } + + needsSpecialHandling(): boolean { + return false; + } +} + +/** + * GPT-specific adapter with aggressive prompting + */ +class GPTModelAdapter extends BaseModelAdapter { + enhanceSystemPrompt( + basePrompt: string, + toolDescriptions: string, + availableToolNames?: string[], + toolMetadata?: ToolMetadata[] + ): string { + const baseSystemPrompt = super.enhanceSystemPrompt( + basePrompt, + toolDescriptions, + availableToolNames, + toolMetadata + ); + + const tools = availableToolNames || []; + const hasComposerTools = tools.includes("writeToFile") || tools.includes("replaceInFile"); + + // Insert GPT-specific instructions after the base prompt + let gptSpecificSection = ` + +CRITICAL FOR GPT MODELS: You MUST ALWAYS include XML tool calls in your response. Do not just describe what you plan to do - you MUST include the actual XML tool call blocks.`; + + if (hasComposerTools) { + gptSpecificSection += ` + +🚨 FILE EDITING WITH COMPOSER TOOLS - GPT SPECIFIC EXAMPLES 🚨 + +When user asks you to edit or modify a file, you MUST: +1. Determine if it's a small edit (use replaceInFile) or major rewrite (use writeToFile) +2. Include the tool call immediately in your response + +EXAMPLE 1 - User says "fix the typo 'teh' to 'the' in my note": +✅ CORRECT RESPONSE: +"I'll fix the typo in your note. + + +replaceInFile +path/to/note.md +\`\`\` +------- SEARCH +teh +======= +the ++++++++ REPLACE +\`\`\` +" + +EXAMPLE 2 - User says "add item 4 to the list": +✅ CORRECT RESPONSE: +"I'll add item 4 to your list. + + +replaceInFile +path/to/file.md +\`\`\` +------- SEARCH +- Item 1 +- Item 2 +- Item 3 +======= +- Item 1 +- Item 2 +- Item 3 +- Item 4 ++++++++ REPLACE +\`\`\` +" + +❌ WRONG (DO NOT DO THIS): +"I'll help you add item 4 to the list. Let me update that for you." +[No tool call = FAILURE] + +CRITICAL: The diff parameter MUST contain the SEARCH/REPLACE blocks wrapped in triple backticks EXACTLY as shown above.`; + } + + gptSpecificSection += ` + +FINAL REMINDER FOR GPT MODELS: If the user asks you to search, find, edit, or modify anything, you MUST include the appropriate XML block in your very next response. Do not wait for another turn.`; + + return baseSystemPrompt + gptSpecificSection; + } + + enhanceUserMessage(message: string, requiresTools: boolean): string { + if (requiresTools) { + const lowerMessage = message.toLowerCase(); + const requiresSearch = + lowerMessage.includes("find") || + lowerMessage.includes("search") || + lowerMessage.includes("my notes"); + + const requiresFileEdit = + lowerMessage.includes("edit") || + lowerMessage.includes("modify") || + lowerMessage.includes("update") || + lowerMessage.includes("change") || + lowerMessage.includes("fix") || + lowerMessage.includes("add") || + lowerMessage.includes("typo"); + + if (requiresSearch) { + return `${message}\n\nREMINDER: Use the XML format to call the localSearch tool.`; + } + + if (requiresFileEdit) { + return `${message}\n\n🚨 GPT REMINDER: Use replaceInFile for small edits (with SEARCH/REPLACE blocks in diff parameter). The diff parameter MUST contain triple backticks around the SEARCH/REPLACE blocks. Check the examples in your system prompt.`; + } + } + return message; + } + + needsSpecialHandling(): boolean { + return true; + } +} + +/** + * Claude adapter with special handling for thinking models + */ +class ClaudeModelAdapter extends BaseModelAdapter { + /** + * Check if this is a Claude thinking model (3.7 Sonnet or Claude 4) + * @returns True if the model supports thinking/reasoning modes + */ + private isThinkingModel(): boolean { + return ( + this.modelName.includes("claude-3-7-sonnet") || + this.modelName.includes("claude-sonnet-4") || + this.modelName.includes("claude-3.7-sonnet") || + this.modelName.includes("claude-4-sonnet") + ); + } + + /** + * Check if this is specifically Claude Sonnet 4 (has hallucination issues) + * @returns True if the model is Claude Sonnet 4 variants + */ + private isClaudeSonnet4(): boolean { + return ( + this.modelName.includes("claude-sonnet-4") || + this.modelName.includes("claude-4-sonnet") || + this.modelName.includes("claude-sonnet-4-20250514") + ); + } + + enhanceSystemPrompt( + basePrompt: string, + toolDescriptions: string, + availableToolNames?: string[], + toolMetadata?: ToolMetadata[] + ): string { + const baseSystemPrompt = super.enhanceSystemPrompt( + basePrompt, + toolDescriptions, + availableToolNames, + toolMetadata + ); + + // Add specific instructions for thinking models + if (this.isThinkingModel()) { + let thinkingModelSection = ` + +IMPORTANT FOR CLAUDE THINKING MODELS: +- You are a thinking model with internal reasoning capability +- Your thinking process will be automatically wrapped in tags - do not manually add thinking tags +- Place ALL tool calls AFTER your thinking is complete, in the main response body +- Tool calls must be in the main response body, NOT inside thinking sections +- Format tool calls exactly as shown in the examples above +- Do not provide final answers before using tools - use tools first, then provide your response based on the results +- If you need to use tools, include them immediately after your thinking, before any final response + +CORRECT FLOW FOR THINKING MODELS: +1. Think through the problem (this happens automatically) +2. Use tools to gather information (place tool calls in main response) +3. Wait for tool results +4. Provide final response based on gathered information + +INCORRECT: Providing a final answer before using tools +CORRECT: Using tools first, then providing answer based on results`; + + // Add even stronger instructions specifically for Claude Sonnet 4 + if (this.isClaudeSonnet4()) { + thinkingModelSection += ` + +🚨 CRITICAL INSTRUCTIONS FOR CLAUDE SONNET 4 - AUTONOMOUS AGENT MODE 🚨 + +⚠️ WARNING: You have a specific tendency to write complete responses immediately after tool calls. This BREAKS the autonomous agent pattern! + +🔄 CORRECT AUTONOMOUS AGENT ITERATION PATTERN: +1. User asks question +2. Brief sentence about what you'll do (1 sentence max) +3. Use tools to gather information: ... +4. ✋ STOP after tool calls - Do not write anything else +5. Wait for tool results (system provides them) +6. Evaluate results and either: Use more tools OR provide final answer + +✅ IDEAL RESPONSE FLOW: +- Brief action statement (1 sentence) +- Tool calls +- STOP (wait for results) +- Brief transition statement (1 sentence) +- More tool calls OR final answer + +🎯 CORRECT FIRST RESPONSE PATTERN (when tools needed): +I'll search your vault for piano practice information. + + +localSearch +piano practice +["piano", "practice"] + + +🌐 MULTILINGUAL EXAMPLE (PRESERVE ORIGINAL LANGUAGE): + +localSearch +ピアノの練習方法 +["ピアノ", "練習", "方法"] + + + +webSearch +piano techniques +[] + + +[RESPONSE ENDS HERE - NO MORE TEXT] + +🎯 CORRECT FOLLOW-UP RESPONSE PATTERN: +Let me gather more specific information about practice schedules. + + +localSearch +practice schedule +["practice", "schedule"] + + +[RESPONSE ENDS HERE - NO MORE TEXT] + +❌ WRONG PATTERN (DO NOT DO THIS): +... + +Based on the search results, here's a complete practice plan... +[This is FORBIDDEN - you haven't received results yet!] + +🔑 KEY UNDERSTANDING FOR CLAUDE 4: +- Brief 1-sentence explanations BEFORE tool calls are good +- Each response is ONE STEP in a multi-step process +- After tool calls, STOP and wait for the system to provide results +- Your thinking is automatically handled in blocks + +⚡ AUTONOMOUS AGENT RULES FOR CLAUDE 4: +1. If you need tools: Brief sentence + tool calls, then STOP +2. If you receive tool results: Evaluate if you need more tools +3. If you need more tools: Brief sentence + more tool calls, then STOP +4. If you have enough info: THEN provide your final response + +REMEMBER: One brief sentence before tools is perfect. Nothing after tool calls.`; + } + + return baseSystemPrompt + thinkingModelSection; + } + + return baseSystemPrompt; + } + + needsSpecialHandling(): boolean { + return this.isThinkingModel(); + } + + /** + * Detect premature responses in Claude models, especially Claude 4 + * Checks for content before tool calls (allowed in small amounts) and after tool calls (not allowed) + * @param response - The model's response text to analyze + * @returns Object indicating if premature content was detected and its type + */ + detectPrematureResponse(response: string): { + hasPremature: boolean; + type: "before" | "after" | null; + } { + const firstToolCallIndex = response.indexOf(""); + + if (firstToolCallIndex === -1) { + // No tool calls at all, so it's either a final response or a problem + return { hasPremature: false, type: null }; + } + + // Check content before tool calls - allow brief sentences (1-2 sentences max) + const contentBeforeTools = response.substring(0, firstToolCallIndex).trim(); + const contentBeforeWithoutThinking = contentBeforeTools + .replace(/[\s\S]*?<\/think>/g, "") + .trim(); + + // Allow brief action statements (up to 2 sentences, 200 characters) + const sentences = contentBeforeWithoutThinking + .split(/[.!?]+/) + .filter((s) => s.trim().length > 0); + const BRIEF_STATEMENT_THRESHOLD = 200; // characters + + if (sentences.length > 2 || contentBeforeWithoutThinking.length > BRIEF_STATEMENT_THRESHOLD) { + return { hasPremature: true, type: "before" }; + } + + // Check content after last tool call (main Claude 4 issue) + const lastToolCallEndIndex = response.lastIndexOf(""); + if (lastToolCallEndIndex !== -1) { + const contentAfterTools = response + .substring(lastToolCallEndIndex + "".length) + .trim(); + + // Remove think blocks to analyze only non-thinking content + const contentAfterWithoutThinking = contentAfterTools + .replace(/[\s\S]*?<\/think>/g, "") + .trim(); + + // Simple threshold: any substantial non-thinking content after tools is premature + const SUBSTANTIAL_CONTENT_THRESHOLD = 100; // characters + + if (contentAfterWithoutThinking.length > SUBSTANTIAL_CONTENT_THRESHOLD) { + return { hasPremature: true, type: "after" }; + } + } + + return { hasPremature: false, type: null }; + } + + /** + * Sanitize Claude 4 responses by removing premature content after tool calls + * Preserves all think blocks while removing substantial non-thinking content + * @param response - The model's response text to sanitize + * @param iteration - The current iteration number (only sanitizes first iteration) + * @returns The sanitized response text + */ + sanitizeResponse(response: string, iteration: number): string { + if (!this.isClaudeSonnet4() || iteration !== 1) { + return response; + } + + const prematureResult = this.detectPrematureResponse(response); + if (!prematureResult.hasPremature) { + return response; + } + + if (prematureResult.type === "after") { + // Simple approach: preserve ALL think blocks, remove substantial non-thinking content + const lastToolCallEndIndex = response.lastIndexOf(""); + if (lastToolCallEndIndex !== -1) { + const baseResponse = response.substring(0, lastToolCallEndIndex + "".length); + const contentAfterTools = response.substring(lastToolCallEndIndex + "".length); + + // Extract and preserve ALL think blocks + const thinkBlockRegex = /[\s\S]*?<\/think>/g; + const thinkBlocks = contentAfterTools.match(thinkBlockRegex) || []; + + // Return base response + all think blocks (remove everything else) + return baseResponse + (thinkBlocks.length > 0 ? "\n" + thinkBlocks.join("\n") : ""); + } + } + + return response; + } + + /** + * Check if Claude 4 streaming should be truncated to prevent hallucinated content + * @param partialResponse - The partial response text received during streaming + * @returns True if streaming should be truncated at this point + */ + shouldTruncateStreaming(partialResponse: string): boolean { + if (!this.isClaudeSonnet4()) { + return false; + } + + // Check if we have tool calls and substantial non-thinking content after them + const lastToolCallEndIndex = partialResponse.lastIndexOf(""); + if (lastToolCallEndIndex === -1) { + return false; + } + + const contentAfterTools = partialResponse + .substring(lastToolCallEndIndex + "".length) + .trim(); + + // Remove think blocks to analyze only non-thinking content + const contentAfterWithoutThinking = contentAfterTools + .replace(/[\s\S]*?<\/think>/g, "") + .trim(); + + // Simple threshold: if there's substantial non-thinking content, truncate + const STREAMING_TRUNCATE_THRESHOLD = 50; + return contentAfterWithoutThinking.length > STREAMING_TRUNCATE_THRESHOLD; + } +} + +/** + * Gemini adapter with aggressive tool calling prompts + */ +class GeminiModelAdapter extends BaseModelAdapter { + enhanceSystemPrompt( + basePrompt: string, + toolDescriptions: string, + availableToolNames?: string[], + toolMetadata?: ToolMetadata[] + ): string { + const baseSystemPrompt = super.enhanceSystemPrompt( + basePrompt, + toolDescriptions, + availableToolNames, + toolMetadata + ); + + // Gemini needs very explicit instructions about tool usage + const tools = availableToolNames || []; + const hasLocalSearch = tools.includes("localSearch"); + + const geminiSpecificSection = ` + +🚨 CRITICAL INSTRUCTIONS FOR GEMINI - AUTONOMOUS AGENT MODE 🚨 + +You MUST use tools to complete tasks. DO NOT ask the user questions about how to proceed. +${ + hasLocalSearch + ? ` +When the user mentions "my notes" or "my vault", use the localSearch tool. + +❌ WRONG: +"Let's start by searching your notes. What kind of information should I look for?" + +✅ CORRECT: + +localSearch +piano +["piano"] + +` + : "" +} +GEMINI SPECIFIC RULES: +1. When user mentions "my notes" about X → use localSearch with query "X" +2. DO NOT ask clarifying questions about search terms +3. DO NOT wait for permission to use tools +4. Use tools based on the user's request + +PATTERN FOR MULTI-STEP REQUESTS: +User: "based on my project roadmap notes and create summary" +Your response: + +localSearch +project roadmap +["project", "roadmap"] + + +Remember: The user has already told you what to do. Execute it NOW with the available tools.`; + + return baseSystemPrompt + geminiSpecificSection; + } + + enhanceUserMessage(message: string, requiresTools: boolean): string { + if (requiresTools) { + // Add explicit reminder for Gemini + return `${message}\n\nREMINDER: Use the tools immediately. Do not ask questions. For "my notes", use localSearch.`; + } + return message; + } + + needsSpecialHandling(): boolean { + return true; + } +} + +/** + * Factory to create appropriate adapter based on model + */ +export class ModelAdapterFactory { + static createAdapter(model: BaseChatModel): ModelAdapter { + const modelName = ((model as any).modelName || (model as any).model || "").toLowerCase(); + + logInfo(`Creating model adapter for: ${modelName}`); + + // GPT models need special handling + if (modelName.includes("gpt")) { + logInfo("Using GPTModelAdapter"); + return new GPTModelAdapter(modelName); + } + + // Claude models + if (modelName.includes("claude")) { + logInfo("Using ClaudeModelAdapter"); + return new ClaudeModelAdapter(modelName); + } + + // Gemini models (check for both "gemini" and "google" prefixes) + if (modelName.includes("gemini") || modelName.includes("google/gemini")) { + logInfo("Using GeminiModelAdapter"); + return new GeminiModelAdapter(modelName); + } + + // Copilot Plus models + if (modelName.includes("copilot-plus")) { + logInfo("Using BaseModelAdapter for Copilot Plus"); + return new BaseModelAdapter(modelName); + } + + // Default adapter for unknown models + logInfo("Using BaseModelAdapter (default)"); + return new BaseModelAdapter(modelName); + } +} + +/** + * Helper to detect if user message likely requires tools + * @param message - The user's message to analyze + * @returns True if the message likely requires tool usage + */ +export function messageRequiresTools(message: string): boolean { + const toolIndicators = [ + "find", + "search", + "look for", + "look up", + "my notes", + "in my vault", + "from my vault", + "check the web", + "search online", + "from the internet", + "current time", + "what time", + "timer", + "youtube", + "video", + "transcript", + ]; + + const lowerMessage = message.toLowerCase(); + return toolIndicators.some((indicator) => lowerMessage.includes(indicator)); +} diff --git a/src/LLMProviders/chainRunner/utils/toolCallParser.ts b/src/LLMProviders/chainRunner/utils/toolCallParser.ts new file mode 100644 index 00000000..db36b88a --- /dev/null +++ b/src/LLMProviders/chainRunner/utils/toolCallParser.ts @@ -0,0 +1,122 @@ +export interface ToolCallMarker { + id: string; + toolName: string; + displayName: string; + emoji: string; + confirmationMessage?: string; + isExecuting: boolean; + result?: string; + startIndex: number; + endIndex: number; +} + +export interface ParsedMessage { + segments: Array<{ + type: "text" | "toolCall"; + content: string; + toolCall?: ToolCallMarker; + }>; +} + +/** + * Parse tool call markers from a message + * Format: content + */ +export function parseToolCallMarkers(message: string): ParsedMessage { + const segments: ParsedMessage["segments"] = []; + // Use [\s\S] instead of . with 's' flag for compatibility with ES6 + const toolCallRegex = + /([\s\S]*?)/g; + + let lastIndex = 0; + let match; + + while ((match = toolCallRegex.exec(message)) !== null) { + // Add text before the tool call + if (match.index > lastIndex) { + segments.push({ + type: "text", + content: message.slice(lastIndex, match.index), + }); + } + + // Parse the tool call + const [ + fullMatch, + id, + toolName, + displayName, + emoji, + confirmationMessage, + isExecuting, + content, + result, + ] = match; + + segments.push({ + type: "toolCall", + content: content, + toolCall: { + id, + toolName, + displayName, + emoji, + confirmationMessage: confirmationMessage || undefined, + isExecuting: isExecuting === "true", + result: result || undefined, + startIndex: match.index, + endIndex: match.index + fullMatch.length, + }, + }); + + lastIndex = match.index + fullMatch.length; + } + + // Add any remaining text + if (lastIndex < message.length) { + segments.push({ + type: "text", + content: message.slice(lastIndex), + }); + } + + // If no tool calls found, return the entire message as text + if (segments.length === 0) { + segments.push({ + type: "text", + content: message, + }); + } + + return { segments }; +} + +/** + * Create a tool call marker + */ +export function createToolCallMarker( + id: string, + toolName: string, + displayName: string, + emoji: string, + confirmationMessage: string = "", + isExecuting: boolean = true, + content: string = "", + result: string = "" +): string { + return `${content}`; +} + +/** + * Update a tool call marker with result + */ +export function updateToolCallMarker(message: string, id: string, result: string): string { + // Escape the id to prevent regex injection + const escapedId = id.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); + const regex = new RegExp( + `([\\s\\S]*?`, + "g" + ); + + return message.replace(regex, `$1false$2${result}-->`); +} diff --git a/src/LLMProviders/chainRunner/utils/toolExecution.test.ts b/src/LLMProviders/chainRunner/utils/toolExecution.test.ts new file mode 100644 index 00000000..5b936953 --- /dev/null +++ b/src/LLMProviders/chainRunner/utils/toolExecution.test.ts @@ -0,0 +1,122 @@ +import { executeSequentialToolCall } from "./toolExecution"; +import { createTool } from "@/tools/SimpleTool"; +import { z } from "zod"; + +// Mock dependencies +jest.mock("@/plusUtils", () => ({ + checkIsPlusUser: jest.fn(), +})); + +jest.mock("@/logger", () => ({ + logError: jest.fn(), + logInfo: jest.fn(), + logWarn: jest.fn(), +})); + +jest.mock("@/tools/toolManager", () => ({ + ToolManager: { + callTool: jest.fn(), + }, +})); + +import { checkIsPlusUser } from "@/plusUtils"; +import { ToolManager } from "@/tools/toolManager"; + +describe("toolExecution", () => { + const mockCheckIsPlusUser = checkIsPlusUser as jest.MockedFunction; + const mockCallTool = ToolManager.callTool as jest.MockedFunction; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe("executeSequentialToolCall", () => { + it("should execute tools without isPlusOnly flag", async () => { + const testTool = createTool({ + name: "testTool", + description: "Test tool", + schema: z.object({ input: z.string() }), + handler: async ({ input }) => `Result: ${input}`, + }); + + mockCallTool.mockResolvedValueOnce("Tool executed successfully"); + + const result = await executeSequentialToolCall( + { name: "testTool", args: { input: "test" } }, + [testTool] + ); + + expect(result).toEqual({ + toolName: "testTool", + result: "Tool executed successfully", + success: true, + }); + expect(mockCheckIsPlusUser).not.toHaveBeenCalled(); + }); + + it("should block plus-only tools for non-plus users", async () => { + const plusTool = createTool({ + name: "plusTool", + description: "Plus-only tool", + schema: z.void(), + handler: async () => "Should not execute", + isPlusOnly: true, + }); + + mockCheckIsPlusUser.mockResolvedValueOnce(false); + + const result = await executeSequentialToolCall({ name: "plusTool", args: {} }, [plusTool]); + + expect(result).toEqual({ + toolName: "plusTool", + result: "Error: plusTool requires a Copilot Plus subscription", + success: false, + }); + expect(mockCallTool).not.toHaveBeenCalled(); + }); + + it("should allow plus-only tools for plus users", async () => { + const plusTool = createTool({ + name: "plusTool", + description: "Plus-only tool", + schema: z.void(), + handler: async () => "Plus tool executed", + isPlusOnly: true, + }); + + mockCheckIsPlusUser.mockResolvedValueOnce(true); + mockCallTool.mockResolvedValueOnce("Plus tool executed"); + + const result = await executeSequentialToolCall({ name: "plusTool", args: {} }, [plusTool]); + + expect(result).toEqual({ + toolName: "plusTool", + result: "Plus tool executed", + success: true, + }); + expect(mockCheckIsPlusUser).toHaveBeenCalled(); + expect(mockCallTool).toHaveBeenCalled(); + }); + + it("should handle tool not found", async () => { + const result = await executeSequentialToolCall({ name: "unknownTool", args: {} }, []); + + expect(result).toEqual({ + toolName: "unknownTool", + result: + "Error: Tool 'unknownTool' not found. Available tools: . Make sure you have the tool enabled in the Agent settings.", + success: false, + }); + }); + + it("should handle invalid tool call", async () => { + const result = await executeSequentialToolCall(null as any, []); + + expect(result).toEqual({ + toolName: "unknown", + result: "Error: Invalid tool call - missing tool name", + success: false, + }); + }); + }); +}); diff --git a/src/LLMProviders/chainRunner/utils/toolExecution.ts b/src/LLMProviders/chainRunner/utils/toolExecution.ts new file mode 100644 index 00000000..c00252a3 --- /dev/null +++ b/src/LLMProviders/chainRunner/utils/toolExecution.ts @@ -0,0 +1,231 @@ +import { logError, logInfo, logWarn } from "@/logger"; +import { checkIsPlusUser } from "@/plusUtils"; +import { ToolManager } from "@/tools/toolManager"; +import { err2String } from "@/utils"; +import { ToolCall } from "./xmlParsing"; + +export interface ToolExecutionResult { + toolName: string; + result: string; + success: boolean; +} + +/** + * Executes a single tool call with timeout and error handling + */ +export async function executeSequentialToolCall( + toolCall: ToolCall, + availableTools: any[], + originalUserMessage?: string +): Promise { + const DEFAULT_TOOL_TIMEOUT = 30000; // 30 seconds timeout per tool + + try { + // Validate tool call + if (!toolCall || !toolCall.name) { + return { + toolName: toolCall?.name || "unknown", + result: "Error: Invalid tool call - missing tool name", + success: false, + }; + } + + // Find the tool in the existing tool registry + const tool = availableTools.find((t) => t.name === toolCall.name); + + if (!tool) { + const availableToolNames = availableTools.map((t) => t.name).join(", "); + return { + toolName: toolCall.name, + result: `Error: Tool '${toolCall.name}' not found. Available tools: ${availableToolNames}. Make sure you have the tool enabled in the Agent settings.`, + success: false, + }; + } + + // Check if tool requires Plus subscription + if (tool.isPlusOnly) { + const isPlusUser = await checkIsPlusUser(); + if (!isPlusUser) { + return { + toolName: toolCall.name, + result: `Error: ${getToolDisplayName(toolCall.name)} requires a Copilot Plus subscription`, + success: false, + }; + } + } + + // Prepare tool arguments + const toolArgs = { ...toolCall.args }; + + // If tool requires user message content and it's provided, inject it + if (tool.requiresUserMessageContent && originalUserMessage) { + toolArgs._userMessageContent = originalUserMessage; + } + + // Determine timeout for this tool + let timeout = DEFAULT_TOOL_TIMEOUT; + if (typeof tool.timeoutMs === "number") { + timeout = tool.timeoutMs; + } + + let result; + if (!timeout || timeout === Infinity) { + // No timeout for this tool + result = await ToolManager.callTool(tool, toolArgs); + } else { + // Use timeout + result = await Promise.race([ + ToolManager.callTool(tool, toolArgs), + new Promise((_, reject) => + setTimeout( + () => reject(new Error(`Tool execution timed out after ${timeout}ms`)), + timeout + ) + ), + ]); + } + + // Validate result + if (result === null || result === undefined) { + logWarn(`Tool ${toolCall.name} returned null/undefined result`); + return { + toolName: toolCall.name, + result: "Tool executed but returned no result", + success: true, + }; + } + + return { + toolName: toolCall.name, + result: typeof result === "string" ? result : JSON.stringify(result), + success: true, + }; + } catch (error) { + logError(`Error executing tool ${toolCall.name}:`, error); + return { + toolName: toolCall.name, + result: `Error: ${err2String(error)}`, + success: false, + }; + } +} + +/** + * Get display name for tool (user-friendly version) + */ +export function getToolDisplayName(toolName: string): string { + const displayNameMap: Record = { + localSearch: "vault search", + webSearch: "web search", + getFileTree: "file tree", + getCurrentTime: "current time", + getTimeRangeMs: "time range", + getTimeInfoByEpoch: "time info", + convertTimeBetweenTimezones: "timezone converter", + startPomodoro: "pomodoro timer", + pomodoroTool: "pomodoro timer", + simpleYoutubeTranscriptionTool: "YouTube transcription", + youtubeTranscription: "YouTube transcription", + indexVault: "vault indexing", + indexTool: "index", + writeToFile: "file editor", + replaceInFile: "file editor", + }; + + return displayNameMap[toolName] || toolName; +} + +/** + * Get emoji for tool display + */ +export function getToolEmoji(toolName: string): string { + const emojiMap: Record = { + localSearch: "🔍", + webSearch: "🌐", + getFileTree: "📁", + getCurrentTime: "🕒", + getTimeRangeMs: "📅", + getTimeInfoByEpoch: "🕰️", + convertTimeBetweenTimezones: "🌍", + startPomodoro: "⏱️", + pomodoroTool: "⏱️", + simpleYoutubeTranscriptionTool: "📺", + youtubeTranscription: "📺", + indexVault: "📚", + indexTool: "📚", + writeToFile: "✏️", + replaceInFile: "🔄", + }; + + return emojiMap[toolName] || "🔧"; +} + +/** + * Get user confirmation message for tool call + */ +export function getToolConfirmtionMessage(toolName: string): string | null { + if (toolName == "writeToFile" || toolName == "replaceInFile") { + return "Accept / reject in the Preview"; + } + return null; +} + +/** + * Log tool call details for debugging + */ +export function logToolCall(toolCall: ToolCall, iteration: number): void { + const displayName = getToolDisplayName(toolCall.name); + const emoji = getToolEmoji(toolCall.name); + + // Create clean parameter display + const paramDisplay = + Object.keys(toolCall.args).length > 0 + ? JSON.stringify(toolCall.args, null, 2) + : "(no parameters)"; + + logInfo(`${emoji} [Iteration ${iteration}] ${displayName.toUpperCase()}`); + logInfo(`Parameters:`, paramDisplay); + logInfo("---"); +} + +/** + * Log tool execution result + */ +export function logToolResult(toolName: string, result: ToolExecutionResult): void { + const displayName = getToolDisplayName(toolName); + const emoji = getToolEmoji(toolName); + const status = result.success ? "✅ SUCCESS" : "❌ FAILED"; + + logInfo(`${emoji} ${displayName.toUpperCase()} RESULT: ${status}`); + + // Log abbreviated result for readability + if (result.result.length > 500) { + logInfo( + `Result: ${result.result.substring(0, 500)}... (truncated, ${result.result.length} chars total)` + ); + } else { + logInfo(`Result:`, result.result); + } + logInfo("---"); +} + +/** + * Deduplicate sources by path, keeping highest score + * If path is not available, falls back to title + */ +export function deduplicateSources( + sources: { title: string; path: string; score: number }[] +): { title: string; path: string; score: number }[] { + const uniqueSources = new Map(); + + for (const source of sources) { + // Use path as the unique key, falling back to title if path is not available + const key = source.path || source.title; + const existing = uniqueSources.get(key); + if (!existing || source.score > existing.score) { + uniqueSources.set(key, source); + } + } + + return Array.from(uniqueSources.values()).sort((a, b) => b.score - a.score); +} diff --git a/src/LLMProviders/chainRunner/utils/xmlParsing.test.ts b/src/LLMProviders/chainRunner/utils/xmlParsing.test.ts new file mode 100644 index 00000000..9408220c --- /dev/null +++ b/src/LLMProviders/chainRunner/utils/xmlParsing.test.ts @@ -0,0 +1,514 @@ +import { parseXMLToolCalls, escapeXml, escapeXmlAttribute, stripToolCallXML } from "./xmlParsing"; + +describe("parseXMLToolCalls", () => { + it("should parse hybrid XML tool calls with JSON arrays", () => { + const text = ` +I'll search your vault for piano learning notes. + + +localSearch +piano learning +["piano", "learning", "practice", "music"] + + +Let me also search the web. + + +webSearch +piano techniques +[] + + `; + + const toolCalls = parseXMLToolCalls(text); + + expect(toolCalls).toHaveLength(2); + + // First tool call + expect(toolCalls[0].name).toBe("localSearch"); + expect(toolCalls[0].args.query).toBe("piano learning"); + expect(toolCalls[0].args.salientTerms).toEqual(["piano", "learning", "practice", "music"]); + + // Second tool call + expect(toolCalls[1].name).toBe("webSearch"); + expect(toolCalls[1].args.query).toBe("piano techniques"); + expect(toolCalls[1].args.chatHistory).toEqual([]); // JSON array format + }); + + it("should handle tool calls with no parameters", () => { + const text = ` + +getFileTree + + `; + + const toolCalls = parseXMLToolCalls(text); + + expect(toolCalls).toHaveLength(1); + expect(toolCalls[0].name).toBe("getFileTree"); + expect(toolCalls[0].args).toEqual({}); + }); + + it("should handle string parameters without JSON parsing", () => { + const text = ` + +simpleYoutubeTranscription +https://youtube.com/watch?v=123 +en + + `; + + const toolCalls = parseXMLToolCalls(text); + + expect(toolCalls).toHaveLength(1); + expect(toolCalls[0].name).toBe("simpleYoutubeTranscription"); + expect(toolCalls[0].args.url).toBe("https://youtube.com/watch?v=123"); + expect(toolCalls[0].args.language).toBe("en"); + }); + + it("should handle hybrid approach with both JSON and XML formats", () => { + const text = ` + +hybridTool +simple string +["item1", "item2"] +{"key": "value", "number": 42} + + xml1 + xml2 + + + `; + + const toolCalls = parseXMLToolCalls(text); + + expect(toolCalls).toHaveLength(1); + expect(toolCalls[0].name).toBe("hybridTool"); + expect(toolCalls[0].args.stringParam).toBe("simple string"); + expect(toolCalls[0].args.jsonArray).toEqual(["item1", "item2"]); // JSON parsed + expect(toolCalls[0].args.jsonObject).toEqual({ key: "value", number: 42 }); // JSON parsed + expect(toolCalls[0].args.xmlArray).toEqual(["xml1", "xml2"]); // XML parsed + }); + + it("should handle nested item tags", () => { + const text = ` + +nestedTool +working string + + + Item 1 + Value 1 + + + Item 2 + Value 2 + + + + `; + + const toolCalls = parseXMLToolCalls(text); + + expect(toolCalls).toHaveLength(1); + expect(toolCalls[0].name).toBe("nestedTool"); + expect(toolCalls[0].args.simpleParam).toBe("working string"); + expect(toolCalls[0].args.nestedArray).toEqual([ + { title: "Item 1", value: "Value 1" }, + { title: "Item 2", value: "Value 2" }, + ]); + }); + + it("should ignore tool calls with empty names", () => { + const text = ` + + +value + + + +validTool +value + + `; + + const toolCalls = parseXMLToolCalls(text); + + expect(toolCalls).toHaveLength(1); + expect(toolCalls[0].name).toBe("validTool"); + }); + + it("should handle whitespace and formatting", () => { + const text = ` + + localSearch + search query + ["term1", "term2"] + + `; + + const toolCalls = parseXMLToolCalls(text); + + expect(toolCalls).toHaveLength(1); + expect(toolCalls[0].name).toBe("localSearch"); + expect(toolCalls[0].args.query).toBe("search query"); + expect(toolCalls[0].args.salientTerms).toEqual(["term1", "term2"]); + }); + + it("should return empty array for text with no tool calls", () => { + const text = "This is just regular text with no tool calls."; + + const toolCalls = parseXMLToolCalls(text); + + expect(toolCalls).toHaveLength(0); + }); + + it("should handle webSearch tool with hybrid chatHistory formats", () => { + const text = ` + +webSearch +piano practice tips +[] + + + +webSearch +follow-up search + + + user + Tell me about piano + + + assistant + Piano is a musical instrument + + + + `; + + const toolCalls = parseXMLToolCalls(text); + + expect(toolCalls).toHaveLength(2); + + // First tool call - JSON empty array + expect(toolCalls[0].name).toBe("webSearch"); + expect(toolCalls[0].args.query).toBe("piano practice tips"); + expect(toolCalls[0].args.chatHistory).toEqual([]); + + // Second tool call - XML item format + expect(toolCalls[1].name).toBe("webSearch"); + expect(toolCalls[1].args.query).toBe("follow-up search"); + expect(toolCalls[1].args.chatHistory).toEqual([ + { role: "user", content: "Tell me about piano" }, + { role: "assistant", content: "Piano is a musical instrument" }, + ]); + }); + + it("should handle malformed JSON gracefully", () => { + const text = ` + +testTool +working string +[invalid json +["valid", "array"] + + `; + + const toolCalls = parseXMLToolCalls(text); + + expect(toolCalls).toHaveLength(1); + expect(toolCalls[0].name).toBe("testTool"); + expect(toolCalls[0].args.goodParam).toBe("working string"); + expect(toolCalls[0].args.badJson).toBe("[invalid json"); // Falls back to string + expect(toolCalls[0].args.goodJson).toEqual(["valid", "array"]); // JSON parsed + }); +}); + +describe("XML Escaping Utilities", () => { + describe("escapeXml", () => { + it("should escape ampersands", () => { + expect(escapeXml("foo & bar")).toBe("foo & bar"); + }); + + it("should escape less than signs", () => { + expect(escapeXml("foo < bar")).toBe("foo < bar"); + }); + + it("should escape greater than signs", () => { + expect(escapeXml("foo > bar")).toBe("foo > bar"); + }); + + it("should escape double quotes", () => { + expect(escapeXml('foo "bar" baz')).toBe("foo "bar" baz"); + }); + + it("should escape single quotes", () => { + expect(escapeXml("foo 'bar' baz")).toBe("foo 'bar' baz"); + }); + + it("should escape multiple special characters", () => { + expect(escapeXml('content & more')).toBe( + "<tag attr="value">content & more</tag>" + ); + }); + + it("should handle empty strings", () => { + expect(escapeXml("")).toBe(""); + }); + + it("should handle strings with no special characters", () => { + expect(escapeXml("hello world")).toBe("hello world"); + }); + + it("should handle non-string inputs", () => { + expect(escapeXml(null as any)).toBe(""); + expect(escapeXml(undefined as any)).toBe(""); + expect(escapeXml(123 as any)).toBe(""); + }); + + it("should escape XML entity references", () => { + expect(escapeXml("<>"'&")).toBe( + "&lt;&gt;&quot;&apos;&amp;" + ); + }); + }); + + describe("escapeXmlAttribute", () => { + it("should escape attribute values the same as escapeXml", () => { + const testString = 'content & more'; + expect(escapeXmlAttribute(testString)).toBe(escapeXml(testString)); + }); + + it("should handle variable names with special characters", () => { + expect(escapeXmlAttribute('my"variable')).toBe("my"variable"); + expect(escapeXmlAttribute("my'variable")).toBe("my'variable"); + expect(escapeXmlAttribute("my")).toBe("my<variable>"); + }); + }); +}); + +describe("stripToolCallXML", () => { + describe("complete tool calls", () => { + it("should remove complete tool call blocks", () => { + const text = `Before tool call. + + +localSearch +test query + + +After tool call.`; + + const result = stripToolCallXML(text); + expect(result).toBe("Before tool call.\n\nAfter tool call."); + }); + + it("should remove multiple complete tool call blocks", () => { + const text = `First text. + + +localSearch +first search + + +Middle text. + + +webSearch +second search + + +Final text.`; + + const result = stripToolCallXML(text); + expect(result).toBe("First text.\n\nMiddle text.\n\nFinal text."); + }); + }); + + describe("partial tool calls", () => { + it("should show calling message for partial tool call at end of text", () => { + const text = `Some text before. + + +localSearch +incomplete`; + + const result = stripToolCallXML(text); + expect(result).toBe("Some text before.\n\nCalling vault search..."); + }); + + it("should show generic calling message for partial tool call with only opening tag", () => { + const text = `Some text before. + +`; + + const result = stripToolCallXML(text); + expect(result).toBe("Some text before.\n\nCalling tool..."); + }); + + it("should show calling message for partial tool call with incomplete parameters", () => { + const text = `Some text before. + + +webSearch +incomplete query +value`; + + const result = stripToolCallXML(text); + expect(result).toBe("Some text before.\n\nCalling web search..."); + }); + + it("should handle mixed complete and partial tool calls", () => { + const text = `Start text. + + +localSearch +complete search + + +Middle text. + + +webSearch +incomplete`; + + const result = stripToolCallXML(text); + expect(result).toBe("Start text.\n\nMiddle text.\n\nCalling web search..."); + }); + + it("should show calling message for partial tool call in middle when followed by text", () => { + const text = `Before text. + + +localSearch +incomplete query without closing + +This text should remain.`; + + const result = stripToolCallXML(text); + expect(result).toBe("Before text.\n\nCalling vault search..."); + }); + + it("should show generic calling message when tool name is not yet available", () => { + const text = `Some text before. + + +`; + + const result = stripToolCallXML(text); + expect(result).toBe("Some text before.\n\nCalling tool..."); + }); + }); + + describe("edge cases", () => { + it("should handle text with no tool calls", () => { + const text = "Just regular text with no tool calls."; + const result = stripToolCallXML(text); + expect(result).toBe("Just regular text with no tool calls."); + }); + + it("should handle empty string", () => { + const result = stripToolCallXML(""); + expect(result).toBe(""); + }); + + it("should handle text with just whitespace", () => { + const result = stripToolCallXML(" \n \n "); + expect(result).toBe(""); + }); + + it("should handle multiple partial tool calls", () => { + const text = `Text before. + + +localSearch + +More text. + + +webSearch +value`; + + const result = stripToolCallXML(text); + expect(result).toBe("Text before.\n\nCalling vault search..."); + }); + + it("should preserve non-tool XML tags", () => { + const text = `Some text with bold and italic tags. + + +localSearch +search + + +More
HTML-like
content.`; + + const result = stripToolCallXML(text); + expect(result).toBe( + "Some text with bold and italic tags.\n\nMore
HTML-like
content." + ); + }); + + it("should handle tool calls with complex nested content", () => { + const text = `Before tool. + + +complexTool + + value1 + value2 + +["item1", "item2"] + + +After tool.`; + + const result = stripToolCallXML(text); + expect(result).toBe("Before tool.\n\nAfter tool."); + }); + }); + + describe("code block removal", () => { + it("should remove empty code blocks", () => { + const text = `Some text. + +\`\`\` +\`\`\` + +More text. + +\`\`\`javascript +\`\`\``; + + const result = stripToolCallXML(text); + expect(result).toBe("Some text.\n\nMore text."); + }); + + it("should remove tool_code blocks", () => { + const text = `Some text. + +\`\`\`tool_code +some tool code here +\`\`\` + +More text.`; + + const result = stripToolCallXML(text); + expect(result).toBe("Some text.\n\nMore text."); + }); + + it("should clean up excessive whitespace", () => { + const text = `Text with + + +multiple + + +newlines.`; + + const result = stripToolCallXML(text); + expect(result).toBe("Text with\n\nmultiple\n\nnewlines."); + }); + }); +}); diff --git a/src/LLMProviders/chainRunner/utils/xmlParsing.ts b/src/LLMProviders/chainRunner/utils/xmlParsing.ts new file mode 100644 index 00000000..6a30c0a5 --- /dev/null +++ b/src/LLMProviders/chainRunner/utils/xmlParsing.ts @@ -0,0 +1,208 @@ +import { logError, logWarn } from "@/logger"; +import { getToolDisplayName } from "./toolExecution"; + +/** + * Escapes special XML characters in a string to prevent XML injection + * @param str - The string to escape + * @returns The escaped string safe for XML content + */ +export function escapeXml(str: string): string { + if (typeof str !== "string") { + return ""; + } + + return str + .replace(/&/g, "&") + .replace(//g, ">") + .replace(/"/g, """) + .replace(/'/g, "'"); +} + +/** + * Escapes special XML characters for use in XML attributes + * @param str - The string to escape for attribute use + * @returns The escaped string safe for XML attributes + */ +export function escapeXmlAttribute(str: string): string { + return escapeXml(str); +} + +export interface ToolCall { + name: string; + args: any; +} + +/** + * Parses XML tool call blocks from AI responses using pure XML format + * Format: toolNamevalue1val1val2 + */ +export function parseXMLToolCalls(text: string): ToolCall[] { + const toolCalls: ToolCall[] = []; + + try { + const regex = /([\s\S]*?)<\/use_tool>/g; + + let match; + while ((match = regex.exec(text)) !== null) { + const content = match[1]; + const nameMatch = content.match(/([\s\S]*?)<\/name>/); + + if (nameMatch) { + const name = nameMatch[1].trim(); + + // Validate tool name + if (!name || name.length === 0) { + logWarn("Skipping tool call with empty name"); + continue; + } + + // Parse individual XML parameters using pure XML approach + const args: any = {}; + + // Remove the name tag from content to avoid matching it + const contentWithoutName = content.replace(/[\s\S]*?<\/name>/, ""); + + // Find all parameter tags + const paramRegex = /<([^>]+)>([\s\S]*?)<\/\1>/g; + let paramMatch; + + while ((paramMatch = paramRegex.exec(contentWithoutName)) !== null) { + const paramName = paramMatch[1].trim(); + const paramContent = paramMatch[2].trim(); + + // Skip empty parameter names + if (!paramName) continue; + + // Parse parameter content as pure XML + args[paramName] = parseParameterContent(paramContent, paramName); + } + + toolCalls.push({ name, args }); + } + } + } catch (error) { + logError("Error parsing XML tool calls:", error); + // Return empty array if parsing fails completely + return []; + } + + return toolCalls; +} + +/** + * Parses parameter content using hybrid approach + * - JSON arrays/objects: ["item1", "item2"] or {"key": "value"} + * - Pure XML arrays: value1value2 + * - Pure XML objects: value1value2 + * - Simple strings: just the text value + * - Empty content: returns appropriate empty value based on parameter name + */ +function parseParameterContent(content: string, parameterName?: string): any { + if (!content) { + // Special handling for known array parameters that should default to empty arrays + if (parameterName === "chatHistory" || parameterName === "salientTerms") { + return []; + } + return ""; + } + + // Check if content contains XML tags + const hasXmlTags = /<[^>]+>/.test(content); + + if (!hasXmlTags) { + // Try to parse as JSON if it looks like a JSON array/object + if ( + (content.startsWith("[") && content.endsWith("]")) || + (content.startsWith("{") && content.endsWith("}")) + ) { + try { + return JSON.parse(content); + } catch { + // If JSON parsing fails, use as string + return content; + } + } + // Simple string value + return content; + } + + // Check if it's an array format with tags + const itemMatches = content.match(/([\s\S]*?)<\/item>/g); + if (itemMatches) { + return itemMatches.map((match) => { + const itemContent = match.replace(/<\/?item>/g, "").trim(); + return parseParameterContent(itemContent); // Recursive for nested structures + }); + } + + // Check if it's an object format with key-value pairs + const objectRegex = /<([^>]+)>([\s\S]*?)<\/\1>/g; + const objectEntries: [string, any][] = []; + let objectMatch; + + while ((objectMatch = objectRegex.exec(content)) !== null) { + const key = objectMatch[1].trim(); + const value = objectMatch[2].trim(); + objectEntries.push([key, parseParameterContent(value)]); + } + + if (objectEntries.length > 0) { + return Object.fromEntries(objectEntries); + } + + // Fallback to string if we can't parse as structured data + return content; +} + +/** + * Extracts tool name from a partial tool call block + */ +function extractToolNameFromPartialBlock(partialContent: string): string | null { + const nameMatch = partialContent.match(/([\s\S]*?)<\/name>/); + if (nameMatch) { + const name = nameMatch[1].trim(); + return name || null; + } + return null; +} + +/** + * Strips XML tool call blocks, thinking blocks, and various code blocks from text. + * @param text - The text to clean + * @returns The cleaned text + */ +export function stripToolCallXML(text: string): string { + let cleaned = text; + + // Remove all complete ... blocks + cleaned = cleaned.replace(/[\s\S]*?<\/use_tool>/g, ""); + + // Replace partial tool calls with calling message + cleaned = cleaned.replace(/([\s\S]*)$/g, (match, partialContent) => { + const toolName = extractToolNameFromPartialBlock(partialContent); + if (toolName) { + const displayName = getToolDisplayName(toolName); + return `Calling ${displayName}...`; + } else { + return `Calling tool...`; + } + }); + + // Keep thinking blocks in autonomous agent mode as they provide valuable context + // They are only removed in other contexts where they add noise + + // Remove empty code blocks that might appear + cleaned = cleaned.replace(/```\w*\s*```/g, ""); + + // Remove tool_code blocks (both empty and with content) + cleaned = cleaned.replace(/```tool_code[\s\S]*?```/g, ""); + + // Remove any remaining empty code blocks with various languages + cleaned = cleaned.replace(/```[\w]*[\s\n]*```/g, ""); + + // Clean up excessive whitespace and trim + cleaned = cleaned.replace(/\n\s*\n\s*\n/g, "\n\n").trim(); + + return cleaned; +} diff --git a/src/LLMProviders/chatModelManager.ts b/src/LLMProviders/chatModelManager.ts index b789c3b2..17b6b3d6 100644 --- a/src/LLMProviders/chatModelManager.ts +++ b/src/LLMProviders/chatModelManager.ts @@ -140,7 +140,7 @@ export default class ChatModelManager { fetch: customModel.enableCors ? safeFetch : undefined, }, ...(isThinkingEnabled && { - thinking: { type: "enabled", budget_tokens: 1024 }, + thinking: { type: "enabled", budget_tokens: 2048 }, }), }, [ChatModelProviders.AZURE_OPENAI]: { diff --git a/src/LLMProviders/intentAnalyzer.ts b/src/LLMProviders/intentAnalyzer.ts index 45e5a3c1..01204489 100644 --- a/src/LLMProviders/intentAnalyzer.ts +++ b/src/LLMProviders/intentAnalyzer.ts @@ -3,6 +3,7 @@ import { isProjectMode } from "@/aiParams"; import { createGetFileTreeTool } from "@/tools/FileTreeTools"; import { indexTool, localSearchTool, webSearchTool } from "@/tools/SearchTools"; import { + convertTimeBetweenTimezonesTool, getCurrentTimeTool, getTimeInfoByEpochTool, getTimeRangeMsTool, @@ -30,6 +31,7 @@ export class IntentAnalyzer { if (this.tools.length === 0) { this.tools = [ getCurrentTimeTool, + convertTimeBetweenTimezonesTool, getTimeInfoByEpochTool, getTimeRangeMsTool, localSearchTool, @@ -104,7 +106,7 @@ export class IntentAnalyzer { const { timeRange, salientTerms } = context; // Handle @vault command - if (message.includes("@vault") && (salientTerms.length > 0 || timeRange)) { + if (message.includes("@vault")) { // Remove all @commands from the query const cleanQuery = this.removeAtCommands(originalMessage); diff --git a/src/LLMProviders/projectManager.ts b/src/LLMProviders/projectManager.ts index 676367e6..b9eff45a 100644 --- a/src/LLMProviders/projectManager.ts +++ b/src/LLMProviders/projectManager.ts @@ -569,7 +569,9 @@ ${contextParts.join("\n\n")} // Note: We're only processing markdown files here, other file types // are handled by FileParserManager and stored in the file cache const files = this.app.vault.getFiles().filter((file) => { - return file.extension === "md" && shouldIndexFile(file, inclusionPatterns, exclusionPatterns); + return ( + file.extension === "md" && shouldIndexFile(file, inclusionPatterns, exclusionPatterns, true) + ); }); logInfo(`Found ${files.length} markdown files to process for project context`); diff --git a/src/cache/projectContextCache.ts b/src/cache/projectContextCache.ts index cc21f8d4..62b83e07 100644 --- a/src/cache/projectContextCache.ts +++ b/src/cache/projectContextCache.ts @@ -99,7 +99,7 @@ export class ProjectContextCache { isProject: true, }); - if (shouldIndexFile(file, inclusions, exclusions)) { + if (shouldIndexFile(file, inclusions, exclusions, true)) { // Only invalidate markdown context, keep other contexts await this.invalidateMarkdownContext(project); logInfo( @@ -525,7 +525,7 @@ export class ProjectContextCache { const file = this.vault.getAbstractFileByPath(filePath); // If file no longer exists or doesn't match patterns, remove its reference - if (!(file instanceof TFile) || !shouldIndexFile(file, inclusions, exclusions)) { + if (!(file instanceof TFile) || !shouldIndexFile(file, inclusions, exclusions, true)) { // Note: We don't remove from fileCache to preserve content for future use removedCount++; } else { @@ -571,7 +571,7 @@ export class ProjectContextCache { let addedCount = 0; for (const file of allFiles) { - if (shouldIndexFile(file, inclusions, exclusions)) { + if (shouldIndexFile(file, inclusions, exclusions, true)) { if (contextCacheToUpdate.fileContexts[file.path]) { continue; } diff --git a/src/commands/CustomCommandChatModal.tsx b/src/commands/CustomCommandChatModal.tsx index ba10f8b9..e4b9b0de 100644 --- a/src/commands/CustomCommandChatModal.tsx +++ b/src/commands/CustomCommandChatModal.tsx @@ -21,7 +21,7 @@ import { CustomCommand } from "@/commands/type"; import { useSettingsValue } from "@/settings/model"; // Custom hook for managing chat chain -function useChatChain(selectedModel: CustomModel) { +function useChatChain(selectedModel: CustomModel, systemPrompt?: string) { const [chatMemory] = useState( new BufferMemory({ returnMessages: true, memoryKey: "history" }) ); @@ -32,10 +32,10 @@ function useChatChain(selectedModel: CustomModel) { async function initChatChain() { const chatModel = await ChatModelManager.getInstance().createModelInstance(selectedModel); + const defaultSystemPrompt = + "You are a helpful assistant. You'll help the user with their content editing needs."; const chatPrompt = ChatPromptTemplate.fromMessages([ - SystemMessagePromptTemplate.fromTemplate( - "You are a helpful assistant. You'll help the user with their content editing needs." - ), + SystemMessagePromptTemplate.fromTemplate(systemPrompt || defaultSystemPrompt), new MessagesPlaceholder("history"), HumanMessagePromptTemplate.fromTemplate("{input}"), ]); @@ -57,7 +57,7 @@ function useChatChain(selectedModel: CustomModel) { } initChatChain(); - }, [selectedModel, chatMemory]); + }, [selectedModel, chatMemory, systemPrompt]); return { chatChain, chatMemory }; } @@ -67,6 +67,7 @@ interface CustomCommandChatModalContentProps { command: CustomCommand; onInsert: (message: string) => void; onReplace: (message: string) => void; + systemPrompt?: string; } function CustomCommandChatModalContent({ @@ -74,6 +75,7 @@ function CustomCommandChatModalContent({ command, onInsert, onReplace, + systemPrompt, }: CustomCommandChatModalContentProps) { const [aiCurrentMessage, setAiCurrentMessage] = useState(null); const [processedMessage, setProcessedMessage] = useState(null); @@ -88,7 +90,7 @@ function CustomCommandChatModalContent({ [command.modelKey, modelKey, settings.activeModels] ); - const { chatChain, chatMemory } = useChatChain(selectedModel); + const { chatChain, chatMemory } = useChatChain(selectedModel, systemPrompt); const commandTitle = command.title; @@ -298,12 +300,16 @@ function CustomCommandChatModalContent({
{generating ? ( // When generating, show Stop button - ) : showFollowupSubmit ? ( // When follow-up instruction has content, show Submit button with Enter shortcut - @@ -311,11 +317,12 @@ function CustomCommandChatModalContent({ // Otherwise, show Insert and Replace buttons with shortcut indicators <>