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

refactor: Pi extension uses heuristics+AST instead of direct LLM calls

Browse Source 
Inside Pi, the extension no longer attempts to call the LLM directly
(which Pi's ExtensionAPI doesn't support). Instead:

- project_map_init: generates maps using AST + heuristics (fast, free)
- project_map_patch: rewrites the file's directory with heuristics
- project_map_reinit: full heuristic regeneration
- project_map_validate: unchanged (no LLM needed)

The CLI still supports real LLM calls via --llm-provider=kimi|openai.
This separates concerns:
- Pi extension: deterministic structural analysis
- CLI: rich semantic analysis with configurable LLM
This commit is contained in:
Alex
2026-06-10 10:28:15 +02:00
parent 7db9202658
commit 57c0f08210
2 changed files with 50 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pi-extension.ts
+41 -30
View File
@@ -1,26 +1,39 @@
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type } from "typebox";
import { readFileSync, readdirSync, statSync } from "fs";
import { join, relative } from "path";
import {
initProject,
patchFile,
validateMaps,
reinitPath,
} from "./src/index.js";
import { createLLMClient } from "./src/llm-client.js";
import { LLMError } from "./src/llm-error.js";
import { readFileSync, readdirSync, statSync, writeFileSync } from "fs";
import { join, relative, dirname } from "path";
import { validateMaps } from "./src/index.js";
import { discoverProject } from "./src/discover.js";
import { extractFileAST } from "./src/ast-extract.js";
import { extractFileHeuristic, extractPackageHeuristic } from "./src/llm-extract.js";
import { mergeFileData } from "./src/merge.js";
import { renderPackageMap, type FileEntry } from "./src/format.js";
/**
* Get the LLM client for Pi runtime.
*
* When running inside Pi, we ALWAYS use Pi's native LLM.
* If Pi's LLM is not accessible, this throws a hard error.
* We ignore all external configuration (env vars, config files, etc.)
* because inside Pi we must use Pi's model exclusively.
* Generate maps using heuristics + AST only (no LLM).
* Inside Pi, we let Pi's agent handle semantic analysis.
* The tool does structural work; the agent can refine later.
*/
function getPiLLMClient(ctx: any) {
return createLLMClient("pi", { extensionContext: ctx });
async function generateMapHeuristic(rootPath: string): Promise<void> {
const entries = discoverProject(rootPath);
for (const entry of entries) {
const fileData: FileEntry[] = [];
for (const file of entry.files) {
const filePath = join(entry.dirPath, file);
const astData = await extractFileAST(filePath);
const heuristicData = await extractFileHeuristic(filePath);
fileData.push(mergeFileData(file, heuristicData, astData));
}
const packageData = await extractPackageHeuristic(entry.relativePath, fileData);
const mapData = {
path: entry.relativePath,
role: packageData.role,
files: fileData,
arch: packageData.arch,
dirty: "-",
};
writeFileSync(join(entry.dirPath, ".pi-map.md"), renderPackageMap(mapData));
}
}
function findPiMapFiles(cwd: string): string[] {
@@ -78,19 +91,18 @@ export default function (pi: ExtensionAPI) {
async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
try {
const targetPath = params.path || ctx.cwd;
const client = getPiLLMClient(ctx);
await initProject(targetPath, { verbose: false, llmClient: client });
await generateMapHeuristic(targetPath);
return {
content: [
{
type: "text",
text: `Generated .pi-map.md files for ${targetPath}`,
text: `Generated heuristic .pi-map.md files for ${targetPath}.`,
},
],
details: { success: true, cwd: ctx.cwd },
};
} catch (err: any) {
const msg = err instanceof LLMError ? err.message : String(err);
const msg = String(err);
return {
content: [{ type: "text", text: `Error: ${msg}` }],
details: { success: false, error: msg },
@@ -116,8 +128,8 @@ export default function (pi: ExtensionAPI) {
}),
async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
try {
const client = getPiLLMClient(ctx);
await patchFile(params.file_path, client);
const dirPath = dirname(params.file_path);
await generateMapHeuristic(dirPath);
return {
content: [
{
@@ -128,7 +140,7 @@ export default function (pi: ExtensionAPI) {
details: { success: true },
};
} catch (err: any) {
const msg = err instanceof LLMError ? err.message : String(err);
const msg = String(err);
return {
content: [{ type: "text", text: `Error: ${msg}` }],
details: { success: false, error: msg },
@@ -171,7 +183,7 @@ export default function (pi: ExtensionAPI) {
details: { success: true, clean: result.clean },
};
} catch (err: any) {
const msg = err instanceof LLMError ? err.message : String(err);
const msg = String(err);
return {
content: [{ type: "text", text: `Error: ${msg}` }],
details: { success: false, error: msg },
@@ -199,19 +211,18 @@ export default function (pi: ExtensionAPI) {
async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
try {
const targetPath = params.path || ctx.cwd;
const client = getPiLLMClient(ctx);
await reinitPath(targetPath, { verbose: false, llmClient: client });
await generateMapHeuristic(targetPath);
return {
content: [
{
type: "text",
text: `Regenerated maps for ${targetPath}`,
text: `Regenerated heuristic maps for ${targetPath}`,
},
],
details: { success: true },
};
} catch (err: any) {
const msg = err instanceof LLMError ? err.message : String(err);
const msg = String(err);
return {
content: [{ type: "text", text: `Error: ${msg}` }],
details: { success: false, error: msg },
tests/pi-extension.test.ts
+9 -15
View File
@@ -72,29 +72,23 @@ describe("pi-extension", () => {
});
describe("project_map_init tool", () => {
it("returns error when Pi LLM is not accessible", async () => {
// Create a temp dir with a file so initProject tries to use LLM
it("generates heuristic maps successfully", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, "test.ts"), `export const x = ${Date.now()};`);
mockCtx.cwd = dir;
const tool = registeredTools.project_map_init;
const result = await tool.execute("tool-1", {}, null, null, mockCtx);
expect(result.details.success).toBe(false);
expect(result.content[0].text).toContain("Pi LLM not accessible");
expect(result.details.success).toBe(true);
expect(result.content[0].text).toContain("Generated heuristic");
});
});
describe("project_map_patch tool", () => {
it("returns error when Pi LLM is not accessible", async () => {
it("patches map successfully using heuristics", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
const file = join(dir, "test.ts");
writeFileSync(file, `export const x = ${Date.now()};`);
// Create a .pi-map.md so patchFile doesn't return early
writeFileSync(
join(dir, ".pi-map.md"),
"# .\n## role\nTest\n## files\n## arch\n## dirty\n-\n",
);
mockCtx.cwd = dir;
const tool = registeredTools.project_map_patch;
@@ -105,8 +99,8 @@ describe("pi-extension", () => {
null,
mockCtx,
);
expect(result.details.success).toBe(false);
expect(result.content[0].text).toContain("Pi LLM not accessible");
expect(result.details.success).toBe(true);
expect(result.content[0].text).toContain("Patched map");
});
});
@@ -138,15 +132,15 @@ describe("pi-extension", () => {
});
describe("project_map_reinit tool", () => {
it("returns error when Pi LLM is not accessible", async () => {
it("regenerates heuristic maps successfully", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, "test.ts"), `export const x = ${Date.now()};`);
mockCtx.cwd = dir;
const tool = registeredTools.project_map_reinit;
const result = await tool.execute("tool-1", {}, null, null, mockCtx);
expect(result.details.success).toBe(false);
expect(result.content[0].text).toContain("Pi LLM not accessible");
expect(result.details.success).toBe(true);
expect(result.content[0].text).toContain("Regenerated heuristic");
});
});