diff --git a/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/codebase-search-node.ts b/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/codebase-search-node.ts index 01e0b91..4c9e54c 100644 --- a/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/codebase-search-node.ts +++ b/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/codebase-search-node.ts @@ -6,6 +6,7 @@ import { mergeCodeSnippets } from '@extension/webview-api/chat-context-processor import type { ToolMessage } from '@langchain/core/messages' import { DynamicStructuredTool } from '@langchain/core/tools' import { settledPromiseResults } from '@shared/utils/common' +import { ContextInfoSource } from '@webview/types/chat' import { z } from 'zod' import { @@ -49,7 +50,7 @@ export const createCodebaseSearchTool = async (state: ChatGraphState) => { .map(row => { // eslint-disable-next-line unused-imports/no-unused-vars const { embedding, ...others } = row - return { ...others, code: '' } + return { ...others, code: '', source: ContextInfoSource.ToolNode } }) const mergedCodeSnippets = await mergeCodeSnippets(searchCodeSnippets, { diff --git a/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/doc-retriever-node.ts b/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/doc-retriever-node.ts index 7bffaad..a562c6d 100644 --- a/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/doc-retriever-node.ts +++ b/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/doc-retriever-node.ts @@ -8,6 +8,7 @@ import { docSitesDB } from '@extension/webview-api/lowdb/doc-sites-db' import type { ToolMessage } from '@langchain/core/messages' import { DynamicStructuredTool } from '@langchain/core/tools' import { removeDuplicates, settledPromiseResults } from '@shared/utils/common' +import { ContextInfoSource } from '@webview/types/chat' import { z } from 'zod' import { @@ -31,6 +32,7 @@ export const createDocRetrieverTool = async (state: ChatGraphState) => { const { allowSearchDocSiteNames } = docContext if (!allowSearchDocSiteNames.length) return null + const siteNames = allowSearchDocSiteNames.map(item => item.name) const getRelevantDocs = async ( queryParts: { siteName: string; keywords: string[] }[] @@ -40,7 +42,7 @@ export const createDocRetrieverTool = async (state: ChatGraphState) => { const docPromises = queryParts.map(async ({ siteName, keywords }) => { const docSite = docSites.find(site => site.name === siteName) - if (!docSite?.isIndexed || !allowSearchDocSiteNames.includes(siteName)) { + if (!docSite?.isIndexed || !siteNames.includes(siteName)) { return [] } @@ -63,7 +65,8 @@ export const createDocRetrieverTool = async (state: ChatGraphState) => { const docInfoResults = await settledPromiseResults( searchRows.map(async row => ({ content: await docIndexer.getRowFileContent(row), - path: docSite.url + path: docSite.url, + source: ContextInfoSource.ToolNode })) ) diff --git a/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/web-search-node.ts b/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/web-search-node.ts index eb9ec17..0c0143d 100644 --- a/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/web-search-node.ts +++ b/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/web-search-node.ts @@ -1,5 +1,8 @@ import { createModelProvider } from '@extension/ai/helpers' -import type { WebSearchResult } from '@extension/webview-api/chat-context-processor/types/chat-context' +import { + ContextInfoSource, + type WebSearchResult +} from '@extension/webview-api/chat-context-processor/types/chat-context' import type { LangchainTool } from '@extension/webview-api/chat-context-processor/types/langchain-message' import { findCurrentToolsCallParams } from '@extension/webview-api/chat-context-processor/utils/find-current-tools-call-params' import { searxngSearch } from '@extension/webview-api/chat-context-processor/utils/searxng-search' @@ -141,7 +144,8 @@ export const webSearchNode: ChatGraphNode = async state => { ...lastConversation.attachments!.docContext.relevantDocs, { path: '', - content: result.relevantContent + content: result.relevantContent, + source: ContextInfoSource.ToolNode } ] diff --git a/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/web-visit-node.ts b/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/web-visit-node.ts index c5e096b..9ebd223 100644 --- a/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/web-visit-node.ts +++ b/src/extension/webview-api/chat-context-processor/strategies/chat-strategy/nodes/web-visit-node.ts @@ -4,6 +4,7 @@ import { findCurrentToolsCallParams } from '@extension/webview-api/chat-context- import type { ToolMessage } from '@langchain/core/messages' import { DynamicStructuredTool } from '@langchain/core/tools' import { settledPromiseResults } from '@shared/utils/common' +import { ContextInfoSource } from '@webview/types/chat' import { z } from 'zod' import { @@ -79,7 +80,8 @@ export const webVisitNode: ChatGraphNode = async state => { ...lastConversation.attachments!.docContext.relevantDocs, ...result.contents.map(item => ({ path: item.url, - content: item.content + content: item.content, + source: ContextInfoSource.ToolNode })) ] }) diff --git a/src/extension/webview-api/chat-context-processor/types/chat-context/base-context.ts b/src/extension/webview-api/chat-context-processor/types/chat-context/base-context.ts new file mode 100644 index 0000000..d830e0d --- /dev/null +++ b/src/extension/webview-api/chat-context-processor/types/chat-context/base-context.ts @@ -0,0 +1,13 @@ +export interface BaseToolContext { + enableTool: boolean +} + +export enum ContextInfoSource { + FileSelector = 'file-selector', + Editor = 'editor', + ToolNode = 'tool-node' +} + +export interface BaseContextInfo { + source: ContextInfoSource +} diff --git a/src/extension/webview-api/chat-context-processor/types/chat-context/base-tool-context.ts b/src/extension/webview-api/chat-context-processor/types/chat-context/base-tool-context.ts deleted file mode 100644 index 0f4f3be..0000000 --- a/src/extension/webview-api/chat-context-processor/types/chat-context/base-tool-context.ts +++ /dev/null @@ -1,3 +0,0 @@ -export interface BaseToolContext { - enableTool: boolean -} diff --git a/src/extension/webview-api/chat-context-processor/types/chat-context/code-context.ts b/src/extension/webview-api/chat-context-processor/types/chat-context/code-context.ts index 9ad0b1b..f32bee9 100644 --- a/src/extension/webview-api/chat-context-processor/types/chat-context/code-context.ts +++ b/src/extension/webview-api/chat-context-processor/types/chat-context/code-context.ts @@ -1,4 +1,6 @@ -export interface CodeChunk { +import type { BaseContextInfo } from './base-context' + +export interface CodeChunk extends BaseContextInfo { code: string language: string diff --git a/src/extension/webview-api/chat-context-processor/types/chat-context/codebase-context.ts b/src/extension/webview-api/chat-context-processor/types/chat-context/codebase-context.ts index 8b37547..5149dc2 100644 --- a/src/extension/webview-api/chat-context-processor/types/chat-context/codebase-context.ts +++ b/src/extension/webview-api/chat-context-processor/types/chat-context/codebase-context.ts @@ -1,6 +1,6 @@ -import type { BaseToolContext } from './base-tool-context' +import type { BaseContextInfo, BaseToolContext } from './base-context' -export interface CodeSnippet { +export interface CodeSnippet extends BaseContextInfo { fileHash: string relativePath: string fullPath: string diff --git a/src/extension/webview-api/chat-context-processor/types/chat-context/doc-context.ts b/src/extension/webview-api/chat-context-processor/types/chat-context/doc-context.ts index bf34e55..1c4060a 100644 --- a/src/extension/webview-api/chat-context-processor/types/chat-context/doc-context.ts +++ b/src/extension/webview-api/chat-context-processor/types/chat-context/doc-context.ts @@ -1,9 +1,15 @@ -export interface DocInfo { +import type { BaseContextInfo } from './base-context' + +export interface DocInfo extends BaseContextInfo { content: string path: string // file path or url } +export interface DocSiteName extends BaseContextInfo { + name: string +} + export interface DocContext { - allowSearchDocSiteNames: string[] + allowSearchDocSiteNames: DocSiteName[] relevantDocs: DocInfo[] } diff --git a/src/extension/webview-api/chat-context-processor/types/chat-context/file-context.ts b/src/extension/webview-api/chat-context-processor/types/chat-context/file-context.ts index 5d9995a..8b8f9ab 100644 --- a/src/extension/webview-api/chat-context-processor/types/chat-context/file-context.ts +++ b/src/extension/webview-api/chat-context-processor/types/chat-context/file-context.ts @@ -1,8 +1,13 @@ -import type { FileInfo, FolderInfo } from '@extension/file-utils/traverse-fs' +import type { + FileInfo as IFileInfo, + FolderInfo as IFolderInfo +} from '@extension/file-utils/traverse-fs' -export type { FileInfo, FolderInfo } +import type { BaseContextInfo } from './base-context' -export interface ImageInfo { +export interface FileInfo extends BaseContextInfo, IFileInfo {} +export interface FolderInfo extends BaseContextInfo, IFolderInfo {} +export interface ImageInfo extends BaseContextInfo { url: string } diff --git a/src/extension/webview-api/chat-context-processor/types/chat-context/git-context.ts b/src/extension/webview-api/chat-context-processor/types/chat-context/git-context.ts index 0f751fd..02899e3 100644 --- a/src/extension/webview-api/chat-context-processor/types/chat-context/git-context.ts +++ b/src/extension/webview-api/chat-context-processor/types/chat-context/git-context.ts @@ -1,4 +1,6 @@ -export interface GitDiff { +import type { BaseContextInfo } from './base-context' + +export interface GitDiff extends BaseContextInfo { /** * @example '.github/workflows/ci.yml' */ @@ -29,7 +31,7 @@ export interface GitDiff { }[] } -export interface GitCommit { +export interface GitCommit extends BaseContextInfo { /** * @example '0bc7f06aa2930c2755c751615cfb2331de41ddb1' */ diff --git a/src/extension/webview-api/chat-context-processor/types/chat-context/index.ts b/src/extension/webview-api/chat-context-processor/types/chat-context/index.ts index c31ebb4..3f64a8a 100644 --- a/src/extension/webview-api/chat-context-processor/types/chat-context/index.ts +++ b/src/extension/webview-api/chat-context-processor/types/chat-context/index.ts @@ -17,7 +17,7 @@ export interface ChatContext { settings: SettingsContext } -export * from './base-tool-context' +export * from './base-context' export * from './code-context' export * from './codebase-context' export * from './conversation' diff --git a/src/extension/webview-api/chat-context-processor/types/chat-context/web-context.ts b/src/extension/webview-api/chat-context-processor/types/chat-context/web-context.ts index d1419ad..8f99723 100644 --- a/src/extension/webview-api/chat-context-processor/types/chat-context/web-context.ts +++ b/src/extension/webview-api/chat-context-processor/types/chat-context/web-context.ts @@ -1,6 +1,6 @@ -import type { BaseToolContext } from './base-tool-context' +import type { BaseContextInfo, BaseToolContext } from './base-context' -export interface WebSearchResult { +export interface WebSearchResult extends BaseContextInfo { url: string title: string content: string diff --git a/src/extension/webview-api/chat-context-processor/utils/get-file-content.ts b/src/extension/webview-api/chat-context-processor/utils/get-file-content.ts index 06eefaa..abeb970 100644 --- a/src/extension/webview-api/chat-context-processor/utils/get-file-content.ts +++ b/src/extension/webview-api/chat-context-processor/utils/get-file-content.ts @@ -1,7 +1,6 @@ +import type { FileInfo } from '@extension/file-utils/traverse-fs' import { VsCodeFS } from '@extension/file-utils/vscode-fs' -import type { FileInfo } from '../types/chat-context' - export const getFileContent = async (fileInfo: FileInfo): Promise => { if (fileInfo.content) { return fileInfo.content diff --git a/src/extension/webview-api/chat-context-processor/utils/searxng-search.ts b/src/extension/webview-api/chat-context-processor/utils/searxng-search.ts index 5cdbc91..7bb6b33 100644 --- a/src/extension/webview-api/chat-context-processor/utils/searxng-search.ts +++ b/src/extension/webview-api/chat-context-processor/utils/searxng-search.ts @@ -1,6 +1,6 @@ import * as cheerio from 'cheerio' -import type { WebSearchResult } from '../types/chat-context' +import { ContextInfoSource, type WebSearchResult } from '../types/chat-context' import { getRandomHeaders } from './fake-request-headers' interface SearxngSearchOptions { @@ -42,7 +42,8 @@ const parseHtml = (htmlContent: string): SearxngResults => { .map((_, element) => ({ title: $(element).find('h3').text().trim(), url: $(element).find('a').attr('href') || '', - content: $(element).find('.content').text().trim() + content: $(element).find('.content').text().trim(), + source: ContextInfoSource.ToolNode })) .get() diff --git a/src/extension/webview-api/controllers/git.controller.ts b/src/extension/webview-api/controllers/git.controller.ts index 48cad36..0da1b6d 100644 --- a/src/extension/webview-api/controllers/git.controller.ts +++ b/src/extension/webview-api/controllers/git.controller.ts @@ -4,9 +4,10 @@ import { getWorkspaceFolder } from '@extension/utils' import { settledPromiseResults } from '@shared/utils/common' import simpleGit, { SimpleGit } from 'simple-git' -import type { - GitCommit, - GitDiff +import { + ContextInfoSource, + type GitCommit, + type GitDiff } from '../chat-context-processor/types/chat-context' import { Controller } from '../types' @@ -36,7 +37,8 @@ export class GitController extends Controller { message: commit.message, diff: this.parseDiff(diff), author: commit.author_name, - date: commit.date + date: commit.date, + source: ContextInfoSource.Editor } }) ) @@ -95,7 +97,8 @@ export class GitController extends Controller { content: `@@ ${content}`, lines: lines.filter(line => line.trim() !== '') } - }) + }), + source: ContextInfoSource.Editor }) }) diff --git a/src/shared/utils/common.ts b/src/shared/utils/common.ts index 2cd2b77..5117990 100644 --- a/src/shared/utils/common.ts +++ b/src/shared/utils/common.ts @@ -3,7 +3,8 @@ export const sleep = (ms: number) => export const removeDuplicates = ( arr: T[], - keys?: (keyof T)[] | ((item: T) => any) + keys?: (keyof T)[] | ((item: T) => any), + prioritySelector?: (a: T, b: T) => T ): T[] => { if (!keys) { return Array.from(new Set(arr)) @@ -14,11 +15,20 @@ export const removeDuplicates = ( ? keys : (item: T) => keys.map(k => item[k]).join('|') - const seen = new Set() - return arr.filter(item => { + const uniqueMap = new Map() + + for (const item of arr) { const key = keyFn(item) - return seen.has(key) ? false : seen.add(key) - }) + if (!uniqueMap.has(key)) { + uniqueMap.set(key, item) + } else if (prioritySelector) { + const existingItem = uniqueMap.get(key)! + const priorityItem = prioritySelector(existingItem, item) + uniqueMap.set(key, priorityItem) + } + } + + return Array.from(uniqueMap.values()) } export const tryParseJSON = (jsonString: string) => { diff --git a/src/webview/components/chat/editor/chat-editor.tsx b/src/webview/components/chat/editor/chat-editor.tsx index 60ebd14..e385148 100644 --- a/src/webview/components/chat/editor/chat-editor.tsx +++ b/src/webview/components/chat/editor/chat-editor.tsx @@ -1,4 +1,4 @@ -import { useEffect, useImperativeHandle, type FC, type Ref } from 'react' +import { useEffect, useId, useImperativeHandle, type FC, type Ref } from 'react' import { AutoFocusPlugin } from '@lexical/react/LexicalAutoFocusPlugin' import { LexicalComposer, @@ -66,13 +66,13 @@ export const ChatEditor: FC = ({ initialConfig, placeholder, autoFocus = false, - conversation, onComplete, onChange, ...otherProps }) => { + const id = useId() const finalInitialConfig: InitialConfigType = { - namespace: `TextComponentEditor-${conversation.id}`, + namespace: `TextComponentEditor-${id}`, // theme: normalTheme, onError, editable: true, @@ -87,7 +87,6 @@ export const ChatEditor: FC = ({ className={className} placeholder={placeholder} autoFocus={autoFocus} - conversation={conversation} onComplete={onComplete} onChange={onChange} {...otherProps} @@ -105,10 +104,6 @@ const ChatEditorInner: FC = ({ onComplete, onChange, - // mention plugin props - conversation, - setConversation, - // div props ...otherProps }) => { @@ -218,10 +213,7 @@ const ChatEditorInner: FC = ({ ErrorBoundary={LexicalErrorBoundary} /> - + {autoFocus && } diff --git a/src/webview/components/chat/editor/chat-input.tsx b/src/webview/components/chat/editor/chat-input.tsx index 3265870..d9a2290 100644 --- a/src/webview/components/chat/editor/chat-input.tsx +++ b/src/webview/components/chat/editor/chat-input.tsx @@ -2,11 +2,20 @@ import { useEffect, useRef, type FC } from 'react' import { tryParseJSON, tryStringifyJSON } from '@shared/utils/common' import { convertToLangchainMessageContents } from '@shared/utils/convert-to-langchain-message-contents' import { getAllTextFromLangchainMessageContents } from '@shared/utils/get-all-text-from-langchain-message-contents' -import { getDefaultConversationAttachments } from '@shared/utils/get-default-conversation-attachments' import { mergeLangchainMessageContents } from '@shared/utils/merge-langchain-message-contents' import { Button } from '@webview/components/ui/button' import { useCloneState } from '@webview/hooks/use-clone-state' -import type { ChatContext, Conversation, FileInfo } from '@webview/types/chat' +import { + AttachmentType, + ContextInfoSource, + type ChatContext, + type Conversation, + type FileInfo +} from '@webview/types/chat' +import { + getAttachmentsFromEditorState, + overrideAttachmentItemsBySource +} from '@webview/utils/attachments' import { cn } from '@webview/utils/common' import { $createParagraphNode, @@ -62,11 +71,9 @@ export const ChatInput: FC = ({ _setConversation ) - const selectedFiles = - conversation.attachments?.fileContext?.selectedFiles ?? [] - - const handleEditorChange = (editorState: EditorState) => { + const handleEditorChange = async (editorState: EditorState) => { const newRichText = tryStringifyJSON(editorState.toJSON()) || '' + setConversation(draft => { if (draft.richText !== newRichText) { draft.richText = newRichText @@ -75,6 +82,10 @@ export const ChatInput: FC = ({ editorState.read(() => $getRoot().getTextContent()) ) ) + draft.attachments = getAttachmentsFromEditorState( + editorState, + draft.attachments + ) } }) } @@ -114,13 +125,21 @@ export const ChatInput: FC = ({ handleSend() } + const selectedFiles = + conversation.attachments?.fileContext?.selectedFiles?.filter( + file => file.source === ContextInfoSource.FileSelector + ) || [] + const handleSelectedFiles = (files: FileInfo[]) => { setConversation(draft => { - if (!draft.attachments) { - draft.attachments = getDefaultConversationAttachments() - } - - draft.attachments.fileContext.selectedFiles = files + draft.attachments = overrideAttachmentItemsBySource( + ContextInfoSource.FileSelector, + draft.attachments, + files.map(file => ({ + type: AttachmentType.Files, + data: file + })) + ) }) } @@ -171,8 +190,6 @@ export const ChatInput: FC = ({ onChange={handleEditorChange} placeholder="Type your message here..." autoFocus={autoFocus} - conversation={conversation} - setConversation={setConversation} className={cn( 'min-h-24 max-h-64 my-2 border overflow-y-auto rounded-lg bg-background shadow-none focus-visible:ring-0', [ChatInputMode.MessageReadonly, ChatInputMode.MessageEdit].includes( diff --git a/src/webview/components/chat/selectors/context-selector.tsx b/src/webview/components/chat/selectors/context-selector.tsx index eae6b58..5998745 100644 --- a/src/webview/components/chat/selectors/context-selector.tsx +++ b/src/webview/components/chat/selectors/context-selector.tsx @@ -3,10 +3,11 @@ import React, { useState } from 'react' import { ImageIcon } from '@radix-ui/react-icons' import { getDefaultConversationAttachments } from '@shared/utils/get-default-conversation-attachments' import { Button } from '@webview/components/ui/button' -import type { - ChatContext, - Conversation, - ModelOption +import { + ContextInfoSource, + type ChatContext, + type Conversation, + type ModelOption } from '@webview/types/chat' import type { Updater } from 'use-immer' @@ -60,7 +61,8 @@ export const ContextSelector: React.FC = ({ } draft.attachments.fileContext.selectedImages.push({ - url: 'https://example.com/image.jpg' + url: 'https://example.com/image.jpg', + source: ContextInfoSource.FileSelector }) }) } diff --git a/src/webview/components/chat/selectors/file-selector/index.tsx b/src/webview/components/chat/selectors/file-selector/index.tsx index 28f7867..85f15c9 100644 --- a/src/webview/components/chat/selectors/file-selector/index.tsx +++ b/src/webview/components/chat/selectors/file-selector/index.tsx @@ -20,7 +20,7 @@ import { import { useFilesSearch } from '@webview/hooks/chat/use-files-search' import { useControllableState } from '@webview/hooks/use-controllable-state' import { useKeyboardNavigation } from '@webview/hooks/use-keyboard-navigation' -import type { FileInfo } from '@webview/types/chat' +import { ContextInfoSource, type FileInfo } from '@webview/types/chat' import { useEvent } from 'react-use' import { FileListView } from './file-list-view' @@ -64,6 +64,11 @@ export const FileSelector: React.FC = ({ const inputRef = useRef(null) const tabRefs = useRef<(HTMLButtonElement | null)[]>([]) + const filteredFilesWithSource = filteredFiles.map(file => ({ + ...file, + source: ContextInfoSource.FileSelector + })) + const { handleKeyDown, setFocusedIndex } = useKeyboardNavigation({ itemCount: tabOptions.length, itemRefs: tabRefs, @@ -157,14 +162,14 @@ export const FileSelector: React.FC = ({
void + onSelect: (option: MentionOption) => void open?: boolean onOpenChange?: (open: boolean) => void onCloseWithoutSelect?: () => void @@ -102,12 +96,8 @@ export const MentionSelector: React.FC = ({ const handleSelect = (option: MentionOption) => { if (isFlattened) { - if (option.mentionStrategy) { - onSelect({ - name: option.label, - strategy: option.mentionStrategy, - strategyAddData: option.data || { label: option.label } - }) + if (option.data) { + onSelect(option) } setIsFlattened(false) setIsOpen(false) @@ -119,12 +109,8 @@ export const MentionSelector: React.FC = ({ setOptionsStack(prevStack => [...prevStack, option.children || []]) onCloseWithoutSelect?.() } else { - if (option.mentionStrategy) { - onSelect({ - name: option.label, - strategy: option.mentionStrategy, - strategyAddData: option.data || { label: option.label } - }) + if (option.data) { + onSelect(option) } setIsOpen(false) } diff --git a/src/webview/hooks/chat/use-mention-options.tsx b/src/webview/hooks/chat/use-mention-options.tsx index 8384f61..20415fc 100644 --- a/src/webview/hooks/chat/use-mention-options.tsx +++ b/src/webview/hooks/chat/use-mention-options.tsx @@ -13,16 +13,14 @@ import { import { MentionFilePreview } from '@webview/components/chat/selectors/mention-selector/files/mention-file-preview' import { MentionFolderPreview } from '@webview/components/chat/selectors/mention-selector/folders/mention-folder-preview' import { FileIcon as FileIcon2 } from '@webview/components/file-icon' -import { RelevantCodeSnippetsMentionStrategy } from '@webview/lexical/mentions/codebase/relevant-code-snippets-mention-strategy' -import { AllowSearchDocSiteNamesToolMentionStrategy } from '@webview/lexical/mentions/docs/allow-search-doc-site-names-mention-strategy' -import { SelectedFilesMentionStrategy } from '@webview/lexical/mentions/files/selected-files-mention-strategy' -import { SelectedFoldersMentionStrategy } from '@webview/lexical/mentions/folders/selected-folders-mention-strategy' -import { GitCommitsMentionStrategy } from '@webview/lexical/mentions/git/git-commits-mention-strategy' -import { GitDiffsMentionStrategy } from '@webview/lexical/mentions/git/git-diffs-mention-strategy' -import { EnableWebToolMentionStrategy } from '@webview/lexical/mentions/web/enable-web-tool-mention-strategy' import { - MentionCategory, + AttachmentType, + ContextInfoSource, SearchSortStrategy, + type DocSiteName, + type FileInfo, + type FolderInfo, + type GitCommit, type MentionOption } from '@webview/types/chat' import { getFileNameFromPath } from '@webview/utils/path' @@ -43,11 +41,10 @@ export const useMentionOptions = () => { ({ id: `file#${file.fullPath}`, label: getFileNameFromPath(file.fullPath), - category: MentionCategory.Files, - mentionStrategy: new SelectedFilesMentionStrategy(), + type: AttachmentType.Files, searchKeywords: [file.relativePath], searchSortStrategy: SearchSortStrategy.EndMatch, - data: file, + data: { ...file, source: ContextInfoSource.Editor } satisfies FileInfo, itemLayoutProps: { icon: ( @@ -65,11 +62,13 @@ export const useMentionOptions = () => { ({ id: `folder#${folder.fullPath}`, label: getFileNameFromPath(folder.fullPath), - category: MentionCategory.Folders, - mentionStrategy: new SelectedFoldersMentionStrategy(), + type: AttachmentType.Folders, searchKeywords: [folder.relativePath], searchSortStrategy: SearchSortStrategy.EndMatch, - data: folder, + data: { + ...folder, + source: ContextInfoSource.Editor + } satisfies FolderInfo, itemLayoutProps: { icon: ( <> @@ -94,10 +93,12 @@ export const useMentionOptions = () => { ({ id: `git-commit#${commit.sha}`, label: commit.message, - category: MentionCategory.Git, - mentionStrategy: new GitCommitsMentionStrategy(), + type: AttachmentType.GitCommit, searchKeywords: [commit.sha, commit.message], - data: commit, + data: { + ...commit, + source: ContextInfoSource.Editor + } satisfies GitCommit, itemLayoutProps: { icon: , label: commit.message, @@ -106,13 +107,15 @@ export const useMentionOptions = () => { }) satisfies MentionOption ) - const docSitesMentionOptions: MentionOption[] = docSites.map(site => ({ + const docSiteNamesMentionOptions: MentionOption[] = docSites.map(site => ({ id: `doc-site#${site.id}`, label: site.name, - category: MentionCategory.Docs, - mentionStrategy: new AllowSearchDocSiteNamesToolMentionStrategy(), + type: AttachmentType.Docs, searchKeywords: [site.name, site.url], - data: site, + data: { + name: site.name, + source: ContextInfoSource.Editor + } satisfies DocSiteName, itemLayoutProps: { icon: , label: site.name, @@ -124,7 +127,7 @@ export const useMentionOptions = () => { { id: 'files', label: 'Files', - category: MentionCategory.Files, + type: AttachmentType.Files, searchKeywords: ['files'], children: filesMentionOptions, itemLayoutProps: { @@ -135,7 +138,7 @@ export const useMentionOptions = () => { { id: 'folders', label: 'Folders', - category: MentionCategory.Folders, + type: AttachmentType.Folders, searchKeywords: ['folders'], children: foldersMentionOptions, itemLayoutProps: { @@ -146,20 +149,18 @@ export const useMentionOptions = () => { { id: 'code', label: 'Code', - category: MentionCategory.Code, + type: AttachmentType.Code, searchKeywords: ['code'], itemLayoutProps: { icon: , label: 'Code' } - // mentionStrategy: new CodeChunksMentionStrategy() }, { id: 'web', label: 'Web', - category: MentionCategory.Web, + type: AttachmentType.Web, searchKeywords: ['web'], - mentionStrategy: new EnableWebToolMentionStrategy(), itemLayoutProps: { icon: , label: 'Web' @@ -168,18 +169,17 @@ export const useMentionOptions = () => { { id: 'docs', label: 'Docs', - category: MentionCategory.Docs, + type: AttachmentType.Docs, searchKeywords: ['docs'], itemLayoutProps: { icon: , label: 'Docs' }, - children: docSitesMentionOptions + children: docSiteNamesMentionOptions }, { id: 'git', label: 'Git', - category: MentionCategory.Git, searchKeywords: ['git'], itemLayoutProps: { icon: , @@ -189,9 +189,8 @@ export const useMentionOptions = () => { { id: 'git#diff', label: 'Diff (Diff of Working State)', - category: MentionCategory.Git, + type: AttachmentType.GitDiff, searchKeywords: ['diff'], - mentionStrategy: new GitDiffsMentionStrategy(), itemLayoutProps: { icon: , label: 'Diff (Diff of Working State)' @@ -200,9 +199,8 @@ export const useMentionOptions = () => { { id: 'git#pull-request', label: 'PR (Diff with Main Branch)', - category: MentionCategory.Git, + type: AttachmentType.GitPr, searchKeywords: ['pull request', 'pr', 'diff'], - mentionStrategy: new GitDiffsMentionStrategy(), itemLayoutProps: { icon: , label: 'PR (Diff with Main Branch)' @@ -214,9 +212,8 @@ export const useMentionOptions = () => { { id: 'codebase', label: 'Codebase', - category: MentionCategory.Codebase, + type: AttachmentType.Codebase, searchKeywords: ['codebase'], - mentionStrategy: new RelevantCodeSnippetsMentionStrategy(), itemLayoutProps: { icon: , label: 'Codebase' diff --git a/src/webview/lexical/hooks/use-mention-manager.ts b/src/webview/lexical/hooks/use-mention-manager.ts deleted file mode 100644 index e2d0fb1..0000000 --- a/src/webview/lexical/hooks/use-mention-manager.ts +++ /dev/null @@ -1,58 +0,0 @@ -import { getDefaultConversationAttachments } from '@shared/utils/get-default-conversation-attachments' -import type { - Attachments, - Conversation, - IMentionStrategy -} from '@webview/types/chat' -import { logger } from '@webview/utils/logger' -import type { Updater } from 'use-immer' - -export interface UseMentionManagerProps { - conversation: Conversation - setConversation: Updater -} - -export function useMentionManager(props: UseMentionManagerProps) { - const { conversation, setConversation } = props - - const currentAttachments = - conversation.attachments || getDefaultConversationAttachments() - - const updateCurrentAttachments = (attachments: Partial) => { - setConversation(draft => { - if (!draft.attachments) { - draft.attachments = getDefaultConversationAttachments() - } - - draft.attachments = { - ...draft.attachments, - ...attachments - } - }) - } - - const addMention = async ({ - strategy, - strategyAddData - }: { - strategy: IMentionStrategy - strategyAddData: any - }) => { - try { - if (strategy) { - const updatedAttachments = - await strategy.buildNewAttachmentsAfterAddMention( - strategyAddData, - currentAttachments - ) - updateCurrentAttachments(updatedAttachments) - } - } catch (error) { - logger.warn('Error adding mention:', error) - } - } - - return { - addMention - } -} diff --git a/src/webview/lexical/mentions/code/code-chunks-mention-strategy.ts b/src/webview/lexical/mentions/code/code-chunks-mention-strategy.ts deleted file mode 100644 index 5a3bb06..0000000 --- a/src/webview/lexical/mentions/code/code-chunks-mention-strategy.ts +++ /dev/null @@ -1,33 +0,0 @@ -import { removeDuplicates } from '@shared/utils/common' -import { - MentionCategory, - type Attachments, - type CodeChunk, - type IMentionStrategy -} from '@webview/types/chat' - -export class CodeChunksMentionStrategy implements IMentionStrategy { - category = MentionCategory.Code as const - - name = 'CodeChunksMentionStrategy' as const - - async buildNewAttachmentsAfterAddMention( - data: CodeChunk | CodeChunk[], - currentAttachments: Attachments - ): Promise> { - const codeChunks = Array.isArray(data) ? data : [data] - - return { - codeContext: { - ...currentAttachments.codeContext, - codeChunks: removeDuplicates( - [ - ...(currentAttachments.codeContext?.codeChunks || []), - ...codeChunks - ], - ['relativePath', 'code'] - ) - } - } - } -} diff --git a/src/webview/lexical/mentions/codebase/relevant-code-snippets-mention-strategy.ts b/src/webview/lexical/mentions/codebase/relevant-code-snippets-mention-strategy.ts deleted file mode 100644 index fd2e01d..0000000 --- a/src/webview/lexical/mentions/codebase/relevant-code-snippets-mention-strategy.ts +++ /dev/null @@ -1,23 +0,0 @@ -import { - MentionCategory, - type Attachments, - type IMentionStrategy -} from '@webview/types/chat' - -export class RelevantCodeSnippetsMentionStrategy implements IMentionStrategy { - category = MentionCategory.Codebase as const - - name = 'RelevantCodeSnippetsMentionStrategy' as const - - async buildNewAttachmentsAfterAddMention( - data: undefined, - currentAttachments: Attachments - ): Promise> { - return { - codebaseContext: { - ...currentAttachments.codebaseContext, - enableTool: true - } - } - } -} diff --git a/src/webview/lexical/mentions/docs/allow-search-doc-site-names-mention-strategy.ts b/src/webview/lexical/mentions/docs/allow-search-doc-site-names-mention-strategy.ts deleted file mode 100644 index ec20db2..0000000 --- a/src/webview/lexical/mentions/docs/allow-search-doc-site-names-mention-strategy.ts +++ /dev/null @@ -1,32 +0,0 @@ -import { removeDuplicates } from '@shared/utils/common' -import { - IMentionStrategy, - MentionCategory, - type Attachments, - type DocSite -} from '@webview/types/chat' - -export class AllowSearchDocSiteNamesToolMentionStrategy - implements IMentionStrategy -{ - category = MentionCategory.Docs as const - - name = 'AllowSearchDocSiteNamesToolMentionStrategy' as const - - async buildNewAttachmentsAfterAddMention( - data: DocSite | DocSite[], - currentAttachments: Attachments - ): Promise> { - const sites = Array.isArray(data) ? data : [data] - - return { - docContext: { - ...currentAttachments.docContext, - allowSearchDocSiteNames: removeDuplicates([ - ...(currentAttachments.docContext?.allowSearchDocSiteNames || []), - ...sites.map(site => site.name) - ]) - } - } - } -} diff --git a/src/webview/lexical/mentions/files/selected-files-mention-strategy.ts b/src/webview/lexical/mentions/files/selected-files-mention-strategy.ts deleted file mode 100644 index d39e137..0000000 --- a/src/webview/lexical/mentions/files/selected-files-mention-strategy.ts +++ /dev/null @@ -1,30 +0,0 @@ -import { removeDuplicates } from '@shared/utils/common' -import { - IMentionStrategy, - MentionCategory, - type Attachments, - type FileInfo -} from '@webview/types/chat' - -export class SelectedFilesMentionStrategy implements IMentionStrategy { - category = MentionCategory.Files as const - - name = 'SelectedFilesMentionStrategy' as const - - async buildNewAttachmentsAfterAddMention( - data: FileInfo | FileInfo[], - currentAttachments: Attachments - ): Promise> { - const files = Array.isArray(data) ? data : [data] - - return { - fileContext: { - ...currentAttachments.fileContext, - selectedFiles: removeDuplicates( - [...(currentAttachments.fileContext?.selectedFiles || []), ...files], - ['fullPath'] - ) - } - } - } -} diff --git a/src/webview/lexical/mentions/files/selected-images-mention-strategy.ts b/src/webview/lexical/mentions/files/selected-images-mention-strategy.ts deleted file mode 100644 index 0917762..0000000 --- a/src/webview/lexical/mentions/files/selected-images-mention-strategy.ts +++ /dev/null @@ -1,30 +0,0 @@ -import { removeDuplicates } from '@shared/utils/common' -import { - IMentionStrategy, - MentionCategory, - type Attachments, - type ImageInfo -} from '@webview/types/chat' - -export class SelectedImagesMentionStrategy implements IMentionStrategy { - category = MentionCategory.Files as const - - name = 'SelectedImagesMentionStrategy' as const - - async buildNewAttachmentsAfterAddMention( - data: ImageInfo | ImageInfo[], - currentAttachments: Attachments - ): Promise> { - const imgs = Array.isArray(data) ? data : [data] - - return { - fileContext: { - ...currentAttachments.fileContext, - selectedImages: removeDuplicates( - [...(currentAttachments.fileContext?.selectedImages || []), ...imgs], - ['url'] - ) - } - } - } -} diff --git a/src/webview/lexical/mentions/folders/selected-folders-mention-strategy.ts b/src/webview/lexical/mentions/folders/selected-folders-mention-strategy.ts deleted file mode 100644 index ab3b351..0000000 --- a/src/webview/lexical/mentions/folders/selected-folders-mention-strategy.ts +++ /dev/null @@ -1,33 +0,0 @@ -import { removeDuplicates } from '@shared/utils/common' -import { - IMentionStrategy, - MentionCategory, - type Attachments, - type FolderInfo -} from '@webview/types/chat' - -export class SelectedFoldersMentionStrategy implements IMentionStrategy { - category = MentionCategory.Folders as const - - name = 'SelectedFoldersMentionStrategy' as const - - async buildNewAttachmentsAfterAddMention( - data: FolderInfo | FolderInfo[], - currentAttachments: Attachments - ): Promise> { - const folders = Array.isArray(data) ? data : [data] - - return { - fileContext: { - ...currentAttachments.fileContext, - selectedFolders: removeDuplicates( - [ - ...(currentAttachments.fileContext?.selectedFolders || []), - ...folders - ], - ['fullPath'] - ) - } - } - } -} diff --git a/src/webview/lexical/mentions/git/git-commits-mention-strategy.ts b/src/webview/lexical/mentions/git/git-commits-mention-strategy.ts deleted file mode 100644 index 1db2d63..0000000 --- a/src/webview/lexical/mentions/git/git-commits-mention-strategy.ts +++ /dev/null @@ -1,30 +0,0 @@ -import { removeDuplicates } from '@shared/utils/common' -import { - MentionCategory, - type Attachments, - type GitCommit, - type IMentionStrategy -} from '@webview/types/chat' - -export class GitCommitsMentionStrategy implements IMentionStrategy { - category = MentionCategory.Git as const - - name = 'GitCommitsMentionStrategy' as const - - async buildNewAttachmentsAfterAddMention( - data: GitCommit | GitCommit[], - currentAttachments: Attachments - ): Promise> { - const commits = Array.isArray(data) ? data : [data] - - return { - gitContext: { - ...currentAttachments.gitContext, - gitCommits: removeDuplicates( - [...(currentAttachments.gitContext?.gitCommits || []), ...commits], - ['sha'] - ) - } - } - } -} diff --git a/src/webview/lexical/mentions/git/git-diffs-mention-strategy.ts b/src/webview/lexical/mentions/git/git-diffs-mention-strategy.ts deleted file mode 100644 index 0b9b89f..0000000 --- a/src/webview/lexical/mentions/git/git-diffs-mention-strategy.ts +++ /dev/null @@ -1,31 +0,0 @@ -import { removeDuplicates } from '@shared/utils/common' -import { - MentionCategory, - type Attachments, - type GitDiff, - type IMentionStrategy -} from '@webview/types/chat' - -export class GitDiffsMentionStrategy implements IMentionStrategy { - category = MentionCategory.Git as const - - name = 'GitDiffsMentionStrategy' as const - - async buildNewAttachmentsAfterAddMention( - data: GitDiff | GitDiff[], - currentAttachments: Attachments - ): Promise> { - const diffs = Array.isArray(data) ? data : [data] - - return { - gitContext: { - ...currentAttachments.gitContext, - gitDiffs: removeDuplicates( - [...(currentAttachments.gitContext?.gitDiffs || []), ...diffs], - diff => - `${diff.from}|${diff.to}|${diff.chunks.map(chunk => chunk.content).join('|')}` - ) - } - } - } -} diff --git a/src/webview/lexical/mentions/web/enable-web-tool-mention-strategy.ts b/src/webview/lexical/mentions/web/enable-web-tool-mention-strategy.ts deleted file mode 100644 index 9822773..0000000 --- a/src/webview/lexical/mentions/web/enable-web-tool-mention-strategy.ts +++ /dev/null @@ -1,23 +0,0 @@ -import { - IMentionStrategy, - MentionCategory, - type Attachments -} from '@webview/types/chat' - -export class EnableWebToolMentionStrategy implements IMentionStrategy { - category = MentionCategory.Web as const - - name = 'EnableWebToolMentionStrategy' as const - - async buildNewAttachmentsAfterAddMention( - data: undefined, - currentAttachments: Attachments - ): Promise> { - return { - webContext: { - ...currentAttachments.webContext, - enableTool: true - } - } - } -} diff --git a/src/webview/lexical/nodes/mention-node.tsx b/src/webview/lexical/nodes/mention-node.tsx index 4318c8c..1642f84 100644 --- a/src/webview/lexical/nodes/mention-node.tsx +++ b/src/webview/lexical/nodes/mention-node.tsx @@ -1,5 +1,6 @@ /* eslint-disable unused-imports/no-unused-vars */ import React from 'react' +import type { AttachmentType } from '@webview/types/chat' import { $applyNodeReplacement, $createTextNode, @@ -18,7 +19,7 @@ import { export type SerializedMentionNode = Spread< { - mentionType: string + mentionType: AttachmentType mentionData: any text: string }, @@ -28,7 +29,9 @@ export type SerializedMentionNode = Spread< function convertMentionElement( domNode: HTMLElement ): DOMConversionOutput | null { - const mentionType = domNode.getAttribute('data-lexical-mention-type') + const mentionType = domNode.getAttribute( + 'data-lexical-mention-type' + ) as AttachmentType const mentionData = domNode.getAttribute('data-lexical-mention-data') const text = domNode.textContent @@ -44,7 +47,7 @@ function convertMentionElement( } export class MentionNode extends DecoratorNode { - __mentionType: string + __mentionType: AttachmentType __mentionData: any @@ -64,7 +67,7 @@ export class MentionNode extends DecoratorNode { } constructor( - mentionType: string, + mentionType: AttachmentType, mentionData: any, text: string, key?: NodeKey @@ -199,7 +202,7 @@ export class MentionNode extends DecoratorNode { } export function $createMentionNode( - mentionType: string, + mentionType: AttachmentType, mentionData: any, text: string ): MentionNode { diff --git a/src/webview/lexical/plugins/mention-plugin.tsx b/src/webview/lexical/plugins/mention-plugin.tsx index ae16b03..61afd87 100644 --- a/src/webview/lexical/plugins/mention-plugin.tsx +++ b/src/webview/lexical/plugins/mention-plugin.tsx @@ -1,11 +1,8 @@ import React, { useState, type FC } from 'react' import { useLexicalComposerContext } from '@lexical/react/LexicalComposerContext' -import { - MentionSelector, - type SelectedMentionStrategy -} from '@webview/components/chat/selectors/mention-selector/mention-selector' +import { MentionSelector } from '@webview/components/chat/selectors/mention-selector/mention-selector' import { useMentionOptions } from '@webview/hooks/chat/use-mention-options' -import type { IMentionStrategy } from '@webview/types/chat' +import type { MentionOption } from '@webview/types/chat' import { $createTextNode, $getSelection, @@ -14,19 +11,14 @@ import { } from 'lexical' import { useEditorCommands } from '../hooks/use-editor-commands' -import { - useMentionManager, - type UseMentionManagerProps -} from '../hooks/use-mention-manager' import { useMentionSearch } from '../hooks/use-mention-search' import { useNearestMentionPosition } from '../hooks/use-nearest-mention-position' import { $createMentionNode } from '../nodes/mention-node' -export interface MentionPluginProps extends UseMentionManagerProps {} +export interface MentionPluginProps {} export const MentionPlugin: FC = props => { const [editor] = useLexicalComposerContext() - const { addMention } = useMentionManager(props) const [isOpen, setIsOpen] = useState(false) const mentionPosition = useNearestMentionPosition(editor) const mentionOptions = useMentionOptions() @@ -38,11 +30,7 @@ export const MentionPlugin: FC = props => { useEditorCommands(editor, isOpen, setIsOpen) - const handleMentionSelect = ({ - name, - strategy, - strategyAddData - }: SelectedMentionStrategy) => { + const handleMentionSelect = (option: MentionOption) => { setIsOpen(false) setSearchQuery('') @@ -50,15 +38,12 @@ export const MentionPlugin: FC = props => { const selection = $getSelection() if ($isRangeSelection(selection)) { insertMention({ - name, selection, - strategy, - strategyAddData, + option, searchQuery }) } }) - addMention({ strategy, strategyAddData }) } const handleCloseWithoutSelect = () => { @@ -93,16 +78,12 @@ export const MentionPlugin: FC = props => { } const insertMention = ({ - name, selection, - strategy, - strategyAddData, + option, searchQuery }: { - name: string selection: RangeSelection - strategy: IMentionStrategy - strategyAddData: any + option: MentionOption searchQuery: string }) => { // Delete the @ symbol and the search query @@ -111,13 +92,11 @@ const insertMention = ({ selection.focus.offset = anchorOffset selection.removeText() + if (!option.type) throw new Error('Mention option type is required') + // Create and insert the mention node - const mentionText = `@${name}` - const mentionNode = $createMentionNode( - strategy.category, - strategyAddData, - mentionText - ) + const mentionText = `@${option.label}` + const mentionNode = $createMentionNode(option.type, option.data, mentionText) selection.insertNodes([mentionNode]) // Insert a space after the mention node diff --git a/src/webview/types/chat.ts b/src/webview/types/chat.ts index 77ca2cf..a4a4360 100644 --- a/src/webview/types/chat.ts +++ b/src/webview/types/chat.ts @@ -23,8 +23,7 @@ export enum SearchSortStrategy { export interface MentionOption { id: string label: string - category: MentionCategory - mentionStrategy?: IMentionStrategy + type?: AttachmentType searchKeywords?: string[] searchSortStrategy?: SearchSortStrategy children?: MentionOption[] @@ -34,18 +33,26 @@ export interface MentionOption { customRenderPreview?: FC } -export enum MentionCategory { +export enum AttachmentType { Files = 'Files', Folders = 'Folders', + Images = 'Images', Code = 'Code', Web = 'Web', Docs = 'Docs', - Git = 'Git', + GitDiff = 'GitDiff', + GitCommit = 'GitCommit', + GitPr = 'GitPr', Codebase = 'Codebase' } +export type AttachmentItem = { + type: AttachmentType + data: any +} + export interface IMentionStrategy { - readonly category: MentionCategory + readonly category: AttachmentType readonly name: string buildLexicalNodeAfterAddMention?: ( diff --git a/src/webview/utils/attachments.ts b/src/webview/utils/attachments.ts new file mode 100644 index 0000000..d53008e --- /dev/null +++ b/src/webview/utils/attachments.ts @@ -0,0 +1,330 @@ +import { removeDuplicates } from '@shared/utils/common' +import { getDefaultConversationAttachments } from '@shared/utils/get-default-conversation-attachments' +import { $isMentionNode } from '@webview/lexical/nodes/mention-node' +import { + AttachmentItem, + Attachments, + AttachmentType, + ContextInfoSource, + type BaseContextInfo, + type GitDiff +} from '@webview/types/chat' +import { produce } from 'immer' +import { + $getRoot, + $isElementNode, + type EditorState, + type LexicalNode +} from 'lexical' + +const deduplicateAttachments = ( + attachments: Attachments, + onlyDeduplicateTypes?: Set, + priorityRemoveSource: Set = new Set([ + ContextInfoSource.Editor + ]) +): Attachments => + produce(attachments, draft => { + const prioritySelector = (a: T, b: T) => + priorityRemoveSource?.has(a.source) ? b : a + + if ( + !onlyDeduplicateTypes || + onlyDeduplicateTypes?.has(AttachmentType.Files) + ) { + draft.fileContext.selectedFiles = removeDuplicates( + draft.fileContext.selectedFiles, + ['fullPath'], + prioritySelector + ) + } + + if ( + !onlyDeduplicateTypes || + onlyDeduplicateTypes?.has(AttachmentType.Folders) + ) { + draft.fileContext.selectedFolders = removeDuplicates( + draft.fileContext.selectedFolders, + ['fullPath'], + prioritySelector + ) + } + + if ( + !onlyDeduplicateTypes || + onlyDeduplicateTypes?.has(AttachmentType.Images) + ) { + draft.fileContext.selectedImages = removeDuplicates( + draft.fileContext.selectedImages, + ['url'], + prioritySelector + ) + } + + if ( + !onlyDeduplicateTypes || + onlyDeduplicateTypes?.has(AttachmentType.GitCommit) + ) { + draft.gitContext.gitCommits = removeDuplicates( + draft.gitContext.gitCommits, + ['sha'], + prioritySelector + ) + } + + if ( + !onlyDeduplicateTypes || + onlyDeduplicateTypes?.has(AttachmentType.GitDiff) + ) { + draft.gitContext.gitDiffs = removeDuplicates( + draft.gitContext.gitDiffs, + diff => + `${diff.from}|${diff.to}|${diff.chunks.map(chunk => chunk.content).join('|')}`, + prioritySelector + ) + } + + if ( + !onlyDeduplicateTypes || + onlyDeduplicateTypes?.has(AttachmentType.Docs) + ) { + draft.docContext.allowSearchDocSiteNames = removeDuplicates( + draft.docContext.allowSearchDocSiteNames, + ['name'], + prioritySelector + ) + } + + if ( + !onlyDeduplicateTypes || + onlyDeduplicateTypes?.has(AttachmentType.Code) + ) { + draft.codeContext.codeChunks = removeDuplicates( + draft.codeContext.codeChunks, + ['relativePath', 'code'], + prioritySelector + ) + } + }) + +export const addAttachmentItems = ( + currentAttachments: Attachments | undefined, + newItems: AttachmentItem[], + priorityRemoveSource?: Set +): Attachments => { + const deduplicateAttachmentTypes = new Set() + + const result = produce( + currentAttachments || getDefaultConversationAttachments(), + draft => { + newItems.forEach(item => { + switch (item.type) { + case AttachmentType.Files: + deduplicateAttachmentTypes.add(AttachmentType.Files) + draft.fileContext.selectedFiles.push(item.data) + break + case AttachmentType.Folders: + deduplicateAttachmentTypes.add(AttachmentType.Folders) + draft.fileContext.selectedFolders.push(item.data) + break + case AttachmentType.Images: + deduplicateAttachmentTypes.add(AttachmentType.Images) + draft.fileContext.selectedImages.push(item.data) + break + case AttachmentType.GitCommit: + deduplicateAttachmentTypes.add(AttachmentType.GitCommit) + draft.gitContext.gitCommits.push(item.data) + break + case AttachmentType.GitDiff: + case AttachmentType.GitPr: + deduplicateAttachmentTypes.add(AttachmentType.GitDiff) + draft.gitContext.gitDiffs.push(item.data) + break + case AttachmentType.Web: + draft.webContext.enableTool = true + break + case AttachmentType.Docs: + deduplicateAttachmentTypes.add(AttachmentType.Docs) + draft.docContext.allowSearchDocSiteNames.push(item.data) + break + case AttachmentType.Code: + deduplicateAttachmentTypes.add(AttachmentType.Code) + draft.codeContext.codeChunks.push(item.data) + break + case AttachmentType.Codebase: + draft.codebaseContext.enableTool = true + break + default: + break + } + }) + } + ) + + return deduplicateAttachments( + result, + deduplicateAttachmentTypes, + priorityRemoveSource + ) +} + +export const removeAttachmentItems = ( + currentAttachments: Attachments | undefined, + itemsToRemove: AttachmentItem[], + priorityRemoveSource?: Set +): Attachments => { + if (!currentAttachments) return getDefaultConversationAttachments() + + const deduplicateAttachmentTypes = new Set() + + const result = produce(currentAttachments, draft => { + itemsToRemove.forEach(item => { + switch (item.type) { + case AttachmentType.Files: + deduplicateAttachmentTypes.add(AttachmentType.Files) + draft.fileContext.selectedFiles = + draft.fileContext.selectedFiles.filter( + file => file.fullPath !== item.data.fullPath + ) + break + case AttachmentType.Folders: + deduplicateAttachmentTypes.add(AttachmentType.Folders) + draft.fileContext.selectedFolders = + draft.fileContext.selectedFolders.filter( + folder => folder.fullPath !== item.data.fullPath + ) + break + case AttachmentType.Images: + deduplicateAttachmentTypes.add(AttachmentType.Images) + draft.fileContext.selectedImages = + draft.fileContext.selectedImages.filter( + image => image.url !== item.data.url + ) + break + case AttachmentType.GitCommit: + deduplicateAttachmentTypes.add(AttachmentType.GitCommit) + draft.gitContext.gitCommits = draft.gitContext.gitCommits.filter( + commit => commit.sha !== item.data.sha + ) + break + case AttachmentType.GitDiff: + case AttachmentType.GitPr: + deduplicateAttachmentTypes.add(AttachmentType.GitDiff) + draft.gitContext.gitDiffs = draft.gitContext.gitDiffs.filter( + diff => + `${diff.from}|${diff.to}|${diff.chunks.map(chunk => chunk.content).join('|')}` !== + `${item.data.from}|${item.data.to}|${(item.data as GitDiff).chunks.map(chunk => chunk.content).join('|')}` + ) + break + case AttachmentType.Docs: + deduplicateAttachmentTypes.add(AttachmentType.Docs) + draft.docContext.allowSearchDocSiteNames = + draft.docContext.allowSearchDocSiteNames.filter( + name => name !== item.data + ) + break + case AttachmentType.Code: + deduplicateAttachmentTypes.add(AttachmentType.Code) + draft.codeContext.codeChunks = draft.codeContext.codeChunks.filter( + chunk => + !( + chunk.relativePath === item.data.relativePath && + chunk.code === item.data.code + ) + ) + break + default: + break + } + }) + }) + + return deduplicateAttachments( + result, + deduplicateAttachmentTypes, + priorityRemoveSource + ) +} + +export const overrideAttachmentItemsBySource = ( + activeSource: ContextInfoSource, + currentAttachments: Attachments | undefined, + newItems: AttachmentItem[], + priorityRemoveSource?: Set +): Attachments => { + let cleanedCurrentAttachments = currentAttachments + + if (!currentAttachments) { + cleanedCurrentAttachments = getDefaultConversationAttachments() + } else { + cleanedCurrentAttachments = produce(currentAttachments, draft => { + // remove the source items + draft.fileContext.selectedFiles = draft.fileContext.selectedFiles.filter( + file => file.source !== activeSource + ) + + draft.fileContext.selectedFolders = + draft.fileContext.selectedFolders.filter( + folder => folder.source !== activeSource + ) + + draft.fileContext.selectedImages = + draft.fileContext.selectedImages.filter( + image => image.source !== activeSource + ) + + draft.gitContext.gitCommits = draft.gitContext.gitCommits.filter( + commit => commit.source !== activeSource + ) + + draft.gitContext.gitDiffs = draft.gitContext.gitDiffs.filter( + diff => diff.source !== activeSource + ) + + draft.webContext.enableTool = false + + draft.docContext.allowSearchDocSiteNames = + draft.docContext.allowSearchDocSiteNames.filter( + name => name.source !== activeSource + ) + + draft.codeContext.codeChunks = draft.codeContext.codeChunks.filter( + chunk => chunk.source !== activeSource + ) + + draft.codebaseContext.enableTool = false + }) + } + + return addAttachmentItems( + cleanedCurrentAttachments, + newItems, + priorityRemoveSource + ) +} + +export const getAttachmentsFromEditorState = ( + editorState: EditorState, + currentAttachments: Attachments | undefined +): Attachments => + editorState.read(() => { + const root = $getRoot() + const attachmentItems: AttachmentItem[] = [] + + const traverseNodes = (node: LexicalNode) => { + if ($isMentionNode(node)) { + const { mentionType, mentionData } = node.exportJSON() + attachmentItems.push({ type: mentionType, data: mentionData }) + } else if ($isElementNode(node)) { + node.getChildren().forEach(traverseNodes) + } + } + + traverseNodes(root) + + return overrideAttachmentItemsBySource( + ContextInfoSource.Editor, + currentAttachments, + attachmentItems + ) + })