From: Cameron Otsuka Date: Sat, 2 May 2026 01:26:14 +0000 (-0700) Subject: local models, basic memory extension X-Git-Url: https://git.otsuka.systems/?a=commitdiff_plain;h=e8b6b76371d44aea5b99cca3cb3399643f1e5b09;p=dotfiles local models, basic memory extension --- diff --git a/dot_pi/agent/extensions/memory.ts b/dot_pi/agent/extensions/memory.ts new file mode 100644 index 0000000..4c2c63c --- /dev/null +++ b/dot_pi/agent/extensions/memory.ts @@ -0,0 +1,463 @@ +import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; +import { spawn } from "node:child_process"; +import fs from "node:fs/promises"; +import path from "node:path"; + +const TOP_FILES = 8; +const TOP_RANGES = 12; +const MAX_RG_EVENTS = 600; +const SNIPPET_CONTEXT_LINES = 2; +const READ_WINDOW_LINES = 140; +const RG_TIMEOUT_MS = 15_000; + +interface MemoryMatch { + file: string; + line: number; + text: string; +} + +interface CandidateFile { + file: string; + matches: MemoryMatch[]; + score: number; + keywords: string[]; +} + +interface CandidateRange { + file: string; + start: number; + end: number; + offset: number; + limit: number; + heading?: string; + headingLine?: number; + matchLines: number[]; + keywords: string[]; + snippet: string; + score: number; +} + +interface ActiveMemorySearch { + directory: string; + keywords: string[]; + allowed: Map>; +} + +function parseArgs(input: string): string[] { + const args: string[] = []; + let current = ""; + let quote: "'" | '"' | undefined; + let escaping = false; + + for (const char of input) { + if (escaping) { + current += char; + escaping = false; + continue; + } + if (char === "\\") { + escaping = true; + continue; + } + if (quote) { + if (char === quote) { + quote = undefined; + } else { + current += char; + } + continue; + } + if (char === "'" || char === '"') { + quote = char; + continue; + } + if (/\s/.test(char)) { + if (current.length > 0) { + args.push(current); + current = ""; + } + continue; + } + current += char; + } + + if (escaping) current += "\\"; + if (quote) throw new Error(`Unclosed ${quote} quote`); + if (current.length > 0) args.push(current); + return args; +} + +function normalizePath(cwd: string, inputPath: string): string { + const trimmed = inputPath.startsWith("@") ? inputPath.slice(1) : inputPath; + return path.resolve(cwd, trimmed); +} + +function isMarkdownFile(filePath: string): boolean { + const ext = path.extname(filePath).toLowerCase(); + return ext === ".md" || ext === ".markdown" || ext === ".mdown" || ext === ".mkd"; +} + +function isInside(parent: string, child: string): boolean { + const relative = path.relative(parent, child); + return relative === "" || (!relative.startsWith("..") && !path.isAbsolute(relative)); +} + +function displayPath(cwd: string, filePath: string): string { + const relative = path.relative(cwd, filePath); + if (relative && !relative.startsWith("..") && !path.isAbsolute(relative)) return relative; + return filePath; +} + +function truncateLine(line: string, max = 220): string { + return line.length <= max ? line : `${line.slice(0, max - 1)}…`; +} + +function keywordHits(text: string, keywords: string[]): string[] { + const lower = text.toLowerCase(); + return keywords.filter((keyword) => lower.includes(keyword.toLowerCase())); +} + +function scoreMatches(matches: MemoryMatch[], keywords: string[]): { score: number; hits: string[] } { + const hits = new Set(); + let headingMatches = 0; + for (const match of matches) { + for (const keyword of keywordHits(match.text, keywords)) hits.add(keyword); + if (/^\s*#{1,6}\s+/.test(match.text)) headingMatches += 1; + } + return { + score: matches.length * 3 + hits.size * 10 + headingMatches * 8, + hits: [...hits], + }; +} + +async function runRipgrep(directory: string, keywords: string[]): Promise { + const args = [ + "--json", + "--line-number", + "--ignore-case", + "--fixed-strings", + "--glob", + "*.md", + "--glob", + "*.markdown", + "--glob", + "*.mdown", + "--glob", + "*.mkd", + ]; + for (const keyword of keywords) args.push("-e", keyword); + args.push("--", directory); + + return new Promise((resolve, reject) => { + const child = spawn("rg", args, { stdio: ["ignore", "pipe", "pipe"] }); + const timer = setTimeout(() => child.kill("SIGTERM"), RG_TIMEOUT_MS); + let stdout = ""; + let stderr = ""; + + child.stdout.setEncoding("utf8"); + child.stderr.setEncoding("utf8"); + child.stdout.on("data", (chunk) => { + stdout += chunk; + // Keep memory bounded if the corpus is very noisy. + if (stdout.length > 5_000_000) child.kill("SIGTERM"); + }); + child.stderr.on("data", (chunk) => { + stderr += chunk; + }); + child.on("error", (error) => { + clearTimeout(timer); + reject(error); + }); + child.on("close", (code, signal) => { + clearTimeout(timer); + if (signal && signal !== "SIGTERM") { + reject(new Error(`rg stopped by signal ${signal}`)); + return; + } + if (code !== 0 && code !== 1 && signal !== "SIGTERM") { + reject(new Error(stderr.trim() || `rg exited with code ${code}`)); + return; + } + + const matches: MemoryMatch[] = []; + for (const line of stdout.split("\n")) { + if (!line.trim()) continue; + let event: any; + try { + event = JSON.parse(line); + } catch { + continue; + } + if (event.type !== "match") continue; + const fileText = event.data?.path?.text; + const lineNumber = event.data?.line_number; + const lineText = event.data?.lines?.text; + if (typeof fileText !== "string" || typeof lineNumber !== "number" || typeof lineText !== "string") { + continue; + } + matches.push({ file: path.resolve(fileText), line: lineNumber, text: lineText.replace(/\r?\n$/, "") }); + if (matches.length >= MAX_RG_EVENTS) break; + } + resolve(matches); + }); + }); +} + +function findNearestHeading(lines: string[], beforeLine: number): { line: number; text: string } | undefined { + for (let i = Math.min(beforeLine, lines.length) - 1; i >= 0; i -= 1) { + const text = lines[i] ?? ""; + if (/^\s*#{1,6}\s+/.test(text)) return { line: i + 1, text: text.trim() }; + } + return undefined; +} + +function findNextHeading(lines: string[], afterLine: number): number | undefined { + for (let i = Math.max(0, afterLine); i < lines.length; i += 1) { + if (/^\s*#{1,6}\s+/.test(lines[i] ?? "")) return i + 1; + } + return undefined; +} + +function clusterLines(lineNumbers: number[]): number[][] { + const sorted = [...new Set(lineNumbers)].sort((a, b) => a - b); + const clusters: number[][] = []; + for (const line of sorted) { + const last = clusters[clusters.length - 1]; + if (!last || line - last[last.length - 1] > 40) { + clusters.push([line]); + } else { + last.push(line); + } + } + return clusters; +} + +async function buildRanges(cwd: string, candidates: CandidateFile[], keywords: string[]): Promise { + const ranges: CandidateRange[] = []; + for (const candidate of candidates) { + let content: string; + try { + content = await fs.readFile(candidate.file, "utf8"); + } catch { + continue; + } + const lines = content.split(/\r?\n/); + const clusters = clusterLines(candidate.matches.map((match) => match.line)); + for (const cluster of clusters) { + const firstMatch = cluster[0]; + const lastMatch = cluster[cluster.length - 1]; + const heading = findNearestHeading(lines, firstMatch); + const nextHeading = findNextHeading(lines, lastMatch + 1); + + let start = Math.max(1, firstMatch - 20); + if (heading && firstMatch - heading.line <= 80) start = heading.line; + + let end = Math.min(lines.length, Math.max(lastMatch + 60, start + 80)); + if (nextHeading && nextHeading > lastMatch && nextHeading - 1 >= start) { + end = Math.min(end, nextHeading - 1); + } + end = Math.min(lines.length, start + READ_WINDOW_LINES - 1, end); + + const snippetStart = Math.max(1, firstMatch - SNIPPET_CONTEXT_LINES); + const snippetEnd = Math.min(lines.length, lastMatch + SNIPPET_CONTEXT_LINES); + const snippetLines: string[] = []; + for (let line = snippetStart; line <= snippetEnd; line += 1) { + const marker = cluster.includes(line) ? ">" : " "; + snippetLines.push(`${marker} ${String(line).padStart(5, " ")}: ${truncateLine(lines[line - 1] ?? "")}`); + } + + const clusterKeywords = new Set(); + for (const line of cluster) { + for (const keyword of keywordHits(lines[line - 1] ?? "", keywords)) clusterKeywords.add(keyword); + } + + ranges.push({ + file: candidate.file, + start, + end, + offset: start, + limit: end - start + 1, + heading: heading?.text, + headingLine: heading?.line, + matchLines: cluster, + keywords: [...clusterKeywords], + snippet: snippetLines.join("\n"), + score: candidate.score + cluster.length * 5 + clusterKeywords.size * 5, + }); + } + } + return ranges.sort((a, b) => b.score - a.score).slice(0, TOP_RANGES); +} + +function groupCandidates(matches: MemoryMatch[], keywords: string[]): CandidateFile[] { + const byFile = new Map(); + for (const match of matches) { + const list = byFile.get(match.file) ?? []; + list.push(match); + byFile.set(match.file, list); + } + + const candidates: CandidateFile[] = []; + for (const [file, fileMatches] of byFile) { + const { score, hits } = scoreMatches(fileMatches, keywords); + candidates.push({ file, matches: fileMatches, score, keywords: hits }); + } + return candidates.sort((a, b) => b.score - a.score).slice(0, TOP_FILES); +} + +function buildPrompt(cwd: string, directory: string, keywords: string[], matches: MemoryMatch[], ranges: CandidateRange[]): string { + const lines: string[] = []; + lines.push(`Memory search results for ${directory}`); + lines.push(`Keywords: ${keywords.map((keyword) => JSON.stringify(keyword)).join(", ")}`); + lines.push(`Search method: ripgrep over Markdown files only (*.md, *.markdown, *.mdown, *.mkd).`); + lines.push(`Total matched lines sampled: ${matches.length}${matches.length >= MAX_RG_EVENTS ? ` (capped at ${MAX_RG_EVENTS})` : ""}.`); + lines.push(""); + lines.push("Instructions for the agent:"); + lines.push("1. Do not bulk-read this Markdown directory."); + lines.push("2. Use Pi's `read` tool only on the suggested ranges below. The memory extension will block other Markdown reads inside this directory for this turn."); + lines.push("3. Start with the top 1-3 ranges. If those are insufficient, read additional listed ranges or ask for a refined `/memory ` search."); + lines.push("4. When answering, cite paths and line numbers from the files you read."); + lines.push(""); + + if (ranges.length === 0) { + lines.push("No candidate ranges were found."); + lines.push("Answer that no matching Markdown memory entries were found, and suggest more specific or alternate keywords."); + return lines.join("\n"); + } + + lines.push("Suggested candidate ranges:"); + for (let i = 0; i < ranges.length; i += 1) { + const range = ranges[i]; + const fileDisplay = displayPath(cwd, range.file); + lines.push(""); + lines.push(`${i + 1}. ${fileDisplay}:${range.start}-${range.end}`); + lines.push(` read arguments: { "path": ${JSON.stringify(range.file)}, "offset": ${range.offset}, "limit": ${range.limit} }`); + if (range.heading) lines.push(` nearest heading: ${range.heading} (line ${range.headingLine})`); + lines.push(` matched lines: ${range.matchLines.join(", ")}`); + if (range.keywords.length) lines.push(` matched keywords: ${range.keywords.map((k) => JSON.stringify(k)).join(", ")}`); + lines.push(" snippet:"); + lines.push("```text"); + lines.push(range.snippet); + lines.push("```"); + } + + lines.push(""); + lines.push("Now read the most relevant listed range(s), then summarize what the memory says about these keywords."); + return lines.join("\n"); +} + +function buildAllowedMap(ranges: CandidateRange[]): Map> { + const allowed = new Map>(); + for (const range of ranges) { + const list = allowed.get(range.file) ?? []; + list.push({ start: range.start, end: range.end }); + allowed.set(range.file, list); + } + return allowed; +} + +export default function memoryExtension(pi: ExtensionAPI) { + let activeMemory: ActiveMemorySearch | undefined; + + pi.registerCommand("memory", { + description: "Search Markdown memory: /memory ", + handler: async (args, ctx) => { + let parsed: string[]; + try { + parsed = parseArgs(args ?? ""); + } catch (error: any) { + ctx.ui.notify(error?.message ?? "Failed to parse /memory arguments", "error"); + return; + } + + const [directoryArg, ...keywordArgs] = parsed; + const keywords = keywordArgs.map((keyword) => keyword.trim()).filter(Boolean); + if (!directoryArg || keywords.length === 0) { + ctx.ui.notify("Usage: /memory ", "error"); + return; + } + + const directory = normalizePath(ctx.cwd, directoryArg); + let stat; + try { + stat = await fs.stat(directory); + } catch { + ctx.ui.notify(`Memory directory not found: ${directory}`, "error"); + return; + } + if (!stat.isDirectory()) { + ctx.ui.notify(`Memory path is not a directory: ${directory}`, "error"); + return; + } + + ctx.ui.notify(`Searching Markdown memory in ${directory}...`, "info"); + + let matches: MemoryMatch[]; + try { + matches = await runRipgrep(directory, keywords); + } catch (error: any) { + ctx.ui.notify(`memory rg failed: ${error?.message ?? String(error)}`, "error"); + return; + } + + const candidates = groupCandidates(matches, keywords); + const ranges = await buildRanges(ctx.cwd, candidates, keywords); + activeMemory = { + directory, + keywords, + allowed: buildAllowedMap(ranges), + }; + + const prompt = buildPrompt(ctx.cwd, directory, keywords, matches, ranges); + pi.sendUserMessage(prompt); + }, + }); + + pi.on("before_agent_start", async (event) => { + if (!activeMemory) return; + return { + systemPrompt: + event.systemPrompt + + "\n\nActive /memory search policy: For this turn, do not bulk-read the Markdown memory directory. Only use the read tool on the exact path/offset/limit ranges listed in the user's /memory search results. If more information is needed, ask the user to run a refined /memory search.", + }; + }); + + pi.on("tool_call", async (event, ctx) => { + if (!activeMemory || event.toolName !== "read") return; + const input = event.input as { path?: string; offset?: number; limit?: number }; + if (!input?.path) return; + + const requestedPath = normalizePath(ctx.cwd, input.path); + if (!isMarkdownFile(requestedPath) || !isInside(activeMemory.directory, requestedPath)) return; + + const allowedRanges = activeMemory.allowed.get(requestedPath); + if (!allowedRanges) { + return { + block: true, + reason: `Blocked by /memory: ${requestedPath} was not one of the top returned Markdown ranges. Run /memory with refined keywords if this file is needed.`, + }; + } + + if (typeof input.offset !== "number" || typeof input.limit !== "number") { + return { + block: true, + reason: "Blocked by /memory: Markdown memory reads must include offset and limit from the suggested ranges.", + }; + } + + const start = input.offset; + const end = input.offset + input.limit - 1; + const allowed = allowedRanges.some((range) => start >= range.start && end <= range.end); + if (!allowed) { + const rangesText = allowedRanges.map((range) => `${range.start}-${range.end}`).join(", "); + return { + block: true, + reason: `Blocked by /memory: requested lines ${start}-${end} are outside the allowed ranges for this file (${rangesText}).`, + }; + } + }); + + pi.on("agent_end", async () => { + activeMemory = undefined; + }); +} diff --git a/dot_pi/agent/models.json b/dot_pi/agent/models.json new file mode 100644 index 0000000..4a0cf2a --- /dev/null +++ b/dot_pi/agent/models.json @@ -0,0 +1,22 @@ +{ + "providers": { + "llama-swap (local)": { + "baseUrl": "http://localhost:12434/v1", + "api": "openai-completions", + "apiKey": "none", + "models": [ + { "id": "qwen3.6" } + ] + }, + "llama-swap (nas)": { + "baseUrl": "http://10.92.3.2:12434/v1", + "api": "openai-completions", + "apiKey": "none", + "models": [ + { "id": "gemma4" }, + { "id": "nemotron-cascade2" }, + { "id": "qwen3.6" } + ] + } + } +}