alex/pi-map
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: remove heuristics, go LLM-only with mocks + restructure

Browse Source 
BREAKING: Heuristic fallback removed — LLM client now required

Changes:
- Remove extractFileHeuristic, extractPackageHeuristic, all helpers
- extractFileLLM / extractPackageLLM now throw LLMError when client missing
- Add hybrid binary detection: extension blacklist + content sniffing
- Increase file size limit: 50KB → 500KB (text files only)
- Restructure src/ into subdirectories:
  - src/llm/ — all LLM clients, extract, batch, cache, error
  - src/ast/ — AST extraction
  - src/cli/ — CLI entry point
- Update all imports across codebase and tests
- Add tests/mock-llm.ts helper for deterministic mock clients
- Update all tests to use mock LLM clients (no heuristics dependency)
- All 52 tests passing (including 8 real LLM integration tests)
This commit is contained in:
Alex
2026-06-10 17:34:37 +02:00
parent d948d88d0a
commit 2f198ea0d2
23 changed files with 382 additions and 507 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.pi-project-map/cache/llm-cache.json
Vendored
+62
View File
@@ -0,0 +1,62 @@
{
"25c2db9d6c4c2c511bec4a736f5adf77423be6011bbedd916fb55003d059b4e3": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452802
},
"889300ab1d0afe082d820fdd9150229625c3ab69b3b84c2ca65d67479be388cc": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452841
},
"585c5da8de102beb6a9906f0e3f07a134b5392bed52fc1a0aaef989e69aa66b6": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452845
},
"a6f92e629a4c341fe1e10f92ef923c54002b8a79ec3bedb0b9079bc7aca4e2e0": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452846
},
"a6c149b54aec4ced46eb38d6faf1f20ec62fd2f9a2f2dcb6a15c4f339dfb5f1e": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452846
},
"4eda95f6452551a12d5b434e284c2e539ce57e0206d5ffa49609f4fa8354a966": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452846
},
"510456d000664badf3cbec9f70c80769a119a590e84b1971e49ac5728882a699": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452847
},
"b0f5449d50fda0b5c0cf47103188df2d4339b047e1e03cab7a15a431f2990185": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452848
},
"eebe9a750792a3d802c6868f32ac264d44ef193a6f1940c4ad5e715b6e31074a": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452848
},
"4958b5d7ec5b4c975bb1e2972519aff5fbd94cb7342be4f8dd854a5c7353bd26": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452849
},
"389a385ef4fc5d1d253253fa3c931c505c7fbcca68db4ebe6b34392b3ca32d65": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452849
},
"879d2eff31eaf101d30de95e0e2df1c1a3a305c4e3b0515a22c85969af49ff25": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452947
},
"c2f0a91350b7f5e11a95848c4d806e49d2017862acbc43afb1dd52a5e73c0b24": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452948
},
"037ecd1db38c230c248787e60fd7bfc0cb0101b187b59535b6e7483be762d350": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452950
},
"b5d546753d33709dff508a6f2c2e547f432267fbb05b68d3da74fa5f7bda9f7a": {
"result": "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
"ts": 1781105452953
}
}
pi-extension.ts
+2 -2
View File
@@ -8,8 +8,8 @@ import {
validateMaps, validateMaps,
reinitPath, reinitPath,
} from "./src/index.js"; } from "./src/index.js";
import { createLLMClient } from "./src/llm-client.js"; import { createLLMClient } from "./src/llm/llm-client.js";
import { LLMError } from "./src/llm-error.js"; import { LLMError } from "./src/llm/llm-error.js";
/** /**
* Get the LLM client for Pi runtime. * Get the LLM client for Pi runtime.
src/ast-extract.ts → src/ast/ast-extract.ts
View File
src/cli.ts → src/cli/cli.ts
+7 -7
View File
@@ -1,11 +1,11 @@
#!/usr/bin/env node #!/usr/bin/env node
import { initProject } from "./init.js"; import { initProject } from "../init.js";
import { patchFile } from "./patch.js"; import { patchFile } from "../patch.js";
import { validateMaps } from "./validate.js"; import { validateMaps } from "../validate.js";
import { reinitPath } from "./init.js"; import { reinitPath } from "../init.js";
import { discoverProject } from "./discover.js"; import { discoverProject } from "../discover.js";
import { createLLMClient, LLMError } from "./llm-client.js"; import { createLLMClient, LLMError } from "../llm/llm-client.js";
import { loadConfig } from "./config.js"; import { loadConfig } from "../config.js";
import pc from "picocolors"; import pc from "picocolors";
const args = process.argv.slice(2); const args = process.argv.slice(2);
src/init.ts
+4 -4
View File
@@ -4,13 +4,13 @@ import {
type PackageMapData, type PackageMapData,
type FileEntry, type FileEntry,
} from "./format.js"; } from "./format.js";
import { extractFileLLM, extractPackageLLM } from "./llm-extract.js"; import { extractFileLLM, extractPackageLLM } from "./llm/llm-extract.js";
import { extractFileAST } from "./ast-extract.js"; import { extractFileAST } from "./ast/ast-extract.js";
import { mergeFileData } from "./merge.js"; import { mergeFileData } from "./merge.js";
import { processFiles } from "./llm-batch.js"; import { processFiles } from "./llm/llm-batch.js";
import { writeFileSync } from "fs"; import { writeFileSync } from "fs";
import { join } from "path"; import { join } from "path";
import type { LLMClient } from "./llm-client.js"; import type { LLMClient } from "./llm/llm-client.js";
export interface InitOptions { export interface InitOptions {
verbose?: boolean; verbose?: boolean;
src/llm-extract.ts
-387
View File
@@ -1,387 +0,0 @@
import { readFileSync, statSync } from "fs";
import { createHash } from "crypto";
import { extname, basename } from "path";
import type { LLMClient } from "./llm-client.js";
import { getCached, setCached } from "./llm-cache.js";
import { LLMError } from "./llm-error.js";
interface LLMFileData {
purpose: string;
exports: string[];
deps: string[];
concepts: string[];
}
interface LLMPackageData {
role: string;
arch: string;
}
const MAX_FILE_SIZE = 50 * 1024; // 50KB
const CONTEXT_BUDGET = 4000; // tokens
const CHARS_PER_TOKEN = 4; // approximate for ASCII
// Heuristic patterns for common file types (used as fallback + for tests)
const FILE_TYPE_PURPOSES: Record<string, string> = {
".ts": "TypeScript module",
".tsx": "React component",
".js": "JavaScript module",
".jsx": "React component",
".py": "Python module",
".go": "Go module",
".rs": "Rust module",
".java": "Java class",
".kt": "Kotlin class",
".swift": "Swift module",
".c": "C source",
".cpp": "C++ source",
".h": "C/C++ header",
".rb": "Ruby module",
".php": "PHP script",
".sh": "Shell script",
".md": "Documentation",
".json": "Configuration",
".yaml": "Configuration",
".yml": "Configuration",
".toml": "Configuration",
".ini": "Configuration",
".env": "Environment config",
".dockerfile": "Docker image definition",
dockerfile: "Docker image definition",
".sql": "Database schema/queries",
".css": "Stylesheet",
".scss": "SCSS stylesheet",
".less": "LESS stylesheet",
".html": "HTML template",
".vue": "Vue component",
".svelte": "Svelte component",
};
function truncateForContext(content: string, promptLength: number): string {
const maxChars = CONTEXT_BUDGET * CHARS_PER_TOKEN - promptLength;
if (content.length <= maxChars) return content;
return content.slice(0, maxChars - 20) + "\n[...truncated]";
}
function buildFilePrompt(filePath: string, content: string): string {
const name = basename(filePath);
const ext = extname(filePath).toLowerCase();
const typeHint = FILE_TYPE_PURPOSES[ext] || FILE_TYPE_PURPOSES[name.toLowerCase()] || ext || "file";
return `Analyze this ${typeHint} file. Respond in this exact format (one line each):
PURPOSE: <concise one-sentence description of what this file does>
DEPS: <comma-separated list of key dependencies/modules it relies on, or "none">
CONCEPTS: <comma-separated list of key concepts/patterns used, or "none">
File: ${name}
\`\`\`
${content}
\`\`\`
`;
}
function parseFileResponse(response: string): { purpose: string; deps: string[]; concepts: string[] } {
const lines = response.split("\n");
let purpose = "";
let deps: string[] = [];
let concepts: string[] = [];
for (const line of lines) {
const trimmed = line.trim();
if (trimmed.startsWith("PURPOSE:")) {
purpose = trimmed.slice("PURPOSE:".length).trim();
} else if (trimmed.startsWith("DEPS:")) {
const depsStr = trimmed.slice("DEPS:".length).trim();
deps = depsStr === "none" ? [] : depsStr.split(",").map((s) => s.trim()).filter(Boolean);
} else if (trimmed.startsWith("CONCEPTS:")) {
const conceptsStr = trimmed.slice("CONCEPTS:".length).trim();
concepts = conceptsStr === "none" ? [] : conceptsStr.split(",").map((s) => s.trim()).filter(Boolean);
}
}
return { purpose, deps, concepts };
}
function buildPackagePrompt(relativePath: string, fileSummaries: { name: string; purpose: string }[]): string {
const filesList = fileSummaries.map((f) => `- ${f.name}: ${f.purpose}`).join("\n");
return `Analyze this code package/directory. Respond in this exact format (one line each):
ROLE: <concise one-sentence description of this package's role in the project>
ARCH: <concise description of architecture/patterns used in this package>
Directory: ${relativePath}
Files:
${filesList}
`;
}
function parsePackageResponse(response: string): { role: string; arch: string } {
const lines = response.split("\n");
let role = "";
let arch = "";
for (const line of lines) {
const trimmed = line.trim();
if (trimmed.startsWith("ROLE:")) {
role = trimmed.slice("ROLE:".length).trim();
} else if (trimmed.startsWith("ARCH:")) {
arch = trimmed.slice("ARCH:".length).trim();
}
}
return { role, arch };
}
// ============================================================================
// PRODUCTION: Real LLM calls
// ============================================================================
export async function extractFileLLM(
filePath: string,
client?: LLMClient,
cacheDir?: string,
): Promise<LLMFileData> {
const content = readFileSync(filePath, "utf8");
const hash = createHash("sha256").update(content).digest("hex");
// Check disk cache
const cached = getCached(hash, cacheDir);
if (cached) {
const parsed = parseFileResponse(cached);
return {
purpose: parsed.purpose,
exports: [], // AST provides precise exports
deps: parsed.deps,
concepts: parsed.concepts,
};
}
// Skip very large files
const size = statSync(filePath).size;
if (size > MAX_FILE_SIZE) {
return {
purpose: "Large/generated file",
exports: [],
deps: [],
concepts: [],
};
}
// If no LLM client provided, fall back to heuristics (for backward compat / tests)
if (!client) {
return extractFileHeuristic(filePath, content);
}
const prompt = buildFilePrompt(filePath, truncateForContext(content, 200));
const response = await client.complete(prompt);
setCached(hash, response, cacheDir);
const parsed = parseFileResponse(response);
return {
purpose: parsed.purpose,
exports: [], // AST provides precise exports
deps: parsed.deps,
concepts: parsed.concepts,
};
}
export async function extractPackageLLM(
relativePath: string,
fileData: { name: string; purpose: string }[],
client?: LLMClient,
_cacheDir?: string,
): Promise<LLMPackageData> {
if (!client) {
return extractPackageHeuristic(relativePath, fileData);
}
const prompt = buildPackagePrompt(relativePath, fileData);
const response = await client.complete(prompt);
const parsed = parsePackageResponse(response);
return {
role: parsed.role || dirNameToRole(relativePath),
arch: parsed.arch || `Contains ${fileData.length} files.`,
};
}
// ============================================================================
// HEURISTIC FALLBACK (for tests / no-LLM mode)
// ============================================================================
export async function extractFileHeuristic(
filePath: string,
content?: string,
): Promise<LLMFileData> {
const fileContent = content ?? readFileSync(filePath, "utf8");
const ext = extname(filePath).toLowerCase();
const name = basename(filePath);
const baseName = basename(filePath, ext);
const exports = extractExportsHeuristic(fileContent, ext, name);
const deps = extractDepsHeuristic(fileContent, ext);
const purpose = generatePurpose(name, ext, baseName, exports);
return { purpose, exports, deps, concepts: [] };
}
export async function extractPackageHeuristic(
relativePath: string,
fileData: { name: string; purpose: string }[],
): Promise<LLMPackageData> {
const dirName = basename(relativePath);
const role = dirNameToRole(dirName);
const purposes = fileData.map((f) => f.purpose);
const hasTests = purposes.some((p) => p.includes("Test"));
const hasTypes = purposes.some((p) => p.includes("Type"));
const hasComponents = purposes.some(
(p) => p.includes("component") || p.includes("Component"),
);
const hasUtils = purposes.some((p) => p.includes("Utility"));
let arch = "";
if (hasTests) arch += "Contains tests. ";
if (hasTypes) arch += "Defines shared types. ";
if (hasComponents) arch += "Component-based architecture. ";
if (hasUtils) arch += "Shared utilities. ";
if (!arch) arch = `Contains ${fileData.length} files.`;
return { role, arch: arch.trim() };
}
// ============================================================================
// INTERNAL HEURISTIC HELPERS
// ============================================================================
function dirNameToRole(dirName: string): string {
if (dirName === ".") return "Project root";
if (dirName === "src" || dirName === "lib" || dirName === "source") return "Source code";
if (dirName === "test" || dirName === "tests" || dirName === "spec") return "Test suite";
if (dirName === "docs" || dirName === "doc") return "Documentation";
if (dirName === "config" || dirName === "configuration") return "Configuration";
if (dirName === "utils" || dirName === "helpers" || dirName === "util") return "Utility functions";
if (dirName === "types" || dirName === "type") return "Type definitions";
if (dirName === "components" || dirName === "component") return "UI components";
if (dirName === "hooks" || dirName === "hook") return "Custom hooks";
if (dirName === "api" || dirName === "apis") return "API endpoints/handlers";
if (dirName === "db" || dirName === "database" || dirName === "models") return "Database layer";
if (dirName === "auth" || dirName === "authentication") return "Authentication layer";
return `Package ${dirName}`;
}
function extractExportsHeuristic(content: string, ext: string, _filename: string): string[] {
const exports: string[] = [];
if ([".ts", ".tsx", ".js", ".jsx", ".mjs"].includes(ext)) {
const exportRegex = /(?:^|\n)\s*export\s+(?:default\s+)?(?:async\s+)?(?:function\s+|class\s+|const\s+|let\s+|var\s+|interface\s+|type\s+|enum\s+)?([A-Za-z_$][A-Za-z0-9_$]*)/g;
let match: RegExpExecArray | null;
match = exportRegex.exec(content);
while (match !== null) {
exports.push(match[1]);
match = exportRegex.exec(content);
}
const namedExportRegex = /(?:^|\n)\s*export\s*\{\s*([^}]+)\s*\}/g;
match = namedExportRegex.exec(content);
while (match !== null) {
const names = match[1].split(",").map((s) => s.trim().split(/\s+as\s+/)[0].trim());
exports.push(...names);
match = namedExportRegex.exec(content);
}
} else if (ext === ".py") {
const pyRegex = /^(?:async\s+)?def\s+([A-Za-z_][A-Za-z0-9_]*)|class\s+([A-Za-z_][A-Za-z0-9_]*)/gm;
let match: RegExpExecArray | null = pyRegex.exec(content);
while (match !== null) {
exports.push(match[1] || match[2]);
match = pyRegex.exec(content);
}
} else if (ext === ".go") {
const goRegex = /^(?:func|type|var|const)\s+([A-Z][A-Za-z0-9_]*)/gm;
let match: RegExpExecArray | null = goRegex.exec(content);
while (match !== null) {
exports.push(match[1]);
match = goRegex.exec(content);
}
} else if (ext === ".rs") {
const rsRegex = /pub\s+(?:fn|struct|enum|trait|type|const|static|use)\s+([A-Za-z_][A-Za-z0-9_]*)/g;
let match: RegExpExecArray | null = rsRegex.exec(content);
while (match !== null) {
exports.push(match[1]);
match = rsRegex.exec(content);
}
}
return [...new Set(exports)];
}
function extractDepsHeuristic(content: string, ext: string): string[] {
const deps: string[] = [];
if ([".ts", ".tsx", ".js", ".jsx", ".mjs"].includes(ext)) {
const importRegex = /import\s+(?:(?:type\s+)?\{[^}]*\}|\*\s+as\s+\w+|\w+)\s+from\s+['"]([^'"]+)['"]/g;
let match: RegExpExecArray | null = importRegex.exec(content);
while (match !== null) {
deps.push(match[1]);
match = importRegex.exec(content);
}
const requireRegex = /require\s*\(\s*['"]([^'"]+)['"]\s*\)/g;
match = requireRegex.exec(content);
while (match !== null) {
deps.push(match[1]);
match = requireRegex.exec(content);
}
} else if (ext === ".py") {
const pyImportRegex = /^(?:from|import)\s+([A-Za-z_][A-Za-z0-9_.]*)/gm;
let match: RegExpExecArray | null = pyImportRegex.exec(content);
while (match !== null) {
deps.push(match[1]);
match = pyImportRegex.exec(content);
}
} else if (ext === ".go") {
const goImportRegex = /"([^"]+)"/g;
let match: RegExpExecArray | null = goImportRegex.exec(content);
while (match !== null) {
if (match[1].includes("/")) deps.push(match[1]);
match = goImportRegex.exec(content);
}
} else if (ext === ".rs") {
const rsUseRegex = /use\s+([A-Za-z_][A-Za-z0-9_:]*)/g;
let match: RegExpExecArray | null = rsUseRegex.exec(content);
while (match !== null) {
deps.push(match[1]);
match = rsUseRegex.exec(content);
}
}
return [...new Set(deps)];
}
function generatePurpose(name: string, ext: string, _baseName: string, exports: string[]): string {
if (/test|spec/i.test(name) && exports.length === 0) return "Test suite";
if (/config|settings/i.test(name)) return "Configuration";
if (/util|helper/i.test(name)) return "Utility functions";
if (/types?\.d?\.ts$/.test(name)) return "Type definitions";
if (/index\./.test(name)) return "Module entry point";
if (/middleware/.test(name)) return "Middleware";
if (/route/.test(name)) return "Route handlers";
if (/controller/.test(name)) return "Controller";
if (/service/.test(name)) return "Service layer";
if (/model/.test(name)) return "Data model";
if (/schema/.test(name)) return "Data schema";
if (/component/.test(name) || /\.tsx$/.test(name) || /\.vue$/.test(name) || /\.svelte$/.test(name)) {
return "UI component";
}
if (/hook|use[A-Z]/.test(name)) return "React hook";
if (/style|\.css|\.scss|\.less/.test(name)) return "Styling";
if (/docker/i.test(name)) return "Container definition";
if (/\.env/.test(name)) return "Environment variables";
if (/readme/i.test(name)) return "Project documentation";
if (exports.length > 0) {
const firstFew = exports.slice(0, 3).join(", ");
if (exports.length <= 3) return `Exports: ${firstFew}`;
return `Exports ${exports.length} symbols: ${firstFew}...`;
}
return FILE_TYPE_PURPOSES[ext] || (ext ? `${ext.slice(1).toUpperCase()} file` : `${name} file`);
}
src/external-llm-client.ts → src/llm/external-llm-client.ts
View File
src/kimi-llm-client.ts → src/llm/kimi-llm-client.ts
View File
src/llm-batch.ts → src/llm/llm-batch.ts
View File
src/llm-cache.ts → src/llm/llm-cache.ts
View File
src/llm-client.ts → src/llm/llm-client.ts
View File
src/llm-error.ts → src/llm/llm-error.ts
View File
src/llm/llm-extract.ts
+217
View File
@@ -0,0 +1,217 @@
import { readFileSync, statSync } from "fs";
import { createHash } from "crypto";
import { extname, basename } from "path";
import type { LLMClient } from "./llm-client.js";
import { getCached, setCached } from "./llm-cache.js";
import { LLMError } from "./llm-error.js";
interface LLMFileData {
purpose: string;
exports: string[];
deps: string[];
concepts: string[];
}
interface LLMPackageData {
role: string;
arch: string;
}
const MAX_FILE_SIZE = 500 * 1024; // 500KB
const CONTEXT_BUDGET = 4000; // tokens
const CHARS_PER_TOKEN = 4; // approximate for ASCII
// Known binary extensions — skip without reading content
const BINARY_EXTENSIONS = new Set([
".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", ".ico", ".svgz",
".mp3", ".mp4", ".avi", ".mov", ".mkv", ".flv", ".wmv",
".wav", ".ogg", ".flac", ".aac", ".wma",
".zip", ".tar", ".gz", ".bz2", ".xz", ".7z", ".rar",
".exe", ".dll", ".so", ".dylib", ".bin",
".pdf", ".doc", ".docx", ".xls", ".xlsx", ".ppt", ".pptx",
".wasm", ".class", ".jar", ".o", ".a",
".ttf", ".otf", ".woff", ".woff2", ".eot",
".db", ".sqlite", ".sqlite3",
]);
function isBinaryFile(filePath: string): boolean {
// Fast-path: check extension
const ext = extname(filePath).toLowerCase();
if (BINARY_EXTENSIONS.has(ext)) return true;
// Fallback: sniff first 8KB for null bytes or non-printable ratio
try {
const buf = readFileSync(filePath).subarray(0, 8192);
let nonPrintable = 0;
for (let i = 0; i < buf.length; i++) {
const b = buf[i];
if (b === 0) return true; // null byte = definitely binary
if (b < 0x20 && b !== 0x09 && b !== 0x0a && b !== 0x0d) {
nonPrintable++;
}
}
// If >30% non-printable, treat as binary
return nonPrintable / buf.length > 0.3;
} catch {
return false;
}
}
function truncateForContext(content: string, promptLength: number): string {
const maxChars = CONTEXT_BUDGET * CHARS_PER_TOKEN - promptLength;
if (content.length <= maxChars) return content;
return `${content.slice(0, maxChars - 20)}\n[...truncated]`;
}
function buildFilePrompt(filePath: string, content: string): string {
const name = basename(filePath);
return `Analyze this file. Respond in this exact format (one line each):
PURPOSE: <concise one-sentence description of what this file does>
DEPS: <comma-separated list of key dependencies/modules it relies on, or "none">
CONCEPTS: <comma-separated list of key concepts/patterns used, or "none">
File: ${name}
\`\`\`
${content}
\`\`\`
`;
}
function parseFileResponse(response: string): { purpose: string; deps: string[]; concepts: string[] } {
const lines = response.split("\n");
let purpose = "";
let deps: string[] = [];
let concepts: string[] = [];
for (const line of lines) {
const trimmed = line.trim();
if (trimmed.startsWith("PURPOSE:")) {
purpose = trimmed.slice("PURPOSE:".length).trim();
} else if (trimmed.startsWith("DEPS:")) {
const depsStr = trimmed.slice("DEPS:".length).trim();
deps = depsStr === "none" ? [] : depsStr.split(",").map((s) => s.trim()).filter(Boolean);
} else if (trimmed.startsWith("CONCEPTS:")) {
const conceptsStr = trimmed.slice("CONCEPTS:".length).trim();
concepts = conceptsStr === "none" ? [] : conceptsStr.split(",").map((s) => s.trim()).filter(Boolean);
}
}
return { purpose, deps, concepts };
}
function buildPackagePrompt(relativePath: string, fileSummaries: { name: string; purpose: string }[]): string {
const filesList = fileSummaries.map((f) => `- ${f.name}: ${f.purpose}`).join("\n");
return `Analyze this code package/directory. Respond in this exact format (one line each):
ROLE: <concise one-sentence description of this package's role in the project>
ARCH: <concise description of architecture/patterns used in this package>
Directory: ${relativePath}
Files:
${filesList}
`;
}
function parsePackageResponse(response: string): { role: string; arch: string } {
const lines = response.split("\n");
let role = "";
let arch = "";
for (const line of lines) {
const trimmed = line.trim();
if (trimmed.startsWith("ROLE:")) {
role = trimmed.slice("ROLE:".length).trim();
} else if (trimmed.startsWith("ARCH:")) {
arch = trimmed.slice("ARCH:".length).trim();
}
}
return { role, arch };
}
export async function extractFileLLM(
filePath: string,
client?: LLMClient,
cacheDir?: string,
): Promise<LLMFileData> {
// LLM-only: require a client
if (!client) {
throw new LLMError(
"No LLM client configured. " +
"Set OPENAI_API_KEY / KIMI_API_KEY environment variable, " +
"or run inside Pi with a configured model.",
);
}
// Skip binary files
if (isBinaryFile(filePath)) {
return {
purpose: "Binary file",
exports: [],
deps: [],
concepts: [],
};
}
const content = readFileSync(filePath, "utf8");
const hash = createHash("sha256").update(content).digest("hex");
// Check disk cache
const cached = getCached(hash, cacheDir);
if (cached) {
const parsed = parseFileResponse(cached);
return {
purpose: parsed.purpose,
exports: [], // AST provides precise exports
deps: parsed.deps,
concepts: parsed.concepts,
};
}
// Skip very large files
const size = statSync(filePath).size;
if (size > MAX_FILE_SIZE) {
return {
purpose: "Large file",
exports: [],
deps: [],
concepts: [],
};
}
const prompt = buildFilePrompt(filePath, truncateForContext(content, 200));
const response = await client.complete(prompt);
setCached(hash, response, cacheDir);
const parsed = parseFileResponse(response);
return {
purpose: parsed.purpose,
exports: [], // AST provides precise exports
deps: parsed.deps,
concepts: parsed.concepts,
};
}
export async function extractPackageLLM(
relativePath: string,
fileData: { name: string; purpose: string }[],
client?: LLMClient,
_cacheDir?: string,
): Promise<LLMPackageData> {
// LLM-only: require a client
if (!client) {
throw new LLMError(
"No LLM client configured. " +
"Set OPENAI_API_KEY / KIMI_API_KEY environment variable, " +
"or run inside Pi with a configured model.",
);
}
const prompt = buildPackagePrompt(relativePath, fileData);
const response = await client.complete(prompt);
const parsed = parsePackageResponse(response);
return {
role: parsed.role || `Package ${basename(relativePath)}`,
arch: parsed.arch || `Contains ${fileData.length} files.`,
};
}
src/pi-llm-client.ts → src/llm/pi-llm-client.ts
View File
src/patch.ts
+3 -3
View File
@@ -1,12 +1,12 @@
import { dirname, join, basename, relative } from "path"; import { dirname, join, basename, relative } from "path";
import { existsSync, readFileSync, writeFileSync } from "fs"; import { existsSync, readFileSync, writeFileSync } from "fs";
import { parsePackageMap, renderPackageMap } from "./format.js"; import { parsePackageMap, renderPackageMap } from "./format.js";
import { extractFileLLM } from "./llm-extract.js"; import { extractFileLLM } from "./llm/llm-extract.js";
import { extractFileAST } from "./ast-extract.js"; import { extractFileAST } from "./ast/ast-extract.js";
import { mergeFileData } from "./merge.js"; import { mergeFileData } from "./merge.js";
import { generateDirectoryMap } from "./init.js"; import { generateDirectoryMap } from "./init.js";
import { readdirSync, statSync } from "fs"; import { readdirSync, statSync } from "fs";
import type { LLMClient } from "./llm-client.js"; import type { LLMClient } from "./llm/llm-client.js";
const SMALL_PACKAGE_THRESHOLD = 10; const SMALL_PACKAGE_THRESHOLD = 10;
src/validate.ts
+1 -1
View File
@@ -2,7 +2,7 @@ import { discoverProject } from "./discover.js";
import { parsePackageMap } from "./format.js"; import { parsePackageMap } from "./format.js";
import { existsSync, readFileSync } from "fs"; import { existsSync, readFileSync } from "fs";
import { join } from "path"; import { join } from "path";
import { extractFileAST } from "./ast-extract.js"; import { extractFileAST } from "./ast/ast-extract.js";
import { generateDirectoryMap } from "./init.js"; import { generateDirectoryMap } from "./init.js";
export interface ValidationResult { export interface ValidationResult {
tests/ast-extract.test.ts
+1 -1
View File
@@ -1,5 +1,5 @@
import { describe, it, expect } from "vitest"; import { describe, it, expect } from "vitest";
import { extractFileAST } from "../src/ast-extract.js"; import { extractFileAST } from "../src/ast/ast-extract.js";
import { mkdtempSync, writeFileSync } from "fs"; import { mkdtempSync, writeFileSync } from "fs";
import { join } from "path"; import { join } from "path";
import { tmpdir } from "os"; import { tmpdir } from "os";
tests/integration.test.ts
+15 -9
View File
@@ -11,6 +11,7 @@ import { tmpdir } from "os";
import { initProject } from "../src/init.js"; import { initProject } from "../src/init.js";
import { patchFile } from "../src/patch.js"; import { patchFile } from "../src/patch.js";
import { validateMaps } from "../src/validate.js"; import { validateMaps } from "../src/validate.js";
import { createMockFileClient, createMockPackageClient } from "./mock-llm.js";
describe("integration", () => { describe("integration", () => {
let dir: string; let dir: string;
@@ -26,25 +27,26 @@ describe("integration", () => {
it("init creates .pi-map.md files", async () => { it("init creates .pi-map.md files", async () => {
mkdirSync(join(dir, "src")); mkdirSync(join(dir, "src"));
writeFileSync(join(dir, "src", "index.ts"), `export function foo() {}\n`); writeFileSync(join(dir, "src", "index.ts"), `export function foo() {}\n`);
await initProject(dir); const client = createMockFileClient();
await initProject(dir, { llmClient: client, verbose: false });
const map = readFileSync(join(dir, "src", ".pi-map.md"), "utf8"); const map = readFileSync(join(dir, "src", ".pi-map.md"), "utf8");
expect(map).toContain("# src"); expect(map).toContain("# src");
expect(map).toContain("foo");
}); });
it("patch updates a file entry", async () => { it("patch updates a file entry", async () => {
mkdirSync(join(dir, "src")); mkdirSync(join(dir, "src"));
writeFileSync(join(dir, "src", "index.ts"), `export function foo() {}\n`); writeFileSync(join(dir, "src", "index.ts"), `export function foo() {}\n`);
writeFileSync(join(dir, "src", "utils.ts"), `export const bar = 1;\n`); writeFileSync(join(dir, "src", "utils.ts"), `export const bar = 1;\n`);
await initProject(dir); const client = createMockFileClient();
await initProject(dir, { llmClient: client, verbose: false });
// Modify a file // Modify a file
writeFileSync( writeFileSync(
join(dir, "src", "index.ts"), join(dir, "src", "index.ts"),
`export function foo() {}\nexport function baz() {}\n`, `export function foo() {}\nexport function baz() {}\n`,
); );
await patchFile(join(dir, "src", "index.ts")); await patchFile(join(dir, "src", "index.ts"), client, dir);
const map = readFileSync(join(dir, "src", ".pi-map.md"), "utf8"); const map = readFileSync(join(dir, "src", ".pi-map.md"), "utf8");
expect(map).toContain("baz"); expect(map).toContain("baz");
@@ -59,14 +61,15 @@ describe("integration", () => {
`export const x${i} = ${i};\n`, `export const x${i} = ${i};\n`,
); );
} }
await initProject(dir); const client = createMockFileClient();
await initProject(dir, { llmClient: client, verbose: false });
// Modify a file // Modify a file
writeFileSync( writeFileSync(
join(dir, "src", "file0.ts"), join(dir, "src", "file0.ts"),
`export const x0 = 0;\nexport const y = 99;\n`, `export const x0 = 0;\nexport const y = 99;\n`,
); );
await patchFile(join(dir, "src", "file0.ts")); await patchFile(join(dir, "src", "file0.ts"), client, dir);
const map = readFileSync(join(dir, "src", ".pi-map.md"), "utf8"); const map = readFileSync(join(dir, "src", ".pi-map.md"), "utf8");
expect(map).toContain("y"); expect(map).toContain("y");
@@ -77,7 +80,8 @@ describe("integration", () => {
it("validate detects new files", async () => { it("validate detects new files", async () => {
mkdirSync(join(dir, "src")); mkdirSync(join(dir, "src"));
writeFileSync(join(dir, "src", "a.ts"), `export const a = 1;\n`); writeFileSync(join(dir, "src", "a.ts"), `export const a = 1;\n`);
await initProject(dir); const client = createMockFileClient();
await initProject(dir, { llmClient: client, verbose: false });
// Add new file // Add new file
writeFileSync(join(dir, "src", "b.ts"), `export const b = 2;\n`); writeFileSync(join(dir, "src", "b.ts"), `export const b = 2;\n`);
@@ -91,7 +95,8 @@ describe("integration", () => {
mkdirSync(join(dir, "src")); mkdirSync(join(dir, "src"));
writeFileSync(join(dir, "src", "a.ts"), `export const a = 1;\n`); writeFileSync(join(dir, "src", "a.ts"), `export const a = 1;\n`);
writeFileSync(join(dir, "src", "b.ts"), `export const b = 2;\n`); writeFileSync(join(dir, "src", "b.ts"), `export const b = 2;\n`);
await initProject(dir); const client = createMockFileClient();
await initProject(dir, { llmClient: client, verbose: false });
// Delete a file // Delete a file
rmSync(join(dir, "src", "b.ts")); rmSync(join(dir, "src", "b.ts"));
@@ -104,7 +109,8 @@ describe("integration", () => {
it("validate detects changed signatures", async () => { it("validate detects changed signatures", async () => {
mkdirSync(join(dir, "src")); mkdirSync(join(dir, "src"));
writeFileSync(join(dir, "src", "a.ts"), `export const a = 1;\n`); writeFileSync(join(dir, "src", "a.ts"), `export const a = 1;\n`);
await initProject(dir); const client = createMockFileClient();
await initProject(dir, { llmClient: client, verbose: false });
// Change exports // Change exports
writeFileSync( writeFileSync(
tests/llm-batch.test.ts
+1 -1
View File
@@ -1,5 +1,5 @@
import { describe, it, expect } from "vitest"; import { describe, it, expect } from "vitest";
import { withRetry, processFiles } from "../src/llm-batch.js"; import { withRetry, processFiles } from "../src/llm/llm-batch.js";
import { LLMError } from "../src/llm-error.js"; import { LLMError } from "../src/llm-error.js";
describe("withRetry", () => { describe("withRetry", () => {
tests/llm-cache.test.ts
+1 -1
View File
@@ -1,5 +1,5 @@
import { describe, it, expect, beforeEach, afterEach } from "vitest"; import { describe, it, expect, beforeEach, afterEach } from "vitest";
import { getCached, setCached } from "../src/llm-cache.js"; import { getCached, setCached } from "../src/llm/llm-cache.js";
import { existsSync, unlinkSync } from "fs"; import { existsSync, unlinkSync } from "fs";
import { join } from "path"; import { join } from "path";
import { tmpdir } from "os"; import { tmpdir } from "os";
tests/llm-extract.test.ts
+46 -86
View File
@@ -1,67 +1,17 @@
import { describe, it, expect } from "vitest"; import { describe, it, expect } from "vitest";
import { extractFileLLM, extractFileHeuristic } from "../src/llm-extract.js"; import { extractFileLLM } from "../src/llm/llm-extract.js";
import { writeFileSync, mkdtempSync } from "fs"; import { writeFileSync, mkdtempSync } from "fs";
import { join } from "path"; import { join } from "path";
import { tmpdir } from "os"; import { tmpdir } from "os";
import type { LLMClient } from "../src/llm-client.js"; import type { LLMClient } from "../src/llm/llm-client.js";
describe("llm-extract heuristics", () => { function createMockClient(response: string): LLMClient {
it("extracts TypeScript exports", async () => { return {
const dir = mkdtempSync(join(tmpdir(), "pi-map-")); async complete() {
const file = join(dir, "test.ts"); return response;
writeFileSync( },
file, };
`export function foo() {} }
export class Bar {}
export const baz = 1;
export type Qux = string;
export { a, b as c };
`,
);
const result = await extractFileLLM(file);
expect(result.exports).toContain("foo");
expect(result.exports).toContain("Bar");
expect(result.exports).toContain("baz");
expect(result.exports).toContain("Qux");
expect(result.exports).toContain("a");
expect(result.exports).toContain("b");
});
it("extracts TypeScript imports", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
const file = join(dir, "test.ts");
writeFileSync(
file,
`import { foo } from "./bar";
import * as baz from "baz-lib";
import type { Qux } from "qux";
const x = require("legacy");
`,
);
const result = await extractFileLLM(file);
expect(result.deps).toContain("./bar");
expect(result.deps).toContain("baz-lib");
expect(result.deps).toContain("qux");
expect(result.deps).toContain("legacy");
});
it("infers purpose from filename patterns", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
const file = join(dir, "userController.ts");
writeFileSync(file, `export class UserController {}`);
const result = await extractFileLLM(file);
expect(result.purpose).toMatch(/Controller|Exports/);
});
it("handles non-code files", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
const file = join(dir, "Dockerfile");
writeFileSync(file, `FROM node:20\nWORKDIR /app`);
const result = await extractFileLLM(file);
expect(result.purpose).toBe("Container definition");
expect(result.exports).toEqual([]);
});
});
describe("llm-extract with mock client", () => { describe("llm-extract with mock client", () => {
it("uses LLM client when provided", async () => { it("uses LLM client when provided", async () => {
@@ -69,11 +19,9 @@ describe("llm-extract with mock client", () => {
const file = join(dir, "test.ts"); const file = join(dir, "test.ts");
writeFileSync(file, `export const foo = 1;`); writeFileSync(file, `export const foo = 1;`);
const mockClient: LLMClient = { const mockClient = createMockClient(
async complete() { "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing",
return "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing"; );
},
};
const result = await extractFileLLM(file, mockClient, tmpdir()); const result = await extractFileLLM(file, mockClient, tmpdir());
expect(result.purpose).toBe("Test file"); expect(result.purpose).toBe("Test file");
@@ -81,42 +29,54 @@ describe("llm-extract with mock client", () => {
expect(result.concepts).toContain("testing"); expect(result.concepts).toContain("testing");
}); });
it("throws without LLM client", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
const file = join(dir, "test.ts");
writeFileSync(file, `export const foo = 1;`);
await expect(extractFileLLM(file)).rejects.toThrow("No LLM client configured");
});
it("skips large files", async () => { it("skips large files", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-")); const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
const file = join(dir, "big.ts"); const file = join(dir, "big.ts");
writeFileSync(file, "x".repeat(60 * 1024)); writeFileSync(file, "x".repeat(600 * 1024));
const mockClient: LLMClient = { const mockClient = createMockClient(
async complete() { "PURPOSE: Should not call\nDEPS: none\nCONCEPTS: none",
return "PURPOSE: Should not call\nDEPS: none\nCONCEPTS: none"; );
},
};
const result = await extractFileLLM(file, mockClient); const result = await extractFileLLM(file, mockClient);
expect(result.purpose).toBe("Large/generated file"); expect(result.purpose).toBe("Large file");
}); });
it("falls back to heuristics without client", async () => { it("skips binary files", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-")); const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
const file = join(dir, "utils.ts"); const file = join(dir, "image.png");
writeFileSync(file, `export function helper() {}`); // Write some binary-looking content with null bytes
const buf = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]);
writeFileSync(file, buf);
const result = await extractFileLLM(file); const mockClient = createMockClient(
expect(result.purpose).toBe("Utility functions"); "PURPOSE: Should not call\nDEPS: none\nCONCEPTS: none",
expect(result.exports).toContain("helper"); );
const result = await extractFileLLM(file, mockClient);
expect(result.purpose).toBe("Binary file");
}); });
});
describe("extractFileHeuristic", () => { it("parses response with deps and concepts", async () => {
it("returns structured data", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-")); const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
const file = join(dir, "test.ts"); const file = join(dir, "test.ts");
writeFileSync(file, `export const x = 1;`); writeFileSync(file, `import { foo } from "bar";\nexport const x = 1;`);
const result = await extractFileHeuristic(file); const mockClient = createMockClient(
expect(result.purpose).toBeDefined(); "PURPOSE: Config module\nDEPS: bar, baz\nCONCEPTS: constants, config",
expect(Array.isArray(result.exports)).toBe(true); );
expect(Array.isArray(result.deps)).toBe(true);
expect(Array.isArray(result.concepts)).toBe(true); const result = await extractFileLLM(file, mockClient, tmpdir());
expect(result.purpose).toBe("Config module");
expect(result.deps).toEqual(["bar", "baz"]);
expect(result.concepts).toEqual(["constants", "config"]);
}); });
}); });
tests/llm-integration.test.ts
+5 -5
View File
@@ -2,9 +2,9 @@ import { describe, it, expect } from "vitest";
import { writeFileSync, mkdtempSync, readFileSync } from "fs"; import { writeFileSync, mkdtempSync, readFileSync } from "fs";
import { join } from "path"; import { join } from "path";
import { tmpdir } from "os"; import { tmpdir } from "os";
import { createLLMClient } from "../src/llm-client.js"; import { createLLMClient } from "../src/llm/llm-client.js";
import { extractFileLLM, extractPackageLLM } from "../src/llm-extract.js"; import { extractFileLLM, extractPackageLLM } from "../src/llm/llm-extract.js";
import { processFiles } from "../src/llm-batch.js"; import { processFiles } from "../src/llm/llm-batch.js";
// Load .env file manually (no dotenv dependency needed) // Load .env file manually (no dotenv dependency needed)
function loadEnv(): Record<string, string> { function loadEnv(): Record<string, string> {
@@ -135,7 +135,7 @@ describe.skipIf(!hasKimiKey)("LLM integration with Kimi", () => {
it("skips large files without calling LLM", async () => { it("skips large files without calling LLM", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-llm-")); const dir = mkdtempSync(join(tmpdir(), "pi-map-llm-"));
const file = join(dir, "big.ts"); const file = join(dir, "big.ts");
writeFileSync(file, "x".repeat(60 * 1024)); writeFileSync(file, "x".repeat(600 * 1024));
let calls = 0; let calls = 0;
const trackingClient = createLLMClient("kimi", { model: kimiModel }); const trackingClient = createLLMClient("kimi", { model: kimiModel });
@@ -146,7 +146,7 @@ describe.skipIf(!hasKimiKey)("LLM integration with Kimi", () => {
}; };
const result = await extractFileLLM(file, trackingClient, dir); const result = await extractFileLLM(file, trackingClient, dir);
expect(result.purpose).toBe("Large/generated file"); expect(result.purpose).toBe("Large file");
expect(calls).toBe(0); // Should never call LLM for large files expect(calls).toBe(0); // Should never call LLM for large files
}); });
}); });
tests/mock-llm.ts
+17
View File
@@ -0,0 +1,17 @@
import type { LLMClient } from "../src/llm/llm-client.js";
export function createMockFileClient(purpose = "Test file"): LLMClient {
return {
async complete() {
return `PURPOSE: ${purpose}\nDEPS: none\nCONCEPTS: testing`;
},
};
}
export function createMockPackageClient(): LLMClient {
return {
async complete() {
return "ROLE: Test package\nARCH: Test architecture";
},
};
}