feat(ui): add repository select in Answer Engine (#3619)

* feat(ui): add repository select in Answer Engine * update: code query * [autofix.ci] apply automated fixes * update * update: rename * update * update * update: rename * update: rename * [autofix.ci] apply automated fixes * update --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
TabbyML · Dec 27, 2024 · fcb1783 · fcb1783
1 parent 84ab6f1
commit fcb1783
Show file tree

Hide file tree

Showing 18 changed files with 532 additions and 256 deletions.
diff --git a/ee/tabby-ui/app/(home)/page.tsx b/ee/tabby-ui/app/(home)/page.tsx
@@ -10,12 +10,16 @@ import { useStore } from 'zustand'
 import { SESSION_STORAGE_KEY } from '@/lib/constants'
 import { useMe } from '@/lib/hooks/use-me'
 import { useSelectedModel } from '@/lib/hooks/use-models'
+import { useSelectedRepository } from '@/lib/hooks/use-repositories'
 import {
   useIsChatEnabled,
   useIsFetchingServerInfo
 } from '@/lib/hooks/use-server-info'
 import { setThreadsPageNo } from '@/lib/stores/answer-engine-store'
-import { updateSelectedModel } from '@/lib/stores/chat-actions'
+import {
+  updateSelectedModel,
+  updateSelectedRepoSourceId
+} from '@/lib/stores/chat-actions'
 import {
   clearHomeScrollPosition,
   setHomeScrollPosition,
@@ -55,7 +59,8 @@ function MainPanel() {
   })
   const scrollY = useStore(useScrollStore, state => state.homePage)
 
-  const { selectedModel, isModelLoading, models } = useSelectedModel()
+  const { selectedModel, isFetchingModels, models } = useSelectedModel()
+  const { selectedRepository, isFetchingRepositories } = useSelectedRepository()
 
   const showMainSection = !!data?.me || !isFetchingServerInfo
 
@@ -85,6 +90,10 @@ function MainPanel() {
     updateSelectedModel(model)
   }
 
+  const onSelectedRepo = (sourceId: string | undefined) => {
+    updateSelectedRepoSourceId(sourceId)
+  }
+
   const onSearch = (question: string, ctx?: ThreadRunContexts) => {
     setIsLoading(true)
     sessionStorage.setItem(SESSION_STORAGE_KEY.SEARCH_INITIAL_MSG, question)
@@ -157,8 +166,12 @@ function MainPanel() {
                   contextInfo={contextInfoData?.contextInfo}
                   fetchingContextInfo={fetchingContextInfo}
                   modelName={selectedModel}
-                  onModelSelect={handleSelectModel}
-                  isModelLoading={isModelLoading}
+                  onSelectModel={handleSelectModel}
+                  repoSourceId={selectedRepository?.sourceId}
+                  onSelectRepo={onSelectedRepo}
+                  isInitializingResources={
+                    isFetchingModels || isFetchingRepositories
+                  }
                   models={models}
                 />
               </AnimationWrapper>

diff --git a/ee/tabby-ui/app/search/components/assistant-message-section.tsx b/ee/tabby-ui/app/search/components/assistant-message-section.tsx
@@ -77,7 +77,9 @@ import { DocDetailView } from '@/components/message-markdown/doc-detail-view'
 import { SiteFavicon } from '@/components/site-favicon'
 import { UserAvatar } from '@/components/user-avatar'
 
-import { ConversationMessage, SearchContext, SOURCE_CARD_STYLE } from './search'
+import { SOURCE_CARD_STYLE } from './search'
+import { SearchContext } from './search-context'
+import { ConversationMessage } from './types'
 
 export function AssistantMessageSection({
   className,
@@ -106,7 +108,8 @@ export function AssistantMessageSection({
     fetchingContextInfo,
     onDeleteMessage,
     isThreadOwner,
-    onUpdateMessage
+    onUpdateMessage,
+    repositories
   } = useContext(SearchContext)
 
   const { supportsOnApplyInEditorV2 } = useContext(ChatContext)
@@ -147,7 +150,15 @@ export function AssistantMessageSection({
 
   const IconAnswer = isLoading ? IconSpinner : IconSparkles
 
-  const relevantCodeGitURL = message?.attachment?.code?.[0]?.gitUrl || ''
+  // match gitUrl for clientCode with codeSourceId
+  const clientCodeGitUrl = useMemo(() => {
+    if (!message.codeSourceId || !repositories?.length) return ''
+
+    const target = repositories.find(
+      info => info.sourceId === message.codeSourceId
+    )
+    return target?.gitUrl ?? ''
+  }, [message.codeSourceId, repositories])
 
   const clientCodeContexts: RelevantCodeContext[] = useMemo(() => {
     if (!clientCode?.length) return []
@@ -158,11 +169,11 @@ export function AssistantMessageSection({
           range: getRangeFromAttachmentCode(code),
           filepath: code.filepath || '',
           content: code.content,
-          git_url: relevantCodeGitURL
+          git_url: clientCodeGitUrl
         }
       }) ?? []
     )
-  }, [clientCode, relevantCodeGitURL])
+  }, [clientCode, clientCodeGitUrl])
 
   const serverCodeContexts: RelevantCodeContext[] = useMemo(() => {
     return (
@@ -185,9 +196,9 @@ export function AssistantMessageSection({
   const messageAttachmentClientCode = useMemo(() => {
     return clientCode?.map(o => ({
       ...o,
-      gitUrl: relevantCodeGitURL
+      gitUrl: clientCodeGitUrl
     }))
-  }, [clientCode, relevantCodeGitURL])
+  }, [clientCode, clientCodeGitUrl])
 
   const messageAttachmentDocs = message?.attachment?.doc
   const messageAttachmentCodeLen =

diff --git a/ee/tabby-ui/app/search/components/header.tsx b/ee/tabby-ui/app/search/components/header.tsx
@@ -32,7 +32,7 @@ import { ThemeToggle } from '@/components/theme-toggle'
 import { MyAvatar } from '@/components/user-avatar'
 import UserPanel from '@/components/user-panel'
 
-import { SearchContext } from './search'
+import { SearchContext } from './search-context'
 
 const deleteThreadMutation = graphql(/* GraphQL */ `
   mutation DeleteThread($id: ID!) {

diff --git a/ee/tabby-ui/app/search/components/search-context.ts b/ee/tabby-ui/app/search/components/search-context.ts
@@ -0,0 +1,32 @@
+import { createContext } from 'react'
+
+import {
+  ContextInfo,
+  RepositorySourceListQuery
+} from '@/lib/gql/generates/graphql'
+import { ExtendedCombinedError } from '@/lib/types'
+
+import { ConversationMessage } from './types'
+
+type SearchContextValue = {
+  // flag for initialize the pathname
+  isPathnameInitialized: boolean
+  isLoading: boolean
+  onRegenerateResponse: (id: string) => void
+  onSubmitSearch: (question: string) => void
+  setDevPanelOpen: (v: boolean) => void
+  setConversationIdForDev: (v: string | undefined) => void
+  enableDeveloperMode: boolean
+  contextInfo: ContextInfo | undefined
+  fetchingContextInfo: boolean
+  onDeleteMessage: (id: string) => void
+  isThreadOwner: boolean
+  onUpdateMessage: (
+    message: ConversationMessage
+  ) => Promise<ExtendedCombinedError | undefined>
+  repositories: RepositorySourceListQuery['repositoryList'] | undefined
+}
+
+export const SearchContext = createContext<SearchContextValue>(
+  {} as SearchContextValue
+)
diff --git a/ee/tabby-ui/app/search/components/search.tsx b/ee/tabby-ui/app/search/components/search.tsx
@@ -1,7 +1,6 @@
 'use client'
 
 import {
-  createContext,
   CSSProperties,
   Fragment,
   useEffect,
@@ -27,12 +26,8 @@ import { useEnableDeveloperMode } from '@/lib/experiment-flags'
 import { graphql } from '@/lib/gql/generates'
 import {
   CodeQueryInput,
-  ContextInfo,
   DocQueryInput,
   InputMaybe,
-  Maybe,
-  Message,
-  MessageAttachmentClientCode,
   Role
 } from '@/lib/gql/generates/graphql'
 import { useCopyToClipboard } from '@/lib/hooks/use-copy-to-clipboard'
@@ -41,10 +36,14 @@ import { useDebounceValue } from '@/lib/hooks/use-debounce'
 import { useLatest } from '@/lib/hooks/use-latest'
 import { useMe } from '@/lib/hooks/use-me'
 import { useSelectedModel } from '@/lib/hooks/use-models'
+import { useSelectedRepository } from '@/lib/hooks/use-repositories'
 import useRouterStuff from '@/lib/hooks/use-router-stuff'
 import { useIsChatEnabled } from '@/lib/hooks/use-server-info'
 import { useThreadRun } from '@/lib/hooks/use-thread-run'
-import { updateSelectedModel } from '@/lib/stores/chat-actions'
+import {
+  updateSelectedModel,
+  updateSelectedRepoSourceId
+} from '@/lib/stores/chat-actions'
 import { clearHomeScrollPosition } from '@/lib/stores/scroll-store'
 import { useMutation } from '@/lib/tabby/gql'
 import {
@@ -53,12 +52,7 @@ import {
   listThreads,
   setThreadPersistedMutation
 } from '@/lib/tabby/query'
-import {
-  AttachmentCodeItem,
-  AttachmentDocItem,
-  ExtendedCombinedError,
-  ThreadRunContexts
-} from '@/lib/types'
+import { ExtendedCombinedError, ThreadRunContexts } from '@/lib/types'
 import {
   cn,
   getMentionsFromText,
@@ -95,49 +89,10 @@ import { AssistantMessageSection } from './assistant-message-section'
 import { DevPanel } from './dev-panel'
 import { Header } from './header'
 import { MessagesSkeleton } from './messages-skeleton'
+import { SearchContext } from './search-context'
+import { ConversationMessage, ConversationPair } from './types'
 import { UserMessageSection } from './user-message-section'
 
-export type ConversationMessage = Omit<
-  Message,
-  '__typename' | 'updatedAt' | 'createdAt' | 'attachment' | 'threadId'
-> & {
-  threadId?: string
-  threadRelevantQuestions?: Maybe<string[]>
-  error?: string
-  attachment?: {
-    clientCode?: Maybe<Array<MessageAttachmentClientCode>> | undefined
-    code: Maybe<Array<AttachmentCodeItem>> | undefined
-    doc: Maybe<Array<AttachmentDocItem>> | undefined
-  }
-}
-
-type ConversationPair = {
-  question: ConversationMessage | null
-  answer: ConversationMessage | null
-}
-
-type SearchContextValue = {
-  // flag for initialize the pathname
-  isPathnameInitialized: boolean
-  isLoading: boolean
-  onRegenerateResponse: (id: string) => void
-  onSubmitSearch: (question: string) => void
-  setDevPanelOpen: (v: boolean) => void
-  setConversationIdForDev: (v: string | undefined) => void
-  enableDeveloperMode: boolean
-  contextInfo: ContextInfo | undefined
-  fetchingContextInfo: boolean
-  onDeleteMessage: (id: string) => void
-  isThreadOwner: boolean
-  onUpdateMessage: (
-    message: ConversationMessage
-  ) => Promise<ExtendedCombinedError | undefined>
-}
-
-export const SearchContext = createContext<SearchContextValue>(
-  {} as SearchContextValue
-)
-
 export const SOURCE_CARD_STYLE = {
   compress: 5.3,
   expand: 6.3
@@ -335,8 +290,9 @@ export function Search() {
 
   const isLoadingRef = useLatest(isLoading)
 
-  const { selectedModel, isModelLoading, models } = useSelectedModel()
-
+  const { selectedModel, isFetchingModels, models } = useSelectedModel()
+  const { selectedRepository, isFetchingRepositories, repos } =
+    useSelectedRepository()
   const currentMessageForDev = useMemo(() => {
     return messages.find(item => item.id === messageIdForDev)
   }, [messageIdForDev, messages])
@@ -575,7 +531,7 @@ export function Search() {
     }
 
     const { sourceIdForCodeQuery, sourceIdsForDocQuery, searchPublic } =
-      getSourceInputs(ctx)
+      getSourceInputs(selectedRepository?.sourceId, ctx)
 
     const codeQuery: InputMaybe<CodeQueryInput> = sourceIdForCodeQuery
       ? { sourceId: sourceIdForCodeQuery, content: question }
@@ -641,7 +597,10 @@ export function Search() {
     )
 
     const { sourceIdForCodeQuery, sourceIdsForDocQuery, searchPublic } =
-      getSourceInputs(getThreadRunContextsFromMentions(mentions))
+      getSourceInputs(
+        selectedRepository?.sourceId,
+        getThreadRunContextsFromMentions(mentions)
+      )
 
     const codeQuery: InputMaybe<CodeQueryInput> = sourceIdForCodeQuery
       ? { sourceId: sourceIdForCodeQuery, content: newUserMessage.content }
@@ -726,10 +685,14 @@ export function Search() {
     )
   }
 
-  const onModelSelect = (model: string) => {
+  const onSelectModel = (model: string) => {
     updateSelectedModel(model)
   }
 
+  const onSelectedRepo = (sourceId: string | undefined) => {
+    updateSelectedRepoSourceId(sourceId)
+  }
+
   const formatedThreadError: ExtendedCombinedError | undefined = useMemo(() => {
     if (!isReady || fetchingThread || !threadIdFromURL) return undefined
     if (threadError || !threadData?.threads?.edges?.length) {
@@ -806,7 +769,8 @@ export function Search() {
         fetchingContextInfo,
         onDeleteMessage,
         isThreadOwner,
-        onUpdateMessage
+        onUpdateMessage,
+        repositories: repos
       }}
     >
       <div className="transition-all" style={style}>
@@ -944,8 +908,12 @@ export function Search() {
                       contextInfo={contextInfoData?.contextInfo}
                       fetchingContextInfo={fetchingContextInfo}
                       modelName={selectedModel}
-                      onModelSelect={onModelSelect}
-                      isModelLoading={isModelLoading}
+                      onSelectModel={onSelectModel}
+                      repoSourceId={selectedRepository?.sourceId}
+                      onSelectRepo={onSelectedRepo}
+                      isInitializingResources={
+                        isFetchingModels || isFetchingRepositories
+                      }
                       models={models}
                     />
                   </div>
@@ -1026,17 +994,24 @@ function ThreadMessagesErrorView({
   )
 }
 
-function getSourceInputs(ctx: ThreadRunContexts | undefined) {
+function getSourceInputs(
+  repositorySourceId: string | undefined,
+  ctx: ThreadRunContexts | undefined
+) {
   let sourceIdsForDocQuery: string[] = []
   let sourceIdForCodeQuery: string | undefined
   let searchPublic = false
 
   if (ctx) {
     sourceIdsForDocQuery = uniq(
-      compact([ctx?.codeSourceIds?.[0]].concat(ctx.docSourceIds))
+      // Compatible with existing user messages
+      compact(
+        [repositorySourceId, ctx?.codeSourceIds?.[0]].concat(ctx.docSourceIds)
+      )
     )
     searchPublic = ctx.searchPublic ?? false
-    sourceIdForCodeQuery = ctx.codeSourceIds?.[0] ?? undefined
+    sourceIdForCodeQuery =
+      repositorySourceId || ctx.codeSourceIds?.[0] || undefined
   }
   return {
     sourceIdsForDocQuery,

diff --git a/ee/tabby-ui/app/search/components/types.ts b/ee/tabby-ui/app/search/components/types.ts
@@ -0,0 +1,24 @@
+import {
+  Maybe,
+  Message,
+  MessageAttachmentClientCode
+} from '@/lib/gql/generates/graphql'
+import { AttachmentCodeItem, AttachmentDocItem } from '@/lib/types'
+
+export type ConversationMessage = Omit<
+  Message,
+  '__typename' | 'updatedAt' | 'createdAt' | 'attachment' | 'threadId'
+> & {
+  threadId?: string
+  threadRelevantQuestions?: Maybe<string[]>
+  error?: string
+  attachment?: {
+    clientCode?: Maybe<Array<MessageAttachmentClientCode>> | undefined
+    code: Maybe<Array<AttachmentCodeItem>> | undefined
+    doc: Maybe<Array<AttachmentDocItem>> | undefined
+  }
+}
+export type ConversationPair = {
+  question: ConversationMessage | null
+  answer: ConversationMessage | null
+}
diff --git a/ee/tabby-ui/app/search/components/user-message-section.tsx b/ee/tabby-ui/app/search/components/user-message-section.tsx
@@ -4,7 +4,8 @@ import { cn } from '@/lib/utils'
 import { ChatContext } from '@/components/chat/chat'
 import { MessageMarkdown } from '@/components/message-markdown'
 
-import { ConversationMessage, SearchContext } from './search'
+import { SearchContext } from './search-context'
+import { ConversationMessage } from './types'
 
 interface QuestionBlockProps extends HTMLAttributes<HTMLDivElement> {
   message: ConversationMessage