From bc2d7f660af1ad58af78510d087353e3f2e427df Mon Sep 17 00:00:00 2001 From: Logan Yang Date: Mon, 4 Aug 2025 09:51:18 -0700 Subject: [PATCH] Implement autonomous agent (#1689) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Implement autonomous agent mode and new message architecture - Implement autonomous agent mode with sequential thinking capabilities - Migrate from sharedState to new clean message management architecture - Add ChatManager, MessageRepository, and ContextManager for better separation of concerns - Implement project-based chat isolation with separate message repositories per project - Add ChatPersistenceManager for saving/loading chat history - Refactor chain runner architecture with specialized runners (Autonomous, Project, VaultQA) - Add XML tool call handling for better compatibility across LLM providers - Introduce vault and web search toggles in settings - Enhance model adapters with explicit tool usage instructions - Add comprehensive test coverage for new architecture components - Update UI components to use new ChatUIState instead of SharedState 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude * Enhance context processing and documentation - Added detailed XML structure for context notes, including metadata (creation and modification times) for notes and selected text. - Introduced new tests for message lifecycle and XML tag formatting to ensure proper context handling. - Created TODO.md to track technical debt and future improvements related to context processing and file handling. - Updated CLAUDE.md to reference the new TODO.md for technical debt awareness. This commit improves the clarity and functionality of context processing in messages, ensuring rich context is maintained throughout the message lifecycle. * Fix composer instruction in new message implementation * Add writeToFile tool to agent (#1632) * Add writeToFile tool call to agent * Update src/LLMProviders/chainRunner/AutonomousAgentChainRunner.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update src/tools/ComposerTools.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Remove unuseful comment --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Fix tests * Enhance context processing with new XML tags for notes and selected text (#1633) * Enhance context processing with new XML tags for notes and selected text - Introduced new constants for XML tags: SELECTED_TEXT_TAG, VARIABLE_TAG, VARIABLE_NOTE_TAG, EMBEDDED_PDF_TAG, VAULT_NOTE_TAG, and RETRIEVED_DOCUMENT_TAG. - Updated context processing to format embedded PDFs and notes with XML structure. - Modified prompt processing to include XML tags for selected text and variables, improving clarity and structure in generated outputs. - Enhanced tests to validate the new XML formatting in processed prompts. * Implement XML escaping utility functions and integrate into context processing - Added `escapeXml` and `escapeXmlAttribute` functions to handle XML special character escaping, enhancing security against XML injection. - Updated context processing to utilize these functions for escaping note titles, paths, and content, ensuring proper XML formatting. - Enhanced tests to validate XML escaping functionality across various scenarios, including selected text and variable names. - Introduced new test files for XML utility functions to ensure comprehensive coverage and reliability. * Refactor XML utility functions and update imports - Moved `escapeXml` and `escapeXmlAttribute` functions from `utils/xmlUtils` to `LLMProviders/chainRunner/utils/xmlParsing`, enhancing modularity. - Updated all relevant imports across the codebase to reflect the new location of XML utility functions. - Removed obsolete `xmlUtils` file and its associated tests, consolidating XML handling in a single module for better maintainability. * Add message creation callback in ChatManager for immediate UI updates (#1634) - Introduced a new method `setOnMessageCreatedCallback` in ChatManager to allow setting a callback that triggers when a message is created. - Updated ChatUIState to utilize this callback for notifying listeners, ensuring the UI updates immediately upon message creation. - This enhancement improves the responsiveness of the chat interface by synchronizing message creation events with UI state updates. * Enhance TimeTools tests for invalid expressions - Added setup to mock console.warn in tests for invalid time expressions. - Verified that the appropriate warning message is logged when parsing fails. - Improved test coverage for handling various invalid time inputs. * Refactor local search handling and restore sources display - Unified tool handling in CopilotPlusChainRunner to fix websearch results being ignored when localSearch has results - All tool outputs now go through the same prepareEnhancedUserMessage method - Preserved QA format and instruction when only localSearch is used - Added proper validation for localSearch results (non-empty documents array) - Maintained time expression support for temporal queries - Fixed source tracking and display in ChatUIState - Enhanced error handling for JSON parsing of search results * Enhance AutonomousAgentChainRunner to track tool calls and results (#1639) - Introduced a new property to store LLM-formatted messages for memory management. - Reset LLM messages at the start of each run to ensure fresh context. - Updated response handling to include LLM-formatted outputs, enhancing memory context for future iterations. - Modified the base class to accept an optional LLM-formatted output parameter for improved context saving. - Enhanced the writeToFile tool's response message for clarity on user actions. * Quick Command (cmd+k) (#1640) * Quick Command (cmd+k) * Update custom command modal style * Improve prompt * Use submenu for commands * Fix cursor comment * Align selected_text variable name * Fix embedded image wikilinks (#1641) - Updated the `extractEmbeddedImages` method to accept an optional `sourcePath` parameter for resolving wikilinks. - Implemented logic to determine the source path from context notes or fallback to the active file. - Enhanced logging for unresolved images to improve debugging and user feedback. - Ensured that embedded images are processed correctly based on the new source path handling. * Fix salient term language issue * Dedupe source notes (#1642) * Fix salient term language issue * Refactor source handling to include path in addition to title and score - Updated source structures across multiple classes to include a `path` property alongside `title` and `score`. - Modified methods in `AutonomousAgentChainRunner`, `BaseChainRunner`, and `CopilotPlusChainRunner` to accommodate the new source format. - Enhanced the `deduplicateSources` function to use `path` as the unique key, falling back to `title` if necessary. - Adjusted the `SourcesModal` to display paths correctly and ensure links open using the appropriate source path. - Updated type definitions in `message.ts` to reflect the new source structure. * Update path handling in source mapping for AutonomousAgentChainRunner - Modified the source mapping logic to ensure the `path` property defaults to an empty string if not provided, enhancing robustness in source data handling. - This change improves the consistency of source objects by ensuring that all properties are defined, which aids in downstream processing. * Bump manifest version to v3 * Fix image (#1643) * Fix salient term language issue * Bump manifest version to v3 * Fix image passing - Updated Chat, ChatManager, MessageRepository, and ChatUIState to include an optional `content` parameter in message handling. - Adjusted method signatures and message creation logic to accommodate the new parameter, improving flexibility in message processing. * Fix regenerate (#1644) * Fix image in agent chain (#1646) * Add YouTube transcription processing to AutonomousAgentChainRunner (#1647) - Implemented a new method to extract and process YouTube URLs from user messages. - Fetch transcriptions using the simpleYoutubeTranscriptionTool and handle errors gracefully. - Updated conversation messages to include fetched transcriptions for improved context in responses. * Disable quick command when live-mode is not on (#1648) * Do not call writeToFileTool again for accepted changes as well (#1652) * Trigger youtube tool only for urls in prompt not context (#1653) * Stop calling youtube tool in context notes in agent mode - Added a new instruction for handling YouTube URLs in the DEFAULT_SYSTEM_PROMPT. - Removed the direct usage of simpleYoutubeTranscriptionTool in the AutonomousAgentChainRunner, as YouTube transcriptions are now automatically processed. - Updated parameter naming conventions in modelAdapter for clarity and consistency. - Cleaned up the simpleYoutubeTranscriptionTool definition by removing unnecessary parameters. * Add agent loop limit message (#1655) * Add composer toggle and implement auto-preview (#1651) * Update composer output format to XML and refactor related components - Changed the output format for composer instructions from JSON to XML in `constants.ts`. - Updated `CopilotPlusChainRunner` to utilize the new XML format and adjusted streaming logic accordingly. - Added a toggle for the composer feature in `ChatInput.tsx` and updated UI elements to reflect this change. - Removed the `ComposerCodeBlock` component as it is no longer needed with the new implementation. - Cleaned up unused code and comments in `ChatSingleMessage.tsx` related to the previous composer handling. * Update src/constants.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update src/constants.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Remove redundant tests * Update doc * Refactor composer handling and implement ActionBlockStreamer - Replaced `ComposerBlockStreamer` with `ActionBlockStreamer` to handle `writeToFile` blocks more efficiently. - Updated `constants.ts` to reflect changes in output format for composer instructions. - Enhanced `ChatSingleMessage.tsx` to process `writeToFile` sections and integrate collapsible UI elements. - Removed the deprecated `ComposerBlockStreamer` and its associated tests. - Improved handling of XML codeblocks and unclosed tags during streaming. * Fix reported bugs * Refactor ActionBlockStreamer tests and implementation - Updated `processChunks` to always push content, including null and empty strings, to the output. - Adjusted test expectations to reflect changes in output handling for chunks with null content. - Simplified regex in `findCompleteBlock` method by removing XML format handling. * Fix reported issue --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Fix tool description generation (#1659) * Access tool description from zod.shape when zod.properties is missing. * Remove the access to schema.properties * Remove unused print * Also support schema.propertes * Implement SimpleTool and remove Langchain tool calling (#1662) * Make web search tool not overly eager * Use SimpleTool interface instead of langchain tool call - Replaced LangChain tool definitions with a new SimpleTool interface for better type safety and validation. - Updated existing tools (e.g., writeToFile, localSearch, webSearch, and YouTube transcription) to utilize the new createTool function. - Enhanced parameter extraction from Zod schemas for improved tool description generation. - Cleaned up related code and documentation to reflect the new tool structure and ensure consistency across the codebase. - Add tests * Update ChatManager tests to handle undefined context cases - Modified test cases in ChatManager to include handling of undefined context parameters. - Ensured that the tests accurately reflect scenarios where no active file exists and when context is undefined. - Improved overall test coverage for message processing in various contexts. * Add validation tests for tool schemas and improve schema definitions - Introduced comprehensive validation tests for tool schemas in `allTools.validation.test.ts` and `SearchTools.schema.test.ts` to ensure adherence to best practices and proper typing. - Enhanced `SearchTools.ts` by replacing `z.any()` with specific object schemas for `chatHistory` to improve type safety. - Added tests for various tool patterns, including metadata validation and schema best practices, ensuring tools are correctly defined and validated against expected interfaces. * Add back example content for canvas file in ComposerTools schema * Make web search tool not overly eager * Add license key in agent mode (#1663) * Add tool execution tests and implement Plus subscription checks (#1664) - Introduced unit tests for the `executeSequentialToolCall` function to validate tool execution behavior, including handling of Plus-only tools. - Enhanced `executeSequentialToolCall` to check for Plus subscription requirements before executing tools. - Updated `createTool` and `SimpleTool` interfaces to include an `isPlusOnly` flag, indicating tools that require a Copilot Plus subscription. - Marked existing tools (e.g., webSearch and YouTube transcription) as Plus-only where applicable. * Add timezone conversion tools and enhance current time functionality (#1665) * Add timezone conversion tools and enhance current time functionality - Introduced `convertTimeBetweenTimezonesTool` to allow conversion of time between different timezones. - Updated `getCurrentTimeTool` to accept an optional timezone parameter, improving its functionality to return time in specified timezones. - Enhanced documentation and added examples for both tools in the codebase. - Added comprehensive unit tests for the new timezone conversion functionality to ensure accuracy and reliability. * Enhance timezone functionality and add unit tests - Updated `getCurrentTime` and `convertToTimeInfo` functions to utilize Luxon's offset for accurate timezone calculations. - Adjusted timezone offset handling to reflect correct signs. - Added unit tests for Tokyo and New York timezone offsets to ensure accuracy in timezone calculations. * Add unit test for past time conversion in timezone tool - Introduced a new test case to verify that converting a past time (6:00 AM PT) to Tokyo correctly reflects the same day (January 15) without erroneously adding a day. - Updated the `convertTimeBetweenTimezones` function to ensure accurate handling of past times without unintended date shifts. * Enhance timezone handling and update related tests - Added `convertTimeBetweenTimezonesTool` to facilitate time conversion between different UTC offsets. - Updated `getCurrentTimeTool` to accept timezone offsets instead of names, improving accuracy and flexibility. - Enhanced error handling for invalid timezone offset formats and added comprehensive unit tests to validate new functionality. - Updated existing tests to reflect changes in timezone offset handling and ensure consistent behavior across various scenarios. * Refactor timezone handling in convertTimeBetweenTimezones function - Updated the creation of DateTime objects to interpret parsed dates as already being in the source timezone, improving clarity and accuracy in timezone conversions. * Fix image in chat history (#1666) * Refactor chat history handling to preserve multimodal content - Removed the deprecated `extractChatHistory` function and replaced it with direct access to raw history from memory, allowing for the preservation of multimodal content. - Updated `AutonomousAgentChainRunner`, `BaseChainRunner`, and `CopilotPlusChainRunner` to utilize `BaseMessage` objects for chat history, ensuring that both user and AI messages maintain their original structure and content. - Enhanced message processing to accommodate multimodal inputs, storing them appropriately for later use. * Refactor chat history handling to improve message processing - Introduced `addChatHistoryToMessages` utility to streamline the addition of chat history, ensuring safe handling of various message formats. - Updated `AutonomousAgentChainRunner` and `CopilotPlusChainRunner` to utilize the new utility, enhancing clarity and maintainability of the code. - Removed deprecated manual processing of chat history to improve code efficiency and readability. * Refactor chat history processing to enhance multimodal support - Removed the deprecated `extractChatHistory` function and replaced it with `processRawChatHistory` to ensure consistent handling of chat history. - Introduced `processedMessagesToTextOnly` function to extract text-only content from multimodal messages, improving clarity for question condensing. - Updated `CopilotPlusChainRunner` to utilize the new processing functions, ensuring that both LLM and question condensing use the same data format. - Added comprehensive unit tests for `processedMessagesToTextOnly` to validate its functionality across various message formats. * Refactor memory management in BaseChainRunner to enhance context saving - Removed the creation of `HumanMessage` and `AIMessage` objects for chat history, simplifying the process. - Implemented `saveContext` for atomic operations, ensuring proper memory management while preserving input and output data. - Noted that LangChain's memory now expects text content, which limits the saving of multimodal content. * Remove multimodal content storage from userMessage in AutonomousAgentChainRunner and CopilotPlusChainRunner to streamline message processing. This change aligns with recent refactors to enhance memory management and improve clarity in handling chat history. * Enhance ChatInput component with autonomous agent toggle button (#1669) - Added a toggle for the autonomous agent feature, allowing users to enable or disable it within the ChatInput component. - Updated the state management to synchronize the autonomous agent toggle with user settings. - Introduced a button for toggling the autonomous agent mode, which is only visible in Copilot Plus mode. - Refactored tool call logic to accommodate the new autonomous agent toggle, ensuring proper handling of tool calls based on the toggle state. * Implement custom topic in conversation history (#1670) * Refactor LoadChatHistoryModal and ChatPersistenceManager for enhanced topic handling - Updated LoadChatHistoryModal to prioritize custom topics from file frontmatter, falling back to filename extraction if not present. - Enhanced ChatPersistenceManager to generate AI topics for new chat files and preserve existing topics when updating files. - Introduced new methods for finding files by epoch and generating AI topics based on conversation content, improving chat file management and organization. * Refactor LoadChatHistoryModal and ChatPersistenceManager for improved topic handling - Updated LoadChatHistoryModal to ensure custom topics are trimmed and validated before use. - Enhanced ChatPersistenceManager to utilize reduce for efficient conversation summary generation, introducing constants for message and character limits to improve readability and maintainability. * Implement tool call UI banner (#1671) * Only add tool call to the user message when agent is off - Updated the tool call logic to only add tool calls when the autonomous agent toggle is off, ensuring that the autonomous agent manages all tools internally when enabled. - Improved code clarity by adding comments to explain the new logic. * Implement tool call collapsible UI element - Introduced a new utility for managing tool call markers, allowing for structured display and tracking of tool calls during execution. - Updated the AutonomousAgentChainRunner to utilize the new tool call markers, improving clarity in the streaming response and tool execution process. - Added a ToolCallBanner component for better visualization of tool calls in the chat interface, including execution status and results. - Refactored ChatSingleMessage to handle tool call updates dynamically, ensuring that the UI reflects the current state of tool calls. - Enhanced CSS for tool call containers and added animations for improved user experience. * Fix disappearing tool call in final response * Implement tool call result formatting - Updated ChatSingleMessage to use a stable ID for message roots, enhancing memory management and preventing leaks. - Implemented cleanup logic for old message roots to optimize performance. - Introduced ToolResultFormatter to format tool results for better display in the UI, ensuring user-friendly output for various tool types. - Enhanced ToolCallBanner to utilize the new formatter, improving the presentation of tool call results in the chat interface. * Add YouTube transcription tool in agent and avoid mistrigger from context - Replaced the deprecated simpleYoutubeTranscriptionTool with a new youtubeTranscriptionTool that automatically extracts YouTube URLs from user messages. - Removed the processYouTubeUrls method to streamline the handling of YouTube transcriptions. - Updated tool execution logic to pass the original user message to tools that require it, enhancing flexibility in tool usage. - Improved the ToolResultFormatter to support multi-URL responses and provide better error handling for transcription failures. - Enhanced documentation for YouTube transcription usage in the codebase. * Enhance ToolCallBanner and ToolResultFormatter for improved user experience - Added animation constants to ToolCallBanner for better visual feedback. - Updated ToolResultFormatter to robustly handle JSON parsing, improving error handling for both single objects and arrays. - Implemented input validation in YouTube transcription tool to prevent excessive input lengths and ensure proper data types. * Refactor tool execution messages and enhance JSON parsing in ToolResultFormatter - Updated the confirmation message for the writeToFile tool to be more concise. - Added a private method in ToolResultFormatter to robustly parse JSON strings, improving error handling and ensuring consistent results. - Streamlined the handling of search results by separating JSON parsing and regex extraction logic in ToolResultFormatter. - Enforced required user message content in the YouTube transcription tool to prevent potential errors. * Truncate long tool results for compact memory (#1672) * Truncate long tool results for compact memory - Introduced `processToolResults` utility to format tool results for different contexts, allowing for both truncated and full results. - Updated `AutonomousAgentChainRunner` to utilize the new utility for improved memory management and clarity in tool result presentation. - Added unit tests for `toolResultUtils` to ensure proper functionality of truncation and formatting methods, enhancing reliability and maintainability of tool result processing. * Enhance tool results handling in AutonomousAgentChainRunner - Added a conditional check to ensure that only non-null tool results are pushed to the llmFormattedMessages array, improving the robustness of message formatting and preventing potential errors from undefined values. * Add Agent settings section (#1676) * Add autonomous agent settings and tool configurations - Introduced settings for enabling the autonomous agent mode, including maximum iterations and available tools. - Updated the CopilotPlusSettings component to allow users to toggle tools like local search, web search, and YouTube transcription. - Enhanced the model adapter to utilize the new settings for tool management. * Enhance tool management by adding always available tools and updating tool usage guidelines in autonomous agent mode * Make filetree tool always available * Refactor tool management system to use a centralized ToolRegistry, enabling dynamic tool registration and configuration. Update settings structure to support enabled tool IDs and enhance UI for tool selection in the Copilot settings. * Initialize built-in tools with vault access and improve tool registration logic * Enhance documentation for initializeBuiltinTools function to clarify tool registration process and dynamic filtering of user-enabled tools. * fix: Fix the shouldIndexFile method to return all files when inclusions is empty in project mode. (#1667) * feat: Display all category items by default in the context-manage-modal. (#1657) * Remove escapingXML for context data (#1678) * Add apply command (#1677) * Implement ReplaceInFile tool (#1661) * Relevant note improvements (#1681) * Change message scrolling behavior (#1680) * Fix timetools and tests (#1683) * Refactor time tool descriptions for clarity and examples * Update XML escaping tests to reflect correct behavior and remove unnecessary settings mock * Strip partial tool call tag (#1682) * Strip partial tool call tag * Show a placeholder message * Remove perserveToolIndicators option * Add tool dynamic prompt (#1679) * feat: Enhance tool execution error messages and update custom prompt instructions for built-in tools * feat: Add unit tests for ModelAdapter and enhance tool instructions for time-based queries * feat: Implement detailed formatting for replaceInFile tool results * feat: Enhance GPT model instructions and refine replaceInFile tool usage guidelines * feat: Expand documentation on schema descriptions and custom instructions for tool usage * feat: Refine tool documentation by clarifying schema descriptions and custom prompt instructions * feat: Update GitHub Actions workflow to run on all pull requests regardless of target branch * Fix tests * feat: Update integration tests to validate writeToFile blocks and enhance replaceInFile usage example * Enhance image extraction logic to support both wiki-style and markdown image syntax (#1685) * Turn off agent when in Projects mode (#1686) * Remove think and action blocks from copy and insert (#1687) * Remove think and action blocks from copy, and fix composer for thinking models * Refactor ToolSettingsSection layout for improved readability and consistency --------- Co-authored-by: Claude Co-authored-by: wenzhengjiang Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Zero Liu Co-authored-by: Emt-lin <41323133+Emt-lin@users.noreply.github.com> --- .github/workflows/node.js.yml | 2 +- CLAUDE.md | 1 + MESSAGE_ARCHITECTURE.md | 77 ++ TODO.md | 51 + __mocks__/obsidian.js | 43 + manifest.json | 2 +- src/LLMProviders/chainManager.ts | 9 +- src/LLMProviders/chainRunner.ts | 983 ------------------ .../chainRunner/AutonomousAgentChainRunner.ts | 519 +++++++++ .../chainRunner/BaseChainRunner.ts | 142 +++ .../chainRunner/CopilotPlusChainRunner.ts | 683 ++++++++++++ .../chainRunner/LLMChainRunner.ts | 97 ++ .../chainRunner/ProjectChainRunner.ts | 25 + src/LLMProviders/chainRunner/README.md | 640 ++++++++++++ .../chainRunner/VaultQAChainRunner.ts | 161 +++ src/LLMProviders/chainRunner/index.ts | 22 + .../utils/ActionBlockStreamer.test.ts | 232 +++++ .../chainRunner/utils/ActionBlockStreamer.ts | 93 ++ .../chainRunner/utils/ThinkBlockStreamer.ts | 124 +++ .../utils/chatHistoryUtils.test.ts | 245 +++++ .../chainRunner/utils/chatHistoryUtils.ts | 112 ++ .../chainRunner/utils/imageExtraction.test.ts | 342 ++++++ .../chainRunner/utils/modelAdapter.test.ts | 155 +++ .../chainRunner/utils/modelAdapter.ts | 697 +++++++++++++ .../chainRunner/utils/toolCallParser.ts | 122 +++ .../chainRunner/utils/toolExecution.test.ts | 122 +++ .../chainRunner/utils/toolExecution.ts | 231 ++++ .../chainRunner/utils/xmlParsing.test.ts | 514 +++++++++ .../chainRunner/utils/xmlParsing.ts | 208 ++++ src/LLMProviders/chatModelManager.ts | 2 +- src/LLMProviders/intentAnalyzer.ts | 4 +- src/LLMProviders/projectManager.ts | 4 +- src/cache/projectContextCache.ts | 6 +- src/commands/CustomCommandChatModal.tsx | 32 +- src/commands/constants.ts | 1 + src/commands/contextMenu.ts | 58 +- src/commands/customCommandRegister.ts | 5 +- src/commands/customCommandUtils.test.ts | 63 +- src/commands/customCommandUtils.ts | 119 ++- .../customCommandUtils.xmlescape.test.ts | 148 +++ src/commands/index.ts | 62 +- src/components/Chat.tsx | 91 +- src/components/QuickCommand.tsx | 199 ++++ src/components/chat-components/ChatInput.tsx | 258 ++--- .../chat-components/ChatMessages.tsx | 91 +- .../chat-components/ChatSingleMessage.tsx | 423 +++++--- .../chat-components/ComposerCodeBlock.tsx | 133 --- .../chat-components/RelevantNotes.tsx | 1 + .../chat-components/ToolCallBanner.tsx | 133 +++ src/components/composer/ApplyView.tsx | 57 +- src/components/modals/AdhocPromptModal.tsx | 62 -- .../modals/ApplyCustomCommandModal.tsx | 66 ++ .../modals/LoadChatHistoryModal.tsx | 31 +- src/components/modals/SourcesModal.tsx | 22 +- .../modals/project/context-manage-modal.tsx | 199 +++- src/components/ui/ModelSelector.tsx | 95 ++ src/constants.ts | 107 +- src/contextProcessor.ts | 25 +- src/core/ChatManager.test.ts | 39 +- src/core/ChatManager.ts | 35 +- src/core/ChatPersistenceManager.ts | 131 ++- src/core/MessageLifecycle.test.ts | 283 +++++ src/core/MessageLifecycle.xmltags.test.ts | 268 +++++ src/core/MessageRepository.ts | 4 +- src/hooks/useChatScrolling.ts | 149 +++ src/integration_tests/composerPrompt.test.ts | 61 +- src/main.ts | 18 + src/mentions/Mention.ts | 2 +- src/search/findRelevantNotes.ts | 2 +- src/search/searchUtils.ts | 10 +- src/settings/model.ts | 28 +- .../v2/components/CopilotPlusSettings.tsx | 19 + .../v2/components/ToolSettingsSection.tsx | 79 ++ src/state/ChatUIState.ts | 24 +- src/styles/tailwind.css | 19 + src/tests/context-manage-modal.test.ts | 2 +- src/tools/ComposerTools.test.ts | 475 +++++++++ src/tools/ComposerTools.ts | 321 ++++++ src/tools/FileTreeTools.ts | 19 +- src/tools/README.md | 381 +++++++ src/tools/SearchTools.schema.test.ts | 301 ++++++ src/tools/SearchTools.ts | 101 +- src/tools/SimpleTool.test.ts | 322 ++++++ src/tools/SimpleTool.ts | 170 +++ src/tools/TimeTools.test.ts | 18 +- src/tools/TimeTools.timezone.test.ts | 192 ++++ src/tools/TimeTools.ts | 296 ++++-- src/tools/ToolRegistry.ts | 128 +++ src/tools/ToolResultFormatter.ts | 401 +++++++ src/tools/YoutubeTools.ts | 114 +- src/tools/allTools.validation.test.ts | 325 ++++++ src/tools/builtinTools.ts | 318 ++++++ src/types.ts | 2 + src/types/message.ts | 4 +- src/utils.cleanMessageForCopy.test.ts | 110 ++ src/utils.ts | 67 +- src/utils/toolResultUtils.test.ts | 67 ++ src/utils/toolResultUtils.ts | 68 ++ 98 files changed, 12499 insertions(+), 2025 deletions(-) create mode 100644 TODO.md delete mode 100644 src/LLMProviders/chainRunner.ts create mode 100644 src/LLMProviders/chainRunner/AutonomousAgentChainRunner.ts create mode 100644 src/LLMProviders/chainRunner/BaseChainRunner.ts create mode 100644 src/LLMProviders/chainRunner/CopilotPlusChainRunner.ts create mode 100644 src/LLMProviders/chainRunner/LLMChainRunner.ts create mode 100644 src/LLMProviders/chainRunner/ProjectChainRunner.ts create mode 100644 src/LLMProviders/chainRunner/README.md create mode 100644 src/LLMProviders/chainRunner/VaultQAChainRunner.ts create mode 100644 src/LLMProviders/chainRunner/index.ts create mode 100644 src/LLMProviders/chainRunner/utils/ActionBlockStreamer.test.ts create mode 100644 src/LLMProviders/chainRunner/utils/ActionBlockStreamer.ts create mode 100644 src/LLMProviders/chainRunner/utils/ThinkBlockStreamer.ts create mode 100644 src/LLMProviders/chainRunner/utils/chatHistoryUtils.test.ts create mode 100644 src/LLMProviders/chainRunner/utils/chatHistoryUtils.ts create mode 100644 src/LLMProviders/chainRunner/utils/imageExtraction.test.ts create mode 100644 src/LLMProviders/chainRunner/utils/modelAdapter.test.ts create mode 100644 src/LLMProviders/chainRunner/utils/modelAdapter.ts create mode 100644 src/LLMProviders/chainRunner/utils/toolCallParser.ts create mode 100644 src/LLMProviders/chainRunner/utils/toolExecution.test.ts create mode 100644 src/LLMProviders/chainRunner/utils/toolExecution.ts create mode 100644 src/LLMProviders/chainRunner/utils/xmlParsing.test.ts create mode 100644 src/LLMProviders/chainRunner/utils/xmlParsing.ts create mode 100644 src/commands/customCommandUtils.xmlescape.test.ts create mode 100644 src/components/QuickCommand.tsx delete mode 100644 src/components/chat-components/ComposerCodeBlock.tsx create mode 100644 src/components/chat-components/ToolCallBanner.tsx delete mode 100644 src/components/modals/AdhocPromptModal.tsx create mode 100644 src/components/modals/ApplyCustomCommandModal.tsx create mode 100644 src/components/ui/ModelSelector.tsx create mode 100644 src/core/MessageLifecycle.test.ts create mode 100644 src/core/MessageLifecycle.xmltags.test.ts create mode 100644 src/hooks/useChatScrolling.ts create mode 100644 src/settings/v2/components/ToolSettingsSection.tsx create mode 100644 src/tools/ComposerTools.test.ts create mode 100644 src/tools/ComposerTools.ts create mode 100644 src/tools/README.md create mode 100644 src/tools/SearchTools.schema.test.ts create mode 100644 src/tools/SimpleTool.test.ts create mode 100644 src/tools/SimpleTool.ts create mode 100644 src/tools/TimeTools.timezone.test.ts create mode 100644 src/tools/ToolRegistry.ts create mode 100644 src/tools/ToolResultFormatter.ts create mode 100644 src/tools/allTools.validation.test.ts create mode 100644 src/tools/builtinTools.ts create mode 100644 src/utils.cleanMessageForCopy.test.ts create mode 100644 src/utils/toolResultUtils.test.ts create mode 100644 src/utils/toolResultUtils.ts 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 <>