feat: add source property for context

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, {

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
         }))
       )
 

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
       }
     ]
 

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
       }))
     ]
   })

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
+}

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
 

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

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[]
 }

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
 }
 

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'
    */

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'

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

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<string> => {
   if (fileInfo.content) {
     return fileInfo.content

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()
 

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
       })
     })
 

src/shared/utils/common.ts

@@ -3,7 +3,8 @@ export const sleep = (ms: number) =>
 
 export const removeDuplicates = <T>(
   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 = <T>(
       ? keys
       : (item: T) => keys.map(k => item[k]).join('|')
 
-  const seen = new Set<string>()
-  return arr.filter(item => {
+  const uniqueMap = new Map<string, T>()
+
+  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) => {

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<ChatEditorProps> = ({
   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<ChatEditorProps> = ({
         className={className}
         placeholder={placeholder}
         autoFocus={autoFocus}
-        conversation={conversation}
         onComplete={onComplete}
         onChange={onChange}
         {...otherProps}
@@ -105,10 +104,6 @@ const ChatEditorInner: FC<ChatEditorProps> = ({
   onComplete,
   onChange,
 
-  // mention plugin props
-  conversation,
-  setConversation,
-
   // div props
   ...otherProps
 }) => {
@@ -218,10 +213,7 @@ const ChatEditorInner: FC<ChatEditorProps> = ({
         ErrorBoundary={LexicalErrorBoundary}
       />
       <OnChangePlugin onChange={onChange!} />
-      <MentionPlugin
-        conversation={conversation}
-        setConversation={setConversation}
-      />
+      <MentionPlugin />
       <HistoryPlugin />
       {autoFocus && <AutoFocusPlugin defaultSelection="rootEnd" />}
       <TabIndentationPlugin />