From 3bca939e99c662fcf4efa4ba9d8636ed37de8bbe Mon Sep 17 00:00:00 2001 From: liangfung Date: Fri, 27 Dec 2024 01:31:16 +0800 Subject: [PATCH] update: code query --- .../components/assistant-message-section.tsx | 21 +++++++++---- ee/tabby-ui/app/search/components/search.tsx | 30 +++++++++++++++---- ee/tabby-ui/components/chat/chat.tsx | 17 ++--------- ee/tabby-ui/lib/hooks/use-thread-run.ts | 13 ++++---- ee/tabby-ui/lib/tabby/query.ts | 15 ++++++++++ 5 files changed, 63 insertions(+), 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 index 920f21ad8ad4..e7edaa669e2e 100644 --- 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 index f990e968f613..af1671180951 100644 --- 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 + repositories: RepositorySourceListQuery['repositoryList'] | undefined } export const SearchContext = createContext( @@ -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() @@ -579,7 +586,7 @@ export function Search() { } const { sourceIdForCodeQuery, sourceIdsForDocQuery, searchPublic } = - getSourceInputs(ctx) + getSourceInputs(selectedRepoId, ctx) const codeQuery: InputMaybe = 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 = sourceIdForCodeQuery ? { sourceId: sourceIdForCodeQuery, content: newUserMessage.content } @@ -814,7 +824,8 @@ export function Search() { fetchingContextInfo, onDeleteMessage, isThreadOwner, - onUpdateMessage + onUpdateMessage, + repositories: repositorySourceListData?.repositoryList }} >
@@ -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 index 4294b43752d1..c9c0e5858abe 100644 --- 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(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 index ae8a264ac610..d1a7ea3c6d5f 100644 --- 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 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 index 0e3bfa34e57e..58bc6b35c435 100644 --- 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 + } + } +`)