diff --git a/pi-extension.ts b/pi-extension.ts index 5fc4491..65ec434 100644 --- a/pi-extension.ts +++ b/pi-extension.ts @@ -1,39 +1,20 @@ import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; import { Type } from "typebox"; -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"; +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"; /** - * 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. + * Get the LLM client for Pi runtime. + * + * When running inside Pi, we ALWAYS use Pi's native LLM via + * @mariozechner/pi-ai's complete() function. This respects the + * user's /model selection and /login auth. No external fetch(). */ -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 getPiLLMClient(ctx: any) { + return createLLMClient("pi", { extensionContext: ctx }); } function findPiMapFiles(cwd: string): string[] { @@ -91,18 +72,19 @@ export default function (pi: ExtensionAPI) { async execute(_toolCallId, params, _signal, _onUpdate, ctx) { try { const targetPath = params.path || ctx.cwd; - await generateMapHeuristic(targetPath); + const client = getPiLLMClient(ctx); + await initProject(targetPath, { verbose: false, llmClient: client }); return { content: [ { type: "text", - text: `Generated heuristic .pi-map.md files for ${targetPath}.`, + text: `Generated .pi-map.md files for ${targetPath}`, }, ], details: { success: true, cwd: ctx.cwd }, }; } catch (err: any) { - const msg = String(err); + const msg = err instanceof LLMError ? err.message : String(err); return { content: [{ type: "text", text: `Error: ${msg}` }], details: { success: false, error: msg }, @@ -128,8 +110,8 @@ export default function (pi: ExtensionAPI) { }), async execute(_toolCallId, params, _signal, _onUpdate, ctx) { try { - const dirPath = dirname(params.file_path); - await generateMapHeuristic(dirPath); + const client = getPiLLMClient(ctx); + await patchFile(params.file_path, client); return { content: [ { @@ -140,7 +122,7 @@ export default function (pi: ExtensionAPI) { details: { success: true }, }; } catch (err: any) { - const msg = String(err); + const msg = err instanceof LLMError ? err.message : String(err); return { content: [{ type: "text", text: `Error: ${msg}` }], details: { success: false, error: msg }, @@ -183,7 +165,7 @@ export default function (pi: ExtensionAPI) { details: { success: true, clean: result.clean }, }; } catch (err: any) { - const msg = String(err); + const msg = err instanceof LLMError ? err.message : String(err); return { content: [{ type: "text", text: `Error: ${msg}` }], details: { success: false, error: msg }, @@ -211,18 +193,19 @@ export default function (pi: ExtensionAPI) { async execute(_toolCallId, params, _signal, _onUpdate, ctx) { try { const targetPath = params.path || ctx.cwd; - await generateMapHeuristic(targetPath); + const client = getPiLLMClient(ctx); + await reinitPath(targetPath, { verbose: false, llmClient: client }); return { content: [ { type: "text", - text: `Regenerated heuristic maps for ${targetPath}`, + text: `Regenerated maps for ${targetPath}`, }, ], details: { success: true }, }; } catch (err: any) { - const msg = String(err); + const msg = err instanceof LLMError ? err.message : String(err); return { content: [{ type: "text", text: `Error: ${msg}` }], details: { success: false, error: msg }, diff --git a/src/pi-llm-client.ts b/src/pi-llm-client.ts index 7a1a2ad..e90d80a 100644 --- a/src/pi-llm-client.ts +++ b/src/pi-llm-client.ts @@ -2,10 +2,11 @@ import { LLMError } from "./llm-error.js"; import type { LLMClient } from "./llm-client.js"; /** - * Pi LLM Client — uses Pi's built-in model when running inside the Pi runtime. + * Pi LLM Client — uses Pi's built-in `complete()` from `@mariozechner/pi-ai`. * - * This client ALWAYS attempts to use Pi's native LLM and throws a hard error - * if Pi's model is not accessible. It never falls back to external APIs. + * This runs inside the Pi extension runtime and calls Pi's internal AI layer, + * which respects the user's configured model (/model) and auth (/login). + * No external fetch() — Pi handles transport, retries, and token accounting. */ export class PiLLMClient implements LLMClient { constructor(private extensionContext?: unknown) {} @@ -16,57 +17,66 @@ export class PiLLMClient implements LLMClient { if (!ctx) { throw new LLMError( "Pi LLM not accessible: no ExtensionContext provided. " + - "This extension must run inside Pi.", + "This tool must run inside Pi.", ); } - // Try to access Pi's model or modelRegistry + // Get the active model from Pi's runtime const model = ctx.model ?? ctx.modelRegistry?.get?.(); if (!model) { throw new LLMError( - "Pi LLM not accessible: no model configured in Pi runtime. " + - "Configure a model (e.g., via /model) before using project-map tools.", + "Pi LLM not accessible: no model configured. " + + "Set a model via /model before using project-map tools.", ); } - // Attempt to call the model via Pi's provider interface. - // Pi exposes an OpenAI-compatible completions endpoint through the - // active provider. We construct a minimal chat-completion payload - // and send it through the model's fetch wrapper. try { - const provider = ctx.modelRegistry?.getProvider?.(model.provider); - if (!provider || typeof provider.complete !== "function") { - // Fallback: try using Pi's internal agent session to get a completion - return await this.callViaAgentSession(ctx, prompt); + // Dynamically import Pi's AI module (available in the Pi runtime) + const { complete } = await import("@mariozechner/pi-ai"); + + const response = await complete( + model, + { + systemPrompt: + "You are a code analysis assistant. Analyze the provided file and respond with concise, structured information.", + messages: [ + { + role: "user", + content: prompt, + timestamp: Date.now(), + }, + ], + }, + { + temperature: 0.1, + maxTokens: 256, + }, + ); + + if (response.errorMessage) { + throw new LLMError(`Pi LLM error: ${response.errorMessage}`); } - return await provider.complete(prompt); + + // Extract text content from AssistantMessage + const text = response.content + .filter((c: any) => c.type === "text") + .map((c: any) => c.text) + .join("") + .trim(); + + return text; } catch (err: any) { if (err instanceof LLMError) throw err; + if (err.code === "MODULE_NOT_FOUND") { + throw new LLMError( + "Pi LLM not accessible: @mariozechner/pi-ai is not available. " + + "This extension must run inside the Pi runtime.", + ); + } throw new LLMError( - `Pi LLM call failed: ${err.message || String(err)}. ` + - "Ensure Pi has a working model configured.", + `Pi LLM call failed: ${err.message || String(err)}`, err, ); } } - - private async callViaAgentSession(ctx: any, prompt: string): Promise { - // Pi's ExtensionContext may expose a session or agent that can - // make model calls. This is the fallback when direct provider - // access is not available. - if (typeof ctx.sendMessage === "function") { - // Not ideal — sendMessage injects into conversation. - // Reserved for future Pi SDK improvement. - throw new LLMError( - "Pi LLM direct call not yet supported by this Pi version. " + - "The extension requires a Pi runtime with exposed model provider API.", - ); - } - - throw new LLMError( - "Pi LLM not accessible: the Pi runtime does not expose a callable model API. " + - "This extension requires Pi to provide either ctx.modelRegistry.getProvider() " + - "or a direct model completion interface.", - ); - } } diff --git a/src/types/pi-ai.d.ts b/src/types/pi-ai.d.ts new file mode 100644 index 0000000..30976e0 --- /dev/null +++ b/src/types/pi-ai.d.ts @@ -0,0 +1,24 @@ +// Pi AI module — available only inside the Pi runtime +// Declared here so TypeScript doesn't error during `npm run build` + +declare module "@mariozechner/pi-ai" { + export function complete( + model: any, + context: { + systemPrompt?: string; + messages: Array<{ + role: "user" | "assistant" | "system" | "toolResult"; + content: string; + timestamp?: number; + }>; + }, + options?: { + temperature?: number; + maxTokens?: number; + }, + ): Promise<{ + role: "assistant"; + content: Array<{ type: string; text?: string }>; + errorMessage?: string; + }>; +} diff --git a/tests/pi-extension.test.ts b/tests/pi-extension.test.ts index a1154b3..b0d4e9b 100644 --- a/tests/pi-extension.test.ts +++ b/tests/pi-extension.test.ts @@ -15,14 +15,17 @@ vi.mock("typebox", () => ({ }, })); +// Mock Pi's AI module so PiLLMClient works in tests +vi.mock("@mariozechner/pi-ai", () => ({ + complete: vi.fn(async (_model: any, _context: any) => ({ + role: "assistant" as const, + content: [{ type: "text", text: "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing" }], + })), +})); + import extension from "../pi-extension.js"; -const CACHE_FILE = join( - homedir(), - ".cache", - "pi-project-map", - "llm-cache.json", -); +const CACHE_FILE = join(homedir(), ".cache", "pi-project-map", "llm-cache.json"); function clearCache() { if (existsSync(CACHE_FILE)) unlinkSync(CACHE_FILE); } @@ -40,7 +43,10 @@ describe("pi-extension", () => { mockNotify = vi.fn(); mockCtx = { cwd: "/home/project", - modelRegistry: {}, + model: { provider: "openai", id: "gpt-4o-mini", api: "openai-completions" }, + modelRegistry: { + getApiKeyAndHeaders: vi.fn(async () => ({ apiKey: "test-key", headers: {} })), + }, ui: { notify: mockNotify }, }; @@ -72,7 +78,7 @@ describe("pi-extension", () => { }); describe("project_map_init tool", () => { - it("generates heuristic maps successfully", async () => { + it("generates maps with mocked Pi LLM", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); writeFileSync(join(dir, "test.ts"), `export const x = ${Date.now()};`); mockCtx.cwd = dir; @@ -80,12 +86,24 @@ describe("pi-extension", () => { const tool = registeredTools.project_map_init; const result = await tool.execute("tool-1", {}, null, null, mockCtx); expect(result.details.success).toBe(true); - expect(result.content[0].text).toContain("Generated heuristic"); + expect(result.content[0].text).toContain("Generated"); + }); + + it("returns error when no Pi model configured", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, "test.ts"), `export const x = ${Date.now()};`); + mockCtx.cwd = dir; + mockCtx.model = null; + + 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("no model configured"); }); }); describe("project_map_patch tool", () => { - it("patches map successfully using heuristics", async () => { + it("patches map with mocked Pi LLM", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); const file = join(dir, "test.ts"); writeFileSync(file, `export const x = ${Date.now()};`); @@ -132,7 +150,7 @@ describe("pi-extension", () => { }); describe("project_map_reinit tool", () => { - it("regenerates heuristic maps successfully", async () => { + it("regenerates maps with mocked Pi LLM", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); writeFileSync(join(dir, "test.ts"), `export const x = ${Date.now()};`); mockCtx.cwd = dir; @@ -140,7 +158,7 @@ describe("pi-extension", () => { const tool = registeredTools.project_map_reinit; const result = await tool.execute("tool-1", {}, null, null, mockCtx); expect(result.details.success).toBe(true); - expect(result.content[0].text).toContain("Regenerated heuristic"); + expect(result.content[0].text).toContain("Regenerated"); }); });