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:
@@ -0,0 +1,4 @@
|
||||
dist/
|
||||
node_modules/
|
||||
.pi-map.md
|
||||
.env
|
||||
@@ -0,0 +1,9 @@
|
||||
# Sample Project
|
||||
|
||||
A small test project for pi-project-map functionality.
|
||||
|
||||
## Structure
|
||||
|
||||
- `src/` — Source code
|
||||
- `tests/` — Test suites
|
||||
- `docs/` — Documentation
|
||||
@@ -0,0 +1,24 @@
|
||||
# API Documentation
|
||||
|
||||
## User Model
|
||||
|
||||
### `createUser(data)`
|
||||
|
||||
Creates a new user with auto-generated ID and timestamp.
|
||||
|
||||
**Parameters:**
|
||||
- `data` — Object with `email` and `name` fields
|
||||
|
||||
**Returns:** `User` object
|
||||
|
||||
**Throws:** `Error` if email is invalid
|
||||
|
||||
### `serializeUser(user)`
|
||||
|
||||
Serializes a user to JSON string.
|
||||
|
||||
## Validation Utilities
|
||||
|
||||
- `validateEmail(email)` — Validates email format
|
||||
- `validateNotEmpty(value)` — Checks string is not empty
|
||||
- `validateMinLength(value, min)` — Checks minimum string length
|
||||
@@ -0,0 +1,10 @@
|
||||
{
|
||||
"name": "sample-project",
|
||||
"version": "1.0.0",
|
||||
"description": "Sample project for testing pi-project-map",
|
||||
"main": "dist/index.js",
|
||||
"scripts": {
|
||||
"build": "tsc",
|
||||
"test": "vitest"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
import React from "react";
|
||||
|
||||
export interface ButtonProps {
|
||||
label: string;
|
||||
variant?: "primary" | "secondary" | "danger";
|
||||
onClick?: () => void;
|
||||
disabled?: boolean;
|
||||
}
|
||||
|
||||
export function Button({
|
||||
label,
|
||||
variant = "primary",
|
||||
onClick,
|
||||
disabled = false,
|
||||
}: ButtonProps): JSX.Element {
|
||||
return (
|
||||
<button
|
||||
className={`btn btn--${variant}`}
|
||||
onClick={onClick}
|
||||
disabled={disabled}
|
||||
>
|
||||
{label}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
export default Button;
|
||||
@@ -0,0 +1,23 @@
|
||||
import React from "react";
|
||||
import type { User } from "../models/user.js";
|
||||
|
||||
export interface UserCardProps {
|
||||
user: User;
|
||||
onEdit?: (user: User) => void;
|
||||
onDelete?: (id: string) => void;
|
||||
}
|
||||
|
||||
export function UserCard({ user, onEdit, onDelete }: UserCardProps): JSX.Element {
|
||||
return (
|
||||
<div className="user-card">
|
||||
<h3>{user.name}</h3>
|
||||
<p>{user.email}</p>
|
||||
<div className="actions">
|
||||
{onEdit && <button onClick={() => onEdit(user)}>Edit</button>}
|
||||
{onDelete && <button onClick={() => onDelete(user.id)}>Delete</button>}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
export default UserCard;
|
||||
@@ -0,0 +1,14 @@
|
||||
import { createUser } from "./models/user.js";
|
||||
import { validateEmail } from "./utils/validation.js";
|
||||
import { logger } from "./utils/logger.js";
|
||||
|
||||
export async function main() {
|
||||
const user = createUser({ email: "test@example.com", name: "Alice" });
|
||||
if (!validateEmail(user.email)) {
|
||||
logger.error("Invalid email");
|
||||
return;
|
||||
}
|
||||
logger.info(`User created: ${user.name}`);
|
||||
}
|
||||
|
||||
export { createUser, validateEmail, logger };
|
||||
@@ -0,0 +1,23 @@
|
||||
import { validateEmail } from "../utils/validation.js";
|
||||
|
||||
export interface User {
|
||||
id: string;
|
||||
email: string;
|
||||
name: string;
|
||||
createdAt: Date;
|
||||
}
|
||||
|
||||
export function createUser(data: Omit<User, "id" | "createdAt">): User {
|
||||
if (!validateEmail(data.email)) {
|
||||
throw new Error("Invalid email");
|
||||
}
|
||||
return {
|
||||
id: crypto.randomUUID(),
|
||||
...data,
|
||||
createdAt: new Date(),
|
||||
};
|
||||
}
|
||||
|
||||
export function serializeUser(user: User): string {
|
||||
return JSON.stringify(user);
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
export type LogLevel = "debug" | "info" | "warn" | "error";
|
||||
|
||||
export function log(level: LogLevel, message: string): void {
|
||||
const timestamp = new Date().toISOString();
|
||||
console.log(`[${timestamp}] ${level.toUpperCase()}: ${message}`);
|
||||
}
|
||||
|
||||
export function debug(message: string): void {
|
||||
log("debug", message);
|
||||
}
|
||||
|
||||
export function info(message: string): void {
|
||||
log("info", message);
|
||||
}
|
||||
|
||||
export function warn(message: string): void {
|
||||
log("warn", message);
|
||||
}
|
||||
|
||||
export function error(message: string): void {
|
||||
log("error", message);
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
const EMAIL_REGEX = /^[^\s@]+@[^\s@]+\.[^\s@]+$/;
|
||||
|
||||
export function validateEmail(email: string): boolean {
|
||||
return EMAIL_REGEX.test(email);
|
||||
}
|
||||
|
||||
export function validateNotEmpty(value: string): boolean {
|
||||
return value.trim().length > 0;
|
||||
}
|
||||
|
||||
export function validateMinLength(value: string, min: number): boolean {
|
||||
return value.length >= min;
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { createUser, serializeUser } from "../src/models/user.js";
|
||||
|
||||
describe("User model", () => {
|
||||
it("creates a user with valid email", () => {
|
||||
const user = createUser({ email: "alice@example.com", name: "Alice" });
|
||||
expect(user.email).toBe("alice@example.com");
|
||||
expect(user.name).toBe("Alice");
|
||||
expect(user.id).toBeDefined();
|
||||
expect(user.createdAt).toBeInstanceOf(Date);
|
||||
});
|
||||
|
||||
it("throws on invalid email", () => {
|
||||
expect(() =>
|
||||
createUser({ email: "not-an-email", name: "Bob" }),
|
||||
).toThrow("Invalid email");
|
||||
});
|
||||
|
||||
it("serializes to JSON", () => {
|
||||
const user = createUser({ email: "charlie@example.com", name: "Charlie" });
|
||||
const json = serializeUser(user);
|
||||
expect(JSON.parse(json).name).toBe("Charlie");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,21 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { validateEmail, validateNotEmpty, validateMinLength } from "../src/utils/validation.js";
|
||||
|
||||
describe("Validation utils", () => {
|
||||
it("validates email format", () => {
|
||||
expect(validateEmail("test@example.com")).toBe(true);
|
||||
expect(validateEmail("invalid")).toBe(false);
|
||||
expect(validateEmail("")).toBe(false);
|
||||
});
|
||||
|
||||
it("checks non-empty strings", () => {
|
||||
expect(validateNotEmpty("hello")).toBe(true);
|
||||
expect(validateNotEmpty(" ")).toBe(false);
|
||||
expect(validateNotEmpty("")).toBe(false);
|
||||
});
|
||||
|
||||
it("checks minimum length", () => {
|
||||
expect(validateMinLength("hello", 3)).toBe(true);
|
||||
expect(validateMinLength("hi", 3)).toBe(false);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,9 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"target": "ES2022",
|
||||
"module": "NodeNext",
|
||||
"strict": true,
|
||||
"outDir": "./dist",
|
||||
"rootDir": "./src"
|
||||
}
|
||||
}
|
||||
Generated
+59
-25
@@ -10,6 +10,8 @@
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"ignore": "^5.3.0",
|
||||
"openai": "^6.42.0",
|
||||
"p-limit": "^7.3.0",
|
||||
"picocolors": "^1.1.1",
|
||||
"tree-sitter": "^0.21.0",
|
||||
"tree-sitter-go": "^0.25.0",
|
||||
@@ -1356,19 +1358,6 @@
|
||||
"url": "https://github.com/sponsors/sindresorhus"
|
||||
}
|
||||
},
|
||||
"node_modules/@vitest/runner/node_modules/yocto-queue": {
|
||||
"version": "1.2.2",
|
||||
"resolved": "https://registry.npmjs.org/yocto-queue/-/yocto-queue-1.2.2.tgz",
|
||||
"integrity": "sha512-4LCcse/U2MHZ63HAJVE+v71o7yOdIe4cZ70Wpf8D/IyjDKYQLV5GD46B+hSTjJsvV5PztjvHoU580EftxjDZFQ==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"engines": {
|
||||
"node": ">=12.20"
|
||||
},
|
||||
"funding": {
|
||||
"url": "https://github.com/sponsors/sindresorhus"
|
||||
}
|
||||
},
|
||||
"node_modules/@vitest/snapshot": {
|
||||
"version": "1.6.1",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/snapshot/-/snapshot-1.6.1.tgz",
|
||||
@@ -2737,6 +2726,24 @@
|
||||
"url": "https://github.com/sponsors/sindresorhus"
|
||||
}
|
||||
},
|
||||
"node_modules/openai": {
|
||||
"version": "6.42.0",
|
||||
"resolved": "https://registry.npmjs.org/openai/-/openai-6.42.0.tgz",
|
||||
"integrity": "sha512-1WFEt/uXMXOLhYRNkgJWo08Y2YNvNwpVU72K7ibrWgWpNOXd4VojXLbe6SQ4bLiUQ3Y8jz4IiyVkylJCL1DtZg==",
|
||||
"license": "Apache-2.0",
|
||||
"peerDependencies": {
|
||||
"ws": "^8.18.0",
|
||||
"zod": "^3.25 || ^4.0"
|
||||
},
|
||||
"peerDependenciesMeta": {
|
||||
"ws": {
|
||||
"optional": true
|
||||
},
|
||||
"zod": {
|
||||
"optional": true
|
||||
}
|
||||
}
|
||||
},
|
||||
"node_modules/optionator": {
|
||||
"version": "0.9.4",
|
||||
"resolved": "https://registry.npmjs.org/optionator/-/optionator-0.9.4.tgz",
|
||||
@@ -2756,6 +2763,37 @@
|
||||
}
|
||||
},
|
||||
"node_modules/p-limit": {
|
||||
"version": "7.3.0",
|
||||
"resolved": "https://registry.npmjs.org/p-limit/-/p-limit-7.3.0.tgz",
|
||||
"integrity": "sha512-7cIXg/Z0M5WZRblrsOla88S4wAK+zOQQWeBYfV3qJuJXMr+LnbYjaadrFaS0JILfEDPVqHyKnZ1Z/1d6J9VVUw==",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"yocto-queue": "^1.2.1"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=20"
|
||||
},
|
||||
"funding": {
|
||||
"url": "https://github.com/sponsors/sindresorhus"
|
||||
}
|
||||
},
|
||||
"node_modules/p-locate": {
|
||||
"version": "5.0.0",
|
||||
"resolved": "https://registry.npmjs.org/p-locate/-/p-locate-5.0.0.tgz",
|
||||
"integrity": "sha512-LaNjtRWUBY++zB5nE/NwcaoMylSPk+S+ZHNB1TzdbMJMny6dynpAGt7X/tl/QYq3TIeE6nxHppbo2LGymrG5Pw==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"p-limit": "^3.0.2"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=10"
|
||||
},
|
||||
"funding": {
|
||||
"url": "https://github.com/sponsors/sindresorhus"
|
||||
}
|
||||
},
|
||||
"node_modules/p-locate/node_modules/p-limit": {
|
||||
"version": "3.1.0",
|
||||
"resolved": "https://registry.npmjs.org/p-limit/-/p-limit-3.1.0.tgz",
|
||||
"integrity": "sha512-TYOanM3wGwNGsZN2cVTYPArw454xnXj5qmWF1bEoAc4+cU/ol7GVh7odevjp1FNHduHc3KZMcFduxU5Xc6uJRQ==",
|
||||
@@ -2771,15 +2809,12 @@
|
||||
"url": "https://github.com/sponsors/sindresorhus"
|
||||
}
|
||||
},
|
||||
"node_modules/p-locate": {
|
||||
"version": "5.0.0",
|
||||
"resolved": "https://registry.npmjs.org/p-locate/-/p-locate-5.0.0.tgz",
|
||||
"integrity": "sha512-LaNjtRWUBY++zB5nE/NwcaoMylSPk+S+ZHNB1TzdbMJMny6dynpAGt7X/tl/QYq3TIeE6nxHppbo2LGymrG5Pw==",
|
||||
"node_modules/p-locate/node_modules/yocto-queue": {
|
||||
"version": "0.1.0",
|
||||
"resolved": "https://registry.npmjs.org/yocto-queue/-/yocto-queue-0.1.0.tgz",
|
||||
"integrity": "sha512-rVksvsnNCdJ/ohGc6xgPwyN8eheCxsiLM8mxuE/t/mOVqJewPuO1miLpTHQiRgTKCLexL4MeAFVagts7HmNZ2Q==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"p-limit": "^3.0.2"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=10"
|
||||
},
|
||||
@@ -4116,13 +4151,12 @@
|
||||
"license": "ISC"
|
||||
},
|
||||
"node_modules/yocto-queue": {
|
||||
"version": "0.1.0",
|
||||
"resolved": "https://registry.npmjs.org/yocto-queue/-/yocto-queue-0.1.0.tgz",
|
||||
"integrity": "sha512-rVksvsnNCdJ/ohGc6xgPwyN8eheCxsiLM8mxuE/t/mOVqJewPuO1miLpTHQiRgTKCLexL4MeAFVagts7HmNZ2Q==",
|
||||
"dev": true,
|
||||
"version": "1.2.2",
|
||||
"resolved": "https://registry.npmjs.org/yocto-queue/-/yocto-queue-1.2.2.tgz",
|
||||
"integrity": "sha512-4LCcse/U2MHZ63HAJVE+v71o7yOdIe4cZ70Wpf8D/IyjDKYQLV5GD46B+hSTjJsvV5PztjvHoU580EftxjDZFQ==",
|
||||
"license": "MIT",
|
||||
"engines": {
|
||||
"node": ">=10"
|
||||
"node": ">=12.20"
|
||||
},
|
||||
"funding": {
|
||||
"url": "https://github.com/sponsors/sindresorhus"
|
||||
|
||||
@@ -42,6 +42,8 @@
|
||||
},
|
||||
"dependencies": {
|
||||
"ignore": "^5.3.0",
|
||||
"openai": "^6.42.0",
|
||||
"p-limit": "^7.3.0",
|
||||
"picocolors": "^1.1.1",
|
||||
"tree-sitter": "^0.21.0",
|
||||
"tree-sitter-go": "^0.25.0",
|
||||
|
||||
+64
-10
@@ -4,6 +4,8 @@ import { patchFile } from "./patch.js";
|
||||
import { validateMaps } from "./validate.js";
|
||||
import { reinitPath } from "./init.js";
|
||||
import { discoverProject } from "./discover.js";
|
||||
import { createLLMClient, LLMError } from "./llm-client.js";
|
||||
import { loadConfig } from "./config.js";
|
||||
import pc from "picocolors";
|
||||
|
||||
const args = process.argv.slice(2);
|
||||
@@ -31,11 +33,22 @@ function printUsage() {
|
||||
console.log(
|
||||
` project-map ${pc.cyan("--version")} Show version\n`,
|
||||
);
|
||||
console.log(`${pc.bold("Options:")}`);
|
||||
console.log(
|
||||
` --llm-provider=openai|kimi LLM provider (default: config or openai)`,
|
||||
);
|
||||
console.log(
|
||||
` --llm-model=<model> LLM model name (or set LLM_MODEL env var)`,
|
||||
);
|
||||
console.log(` --llm-base-url=<url> Custom base URL for LLM API\n`);
|
||||
console.log(`${pc.bold("Examples:")}`);
|
||||
console.log(` project-map init`);
|
||||
console.log(` project-map patch src/components/Button.tsx`);
|
||||
console.log(` project-map validate --fix`);
|
||||
console.log(` project-map reinit`);
|
||||
console.log(
|
||||
` project-map init --llm-provider=kimi --llm-model=kimi-k2-6`,
|
||||
);
|
||||
}
|
||||
|
||||
function printVersion() {
|
||||
@@ -50,17 +63,49 @@ function formatCount(count: number, label: string): string {
|
||||
return `${pc.bold(String(count))} ${count === 1 ? label : plural}`;
|
||||
}
|
||||
|
||||
function parseValidateArgs(args: string[]): { path: string; fix: boolean } {
|
||||
function parseArgs(args: string[]): {
|
||||
path: string;
|
||||
fix: boolean;
|
||||
llmProvider?: string;
|
||||
llmModel?: string;
|
||||
llmBaseUrl?: string;
|
||||
positional: string[];
|
||||
} {
|
||||
let path = ".";
|
||||
let fix = false;
|
||||
let llmProvider: string | undefined;
|
||||
let llmModel: string | undefined;
|
||||
let llmBaseUrl: string | undefined;
|
||||
const positional: string[] = [];
|
||||
|
||||
for (const arg of args.slice(1)) {
|
||||
if (arg === "--fix") {
|
||||
fix = true;
|
||||
} else if (arg.startsWith("--llm-provider=")) {
|
||||
llmProvider = arg.slice("--llm-provider=".length);
|
||||
} else if (arg.startsWith("--llm-model=")) {
|
||||
llmModel = arg.slice("--llm-model=".length);
|
||||
} else if (arg.startsWith("--llm-base-url=")) {
|
||||
llmBaseUrl = arg.slice("--llm-base-url=".length);
|
||||
} else if (!arg.startsWith("-")) {
|
||||
positional.push(arg);
|
||||
path = arg;
|
||||
}
|
||||
}
|
||||
return { path, fix };
|
||||
|
||||
return { path, fix, llmProvider, llmModel, llmBaseUrl, positional };
|
||||
}
|
||||
|
||||
function createClientFromArgs(args: ReturnType<typeof parseArgs>) {
|
||||
const config = loadConfig();
|
||||
const provider = (args.llmProvider || config.llmProvider) as
|
||||
| "openai"
|
||||
| "kimi"
|
||||
| "pi";
|
||||
return createLLMClient(provider, {
|
||||
model: args.llmModel || config.llmModel || process.env.LLM_MODEL,
|
||||
baseUrl: args.llmBaseUrl || config.llmBaseUrl,
|
||||
});
|
||||
}
|
||||
|
||||
async function main() {
|
||||
@@ -74,13 +119,16 @@ async function main() {
|
||||
process.exit(0);
|
||||
}
|
||||
|
||||
const parsed = parseArgs(args);
|
||||
|
||||
switch (command) {
|
||||
case "init": {
|
||||
const targetPath = args[1] || ".";
|
||||
const targetPath = parsed.positional[0] || ".";
|
||||
const start = Date.now();
|
||||
const entries = discoverProject(targetPath);
|
||||
console.log(`Scanning ${formatCount(entries.length, "directory")}...`);
|
||||
await initProject(targetPath, { verbose: false });
|
||||
const client = createClientFromArgs(parsed);
|
||||
await initProject(targetPath, { verbose: false, llmClient: client });
|
||||
const elapsed = ((Date.now() - start) / 1000).toFixed(1);
|
||||
console.log(
|
||||
`${pc.green("✓")} Generated ${formatCount(entries.length, ".pi-map.md file")} in ${elapsed}s`,
|
||||
@@ -88,18 +136,19 @@ async function main() {
|
||||
break;
|
||||
}
|
||||
case "patch": {
|
||||
if (!args[1]) {
|
||||
if (!parsed.positional[0]) {
|
||||
console.error(
|
||||
`${pc.red("Error:")} Missing file path. Usage: project-map patch <file>`,
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
await patchFile(args[1]);
|
||||
const client = createClientFromArgs(parsed);
|
||||
await patchFile(parsed.positional[0], client);
|
||||
console.log(`${pc.green("✓")} Patched`);
|
||||
break;
|
||||
}
|
||||
case "validate": {
|
||||
const { path, fix } = parseValidateArgs(args);
|
||||
const { path, fix } = parsed;
|
||||
const result = await validateMaps(path, { fix, verbose: true });
|
||||
if (result.clean) {
|
||||
console.log(`${pc.green("✓")} All .pi-map.md files are clean.`);
|
||||
@@ -123,13 +172,14 @@ async function main() {
|
||||
break;
|
||||
}
|
||||
case "reinit": {
|
||||
const targetPath = args[1] || ".";
|
||||
const targetPath = parsed.positional[0] || ".";
|
||||
const start = Date.now();
|
||||
const entries = discoverProject(targetPath);
|
||||
console.log(
|
||||
`Regenerating ${formatCount(entries.length, ".pi-map.md file")}...`,
|
||||
);
|
||||
await reinitPath(targetPath, { verbose: false });
|
||||
const client = createClientFromArgs(parsed);
|
||||
await reinitPath(targetPath, { verbose: false, llmClient: client });
|
||||
const elapsed = ((Date.now() - start) / 1000).toFixed(1);
|
||||
console.log(`${pc.green("✓")} Regenerated in ${elapsed}s`);
|
||||
break;
|
||||
@@ -142,6 +192,10 @@ async function main() {
|
||||
}
|
||||
|
||||
main().catch((err) => {
|
||||
console.error(`${pc.red("Error:")} ${err.message}`);
|
||||
if (err instanceof LLMError) {
|
||||
console.error(`${pc.red("LLM Error:")} ${err.message}`);
|
||||
} else {
|
||||
console.error(`${pc.red("Error:")} ${err.message}`);
|
||||
}
|
||||
process.exit(1);
|
||||
});
|
||||
|
||||
@@ -4,7 +4,9 @@ import { join } from "path";
|
||||
export interface SkillConfig {
|
||||
ignorePatterns: string[];
|
||||
smallPackageThreshold: number;
|
||||
llmProvider: "openai" | "kimi" | "pi";
|
||||
llmModel: string;
|
||||
llmBaseUrl?: string;
|
||||
contextBudget: number;
|
||||
autoInjectPrompt: boolean;
|
||||
}
|
||||
@@ -32,6 +34,7 @@ export const DEFAULT_CONFIG: SkillConfig = {
|
||||
".prettiercache",
|
||||
],
|
||||
smallPackageThreshold: 10,
|
||||
llmProvider: "openai",
|
||||
llmModel: "gpt-4o-mini",
|
||||
contextBudget: 4000,
|
||||
autoInjectPrompt: true,
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
import OpenAI from "openai";
|
||||
import { LLMError } from "./llm-error.js";
|
||||
import type { LLMClient, LLMClientOptions } from "./llm-client.js";
|
||||
|
||||
export class ExternalLLMClient implements LLMClient {
|
||||
private client: OpenAI;
|
||||
private model: string;
|
||||
|
||||
constructor(options: LLMClientOptions = {}) {
|
||||
const apiKey = options.apiKey || process.env.OPENAI_API_KEY;
|
||||
if (!apiKey) {
|
||||
throw new LLMError(
|
||||
"No OpenAI API key provided. Set OPENAI_API_KEY environment variable or pass apiKey in options.",
|
||||
);
|
||||
}
|
||||
this.client = new OpenAI({
|
||||
apiKey,
|
||||
baseURL: options.baseUrl,
|
||||
});
|
||||
this.model =
|
||||
options.model ||
|
||||
process.env.OPENAI_MODEL ||
|
||||
process.env.LLM_MODEL ||
|
||||
"gpt-4o-mini";
|
||||
}
|
||||
|
||||
async complete(prompt: string): Promise<string> {
|
||||
try {
|
||||
const response = await this.client.chat.completions.create({
|
||||
model: this.model,
|
||||
messages: [
|
||||
{
|
||||
role: "system",
|
||||
content:
|
||||
"You are a code analysis assistant. Analyze the provided file and respond with concise, structured information.",
|
||||
},
|
||||
{ role: "user", content: prompt },
|
||||
],
|
||||
temperature: 0.1,
|
||||
max_tokens: 256,
|
||||
});
|
||||
return response.choices[0]?.message?.content?.trim() || "";
|
||||
} catch (err: any) {
|
||||
throw new LLMError(
|
||||
`External LLM request failed: ${err.message || String(err)}`,
|
||||
err,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
+17
-6
@@ -9,34 +9,45 @@ import { extractFileAST } from "./ast-extract.js";
|
||||
import { mergeFileData } from "./merge.js";
|
||||
import { writeFileSync } from "fs";
|
||||
import { join } from "path";
|
||||
import type { LLMClient } from "./llm-client.js";
|
||||
|
||||
export interface InitOptions {
|
||||
verbose?: boolean;
|
||||
llmClient?: LLMClient;
|
||||
}
|
||||
|
||||
export async function initProject(
|
||||
rootPath: string,
|
||||
options?: { verbose?: boolean },
|
||||
options: InitOptions = {},
|
||||
): Promise<void> {
|
||||
const entries = discoverProject(rootPath);
|
||||
|
||||
for (const entry of entries) {
|
||||
await generateDirectoryMap(entry);
|
||||
await generateDirectoryMap(entry, options.llmClient);
|
||||
}
|
||||
|
||||
if (options?.verbose !== false) {
|
||||
if (options.verbose !== false) {
|
||||
console.log(`Generated ${entries.length} .pi-map.md files`);
|
||||
}
|
||||
}
|
||||
|
||||
export async function generateDirectoryMap(
|
||||
entry: DirectoryEntry,
|
||||
llmClient?: LLMClient,
|
||||
): Promise<FileEntry[]> {
|
||||
const fileData: FileEntry[] = [];
|
||||
for (const file of entry.files) {
|
||||
const filePath = join(entry.dirPath, file);
|
||||
const llmData = await extractFileLLM(filePath);
|
||||
const llmData = await extractFileLLM(filePath, llmClient);
|
||||
const astData = await extractFileAST(filePath);
|
||||
fileData.push(mergeFileData(file, llmData, astData));
|
||||
}
|
||||
|
||||
const packageData = await extractPackageLLM(entry.relativePath, fileData);
|
||||
const packageData = await extractPackageLLM(
|
||||
entry.relativePath,
|
||||
fileData,
|
||||
llmClient,
|
||||
);
|
||||
|
||||
const mapData: PackageMapData = {
|
||||
path: entry.relativePath,
|
||||
@@ -53,7 +64,7 @@ export async function generateDirectoryMap(
|
||||
|
||||
export async function reinitPath(
|
||||
path: string,
|
||||
options?: { verbose?: boolean },
|
||||
options: InitOptions = {},
|
||||
): Promise<void> {
|
||||
// Full regeneration clears all dirty markers by overwriting every .pi-map.md
|
||||
await initProject(path, options);
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
import { LLMError } from "./llm-error.js";
|
||||
import type { LLMClient, LLMClientOptions } from "./llm-client.js";
|
||||
|
||||
/**
|
||||
* Kimi.com LLM Client.
|
||||
*
|
||||
* Uses the Anthropic-based API at https://api.kimi.com/coding/
|
||||
* Set KIMI_API_KEY environment variable or pass apiKey in options.
|
||||
*/
|
||||
export class KimiLLMClient implements LLMClient {
|
||||
private apiKey: string;
|
||||
private model: string;
|
||||
private baseUrl: string;
|
||||
|
||||
constructor(options: LLMClientOptions = {}) {
|
||||
const apiKey =
|
||||
options.apiKey ||
|
||||
process.env.KIMI_API_KEY ||
|
||||
process.env.KIMI_COM_API_KEY;
|
||||
if (!apiKey) {
|
||||
throw new LLMError(
|
||||
"No Kimi API key provided. Set KIMI_API_KEY environment variable or pass apiKey in options.",
|
||||
);
|
||||
}
|
||||
this.apiKey = apiKey;
|
||||
this.model =
|
||||
options.model ||
|
||||
process.env.KIMI_MODEL ||
|
||||
process.env.LLM_MODEL ||
|
||||
"kimi-k2-6";
|
||||
this.baseUrl = options.baseUrl || "https://api.kimi.com/coding";
|
||||
}
|
||||
|
||||
async complete(prompt: string): Promise<string> {
|
||||
try {
|
||||
const response = await fetch(`${this.baseUrl}/v1/messages`, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
"x-api-key": this.apiKey,
|
||||
"anthropic-version": "2023-06-01",
|
||||
},
|
||||
body: JSON.stringify({
|
||||
model: this.model,
|
||||
max_tokens: 256,
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: prompt,
|
||||
},
|
||||
],
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const text = await response.text();
|
||||
throw new LLMError(`Kimi API error ${response.status}: ${text}`);
|
||||
}
|
||||
|
||||
const data = (await response.json()) as {
|
||||
content?: Array<{ type: string; text?: string }>;
|
||||
error?: { message: string };
|
||||
};
|
||||
|
||||
if (data.error) {
|
||||
throw new LLMError(`Kimi error: ${data.error.message}`);
|
||||
}
|
||||
|
||||
const text = data.content?.[0]?.text?.trim() || "";
|
||||
return text;
|
||||
} catch (err: any) {
|
||||
if (err instanceof LLMError) throw err;
|
||||
throw new LLMError(
|
||||
`Kimi request failed: ${err.message || String(err)}`,
|
||||
err,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
import pLimit from "p-limit";
|
||||
import { LLMError } from "./llm-error.js";
|
||||
|
||||
export interface BatchOptions {
|
||||
concurrency?: number;
|
||||
batchDelayMs?: number;
|
||||
maxRetries?: number;
|
||||
retryDelaysMs?: number[];
|
||||
}
|
||||
|
||||
const DEFAULT_OPTIONS: Required<BatchOptions> = {
|
||||
concurrency: 4,
|
||||
batchDelayMs: 100,
|
||||
maxRetries: 3,
|
||||
retryDelaysMs: [1000, 2000, 4000],
|
||||
};
|
||||
|
||||
function sleep(ms: number): Promise<void> {
|
||||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||
}
|
||||
|
||||
export async function withRetry<T>(
|
||||
fn: () => Promise<T>,
|
||||
options: Pick<BatchOptions, "maxRetries" | "retryDelaysMs"> = {},
|
||||
): Promise<T> {
|
||||
const maxRetries = options.maxRetries ?? DEFAULT_OPTIONS.maxRetries;
|
||||
const delays = options.retryDelaysMs ?? DEFAULT_OPTIONS.retryDelaysMs;
|
||||
|
||||
let lastError: unknown;
|
||||
for (let attempt = 0; attempt <= maxRetries; attempt++) {
|
||||
try {
|
||||
return await fn();
|
||||
} catch (err) {
|
||||
lastError = err;
|
||||
if (attempt < maxRetries) {
|
||||
const delay = delays[attempt] ?? delays[delays.length - 1] ?? 1000;
|
||||
await sleep(delay);
|
||||
}
|
||||
}
|
||||
}
|
||||
throw lastError;
|
||||
}
|
||||
|
||||
export async function processFiles<T, R>(
|
||||
files: T[],
|
||||
processor: (file: T) => Promise<R>,
|
||||
options: BatchOptions = {},
|
||||
): Promise<R[]> {
|
||||
const opts = { ...DEFAULT_OPTIONS, ...options };
|
||||
const limit = pLimit(opts.concurrency);
|
||||
|
||||
const results: R[] = [];
|
||||
let batchCount = 0;
|
||||
|
||||
const tasks = files.map((file, index) =>
|
||||
limit(async () => {
|
||||
// Small delay between batches based on index
|
||||
if (index > 0 && index % opts.concurrency === 0) {
|
||||
batchCount++;
|
||||
await sleep(opts.batchDelayMs);
|
||||
}
|
||||
return withRetry(() => processor(file), opts);
|
||||
}),
|
||||
);
|
||||
|
||||
const settled = await Promise.all(tasks);
|
||||
return settled;
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
import { readFileSync, writeFileSync, existsSync, mkdirSync, renameSync } from "fs";
|
||||
import { join } from "path";
|
||||
import { homedir } from "os";
|
||||
|
||||
const CACHE_DIR = join(homedir(), ".cache", "pi-project-map");
|
||||
const CACHE_FILE = join(CACHE_DIR, "llm-cache.json");
|
||||
|
||||
interface CacheEntry {
|
||||
result: string;
|
||||
ts: number;
|
||||
}
|
||||
|
||||
function ensureCacheDir(): void {
|
||||
if (!existsSync(CACHE_DIR)) {
|
||||
mkdirSync(CACHE_DIR, { recursive: true });
|
||||
}
|
||||
}
|
||||
|
||||
function loadCache(): Record<string, CacheEntry> {
|
||||
if (!existsSync(CACHE_FILE)) return {};
|
||||
try {
|
||||
const raw = readFileSync(CACHE_FILE, "utf8");
|
||||
return JSON.parse(raw) as Record<string, CacheEntry>;
|
||||
} catch {
|
||||
// Corrupted cache — start fresh
|
||||
return {};
|
||||
}
|
||||
}
|
||||
|
||||
function saveCache(cache: Record<string, CacheEntry>): void {
|
||||
ensureCacheDir();
|
||||
const tmp = CACHE_FILE + ".tmp";
|
||||
writeFileSync(tmp, JSON.stringify(cache, null, 2));
|
||||
renameSync(tmp, CACHE_FILE);
|
||||
}
|
||||
|
||||
export function getCached(hash: string): string | undefined {
|
||||
const cache = loadCache();
|
||||
return cache[hash]?.result;
|
||||
}
|
||||
|
||||
export function setCached(hash: string, result: string): void {
|
||||
const cache = loadCache();
|
||||
cache[hash] = { result, ts: Date.now() };
|
||||
saveCache(cache);
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
import { LLMError } from "./llm-error.js";
|
||||
import { ExternalLLMClient } from "./external-llm-client.js";
|
||||
import { KimiLLMClient } from "./kimi-llm-client.js";
|
||||
import { PiLLMClient } from "./pi-llm-client.js";
|
||||
|
||||
export interface LLMClient {
|
||||
complete(prompt: string): Promise<string>;
|
||||
}
|
||||
|
||||
export interface LLMClientOptions {
|
||||
apiKey?: string;
|
||||
model?: string;
|
||||
baseUrl?: string;
|
||||
// Pi-specific
|
||||
extensionContext?: unknown;
|
||||
}
|
||||
|
||||
export function createLLMClient(
|
||||
mode: "pi" | "openai" | "kimi",
|
||||
options: LLMClientOptions = {},
|
||||
): LLMClient {
|
||||
if (mode === "pi") {
|
||||
return new PiLLMClient(options.extensionContext);
|
||||
}
|
||||
if (mode === "kimi") {
|
||||
return new KimiLLMClient(options);
|
||||
}
|
||||
return new ExternalLLMClient(options);
|
||||
}
|
||||
|
||||
export { LLMError };
|
||||
@@ -0,0 +1,9 @@
|
||||
export class LLMError extends Error {
|
||||
constructor(
|
||||
message: string,
|
||||
public readonly cause?: unknown,
|
||||
) {
|
||||
super(message);
|
||||
this.name = "LLMError";
|
||||
}
|
||||
}
|
||||
+300
-215
@@ -1,11 +1,15 @@
|
||||
import { readFileSync } from "fs";
|
||||
import { readFileSync, statSync } from "fs";
|
||||
import { createHash } from "crypto";
|
||||
import { extname, basename } from "path";
|
||||
import type { LLMClient } from "./llm-client.js";
|
||||
import { getCached, setCached } from "./llm-cache.js";
|
||||
import { LLMError } from "./llm-error.js";
|
||||
|
||||
interface LLMFileData {
|
||||
purpose: string;
|
||||
exports: string[];
|
||||
deps: string[];
|
||||
concepts: string[];
|
||||
}
|
||||
|
||||
interface LLMPackageData {
|
||||
@@ -13,9 +17,11 @@ interface LLMPackageData {
|
||||
arch: string;
|
||||
}
|
||||
|
||||
const cache = new Map<string, LLMFileData>();
|
||||
const MAX_FILE_SIZE = 50 * 1024; // 50KB
|
||||
const CONTEXT_BUDGET = 4000; // tokens
|
||||
const CHARS_PER_TOKEN = 4; // approximate for ASCII
|
||||
|
||||
// Heuristic patterns for common file types
|
||||
// Heuristic patterns for common file types (used as fallback + for tests)
|
||||
const FILE_TYPE_PURPOSES: Record<string, string> = {
|
||||
".ts": "TypeScript module",
|
||||
".tsx": "React component",
|
||||
@@ -51,236 +57,178 @@ const FILE_TYPE_PURPOSES: Record<string, string> = {
|
||||
".svelte": "Svelte component",
|
||||
};
|
||||
|
||||
export async function extractFileLLM(filePath: string): Promise<LLMFileData> {
|
||||
function truncateForContext(content: string, promptLength: number): string {
|
||||
const maxChars = CONTEXT_BUDGET * CHARS_PER_TOKEN - promptLength;
|
||||
if (content.length <= maxChars) return content;
|
||||
return content.slice(0, maxChars - 20) + "\n[...truncated]";
|
||||
}
|
||||
|
||||
function buildFilePrompt(filePath: string, content: string): string {
|
||||
const name = basename(filePath);
|
||||
const ext = extname(filePath).toLowerCase();
|
||||
const typeHint = FILE_TYPE_PURPOSES[ext] || FILE_TYPE_PURPOSES[name.toLowerCase()] || ext || "file";
|
||||
|
||||
return `Analyze this ${typeHint} file. Respond in this exact format (one line each):
|
||||
PURPOSE: <concise one-sentence description of what this file does>
|
||||
DEPS: <comma-separated list of key dependencies/modules it relies on, or "none">
|
||||
CONCEPTS: <comma-separated list of key concepts/patterns used, or "none">
|
||||
|
||||
File: ${name}
|
||||
\`\`\`
|
||||
${content}
|
||||
\`\`\`
|
||||
`;
|
||||
}
|
||||
|
||||
function parseFileResponse(response: string): { purpose: string; deps: string[]; concepts: string[] } {
|
||||
const lines = response.split("\n");
|
||||
let purpose = "";
|
||||
let deps: string[] = [];
|
||||
let concepts: string[] = [];
|
||||
|
||||
for (const line of lines) {
|
||||
const trimmed = line.trim();
|
||||
if (trimmed.startsWith("PURPOSE:")) {
|
||||
purpose = trimmed.slice("PURPOSE:".length).trim();
|
||||
} else if (trimmed.startsWith("DEPS:")) {
|
||||
const depsStr = trimmed.slice("DEPS:".length).trim();
|
||||
deps = depsStr === "none" ? [] : depsStr.split(",").map((s) => s.trim()).filter(Boolean);
|
||||
} else if (trimmed.startsWith("CONCEPTS:")) {
|
||||
const conceptsStr = trimmed.slice("CONCEPTS:".length).trim();
|
||||
concepts = conceptsStr === "none" ? [] : conceptsStr.split(",").map((s) => s.trim()).filter(Boolean);
|
||||
}
|
||||
}
|
||||
|
||||
return { purpose, deps, concepts };
|
||||
}
|
||||
|
||||
function buildPackagePrompt(relativePath: string, fileSummaries: { name: string; purpose: string }[]): string {
|
||||
const filesList = fileSummaries.map((f) => `- ${f.name}: ${f.purpose}`).join("\n");
|
||||
return `Analyze this code package/directory. Respond in this exact format (one line each):
|
||||
ROLE: <concise one-sentence description of this package's role in the project>
|
||||
ARCH: <concise description of architecture/patterns used in this package>
|
||||
|
||||
Directory: ${relativePath}
|
||||
Files:
|
||||
${filesList}
|
||||
`;
|
||||
}
|
||||
|
||||
function parsePackageResponse(response: string): { role: string; arch: string } {
|
||||
const lines = response.split("\n");
|
||||
let role = "";
|
||||
let arch = "";
|
||||
|
||||
for (const line of lines) {
|
||||
const trimmed = line.trim();
|
||||
if (trimmed.startsWith("ROLE:")) {
|
||||
role = trimmed.slice("ROLE:".length).trim();
|
||||
} else if (trimmed.startsWith("ARCH:")) {
|
||||
arch = trimmed.slice("ARCH:".length).trim();
|
||||
}
|
||||
}
|
||||
|
||||
return { role, arch };
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// PRODUCTION: Real LLM calls
|
||||
// ============================================================================
|
||||
|
||||
export async function extractFileLLM(
|
||||
filePath: string,
|
||||
client?: LLMClient,
|
||||
): Promise<LLMFileData> {
|
||||
const content = readFileSync(filePath, "utf8");
|
||||
const hash = createHash("sha256").update(content).digest("hex");
|
||||
|
||||
if (cache.has(hash)) {
|
||||
return cache.get(hash)!;
|
||||
// Check disk cache
|
||||
const cached = getCached(hash);
|
||||
if (cached) {
|
||||
const parsed = parseFileResponse(cached);
|
||||
return {
|
||||
purpose: parsed.purpose,
|
||||
exports: [], // AST provides precise exports
|
||||
deps: parsed.deps,
|
||||
concepts: parsed.concepts,
|
||||
};
|
||||
}
|
||||
|
||||
const ext = extname(filePath).toLowerCase();
|
||||
const name = basename(filePath);
|
||||
const baseName = basename(filePath, ext);
|
||||
|
||||
// Extract exports via heuristics
|
||||
const exports = extractExports(content, ext, name);
|
||||
|
||||
// Extract dependencies via heuristics
|
||||
const deps = extractDeps(content, ext);
|
||||
|
||||
// Generate purpose from filename + content heuristics
|
||||
const purpose = generatePurpose(name, ext, baseName, content, exports);
|
||||
|
||||
const result: LLMFileData = { purpose, exports, deps };
|
||||
cache.set(hash, result);
|
||||
return result;
|
||||
}
|
||||
|
||||
function extractExports(
|
||||
content: string,
|
||||
ext: string,
|
||||
_filename: string,
|
||||
): string[] {
|
||||
const exports: string[] = [];
|
||||
|
||||
if ([".ts", ".tsx", ".js", ".jsx", ".mjs"].includes(ext)) {
|
||||
// ES module exports — only match at start of line (after optional whitespace)
|
||||
// Handles: export function foo, export async function foo, export class Foo,
|
||||
// export const foo, export { foo, bar }, export default foo
|
||||
const exportRegex =
|
||||
/(?:^|\n)\s*export\s+(?:default\s+)?(?:async\s+)?(?:function\s+|class\s+|const\s+|let\s+|var\s+|interface\s+|type\s+|enum\s+)?([A-Za-z_$][A-Za-z0-9_$]*)/g;
|
||||
let match: RegExpExecArray | null;
|
||||
match = exportRegex.exec(content);
|
||||
while (match !== null) {
|
||||
exports.push(match[1]);
|
||||
match = exportRegex.exec(content);
|
||||
}
|
||||
|
||||
// Named export destructuring: export { foo, bar }
|
||||
const namedExportRegex = /(?:^|\n)\s*export\s*\{\s*([^}]+)\s*\}/g;
|
||||
match = namedExportRegex.exec(content);
|
||||
while (match !== null) {
|
||||
const names = match[1].split(",").map((s) =>
|
||||
s
|
||||
.trim()
|
||||
.split(/\s+as\s+/)[0]
|
||||
.trim(),
|
||||
);
|
||||
exports.push(...names);
|
||||
match = namedExportRegex.exec(content);
|
||||
}
|
||||
} else if (ext === ".py") {
|
||||
// Python exports (top-level functions/classes)
|
||||
const pyRegex =
|
||||
/^(?:async\s+)?def\s+([A-Za-z_][A-Za-z0-9_]*)|class\s+([A-Za-z_][A-Za-z0-9_]*)/gm;
|
||||
let match: RegExpExecArray | null = pyRegex.exec(content);
|
||||
while (match !== null) {
|
||||
exports.push(match[1] || match[2]);
|
||||
match = pyRegex.exec(content);
|
||||
}
|
||||
} else if (ext === ".go") {
|
||||
// Go exports (capitalized functions/types)
|
||||
const goRegex = /^(?:func|type|var|const)\s+([A-Z][A-Za-z0-9_]*)/gm;
|
||||
let match: RegExpExecArray | null = goRegex.exec(content);
|
||||
while (match !== null) {
|
||||
exports.push(match[1]);
|
||||
match = goRegex.exec(content);
|
||||
}
|
||||
} else if (ext === ".rs") {
|
||||
// Rust exports (pub items)
|
||||
const rsRegex =
|
||||
/pub\s+(?:fn|struct|enum|trait|type|const|static|use)\s+([A-Za-z_][A-Za-z0-9_]*)/g;
|
||||
let match: RegExpExecArray | null = rsRegex.exec(content);
|
||||
while (match !== null) {
|
||||
exports.push(match[1]);
|
||||
match = rsRegex.exec(content);
|
||||
}
|
||||
// Skip very large files
|
||||
const size = statSync(filePath).size;
|
||||
if (size > MAX_FILE_SIZE) {
|
||||
return {
|
||||
purpose: "Large/generated file",
|
||||
exports: [],
|
||||
deps: [],
|
||||
concepts: [],
|
||||
};
|
||||
}
|
||||
|
||||
// Deduplicate while preserving order
|
||||
return [...new Set(exports)];
|
||||
}
|
||||
|
||||
function extractDeps(content: string, ext: string): string[] {
|
||||
const deps: string[] = [];
|
||||
|
||||
if ([".ts", ".tsx", ".js", ".jsx", ".mjs"].includes(ext)) {
|
||||
// ES imports
|
||||
const importRegex =
|
||||
/import\s+(?:(?:type\s+)?\{[^}]*\}|\*\s+as\s+\w+|\w+)\s+from\s+['"]([^'"]+)['"]/g;
|
||||
let match: RegExpExecArray | null = importRegex.exec(content);
|
||||
while (match !== null) {
|
||||
deps.push(match[1]);
|
||||
match = importRegex.exec(content);
|
||||
}
|
||||
// CommonJS requires
|
||||
const requireRegex = /require\s*\(\s*['"]([^'"]+)['"]\s*\)/g;
|
||||
match = requireRegex.exec(content);
|
||||
while (match !== null) {
|
||||
deps.push(match[1]);
|
||||
match = requireRegex.exec(content);
|
||||
}
|
||||
} else if (ext === ".py") {
|
||||
// Python imports
|
||||
const pyImportRegex = /^(?:from|import)\s+([A-Za-z_][A-Za-z0-9_.]*)/gm;
|
||||
let match: RegExpExecArray | null = pyImportRegex.exec(content);
|
||||
while (match !== null) {
|
||||
deps.push(match[1]);
|
||||
match = pyImportRegex.exec(content);
|
||||
}
|
||||
} else if (ext === ".go") {
|
||||
// Go imports
|
||||
const goImportRegex = /"([^"]+)"/g;
|
||||
let match: RegExpExecArray | null = goImportRegex.exec(content);
|
||||
while (match !== null) {
|
||||
if (match[1].includes("/")) deps.push(match[1]);
|
||||
match = goImportRegex.exec(content);
|
||||
}
|
||||
} else if (ext === ".rs") {
|
||||
// Rust use statements
|
||||
const rsUseRegex = /use\s+([A-Za-z_][A-Za-z0-9_:]*)/g;
|
||||
let match: RegExpExecArray | null = rsUseRegex.exec(content);
|
||||
while (match !== null) {
|
||||
deps.push(match[1]);
|
||||
match = rsUseRegex.exec(content);
|
||||
}
|
||||
// If no LLM client provided, fall back to heuristics (for backward compat / tests)
|
||||
if (!client) {
|
||||
return extractFileHeuristic(filePath, content);
|
||||
}
|
||||
|
||||
// Deduplicate
|
||||
return [...new Set(deps)];
|
||||
}
|
||||
const prompt = buildFilePrompt(filePath, truncateForContext(content, 200));
|
||||
const response = await client.complete(prompt);
|
||||
setCached(hash, response);
|
||||
|
||||
function generatePurpose(
|
||||
name: string,
|
||||
ext: string,
|
||||
_baseName: string,
|
||||
_content: string,
|
||||
exports: string[],
|
||||
): string {
|
||||
// Check for specific file patterns
|
||||
if (/test|spec/i.test(name) && exports.length === 0) {
|
||||
return "Test suite";
|
||||
}
|
||||
if (/config|settings/i.test(name)) return "Configuration";
|
||||
if (/util|helper/i.test(name)) return "Utility functions";
|
||||
if (/types?\.d?\.ts$/.test(name)) return "Type definitions";
|
||||
if (/index\./.test(name)) return "Module entry point";
|
||||
if (/middleware/.test(name)) return "Middleware";
|
||||
if (/route/.test(name)) return "Route handlers";
|
||||
if (/controller/.test(name)) return "Controller";
|
||||
if (/service/.test(name)) return "Service layer";
|
||||
if (/model/.test(name)) return "Data model";
|
||||
if (/schema/.test(name)) return "Data schema";
|
||||
if (
|
||||
/component/.test(name) ||
|
||||
/\.tsx$/.test(name) ||
|
||||
/\.vue$/.test(name) ||
|
||||
/\.svelte$/.test(name)
|
||||
) {
|
||||
return "UI component";
|
||||
}
|
||||
if (/hook|use[A-Z]/.test(name)) return "React hook";
|
||||
if (/style|\.css|\.scss|\.less/.test(name)) return "Styling";
|
||||
if (/docker/i.test(name)) return "Container definition";
|
||||
if (/\.env/.test(name)) return "Environment variables";
|
||||
if (/readme/i.test(name)) return "Project documentation";
|
||||
|
||||
// Use exports to infer purpose
|
||||
if (exports.length > 0) {
|
||||
const firstFew = exports.slice(0, 3).join(", ");
|
||||
if (exports.length <= 3) return `Exports: ${firstFew}`;
|
||||
return `Exports ${exports.length} symbols: ${firstFew}...`;
|
||||
}
|
||||
|
||||
// Fallback to file type
|
||||
return (
|
||||
FILE_TYPE_PURPOSES[ext] ||
|
||||
(ext ? `${ext.slice(1).toUpperCase()} file` : `${name} file`)
|
||||
);
|
||||
const parsed = parseFileResponse(response);
|
||||
return {
|
||||
purpose: parsed.purpose,
|
||||
exports: [], // AST provides precise exports
|
||||
deps: parsed.deps,
|
||||
concepts: parsed.concepts,
|
||||
};
|
||||
}
|
||||
|
||||
export async function extractPackageLLM(
|
||||
relativePath: string,
|
||||
fileData: { name: string; purpose: string }[],
|
||||
client?: LLMClient,
|
||||
): Promise<LLMPackageData> {
|
||||
const dirName = basename(relativePath);
|
||||
|
||||
// Infer role from directory name
|
||||
let role = dirName === "." ? "Project root" : `Package ${dirName}`;
|
||||
if (dirName === "src" || dirName === "lib" || dirName === "source") {
|
||||
role = "Source code";
|
||||
} else if (dirName === "test" || dirName === "tests" || dirName === "spec") {
|
||||
role = "Test suite";
|
||||
} else if (dirName === "docs" || dirName === "doc") {
|
||||
role = "Documentation";
|
||||
} else if (dirName === "config" || dirName === "configuration") {
|
||||
role = "Configuration";
|
||||
} else if (
|
||||
dirName === "utils" ||
|
||||
dirName === "helpers" ||
|
||||
dirName === "util"
|
||||
) {
|
||||
role = "Utility functions";
|
||||
} else if (dirName === "types" || dirName === "type") {
|
||||
role = "Type definitions";
|
||||
} else if (dirName === "components" || dirName === "component") {
|
||||
role = "UI components";
|
||||
} else if (dirName === "hooks" || dirName === "hook") {
|
||||
role = "Custom hooks";
|
||||
} else if (dirName === "api" || dirName === "apis") {
|
||||
role = "API endpoints/handlers";
|
||||
} else if (
|
||||
dirName === "db" ||
|
||||
dirName === "database" ||
|
||||
dirName === "models"
|
||||
) {
|
||||
role = "Database layer";
|
||||
} else if (dirName === "auth" || dirName === "authentication") {
|
||||
role = "Authentication layer";
|
||||
if (!client) {
|
||||
return extractPackageHeuristic(relativePath, fileData);
|
||||
}
|
||||
|
||||
// Infer architecture from file patterns
|
||||
const prompt = buildPackagePrompt(relativePath, fileData);
|
||||
const response = await client.complete(prompt);
|
||||
const parsed = parsePackageResponse(response);
|
||||
|
||||
return {
|
||||
role: parsed.role || dirNameToRole(relativePath),
|
||||
arch: parsed.arch || `Contains ${fileData.length} files.`,
|
||||
};
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// HEURISTIC FALLBACK (for tests / no-LLM mode)
|
||||
// ============================================================================
|
||||
|
||||
export async function extractFileHeuristic(
|
||||
filePath: string,
|
||||
content?: string,
|
||||
): Promise<LLMFileData> {
|
||||
const fileContent = content ?? readFileSync(filePath, "utf8");
|
||||
const ext = extname(filePath).toLowerCase();
|
||||
const name = basename(filePath);
|
||||
const baseName = basename(filePath, ext);
|
||||
|
||||
const exports = extractExportsHeuristic(fileContent, ext, name);
|
||||
const deps = extractDepsHeuristic(fileContent, ext);
|
||||
const purpose = generatePurpose(name, ext, baseName, exports);
|
||||
|
||||
return { purpose, exports, deps, concepts: [] };
|
||||
}
|
||||
|
||||
export async function extractPackageHeuristic(
|
||||
relativePath: string,
|
||||
fileData: { name: string; purpose: string }[],
|
||||
): Promise<LLMPackageData> {
|
||||
const dirName = basename(relativePath);
|
||||
const role = dirNameToRole(dirName);
|
||||
|
||||
const purposes = fileData.map((f) => f.purpose);
|
||||
const hasTests = purposes.some((p) => p.includes("Test"));
|
||||
const hasTypes = purposes.some((p) => p.includes("Type"));
|
||||
@@ -298,3 +246,140 @@ export async function extractPackageLLM(
|
||||
|
||||
return { role, arch: arch.trim() };
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// INTERNAL HEURISTIC HELPERS
|
||||
// ============================================================================
|
||||
|
||||
function dirNameToRole(dirName: string): string {
|
||||
if (dirName === ".") return "Project root";
|
||||
if (dirName === "src" || dirName === "lib" || dirName === "source") return "Source code";
|
||||
if (dirName === "test" || dirName === "tests" || dirName === "spec") return "Test suite";
|
||||
if (dirName === "docs" || dirName === "doc") return "Documentation";
|
||||
if (dirName === "config" || dirName === "configuration") return "Configuration";
|
||||
if (dirName === "utils" || dirName === "helpers" || dirName === "util") return "Utility functions";
|
||||
if (dirName === "types" || dirName === "type") return "Type definitions";
|
||||
if (dirName === "components" || dirName === "component") return "UI components";
|
||||
if (dirName === "hooks" || dirName === "hook") return "Custom hooks";
|
||||
if (dirName === "api" || dirName === "apis") return "API endpoints/handlers";
|
||||
if (dirName === "db" || dirName === "database" || dirName === "models") return "Database layer";
|
||||
if (dirName === "auth" || dirName === "authentication") return "Authentication layer";
|
||||
return `Package ${dirName}`;
|
||||
}
|
||||
|
||||
function extractExportsHeuristic(content: string, ext: string, _filename: string): string[] {
|
||||
const exports: string[] = [];
|
||||
|
||||
if ([".ts", ".tsx", ".js", ".jsx", ".mjs"].includes(ext)) {
|
||||
const exportRegex = /(?:^|\n)\s*export\s+(?:default\s+)?(?:async\s+)?(?:function\s+|class\s+|const\s+|let\s+|var\s+|interface\s+|type\s+|enum\s+)?([A-Za-z_$][A-Za-z0-9_$]*)/g;
|
||||
let match: RegExpExecArray | null;
|
||||
match = exportRegex.exec(content);
|
||||
while (match !== null) {
|
||||
exports.push(match[1]);
|
||||
match = exportRegex.exec(content);
|
||||
}
|
||||
|
||||
const namedExportRegex = /(?:^|\n)\s*export\s*\{\s*([^}]+)\s*\}/g;
|
||||
match = namedExportRegex.exec(content);
|
||||
while (match !== null) {
|
||||
const names = match[1].split(",").map((s) => s.trim().split(/\s+as\s+/)[0].trim());
|
||||
exports.push(...names);
|
||||
match = namedExportRegex.exec(content);
|
||||
}
|
||||
} else if (ext === ".py") {
|
||||
const pyRegex = /^(?:async\s+)?def\s+([A-Za-z_][A-Za-z0-9_]*)|class\s+([A-Za-z_][A-Za-z0-9_]*)/gm;
|
||||
let match: RegExpExecArray | null = pyRegex.exec(content);
|
||||
while (match !== null) {
|
||||
exports.push(match[1] || match[2]);
|
||||
match = pyRegex.exec(content);
|
||||
}
|
||||
} else if (ext === ".go") {
|
||||
const goRegex = /^(?:func|type|var|const)\s+([A-Z][A-Za-z0-9_]*)/gm;
|
||||
let match: RegExpExecArray | null = goRegex.exec(content);
|
||||
while (match !== null) {
|
||||
exports.push(match[1]);
|
||||
match = goRegex.exec(content);
|
||||
}
|
||||
} else if (ext === ".rs") {
|
||||
const rsRegex = /pub\s+(?:fn|struct|enum|trait|type|const|static|use)\s+([A-Za-z_][A-Za-z0-9_]*)/g;
|
||||
let match: RegExpExecArray | null = rsRegex.exec(content);
|
||||
while (match !== null) {
|
||||
exports.push(match[1]);
|
||||
match = rsRegex.exec(content);
|
||||
}
|
||||
}
|
||||
|
||||
return [...new Set(exports)];
|
||||
}
|
||||
|
||||
function extractDepsHeuristic(content: string, ext: string): string[] {
|
||||
const deps: string[] = [];
|
||||
|
||||
if ([".ts", ".tsx", ".js", ".jsx", ".mjs"].includes(ext)) {
|
||||
const importRegex = /import\s+(?:(?:type\s+)?\{[^}]*\}|\*\s+as\s+\w+|\w+)\s+from\s+['"]([^'"]+)['"]/g;
|
||||
let match: RegExpExecArray | null = importRegex.exec(content);
|
||||
while (match !== null) {
|
||||
deps.push(match[1]);
|
||||
match = importRegex.exec(content);
|
||||
}
|
||||
const requireRegex = /require\s*\(\s*['"]([^'"]+)['"]\s*\)/g;
|
||||
match = requireRegex.exec(content);
|
||||
while (match !== null) {
|
||||
deps.push(match[1]);
|
||||
match = requireRegex.exec(content);
|
||||
}
|
||||
} else if (ext === ".py") {
|
||||
const pyImportRegex = /^(?:from|import)\s+([A-Za-z_][A-Za-z0-9_.]*)/gm;
|
||||
let match: RegExpExecArray | null = pyImportRegex.exec(content);
|
||||
while (match !== null) {
|
||||
deps.push(match[1]);
|
||||
match = pyImportRegex.exec(content);
|
||||
}
|
||||
} else if (ext === ".go") {
|
||||
const goImportRegex = /"([^"]+)"/g;
|
||||
let match: RegExpExecArray | null = goImportRegex.exec(content);
|
||||
while (match !== null) {
|
||||
if (match[1].includes("/")) deps.push(match[1]);
|
||||
match = goImportRegex.exec(content);
|
||||
}
|
||||
} else if (ext === ".rs") {
|
||||
const rsUseRegex = /use\s+([A-Za-z_][A-Za-z0-9_:]*)/g;
|
||||
let match: RegExpExecArray | null = rsUseRegex.exec(content);
|
||||
while (match !== null) {
|
||||
deps.push(match[1]);
|
||||
match = rsUseRegex.exec(content);
|
||||
}
|
||||
}
|
||||
|
||||
return [...new Set(deps)];
|
||||
}
|
||||
|
||||
function generatePurpose(name: string, ext: string, _baseName: string, exports: string[]): string {
|
||||
if (/test|spec/i.test(name) && exports.length === 0) return "Test suite";
|
||||
if (/config|settings/i.test(name)) return "Configuration";
|
||||
if (/util|helper/i.test(name)) return "Utility functions";
|
||||
if (/types?\.d?\.ts$/.test(name)) return "Type definitions";
|
||||
if (/index\./.test(name)) return "Module entry point";
|
||||
if (/middleware/.test(name)) return "Middleware";
|
||||
if (/route/.test(name)) return "Route handlers";
|
||||
if (/controller/.test(name)) return "Controller";
|
||||
if (/service/.test(name)) return "Service layer";
|
||||
if (/model/.test(name)) return "Data model";
|
||||
if (/schema/.test(name)) return "Data schema";
|
||||
if (/component/.test(name) || /\.tsx$/.test(name) || /\.vue$/.test(name) || /\.svelte$/.test(name)) {
|
||||
return "UI component";
|
||||
}
|
||||
if (/hook|use[A-Z]/.test(name)) return "React hook";
|
||||
if (/style|\.css|\.scss|\.less/.test(name)) return "Styling";
|
||||
if (/docker/i.test(name)) return "Container definition";
|
||||
if (/\.env/.test(name)) return "Environment variables";
|
||||
if (/readme/i.test(name)) return "Project documentation";
|
||||
|
||||
if (exports.length > 0) {
|
||||
const firstFew = exports.slice(0, 3).join(", ");
|
||||
if (exports.length <= 3) return `Exports: ${firstFew}`;
|
||||
return `Exports ${exports.length} symbols: ${firstFew}...`;
|
||||
}
|
||||
|
||||
return FILE_TYPE_PURPOSES[ext] || (ext ? `${ext.slice(1).toUpperCase()} file` : `${name} file`);
|
||||
}
|
||||
|
||||
+14
-7
@@ -6,10 +6,14 @@ import { extractFileAST } from "./ast-extract.js";
|
||||
import { mergeFileData } from "./merge.js";
|
||||
import { generateDirectoryMap } from "./init.js";
|
||||
import { readdirSync, statSync } from "fs";
|
||||
import type { LLMClient } from "./llm-client.js";
|
||||
|
||||
const SMALL_PACKAGE_THRESHOLD = 10;
|
||||
|
||||
export async function patchFile(filePath: string): Promise<void> {
|
||||
export async function patchFile(
|
||||
filePath: string,
|
||||
llmClient?: LLMClient,
|
||||
): Promise<void> {
|
||||
const dirPath = dirname(filePath);
|
||||
const mapPath = join(dirPath, ".pi-map.md");
|
||||
|
||||
@@ -31,18 +35,21 @@ export async function patchFile(filePath: string): Promise<void> {
|
||||
return st.isFile();
|
||||
});
|
||||
const relDir = relative(process.cwd(), dirPath) || ".";
|
||||
await generateDirectoryMap({
|
||||
dirPath,
|
||||
relativePath: relDir,
|
||||
files,
|
||||
});
|
||||
await generateDirectoryMap(
|
||||
{
|
||||
dirPath,
|
||||
relativePath: relDir,
|
||||
files,
|
||||
},
|
||||
llmClient,
|
||||
);
|
||||
console.log(
|
||||
`Full rewrite of ${mapPath} (small package: ${allFiles.length} files)`,
|
||||
);
|
||||
} else {
|
||||
// Section-level patch
|
||||
const existing = parsePackageMap(readFileSync(mapPath, "utf8"));
|
||||
const llmData = await extractFileLLM(filePath);
|
||||
const llmData = await extractFileLLM(filePath, llmClient);
|
||||
const astData = await extractFileAST(filePath);
|
||||
const fileName = basename(filePath);
|
||||
const updatedFile = mergeFileData(fileName, llmData, astData);
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
import { LLMError } from "./llm-error.js";
|
||||
import type { LLMClient } from "./llm-client.js";
|
||||
|
||||
/**
|
||||
* Pi LLM Client — calls Pi's built-in LLM via ExtensionAPI.
|
||||
*
|
||||
* TODO: This is a stub. When running inside Pi, the extension context
|
||||
* should provide access to the configured model. The exact API shape
|
||||
* depends on the Pi runtime version. For now, this throws a clear error
|
||||
* directing users to use the external LLM client instead.
|
||||
*/
|
||||
export class PiLLMClient implements LLMClient {
|
||||
constructor(private _extensionContext?: unknown) {}
|
||||
|
||||
async complete(_prompt: string): Promise<string> {
|
||||
throw new LLMError(
|
||||
"Pi native LLM client is not yet implemented. " +
|
||||
"Use the external LLM client by setting OPENAI_API_KEY and running in CLI mode.",
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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"]);
|
||||
});
|
||||
});
|
||||
@@ -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");
|
||||
});
|
||||
});
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -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({});
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user