diff --git a/pi-extension.ts b/pi-extension.ts index 99b2606..5fc4491 100644 --- a/pi-extension.ts +++ b/pi-extension.ts @@ -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 { + 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 }, diff --git a/tests/pi-extension.test.ts b/tests/pi-extension.test.ts index 6ad2007..a1154b3 100644 --- a/tests/pi-extension.test.ts +++ b/tests/pi-extension.test.ts @@ -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"); }); });