Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(ui): fill repository context into chat sidebar. #1950

Closed
wants to merge 7 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions crates/tabby-common/src/api/code.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ pub struct HitDocument {
pub filepath: String,
pub git_url: String,
pub language: String,
pub start_line: i64,
}

#[derive(Error, Debug)]
Expand Down
3 changes: 3 additions & 0 deletions crates/tabby-common/src/index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ pub struct CodeSearchSchema {
pub field_filepath: Field,
pub field_language: Field,
pub field_body: Field,
pub field_start_line: Field,
}

impl CodeSearchSchema {
Expand All @@ -38,6 +39,7 @@ impl CodeSearchSchema {
let field_filepath = builder.add_text_field("filepath", STRING | STORED);
let field_language = builder.add_text_field("language", STRING | STORED);
let field_body = builder.add_text_field("body", code_options);
let field_start_line = builder.add_i64_field("start_line", STORED);
let schema = builder.build();

Self {
Expand All @@ -46,6 +48,7 @@ impl CodeSearchSchema {
field_filepath,
field_language,
field_body,
field_start_line,
}
}
}
Expand Down
23 changes: 21 additions & 2 deletions crates/tabby-scheduler/src/code/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,26 @@
pub fn chunks<'splitter, 'text: 'splitter>(
&'splitter self,
text: &'text str,
) -> impl Iterator<Item = &'text str> + 'splitter {
self.splitter.chunks(text, 192)
) -> impl Iterator<Item = (i64, &'text str)> + 'splitter {
self.splitter
.chunk_indices(text, 256)
.map(|(offset, chunk)| (line_number_from_byte_offset(text, offset), chunk))

Check warning on line 63 in crates/tabby-scheduler/src/code/mod.rs

View check run for this annotation

Codecov / codecov/patch

crates/tabby-scheduler/src/code/mod.rs#L60-L63

Added lines #L60 - L63 were not covered by tests
}
}

fn line_number_from_byte_offset(s: &str, byte_offset: usize) -> i64 {
let mut line_number = 1; // Start counting from line 1
let mut current_offset = 0;

Check warning on line 69 in crates/tabby-scheduler/src/code/mod.rs

View check run for this annotation

Codecov / codecov/patch

crates/tabby-scheduler/src/code/mod.rs#L67-L69

Added lines #L67 - L69 were not covered by tests

for c in s.chars() {
if c == '\n' {
line_number += 1;
}
current_offset += c.len_utf8();
if current_offset >= byte_offset {
break;
}

Check warning on line 78 in crates/tabby-scheduler/src/code/mod.rs

View check run for this annotation

Codecov / codecov/patch

crates/tabby-scheduler/src/code/mod.rs#L71-L78

Added lines #L71 - L78 were not covered by tests
}

line_number
}

Check warning on line 82 in crates/tabby-scheduler/src/code/mod.rs

View check run for this annotation

Codecov / codecov/patch

crates/tabby-scheduler/src/code/mod.rs#L81-L82

Added lines #L81 - L82 were not covered by tests
3 changes: 2 additions & 1 deletion crates/tabby-scheduler/src/index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,14 +48,15 @@
}
};

for body in intelligence.chunks(&text) {
for (start_line, body) in intelligence.chunks(&text) {

Check warning on line 51 in crates/tabby-scheduler/src/index.rs

View check run for this annotation

Codecov / codecov/patch

crates/tabby-scheduler/src/index.rs#L51

Added line #L51 was not covered by tests
pb.as_mut().map(|b| b.update(body.len())).transpose()?;

writer.add_document(doc!(
code.field_git_url => file.git_url.clone(),
code.field_filepath => file.filepath.clone(),
code.field_language => file.language.clone(),
code.field_body => body,
code.field_start_line => start_line,

Check warning on line 59 in crates/tabby-scheduler/src/index.rs

View check run for this annotation

Codecov / codecov/patch

crates/tabby-scheduler/src/index.rs#L59

Added line #L59 was not covered by tests
))?;
}
}
Expand Down
8 changes: 8 additions & 0 deletions crates/tabby/src/services/code.rs
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@
filepath: get_field(&doc, self.schema.field_filepath),
git_url: get_field(&doc, self.schema.field_git_url),
language: get_field(&doc, self.schema.field_language),
start_line: get_i64_field(&doc, self.schema.field_start_line),

Check warning on line 80 in crates/tabby/src/services/code.rs

View check run for this annotation

Codecov / codecov/patch

crates/tabby/src/services/code.rs#L80

Added line #L80 was not covered by tests
},
id: doc_address.doc_id,
}
Expand Down Expand Up @@ -172,6 +173,13 @@
.to_owned()
}

fn get_i64_field(doc: &Document, field: Field) -> i64 {
doc.get_first(field)
.and_then(|x| x.as_i64())
.expect("Missing field")
.to_owned()
}

Check warning on line 181 in crates/tabby/src/services/code.rs

View check run for this annotation

Codecov / codecov/patch

crates/tabby/src/services/code.rs#L176-L181

Added lines #L176 - L181 were not covered by tests

struct CodeSearchService {
search: Arc<Mutex<Option<CodeSearchImpl>>>,
}
Expand Down
2 changes: 1 addition & 1 deletion crates/tabby/src/services/completion/completion_prompt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@

pub async fn collect(&self, language: &str, segments: &Segments) -> Vec<Snippet> {
let quota_threshold_for_snippets_from_code_search = 256;
let mut max_snippets_chars_in_prompt = 768;
let mut max_snippets_chars_in_prompt = 1024;

Check warning on line 41 in crates/tabby/src/services/completion/completion_prompt.rs

View check run for this annotation

Codecov / codecov/patch

crates/tabby/src/services/completion/completion_prompt.rs#L41

Added line #L41 was not covered by tests
let mut snippets: Vec<Snippet> = vec![];

if let Some((count_characters, snippets_from_segments)) =
Expand Down
92 changes: 85 additions & 7 deletions ee/tabby-ui/app/files/components/chat-side-bar.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@ import React from 'react'

import { useStore } from '@/lib/hooks/use-store'
import { useChatStore } from '@/lib/stores/chat-store'
import fetcher from '@/lib/tabby/fetcher'
import { ISearchHit, SearchReponse } from '@/lib/types'
import { cn } from '@/lib/utils'
import { Button } from '@/components/ui/button'
import { IconClose } from '@/components/ui/icons'
Expand All @@ -22,14 +24,15 @@ export const ChatSideBar: React.FC<ChatSideBarProps> = ({
const activeChatId = useStore(useChatStore, state => state.activeChatId)
const iframeRef = React.useRef<HTMLIFrameElement>(null)

const getPrompt = ({
const getPrompt = async ({
action,
code,
language,
path,
lineFrom,
lineTo
}: QuickActionEventPayload) => {
const contextPrompt = await buildContextPrompt(language, code, path)
let builtInPrompt = ''
switch (action) {
case 'explain':
Expand All @@ -47,18 +50,22 @@ export const ChatSideBar: React.FC<ChatSideBarProps> = ({
const codeBlockMeta = `${
language ?? ''
} is_reference=1 path=${path} line_from=${lineFrom} line_to=${lineTo}`
return `${builtInPrompt}\n${'```'}${codeBlockMeta}\n${code}\n${'```'}\n`
return `${contextPrompt}${builtInPrompt}\n${'```'}${codeBlockMeta}\n${code}\n${'```'}\n`
}

React.useEffect(() => {
async function postPrompt(e: QuickActionEventPayload) {
const contentWindow = iframeRef.current?.contentWindow
contentWindow?.postMessage({
action: 'append',
payload: await getPrompt(e)
})
}

React.useEffect(() => {
if (pendingEvent) {
contentWindow?.postMessage({
action: 'append',
payload: getPrompt(pendingEvent)
postPrompt(pendingEvent).then(() => {
setPendingEvent(undefined)
})
setPendingEvent(undefined)
}
}, [pendingEvent, iframeRef.current?.contentWindow])

Expand Down Expand Up @@ -90,3 +97,74 @@ function Header() {
</div>
)
}

async function buildContextPrompt(
language: string | undefined,
code: string,
path: string | undefined
) {
if (!language || !path) {
return []
}

if (code.length < 128) {
return []
}

const segments = path.split('/');
const repo = segments[0];
path = segments.slice(1).join('/');

const tokens = code.split(/[^\w]/).filter(x => x)

// FIXME(jueliang): restrict query with `git_url` of `repo`.
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@liangfung as we have git_url from repositoryList, please fill it in a query like
git_url:${gitUrl}, and append it to AND ${gitUrlQuery}

const languageQuery = buildLanguageQuery(language)
const bodyQuery = tokens.map(x => `body:${x}`).join(' OR ')
const query = `${languageQuery} AND (${bodyQuery})`

const queryParam = `q=${encodeURIComponent(query)}&limit=20`

const data: SearchReponse = await fetcher(`/v1beta/search?${queryParam}`, {
responseFormat: 'json'
})
const snippets =
data.hits.filter(x => x.score > 30 && path !== x.doc.filepath) || []
return formatContextPrompt(repo, language, snippets.slice(0, 3))
}

function formatContextPrompt(
repo: string,
language: string,
snippets: ISearchHit[]
) {
let prompt = 'Given following relevant code snippets:\n\n'
for (const { doc } of snippets) {
const numLines = doc.body.split(/\r\n|\r|\n/).length
const fromLine = doc.start_line
const toLine = doc.start_line + numLines - 1
const reference = `\`\`\`${language} is_reference=1 path=${repo}/${doc.filepath} line_from=${fromLine} line_to=${toLine}
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
const reference = `\`\`\`${language} is_reference=1 path=${repo}/${doc.filepath} line_from=${fromLine} line_to=${toLine}
const reference = `\`\`\`${language} reference_group=1 path=${repo}/${doc.filepath} line_from=${fromLine} line_to=${toLine}

${doc.body}
\`\`\`
`
prompt += reference
}

if (snippets.length) {
return prompt
} else {
return ''
}
}

function buildLanguageQuery(language: string) {
if (
language == 'javascript' ||
language == 'jsx' ||
language == 'typescript' ||
language == 'tsx'
) {
language = 'javascript-typescript'
}

return `language:${language}`
}
19 changes: 18 additions & 1 deletion ee/tabby-ui/components/prompt-form.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import useSWR from 'swr'

import { useEnterSubmit } from '@/lib/hooks/use-enter-submit'
import fetcher from '@/lib/tabby/fetcher'
import type { ISearchHit, SearchReponse } from '@/lib/types'
import { cn } from '@/lib/utils'
import { Button, buttonVariants } from '@/components/ui/button'
import {
Expand Down Expand Up @@ -349,3 +348,21 @@ function IconForCompletionKind({
return <IconSymbolFunction {...rest} />
}
}

type ISearchHit = {
id: number
score: number
doc: {
body: string
filepath: string
git_url: string
language: string
name: string
kind: string
}
}

type SearchReponse = {
hits: Array<ISearchHit>
num_hits: number
}
19 changes: 10 additions & 9 deletions ee/tabby-ui/lib/types/chat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,18 +12,19 @@ export interface Chat extends Record<string, any> {

export type ISearchHit = {
id: number
doc?: {
body?: string
name?: string
filepath?: string
git_url?: string
kind?: string
language?: string
score: number
doc: {
body: string
filepath: string
git_url: string
language: string
start_line: number
}
}

export type SearchReponse = {
hits?: Array<ISearchHit>
num_hits?: number
hits: Array<ISearchHit>
num_hits: number
}

export type MessageActionType = 'edit' | 'delete' | 'regenerate'
Loading