Skip to content

Commit

Permalink
fix: Return whole lines when looking up source code
Browse files Browse the repository at this point in the history
  • Loading branch information
dividedmind committed Jan 9, 2025
1 parent ec341a8 commit 3182837
Show file tree
Hide file tree
Showing 4 changed files with 72 additions and 11 deletions.
8 changes: 4 additions & 4 deletions packages/cli/src/rpc/explain/buildContext.ts
Original file line number Diff line number Diff line change
Expand Up @@ -70,13 +70,13 @@ export default async function buildContext(

codeSnippetLocations.add(event.location);

const snippets = await lookupSourceCode(result.directory, event.location);
if (snippets) {
const snippet = await lookupSourceCode(result.directory, event.location);
if (snippet) {
codeSnippets.push({
directory: result.directory,
type: ContextV2.ContextItemType.CodeSnippet,
location: event.location,
content: snippets.join('\n'),
location: snippet.location,
content: snippet.content,
score: event.score,
});
}
Expand Down
34 changes: 29 additions & 5 deletions packages/cli/src/rpc/explain/lookupSourceCode.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ const FILE_NAMES = new Set<string>();
export default async function lookupSourceCode(
directory: string,
location: string
): Promise<string[] | undefined> {
): Promise<{ content: string; location: string } | undefined> {
const parsedLocation = parseLocation(location);
if (!parsedLocation) return;

Expand Down Expand Up @@ -77,9 +77,9 @@ export default async function lookupSourceCode(
}

const fileContent = await readFile(fileName, 'utf-8');
if (!lineNo) return [fileContent];
const path = fileName.slice(directory.length + 1);

if (lineNo <= 0) return [fileContent];
if (!lineNo || lineNo <= 0) return { content: fileContent, location: path };

const fileExtension = fileName.slice(fileName.lastIndexOf('.'));
const language = LANGUAGE_BY_FILE_EXTENSION[fileExtension];
Expand All @@ -100,6 +100,30 @@ export default async function lookupSourceCode(
const matches = chunks.filter(
(chunk) => chunk.metadata.loc.lines.from <= lineNo && chunk.metadata.loc.lines.to >= lineNo
);
if (verbose()) warn(chalk.gray(`Obtained ${matches.length} source code chunks for ${location}`));
return matches.map((match) => match.pageContent);

// determine the extent of the snippets and return a single snippet that contains all the complete lines
if (matches.length === 0) {
warn(chalk.gray(`No source code found for ${location}`));
return;
}

const extent = {
from: Math.min(...matches.map((match) => match.metadata.loc.lines.from)),
to: Math.max(...matches.map((match) => match.metadata.loc.lines.to)),
};

if (verbose())
warn(
chalk.gray(
`Found ${matches.length} matches for ${location} (extent: ${extent.from}-${extent.to})`
)
);

return {
content: fileContent
.split('\n')
.slice(extent.from - 1, extent.to)
.join('\n'),
location: `${path}:${extent.from}-${extent.to}`,
};
}
6 changes: 4 additions & 2 deletions packages/cli/tests/unit/rpc/explain/buildContext.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,9 @@ describe('buildContext', () => {

describe('function call', () => {
it('is emitted with source code', async () => {
jest.mocked(lookupSourceCode).mockResolvedValue(['the code']);
jest
.mocked(lookupSourceCode)
.mockResolvedValue({ content: 'the code', location: 'app/models/user.rb:42-47' });

const context = await buildContext([
{
Expand All @@ -53,7 +55,7 @@ describe('buildContext', () => {
).toEqual([
{
directory: 'a',
location: 'app/models/user.rb',
location: 'app/models/user.rb:42-47',
type: 'code-snippet',
score: 1,
content: 'the code',
Expand Down
35 changes: 35 additions & 0 deletions packages/cli/tests/unit/rpc/explain/lookupSourceCode.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
import path from 'node:path';

import { fs, vol } from 'memfs';

import lookupSourceCode from '../../../../src/rpc/explain/lookupSourceCode';

describe('lookupSourceCode', () => {
const directory = '/project';
const fileName = 'example.js';
const sourceCodeSnippet = 'line 1\nline 2\nline 3\nline 4\nline 5';
const filePath = path.join(directory, fileName);

it('returns the correct snippet object', async () => {
const snippet = await lookupSourceCode(directory, `${fileName}:2`);
expect(snippet).toEqual({
content: sourceCodeSnippet,
location: 'example.js:1-5',
});
});

it('returns undefined for non-existent line', async () => {
const snippet = await lookupSourceCode(directory, `${fileName}:10`);
expect(snippet).toBeUndefined();
});

beforeEach(() => {
vol.mkdirSync(directory, { recursive: true });
vol.writeFileSync(filePath, sourceCodeSnippet);
});

afterEach(() => vol.reset());
});

jest.mock('fs', () => fs);
jest.mock('fs/promises', () => fs.promises);

0 comments on commit 3182837

Please sign in to comment.