update: code query

TabbyML · Dec 26, 2024 · 3bca939 · 3bca939
1 parent 178a16c
commit 3bca939
Show file tree

Hide file tree

Showing 5 changed files with 63 additions and 33 deletions.
diff --git a/ee/tabby-ui/app/search/components/assistant-message-section.tsx b/ee/tabby-ui/app/search/components/assistant-message-section.tsx
@@ -105,7 +105,8 @@ export function AssistantMessageSection({
     fetchingContextInfo,
     onDeleteMessage,
     isThreadOwner,
-    onUpdateMessage
+    onUpdateMessage,
+    repositories
   } = useContext(SearchContext)
 
   const { supportsOnApplyInEditorV2 } = useContext(ChatContext)
@@ -146,7 +147,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 []
@@ -157,11 +166,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 (
@@ -183,9 +192,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/search.tsx b/ee/tabby-ui/app/search/components/search.tsx
@@ -33,6 +33,7 @@ import {
   Maybe,
   Message,
   MessageAttachmentClientCode,
+  RepositorySourceListQuery,
   Role
 } from '@/lib/gql/generates/graphql'
 import { useCopyToClipboard } from '@/lib/hooks/use-copy-to-clipboard'
@@ -55,6 +56,7 @@ import {
   contextInfoQuery,
   listThreadMessages,
   listThreads,
+  repositorySourceListQuery,
   setThreadPersistedMutation
 } from '@/lib/tabby/query'
 import {
@@ -136,6 +138,7 @@ type SearchContextValue = {
   onUpdateMessage: (
     message: ConversationMessage
   ) => Promise<ExtendedCombinedError | undefined>
+  repositories: RepositorySourceListQuery['repositoryList'] | undefined
 }
 
 export const SearchContext = createContext<SearchContextValue>(
@@ -223,6 +226,10 @@ export function Search() {
   const [{ data: contextInfoData, fetching: fetchingContextInfo }] = useQuery({
     query: contextInfoQuery
   })
+  const [{ data: repositorySourceListData, fetching: fetchingRepos }] =
+    useQuery({
+      query: repositorySourceListQuery
+    })
 
   const [afterCursor, setAfterCursor] = useState<string | undefined>()
 
@@ -579,7 +586,7 @@ export function Search() {
     }
 
     const { sourceIdForCodeQuery, sourceIdsForDocQuery, searchPublic } =
-      getSourceInputs(ctx)
+      getSourceInputs(selectedRepoId, ctx)
 
     const codeQuery: InputMaybe<CodeQueryInput> = sourceIdForCodeQuery
       ? { sourceId: sourceIdForCodeQuery, content: question }
@@ -645,7 +652,10 @@ export function Search() {
     )
 
     const { sourceIdForCodeQuery, sourceIdsForDocQuery, searchPublic } =
-      getSourceInputs(getThreadRunContextsFromMentions(mentions))
+      getSourceInputs(
+        selectedRepoId,
+        getThreadRunContextsFromMentions(mentions)
+      )
 
     const codeQuery: InputMaybe<CodeQueryInput> = sourceIdForCodeQuery
       ? { sourceId: sourceIdForCodeQuery, content: newUserMessage.content }
@@ -814,7 +824,8 @@ export function Search() {
         fetchingContextInfo,
         onDeleteMessage,
         isThreadOwner,
-        onUpdateMessage
+        onUpdateMessage,
+        repositories: repositorySourceListData?.repositoryList
       }}
     >
       <div className="transition-all" style={style}>
@@ -1036,17 +1047,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/components/chat/chat.tsx b/ee/tabby-ui/components/chat/chat.tsx
@@ -25,6 +25,7 @@ import { useLatest } from '@/lib/hooks/use-latest'
 import { useThreadRun } from '@/lib/hooks/use-thread-run'
 import { filename2prism } from '@/lib/language-utils'
 import { useChatStore } from '@/lib/stores/chat-store'
+import { repositorySourceListQuery } from '@/lib/tabby/query'
 import { ExtendedCombinedError } from '@/lib/types'
 import {
   AssistantMessage,
@@ -49,20 +50,6 @@ import { ChatScrollAnchor } from './chat-scroll-anchor'
 import { EmptyScreen } from './empty-screen'
 import { QuestionAnswerList } from './question-answer'
 
-const repositoryListQuery = graphql(/* GraphQL */ `
-  query RepositorySourceList {
-    repositoryList {
-      id
-      name
-      kind
-      gitUrl
-      sourceId
-      sourceName
-      sourceKind
-    }
-  }
-`)
-
 type ChatContextValue = {
   initialized: boolean
   threadId: string | undefined
@@ -180,7 +167,7 @@ function ChatRenderer(
   const chatPanelRef = React.useRef<ChatPanelRef>(null)
 
   const [{ data: repositoryListData, fetching: fetchingRepos }] = useQuery({
-    query: repositoryListQuery
+    query: repositorySourceListQuery
   })
   const repos = repositoryListData?.repositoryList
 

diff --git a/ee/tabby-ui/lib/hooks/use-thread-run.ts b/ee/tabby-ui/lib/hooks/use-thread-run.ts
@@ -38,6 +38,7 @@ const CreateThreadAndRunSubscription = graphql(/* GraphQL */ `
         questions
       }
       ... on ThreadAssistantMessageAttachmentsCode {
+        codeSourceId
         hits {
           code {
             gitUrl
@@ -115,6 +116,7 @@ const CreateThreadRunSubscription = graphql(/* GraphQL */ `
         questions
       }
       ... on ThreadAssistantMessageAttachmentsCode {
+        codeSourceId
         hits {
           code {
             gitUrl
@@ -188,13 +190,11 @@ const DeleteThreadMessagePairMutation = graphql(/* GraphQL */ `
     )
   }
 `)
-
-type ID = string
-
 export interface AnswerStream {
-  threadId?: ID
-  userMessageId?: ID
-  assistantMessageId?: ID
+  threadId?: string
+  userMessageId?: string
+  assistantMessageId?: string
+  codeSourceId?: string
   relevantQuestions?: Array<string>
   attachmentsCode?: ThreadAssistantMessageAttachmentCodeHits
   attachmentsDoc?: ThreadAssistantMessageAttachmentDocHits
@@ -274,6 +274,7 @@ export function useThreadRun({
         break
       case 'ThreadAssistantMessageAttachmentsCode':
         x.attachmentsCode = data.hits
+        x.codeSourceId = data.codeSourceId
         break
       case 'ThreadAssistantMessageAttachmentsDoc':
         x.attachmentsDoc = data.hits

diff --git a/ee/tabby-ui/lib/tabby/query.ts b/ee/tabby-ui/lib/tabby/query.ts
@@ -397,6 +397,7 @@ export const listThreadMessages = graphql(/* GraphQL */ `
         node {
           id
           threadId
+          codeSourceId
           role
           content
           attachment {
@@ -472,3 +473,17 @@ export const notificationsQuery = graphql(/* GraphQL */ `
     }
   }
 `)
+
+export const repositorySourceListQuery = graphql(/* GraphQL */ `
+  query RepositorySourceList {
+    repositoryList {
+      id
+      name
+      kind
+      gitUrl
+      sourceId
+      sourceName
+      sourceKind
+    }
+  }
+`)