pi-map/tests/llm-integration.test.ts

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;
		}
	});
});