feat: M7 proper LLM integration with dual providers, caching, and parallel batching

- Add LLM client abstraction (src/llm-client.ts) with factory pattern
- Add OpenAI-compatible external client (src/external-llm-client.ts)
- Add Kimi.com client using Anthropic-based API (src/kimi-llm-client.ts)
- Add Pi native LLM stub (src/pi-llm-client.ts) for future ExtensionAPI wiring
- Add SHA-256 disk cache at ~/.cache/pi-project-map/ (src/llm-cache.ts)
- Add parallel batching with p-limit, retry + exponential backoff (src/llm-batch.ts)
- Rewrite llm-extract.ts to use real LLM calls with structured prompts
  - File-level: PURPOSE, DEPS, CONCEPTS
  - Package-level: ROLE, ARCH
  - Context truncation, 50KB skip, cache before LLM call
- Wire CLI with --llm-provider, --llm-model, --llm-base-url flags
- Update config.ts with llmProvider, llmBaseUrl fields
- Update init.ts and patch.ts to accept optional LLMClient
- Add sample project fixture for manual testing
- Add tests: llm-cache (3), llm-batch (5), llm-integration (8 with real Kimi API),
  pi-extension (14 mocked)
- All 56 tests pass
This commit is contained in:
2026-06-09 22:49:34 +02:00
parent 7b67205d43
commit 69d3acda5d
32 changed files with 1565 additions and 264 deletions
+56
View File
@@ -0,0 +1,56 @@
import { describe, it, expect } from "vitest";
import { withRetry, processFiles } from "../src/llm-batch.js";
import { LLMError } from "../src/llm-error.js";
describe("withRetry", () => {
it("returns result on first success", async () => {
const result = await withRetry(async () => "success");
expect(result).toBe("success");
});
it("retries on failure and eventually succeeds", async () => {
let attempts = 0;
const result = await withRetry(async () => {
attempts++;
if (attempts < 3) throw new Error("transient");
return "success";
});
expect(result).toBe("success");
expect(attempts).toBe(3);
});
it("throws after max retries", async () => {
let attempts = 0;
await expect(
withRetry(async () => {
attempts++;
throw new Error("persistent");
}, { maxRetries: 2, retryDelaysMs: [10, 20] }),
).rejects.toThrow("persistent");
expect(attempts).toBe(3); // initial + 2 retries
});
});
describe("processFiles", () => {
it("processes all files in parallel", async () => {
const files = [1, 2, 3, 4, 5];
const results = await processFiles(files, async (n) => n * 2, {
concurrency: 2,
});
expect(results).toEqual([2, 4, 6, 8, 10]);
});
it("retries failed files", async () => {
let attempts = 0;
const results = await processFiles(
[1],
async () => {
attempts++;
if (attempts < 2) throw new Error("fail");
return "ok";
},
{ maxRetries: 2, retryDelaysMs: [10, 20] },
);
expect(results).toEqual(["ok"]);
});
});
+40
View File
@@ -0,0 +1,40 @@
import { describe, it, expect, beforeEach, afterEach } from "vitest";
import { getCached, setCached } from "../src/llm-cache.js";
import { existsSync, unlinkSync, rmdirSync } from "fs";
import { join } from "path";
import { homedir } from "os";
const TEST_CACHE_DIR = join(homedir(), ".cache", "pi-project-map");
const TEST_CACHE_FILE = join(TEST_CACHE_DIR, "llm-cache.json");
describe("llm-cache", () => {
beforeEach(() => {
if (existsSync(TEST_CACHE_FILE)) {
unlinkSync(TEST_CACHE_FILE);
}
});
afterEach(() => {
if (existsSync(TEST_CACHE_FILE)) {
unlinkSync(TEST_CACHE_FILE);
}
});
it("returns undefined for missing entries", () => {
const result = getCached("nonexistent-hash");
expect(result).toBeUndefined();
});
it("stores and retrieves cached results", () => {
setCached("abc123", "PURPOSE: test\nDEPS: none\nCONCEPTS: none");
const result = getCached("abc123");
expect(result).toBe("PURPOSE: test\nDEPS: none\nCONCEPTS: none");
});
it("overwrites existing entries", () => {
setCached("abc123", "old");
setCached("abc123", "new");
const result = getCached("abc123");
expect(result).toBe("new");
});
});
+60 -1
View File
@@ -1,8 +1,9 @@
import { describe, it, expect } from "vitest";
import { extractFileLLM } from "../src/llm-extract.js";
import { extractFileLLM, extractFileHeuristic } from "../src/llm-extract.js";
import { writeFileSync, mkdtempSync } from "fs";
import { join } from "path";
import { tmpdir } from "os";
import type { LLMClient } from "../src/llm-client.js";
describe("llm-extract heuristics", () => {
it("extracts TypeScript exports", async () => {
@@ -61,3 +62,61 @@ const x = require("legacy");
expect(result.exports).toEqual([]);
});
});
describe("llm-extract with mock client", () => {
it("uses LLM client when provided", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
const file = join(dir, "test.ts");
writeFileSync(file, `export const foo = 1;`);
const mockClient: LLMClient = {
async complete() {
return "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing";
},
};
const result = await extractFileLLM(file, mockClient);
expect(result.purpose).toBe("Test file");
expect(result.deps).toEqual([]);
expect(result.concepts).toContain("testing");
});
it("skips large files", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
const file = join(dir, "big.ts");
writeFileSync(file, "x".repeat(60 * 1024));
const mockClient: LLMClient = {
async complete() {
return "PURPOSE: Should not call\nDEPS: none\nCONCEPTS: none";
},
};
const result = await extractFileLLM(file, mockClient);
expect(result.purpose).toBe("Large/generated file");
});
it("falls back to heuristics without client", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
const file = join(dir, "utils.ts");
writeFileSync(file, `export function helper() {}`);
const result = await extractFileLLM(file);
expect(result.purpose).toBe("Utility functions");
expect(result.exports).toContain("helper");
});
});
describe("extractFileHeuristic", () => {
it("returns structured data", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
const file = join(dir, "test.ts");
writeFileSync(file, `export const x = 1;`);
const result = await extractFileHeuristic(file);
expect(result.purpose).toBeDefined();
expect(Array.isArray(result.exports)).toBe(true);
expect(Array.isArray(result.deps)).toBe(true);
expect(Array.isArray(result.concepts)).toBe(true);
});
});
+178
View File
@@ -0,0 +1,178 @@
import { describe, it, expect } from "vitest";
import { writeFileSync, mkdtempSync, readFileSync } from "fs";
import { join } from "path";
import { tmpdir } from "os";
import { createLLMClient } from "../src/llm-client.js";
import { extractFileLLM, extractPackageLLM } from "../src/llm-extract.js";
import { processFiles } from "../src/llm-batch.js";
// Load .env file manually (no dotenv dependency needed)
function loadEnv(): Record<string, string> {
const env: Record<string, string> = {};
try {
const content = readFileSync(".env", "utf8");
for (const line of content.split("\n")) {
const match = line.match(/^([A-Za-z_][A-Za-z0-9_]*)=(.*)$/);
if (match) env[match[1]] = match[2];
}
} catch {
// No .env file
}
return env;
}
const env = loadEnv();
const kimiKey = env.KIMI_API_KEY || process.env.KIMI_API_KEY;
const kimiModel =
env.KIMI_MODEL ||
env.LLM_MODEL ||
process.env.KIMI_MODEL ||
process.env.LLM_MODEL ||
"kimi-k2-6";
const hasKimiKey = !!kimiKey;
// Set env vars so the clients can pick them up
if (env.KIMI_API_KEY) process.env.KIMI_API_KEY = env.KIMI_API_KEY;
if (env.LLM_MODEL) process.env.LLM_MODEL = env.LLM_MODEL;
describe.skipIf(!hasKimiKey)("LLM integration with Kimi", () => {
it("creates kimi client and calls complete", async () => {
const client = createLLMClient("kimi", { model: kimiModel });
const response = await client.complete(
"PURPOSE: test\nAnalyze this: export const x = 1;",
);
expect(typeof response).toBe("string");
expect(response.length).toBeGreaterThan(0);
console.log(" complete() response:", response.slice(0, 120));
}, 30000);
it("extracts file purpose with real LLM", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-llm-"));
const file = join(dir, "config.ts");
writeFileSync(
file,
`export const API_URL = "https://api.example.com";\nexport const TIMEOUT = 5000;`,
);
const client = createLLMClient("kimi", { model: kimiModel });
const result = await extractFileLLM(file, client);
expect(result.purpose).toBeTruthy();
expect(result.purpose.length).toBeGreaterThan(5);
expect(Array.isArray(result.deps)).toBe(true);
expect(Array.isArray(result.concepts)).toBe(true);
console.log(" File purpose:", result.purpose);
console.log(" Concepts:", result.concepts.join(", ") || "none");
}, 30000);
it("extracts package role with real LLM", async () => {
const client = createLLMClient("kimi", { model: kimiModel });
const result = await extractPackageLLM(
"src/utils",
[
{ name: "http.ts", purpose: "HTTP client wrapper" },
{ name: "cache.ts", purpose: "In-memory cache" },
{ name: "retry.ts", purpose: "Retry logic with backoff" },
],
client,
);
expect(result.role).toBeTruthy();
expect(result.role.length).toBeGreaterThan(5);
expect(result.arch).toBeTruthy();
console.log(" Package role:", result.role);
console.log(" Package arch:", result.arch);
}, 30000);
it("caches LLM results on disk", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-llm-"));
const file = join(dir, "test.ts");
writeFileSync(file, `export const version = "1.0.0";`);
const client = createLLMClient("kimi", { model: kimiModel });
const result1 = await extractFileLLM(file, client);
expect(result1.purpose).toBeTruthy();
// Second call should hit cache — much faster
const start = Date.now();
const result2 = await extractFileLLM(file, client);
const elapsed = Date.now() - start;
expect(result2.purpose).toBe(result1.purpose);
expect(elapsed).toBeLessThan(500); // Cache hit should be fast
console.log(" Cache hit time:", elapsed, "ms");
}, 30000);
it("processes multiple files in parallel", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-llm-"));
const files: string[] = [];
for (let i = 0; i < 3; i++) {
const f = join(dir, `file${i}.ts`);
writeFileSync(
f,
`export const val${i} = ${i};\n// Some logic here\nexport function helper${i}() { return val${i}; }`,
);
files.push(f);
}
const client = createLLMClient("kimi", { model: kimiModel });
const start = Date.now();
const results = await processFiles(
files,
async (f) => extractFileLLM(f, client),
{ concurrency: 3, maxRetries: 1, retryDelaysMs: [2000] },
);
const elapsed = Date.now() - start;
expect(results.length).toBe(3);
for (const r of results) {
expect(r.purpose).toBeTruthy();
expect(r.purpose.length).toBeGreaterThan(5);
}
console.log(" Parallel processing:", elapsed, "ms for 3 files");
}, 60000);
it("skips large files without calling LLM", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-map-llm-"));
const file = join(dir, "big.ts");
writeFileSync(file, "x".repeat(60 * 1024));
let calls = 0;
const trackingClient = createLLMClient("kimi", { model: kimiModel });
const originalComplete = trackingClient.complete.bind(trackingClient);
trackingClient.complete = async (...args) => {
calls++;
return originalComplete(...args);
};
const result = await extractFileLLM(file, trackingClient);
expect(result.purpose).toBe("Large/generated file");
expect(calls).toBe(0); // Should never call LLM for large files
});
});
describe("LLM integration without env vars", () => {
it("throws clear error when API key is missing", () => {
const saved = process.env.KIMI_API_KEY;
delete process.env.KIMI_API_KEY;
try {
expect(() => createLLMClient("kimi", { apiKey: undefined })).toThrow(
"No Kimi API key",
);
} finally {
if (saved) process.env.KIMI_API_KEY = saved;
}
});
it("throws clear error for OpenAI without key", () => {
const saved = process.env.OPENAI_API_KEY;
delete process.env.OPENAI_API_KEY;
try {
expect(() => createLLMClient("openai", { apiKey: undefined })).toThrow(
"No OpenAI API key",
);
} finally {
if (saved) process.env.OPENAI_API_KEY = saved;
}
});
});
+245
View File
@@ -0,0 +1,245 @@
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
// Mock child_process before importing the extension
vi.mock("child_process", () => ({
execSync: vi.fn(),
}));
vi.mock("fs", async () => {
const actual = await vi.importActual<typeof import("fs")>("fs");
return {
...actual,
readFileSync: vi.fn(),
};
});
vi.mock("@mariozechner/pi-coding-agent", () => ({
ExtensionAPI: class {},
}));
vi.mock("typebox", () => ({
Type: {
Object: (props: unknown) => props,
Optional: (prop: unknown) => prop,
String: (opts: unknown) => ({ type: "string", ...opts }),
},
}));
import extension from "../pi-extension.js";
import { execSync } from "child_process";
import { readFileSync } from "fs";
describe("pi-extension", () => {
let registeredTools: Record<string, any> = {};
let registeredEvents: Record<string, any> = {};
let mockCtx: any;
let mockNotify: ReturnType<typeof vi.fn>;
beforeEach(() => {
registeredTools = {};
registeredEvents = {};
mockNotify = vi.fn();
mockCtx = {
cwd: "/home/project",
ui: { notify: mockNotify },
};
const mockPi = {
registerTool: vi.fn((tool: any) => {
registeredTools[tool.name] = tool;
}),
on: vi.fn((event: string, handler: any) => {
registeredEvents[event] = handler;
}),
};
extension(mockPi);
vi.clearAllMocks();
});
afterEach(() => {
vi.restoreAllMocks();
});
describe("tool registration", () => {
it("registers 4 tools", () => {
expect(Object.keys(registeredTools)).toHaveLength(4);
expect(registeredTools).toHaveProperty("project_map_init");
expect(registeredTools).toHaveProperty("project_map_patch");
expect(registeredTools).toHaveProperty("project_map_validate");
expect(registeredTools).toHaveProperty("project_map_reinit");
});
it("registers session_start and before_agent_start events", () => {
expect(registeredEvents).toHaveProperty("session_start");
expect(registeredEvents).toHaveProperty("before_agent_start");
});
});
describe("project_map_init tool", () => {
it("calls npx project-map init with default cwd", async () => {
vi.mocked(execSync).mockReturnValue("Generated 5 .pi-map.md files\n");
const tool = registeredTools.project_map_init;
const result = await tool.execute("tool-1", {}, null, null, mockCtx);
expect(execSync).toHaveBeenCalledWith(
expect.stringContaining("npx project-map init"),
expect.objectContaining({ cwd: "/home/project" }),
);
expect(result.details.success).toBe(true);
});
it("calls npx project-map init with custom path", async () => {
vi.mocked(execSync).mockReturnValue("Generated 3 .pi-map.md files\n");
const tool = registeredTools.project_map_init;
const result = await tool.execute(
"tool-1",
{ path: "./src" },
null,
null,
mockCtx,
);
expect(execSync).toHaveBeenCalledWith(
expect.stringContaining("npx project-map init ./src"),
expect.anything(),
);
expect(result.details.success).toBe(true);
});
it("reports failure on error", async () => {
vi.mocked(execSync).mockImplementation(() => {
const err = new Error("Command failed") as any;
err.stdout = "";
err.stderr = "No such file";
throw err;
});
const tool = registeredTools.project_map_init;
const result = await tool.execute("tool-1", {}, null, null, mockCtx);
expect(result.details.success).toBe(false);
});
});
describe("project_map_patch tool", () => {
it("calls npx project-map patch with file path", async () => {
vi.mocked(execSync).mockReturnValue("Patched src/foo.ts\n");
const tool = registeredTools.project_map_patch;
const result = await tool.execute(
"tool-1",
{ file_path: "src/components/Button.tsx" },
null,
null,
mockCtx,
);
expect(execSync).toHaveBeenCalledWith(
expect.stringContaining(
"npx project-map patch src/components/Button.tsx",
),
expect.anything(),
);
expect(result.details.success).toBe(true);
});
});
describe("project_map_validate tool", () => {
it("calls npx project-map validate and reports clean", async () => {
vi.mocked(execSync).mockReturnValue("All .pi-map.md files are clean.\n");
const tool = registeredTools.project_map_validate;
const result = await tool.execute("tool-1", {}, null, null, mockCtx);
expect(execSync).toHaveBeenCalledWith(
expect.stringContaining("npx project-map validate"),
expect.anything(),
);
expect(result.details.clean).toBe(true);
});
it("reports not clean when output lacks 'clean'", async () => {
vi.mocked(execSync).mockReturnValue("Found 2 discrepancies\n");
const tool = registeredTools.project_map_validate;
const result = await tool.execute("tool-1", {}, null, null, mockCtx);
expect(result.details.clean).toBe(false);
});
});
describe("project_map_reinit tool", () => {
it("calls npx project-map reinit with path", async () => {
vi.mocked(execSync).mockReturnValue("Regenerated 7 files\n");
const tool = registeredTools.project_map_reinit;
const result = await tool.execute(
"tool-1",
{ path: "./src" },
null,
null,
mockCtx,
);
expect(execSync).toHaveBeenCalledWith(
expect.stringContaining("npx project-map reinit ./src"),
expect.anything(),
);
expect(result.details.success).toBe(true);
});
});
describe("session_start event", () => {
it("notifies when dirty .pi-map.md files exist", async () => {
vi.mocked(execSync).mockReturnValueOnce(
"./src/.pi-map.md\n./tests/.pi-map.md\n",
); // findPiMapFiles
vi.mocked(readFileSync)
.mockReturnValueOnce("## dirty\n2024-01-01: patched\n") // dirty
.mockReturnValueOnce("## dirty\n-\n"); // clean
const handler = registeredEvents.session_start;
await handler(null, mockCtx);
expect(mockNotify).toHaveBeenCalledWith(
expect.stringContaining("1 dirty packages detected"),
"warning",
);
});
it("does not notify when all clean", async () => {
vi.mocked(execSync).mockReturnValueOnce("./src/.pi-map.md\n");
vi.mocked(readFileSync).mockReturnValueOnce("## dirty\n-\n");
const handler = registeredEvents.session_start;
await handler(null, mockCtx);
expect(mockNotify).not.toHaveBeenCalled();
});
it("does nothing when no maps exist", async () => {
vi.mocked(execSync).mockImplementation(() => {
throw new Error("no maps");
});
const handler = registeredEvents.session_start;
await handler(null, mockCtx);
expect(mockNotify).not.toHaveBeenCalled();
});
});
describe("before_agent_start event", () => {
it("injects hint when .pi-map.md files exist", async () => {
vi.mocked(execSync).mockReturnValueOnce("./src/.pi-map.md\n");
vi.mocked(readFileSync).mockReturnValueOnce("content");
const handler = registeredEvents.before_agent_start;
const result = await handler(null, mockCtx);
expect(result).toHaveProperty("message");
expect(result.message.content).toContain("project_map_patch");
expect(result.message.display).toBe(false);
});
it("returns empty object when no maps exist", async () => {
vi.mocked(execSync).mockImplementation(() => {
throw new Error("no maps");
});
const handler = registeredEvents.before_agent_start;
const result = await handler(null, mockCtx);
expect(result).toEqual({});
});
});
});