feat: 添加 AST RepoMap 代码仓库地图功能
- 实现基于 Tree-sitter 的代码符号提取 (支持 TS/JS/Python) - 实现 PageRank 算法进行符号相关性排序 - 支持个性化权重调整 (提及的标识符、聊天文件等) - 添加磁盘缓存避免重复解析 - 集成 repo_map 工具到工具系统 - 添加 15 个单元测试
This commit is contained in:
@@ -0,0 +1,329 @@
|
||||
/**
|
||||
* RepoMap 测试
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeAll } from 'vitest';
|
||||
import * as path from 'path';
|
||||
import * as fs from 'fs/promises';
|
||||
import * as os from 'os';
|
||||
import { createRepoMap } from '../../src/repomap/repomap.js';
|
||||
import { TagExtractor } from '../../src/repomap/tags/extractor.js';
|
||||
import { Graph, pagerank, distributeRanksToDefinitions } from '../../src/repomap/ranking/index.js';
|
||||
import { DiskCache } from '../../src/repomap/cache/disk-cache.js';
|
||||
|
||||
describe('Graph', () => {
|
||||
it('should add edges and track nodes', () => {
|
||||
const graph = new Graph();
|
||||
|
||||
graph.addEdge({ from: 'a.ts', to: 'b.ts', weight: 1, ident: 'foo' });
|
||||
graph.addEdge({ from: 'a.ts', to: 'c.ts', weight: 2, ident: 'bar' });
|
||||
graph.addEdge({ from: 'b.ts', to: 'c.ts', weight: 1, ident: 'baz' });
|
||||
|
||||
const nodes = graph.getNodes();
|
||||
expect(nodes).toContain('a.ts');
|
||||
expect(nodes).toContain('b.ts');
|
||||
expect(nodes).toContain('c.ts');
|
||||
expect(nodes.length).toBe(3);
|
||||
});
|
||||
|
||||
it('should track in and out edges correctly', () => {
|
||||
const graph = new Graph();
|
||||
|
||||
graph.addEdge({ from: 'a.ts', to: 'b.ts', weight: 1, ident: 'foo' });
|
||||
graph.addEdge({ from: 'a.ts', to: 'c.ts', weight: 2, ident: 'bar' });
|
||||
|
||||
const outEdges = graph.getOutEdges('a.ts');
|
||||
expect(outEdges.length).toBe(2);
|
||||
|
||||
const inEdges = graph.getInEdges('b.ts');
|
||||
expect(inEdges.length).toBe(1);
|
||||
expect(inEdges[0].from).toBe('a.ts');
|
||||
});
|
||||
|
||||
it('should calculate out degree correctly', () => {
|
||||
const graph = new Graph();
|
||||
|
||||
graph.addEdge({ from: 'a.ts', to: 'b.ts', weight: 1, ident: 'foo' });
|
||||
graph.addEdge({ from: 'a.ts', to: 'c.ts', weight: 2, ident: 'bar' });
|
||||
|
||||
expect(graph.getOutDegree('a.ts')).toBe(3); // sum of weights
|
||||
expect(graph.getOutDegree('b.ts')).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('PageRank', () => {
|
||||
it('should compute ranks for a simple graph', () => {
|
||||
const graph = new Graph();
|
||||
|
||||
// 创建一个简单的图: a -> b -> c
|
||||
graph.addEdge({ from: 'a.ts', to: 'b.ts', weight: 1, ident: 'foo' });
|
||||
graph.addEdge({ from: 'b.ts', to: 'c.ts', weight: 1, ident: 'bar' });
|
||||
|
||||
const ranks = pagerank(graph);
|
||||
|
||||
expect(ranks.size).toBe(3);
|
||||
expect(ranks.get('a.ts')).toBeGreaterThan(0);
|
||||
expect(ranks.get('b.ts')).toBeGreaterThan(0);
|
||||
expect(ranks.get('c.ts')).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it('should respect personalization vector', () => {
|
||||
const graph = new Graph();
|
||||
|
||||
graph.addEdge({ from: 'a.ts', to: 'b.ts', weight: 1, ident: 'foo' });
|
||||
graph.addEdge({ from: 'b.ts', to: 'c.ts', weight: 1, ident: 'bar' });
|
||||
graph.addEdge({ from: 'c.ts', to: 'a.ts', weight: 1, ident: 'baz' });
|
||||
|
||||
// 高度偏好 a.ts
|
||||
const personalization = new Map<string, number>();
|
||||
personalization.set('a.ts', 100);
|
||||
personalization.set('b.ts', 1);
|
||||
personalization.set('c.ts', 1);
|
||||
|
||||
const ranks = pagerank(graph, { personalization });
|
||||
|
||||
// a.ts 应该有较高的 rank
|
||||
expect(ranks.get('a.ts')!).toBeGreaterThan(ranks.get('b.ts')!);
|
||||
});
|
||||
|
||||
it('should distribute ranks to definitions', () => {
|
||||
const graph = new Graph();
|
||||
|
||||
graph.addEdge({ from: 'a.ts', to: 'b.ts', weight: 1, ident: 'foo' });
|
||||
graph.addEdge({ from: 'a.ts', to: 'b.ts', weight: 2, ident: 'bar' });
|
||||
|
||||
const nodeRanks = pagerank(graph);
|
||||
const defRanks = distributeRanksToDefinitions(graph, nodeRanks);
|
||||
|
||||
expect(defRanks.has('b.ts:foo')).toBe(true);
|
||||
expect(defRanks.has('b.ts:bar')).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('DiskCache', () => {
|
||||
let cacheDir: string;
|
||||
let cache: DiskCache<{ value: number }>;
|
||||
|
||||
beforeAll(async () => {
|
||||
cacheDir = path.join(os.tmpdir(), `repomap-test-${Date.now()}`);
|
||||
cache = new DiskCache(cacheDir);
|
||||
});
|
||||
|
||||
it('should set and get values', async () => {
|
||||
await cache.set('key1', { value: 42 });
|
||||
const result = await cache.get('key1');
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result?.value).toBe(42);
|
||||
});
|
||||
|
||||
it('should return null for missing keys', async () => {
|
||||
const result = await cache.get('nonexistent');
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
it('should delete values', async () => {
|
||||
await cache.set('key2', { value: 100 });
|
||||
await cache.delete('key2');
|
||||
|
||||
const result = await cache.get('key2');
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
it('should check existence with has', async () => {
|
||||
await cache.set('key3', { value: 200 });
|
||||
|
||||
expect(await cache.has('key3')).toBe(true);
|
||||
expect(await cache.has('nonexistent2')).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('TagExtractor', () => {
|
||||
let tempDir: string;
|
||||
|
||||
beforeAll(async () => {
|
||||
tempDir = path.join(os.tmpdir(), `repomap-extractor-test-${Date.now()}`);
|
||||
await fs.mkdir(tempDir, { recursive: true });
|
||||
});
|
||||
|
||||
it('should extract tags from TypeScript code using regex fallback', async () => {
|
||||
const testFile = path.join(tempDir, 'test.ts');
|
||||
const code = `
|
||||
export function greet(name: string): string {
|
||||
return \`Hello, \${name}!\`;
|
||||
}
|
||||
|
||||
export class Greeter {
|
||||
private name: string;
|
||||
|
||||
constructor(name: string) {
|
||||
this.name = name;
|
||||
}
|
||||
|
||||
greet(): string {
|
||||
return \`Hello, \${this.name}!\`;
|
||||
}
|
||||
}
|
||||
|
||||
export interface IGreeter {
|
||||
greet(): string;
|
||||
}
|
||||
`;
|
||||
|
||||
await fs.writeFile(testFile, code);
|
||||
|
||||
const extractor = new TagExtractor();
|
||||
const tags = await extractor.getTags(testFile, 'test.ts');
|
||||
|
||||
// 检查是否提取到了一些 tags
|
||||
expect(tags.length).toBeGreaterThan(0);
|
||||
|
||||
// 检查是否有定义
|
||||
const defs = tags.filter((t) => t.kind === 'def');
|
||||
expect(defs.length).toBeGreaterThan(0);
|
||||
|
||||
// 检查是否包含函数和类
|
||||
const defNames = defs.map((t) => t.name);
|
||||
expect(defNames).toContain('greet');
|
||||
expect(defNames).toContain('Greeter');
|
||||
});
|
||||
|
||||
it('should extract tags from Python code using regex fallback', async () => {
|
||||
const testFile = path.join(tempDir, 'test.py');
|
||||
const code = `
|
||||
def greet(name: str) -> str:
|
||||
return f"Hello, {name}!"
|
||||
|
||||
class Greeter:
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
|
||||
def greet(self) -> str:
|
||||
return f"Hello, {self.name}!"
|
||||
|
||||
async def async_greet(name: str) -> str:
|
||||
return f"Hello async, {name}!"
|
||||
`;
|
||||
|
||||
await fs.writeFile(testFile, code);
|
||||
|
||||
const extractor = new TagExtractor();
|
||||
const tags = await extractor.getTags(testFile, 'test.py');
|
||||
|
||||
// 检查是否提取到了一些 tags
|
||||
expect(tags.length).toBeGreaterThan(0);
|
||||
|
||||
// 检查是否有定义
|
||||
const defs = tags.filter((t) => t.kind === 'def');
|
||||
expect(defs.length).toBeGreaterThan(0);
|
||||
|
||||
// 检查是否包含函数和类
|
||||
const defNames = defs.map((t) => t.name);
|
||||
expect(defNames).toContain('greet');
|
||||
expect(defNames).toContain('Greeter');
|
||||
expect(defNames).toContain('async_greet');
|
||||
});
|
||||
});
|
||||
|
||||
describe('RepoMap', () => {
|
||||
let tempDir: string;
|
||||
|
||||
beforeAll(async () => {
|
||||
tempDir = path.join(os.tmpdir(), `repomap-test-${Date.now()}`);
|
||||
await fs.mkdir(tempDir, { recursive: true });
|
||||
|
||||
// 创建测试文件
|
||||
await fs.writeFile(
|
||||
path.join(tempDir, 'utils.ts'),
|
||||
`
|
||||
export function add(a: number, b: number): number {
|
||||
return a + b;
|
||||
}
|
||||
|
||||
export function multiply(a: number, b: number): number {
|
||||
return a * b;
|
||||
}
|
||||
`
|
||||
);
|
||||
|
||||
await fs.writeFile(
|
||||
path.join(tempDir, 'calculator.ts'),
|
||||
`
|
||||
import { add, multiply } from './utils';
|
||||
|
||||
export class Calculator {
|
||||
add(a: number, b: number): number {
|
||||
return add(a, b);
|
||||
}
|
||||
|
||||
multiply(a: number, b: number): number {
|
||||
return multiply(a, b);
|
||||
}
|
||||
}
|
||||
`
|
||||
);
|
||||
|
||||
await fs.writeFile(
|
||||
path.join(tempDir, 'main.ts'),
|
||||
`
|
||||
import { Calculator } from './calculator';
|
||||
|
||||
const calc = new Calculator();
|
||||
console.log(calc.add(1, 2));
|
||||
console.log(calc.multiply(3, 4));
|
||||
`
|
||||
);
|
||||
});
|
||||
|
||||
it('should create a repo map', async () => {
|
||||
const repoMap = createRepoMap(tempDir, {
|
||||
mapTokens: 2048,
|
||||
});
|
||||
|
||||
const files = [
|
||||
path.join(tempDir, 'utils.ts'),
|
||||
path.join(tempDir, 'calculator.ts'),
|
||||
path.join(tempDir, 'main.ts'),
|
||||
];
|
||||
|
||||
const map = await repoMap.getRepoMap([], files, new Set(), new Set());
|
||||
|
||||
// 应该生成了一些内容
|
||||
expect(map.length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it('should boost mentioned identifiers', async () => {
|
||||
const repoMap = createRepoMap(tempDir, {
|
||||
mapTokens: 2048,
|
||||
});
|
||||
|
||||
const files = [
|
||||
path.join(tempDir, 'utils.ts'),
|
||||
path.join(tempDir, 'calculator.ts'),
|
||||
path.join(tempDir, 'main.ts'),
|
||||
];
|
||||
|
||||
const mentionedIdents = new Set(['Calculator']);
|
||||
const map = await repoMap.getRepoMap([], files, new Set(), mentionedIdents);
|
||||
|
||||
// 应该包含 Calculator 相关内容
|
||||
expect(map.length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it('should exclude chat files from output', async () => {
|
||||
const repoMap = createRepoMap(tempDir, {
|
||||
mapTokens: 2048,
|
||||
});
|
||||
|
||||
const chatFiles = [path.join(tempDir, 'utils.ts')];
|
||||
const otherFiles = [
|
||||
path.join(tempDir, 'calculator.ts'),
|
||||
path.join(tempDir, 'main.ts'),
|
||||
];
|
||||
|
||||
const map = await repoMap.getRepoMap(chatFiles, otherFiles, new Set(), new Set());
|
||||
|
||||
// utils.ts 不应该出现在输出中
|
||||
expect(map).not.toContain('utils.ts:');
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user