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

feat: Pi extension uses Pi's internal complete() for LLM analysis

Browse Source 
- PiLLMClient now imports @mariozechner/pi-ai's complete() function
- Respects user's /model selection and /login auth
- No external fetch() — Pi handles transport internally
- Added src/types/pi-ai.d.ts for TypeScript declarations
- Extension reverted to real LLM calls (initProject, patchFile, reinitPath)
- Tests mock @mariozechner/pi-ai for verification
This commit is contained in:
Alex
2026-06-10 14:47:47 +02:00
parent 57c0f08210
commit 5c6719817f
4 changed files with 126 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pi-extension.ts
+24 -41
View File
@@ -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<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 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 },
src/pi-llm-client.ts
+48 -38
View File
@@ -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<string> {
// 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.",
);
}
}
src/types/pi-ai.d.ts
Vendored
+24
View File
@@ -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;
}>;
}
tests/pi-extension.test.ts
+30 -12
View File
@@ -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");
});
});