@chendpoc/pi-memory 0.1.0

package/src/trainer/index.ts

@@ -0,0 +1,147 @@
+import { installBundle } from "../bundle/install.js";
+import { defaultBundleRoot, defaultSessionsDir } from "../paths.js";
+
+import { loadSessions } from "./sessionLoader.js";
+import { extractFactsFromSessions, type ExtractFactsOptions } from "./extractFacts.js";
+import { resolveEntities } from "./entityResolver.js";
+import { buildBundle, type BuildBundleResult } from "./bundleBuilder.js";
+import { readMarker, writeMarker } from "./marker.js";
+import { loadExistingBundle } from "./bundleLoader.js";
+import { deltaMerge, type DeltaLog } from "./deltaMerge.js";
+
+export interface TrainBundleConfig {
+  sessionsDir?: string;
+  bundleRoot?: string;
+  /** Ignore marker and rebuild from all sessions. */
+  full?: boolean;
+  /** Show what would be extracted without writing. */
+  dryRun?: boolean;
+  /** Bundle version to stamp (default "0.6.0"). */
+  bundleVersion?: string;
+  /** Extraction options (optional LLM extractor). */
+  extractOpts?: ExtractFactsOptions;
+  /** Skip delta merge — full rebuild even if existing bundle present (default false). */
+  noMerge?: boolean;
+}
+
+export interface TrainBundleResult {
+  sessionsProcessed: number;
+  entityCount: number;
+  relationCount: number;
+  eventCount: number;
+  bundleResult?: BuildBundleResult;
+  installResult?: {
+    bundle_ts: string;
+    bundle_version: string;
+    installed_dir: string;
+    files_copied: number;
+  };
+  dryRun: boolean;
+  /** Delta operation log when merge was performed. */
+  delta?: DeltaLog;
+}
+
+/**
+ * Full training pipeline:
+ *   load existing bundle → extract new → delta merge → build → install → update marker.
+ *
+ * When `noMerge` is true (or no existing bundle), falls back to full rebuild.
+ */
+export async function trainBundle(
+  config: TrainBundleConfig = {},
+): Promise<TrainBundleResult> {
+  const sessionsDir = config.sessionsDir ?? defaultSessionsDir();
+  const bundleRoot = config.bundleRoot ?? defaultBundleRoot();
+
+  let modifiedAfter: Date | null = null;
+  if (!config.full) {
+    modifiedAfter = await readMarker(bundleRoot);
+  }
+
+  const sessions = await loadSessions({ sessionsDir, modifiedAfter });
+
+  if (sessions.length === 0) {
+    return {
+      sessionsProcessed: 0,
+      entityCount: 0,
+      relationCount: 0,
+      eventCount: 0,
+      dryRun: config.dryRun ?? false,
+    };
+  }
+
+  const extracted = await extractFactsFromSessions(sessions, config.extractOpts);
+  const newGraph = resolveEntities(extracted.entities, extracted.relations);
+
+  let finalGraph = newGraph;
+  let finalEvents = extracted.events;
+  let deltaLog: DeltaLog | undefined;
+
+  if (!config.noMerge) {
+    const existingBundle = await loadExistingBundle(bundleRoot);
+    if (existingBundle) {
+      const merged = deltaMerge(
+        existingBundle,
+        { graph: newGraph, events: extracted.events },
+      );
+      finalGraph = merged.graph;
+      finalEvents = merged.events;
+      deltaLog = merged.delta;
+    }
+  }
+
+  if (config.dryRun) {
+    return {
+      sessionsProcessed: sessions.length,
+      entityCount: finalGraph.entities.length,
+      relationCount: finalGraph.relations.length,
+      eventCount: finalEvents.length,
+      dryRun: true,
+      delta: deltaLog,
+    };
+  }
+
+  const bundleResult = await buildBundle(
+    { graph: finalGraph, events: finalEvents },
+    { outputDir: bundleRoot, bundleVersion: config.bundleVersion },
+  );
+
+  const installResult = await installBundle({
+    bundleRoot,
+    sourceDir: bundleResult.bundleDir,
+  });
+
+  const latestMtime = sessions.reduce(
+    (max, s) => (s.modifiedAt > max ? s.modifiedAt : max),
+    sessions[0]!.modifiedAt,
+  );
+  await writeMarker(bundleRoot, latestMtime);
+
+  return {
+    sessionsProcessed: sessions.length,
+    entityCount: bundleResult.stats.entityCount,
+    relationCount: bundleResult.stats.edgeCount,
+    eventCount: bundleResult.stats.eventCount,
+    bundleResult,
+    installResult,
+    dryRun: false,
+    delta: deltaLog,
+  };
+}
+
+export { loadSessions } from "./sessionLoader.js";
+export type { LoadedSession, SessionTurn, SessionLoaderOptions } from "./sessionLoader.js";
+export { extractFacts, extractFactsFromSessions, RELATION_CATALOG, ALL_RELATIONS } from "./extractFacts.js";
+export type {
+  ExtractedEntity, ExtractedRelation, ExtractedEvent,
+  ExtractionResult, EntityType, LLMFactExtractor, ExtractFactsOptions,
+} from "./extractFacts.js";
+export { resolveEntities } from "./entityResolver.js";
+export type { ResolvedEntity, ResolvedRelation, ResolvedGraph } from "./entityResolver.js";
+export { buildBundle } from "./bundleBuilder.js";
+export type { BundleData, BuildBundleOptions, BuildBundleResult } from "./bundleBuilder.js";
+export { readMarker, writeMarker } from "./marker.js";
+export { loadExistingBundle } from "./bundleLoader.js";
+export type { ExistingBundle } from "./bundleLoader.js";
+export { deltaMerge } from "./deltaMerge.js";
+export type { DeltaOp, DeltaLogEntry, DeltaLog, MergeResult } from "./deltaMerge.js";

package/src/trainer/llmExtractor.ts

@@ -0,0 +1,206 @@
+import type { SessionTurn } from "./sessionLoader.js";
+import type {
+  ExtractedEntity,
+  ExtractedRelation,
+  ExtractedEvent,
+  LLMFactExtractor,
+  EntityType,
+} from "./extractFacts.js";
+import { extractFacts } from "./extractFacts.js";
+import type { LoadedSession } from "./sessionLoader.js";
+import { ALL_RELATIONS } from "./extractFacts.js";
+
+/** Provider-agnostic LLM client — any backend that can complete a prompt. */
+export interface LLMClient {
+  complete(prompt: string): Promise<string>;
+}
+
+export interface LLMExtractorOptions {
+  client: LLMClient;
+  /** How many turns to send per LLM call (default 10). */
+  batchSize?: number;
+}
+
+const VALID_ENTITY_TYPES = new Set<string>([
+  "Person", "Project", "Tool", "Company", "Organization",
+  "Location", "Document",
+]);
+
+const VALID_ENTITY_TYPES_FULL = new Set<string>([
+  ...VALID_ENTITY_TYPES,
+  "Language", "Concept", "File", "Product", "Website", "Platform", "Unknown",
+]);
+
+const RELATION_SET: Set<string> = new Set(ALL_RELATIONS);
+
+function buildPrompt(turns: SessionTurn[], sessionId: string): string {
+  const turnsText = turns
+    .map((t) => `[${t.role} #${t.turnIndex}]: ${t.content}`)
+    .join("\n\n");
+
+  return `You are a structured fact extractor. Analyze the following conversation turns and extract:
+
+1. **Entities** — people, projects, tools, companies, organizations, locations, documents mentioned.
+   Each entity: { "name": string, "type": one of "Person"|"Project"|"Tool"|"Company"|"Organization"|"Location"|"Document" }
+
+2. **Relations** — connections between entities. Use ONLY these relation types:
+   ${ALL_RELATIONS.join(", ")}
+   Each relation: { "head": string, "relation": string, "tail": string, "turn_index": number, "evidence": string, "negated": boolean }
+   Set negated=true only when the text explicitly says something is no longer true (e.g. "no longer uses X").
+
+3. **Events** — decisions, milestones, deadlines, deployments.
+   Each event: { "description": string, "turn_index": number }
+
+Session ID: ${sessionId}
+
+Conversation:
+${turnsText}
+
+Respond with ONLY a JSON object (no markdown fences, no explanation):
+{
+  "entities": [...],
+  "relations": [...],
+  "events": [...]
+}`;
+}
+
+function parseEntityType(raw: string): EntityType {
+  if (VALID_ENTITY_TYPES_FULL.has(raw)) return raw as EntityType;
+  return "Unknown";
+}
+
+interface RawLLMEntity {
+  name?: string;
+  type?: string;
+}
+
+interface RawLLMRelation {
+  head?: string;
+  relation?: string;
+  tail?: string;
+  turn_index?: number;
+  evidence?: string;
+  negated?: boolean;
+}
+
+interface RawLLMEvent {
+  description?: string;
+  turn_index?: number;
+}
+
+interface RawLLMResponse {
+  entities?: RawLLMEntity[];
+  relations?: RawLLMRelation[];
+  events?: RawLLMEvent[];
+}
+
+function parseLLMResponse(
+  raw: string,
+  sessionId: string,
+  turns: SessionTurn[],
+): { entities: ExtractedEntity[]; relations: ExtractedRelation[]; events: ExtractedEvent[] } {
+  const cleaned = raw.replace(/^```(?:json)?\s*/m, "").replace(/\s*```\s*$/m, "").trim();
+  const parsed: RawLLMResponse = JSON.parse(cleaned);
+
+  const entities: ExtractedEntity[] = [];
+  if (Array.isArray(parsed.entities)) {
+    for (const e of parsed.entities) {
+      const name = typeof e.name === "string" ? e.name.trim() : "";
+      if (!name || name.length > 100) continue;
+      entities.push({
+        name,
+        type: parseEntityType(typeof e.type === "string" ? e.type : "Unknown"),
+        mentions: [{
+          sessionId,
+          turnIndex: 0,
+          snippet: name,
+        }],
+      });
+    }
+  }
+
+  const relations: ExtractedRelation[] = [];
+  if (Array.isArray(parsed.relations)) {
+    for (const r of parsed.relations) {
+      const head = typeof r.head === "string" ? r.head.trim() : "";
+      const tail = typeof r.tail === "string" ? r.tail.trim() : "";
+      const relation = typeof r.relation === "string" ? r.relation.trim() : "";
+      if (!head || !tail || !relation) continue;
+      if (!RELATION_SET.has(relation)) continue;
+      const turnIdx = typeof r.turn_index === "number" ? r.turn_index : 0;
+      relations.push({
+        headName: head,
+        relation,
+        tailName: tail,
+        sessionId,
+        turnIndex: turnIdx,
+        evidence: typeof r.evidence === "string" ? r.evidence.slice(0, 200) : "",
+        negated: r.negated === true ? true : undefined,
+      });
+    }
+  }
+
+  const events: ExtractedEvent[] = [];
+  if (Array.isArray(parsed.events)) {
+    for (const ev of parsed.events) {
+      const desc = typeof ev.description === "string" ? ev.description.trim() : "";
+      if (!desc) continue;
+      const turnIdx = typeof ev.turn_index === "number" ? ev.turn_index : 0;
+      const matchingTurn = turns.find((t) => t.turnIndex === turnIdx);
+      events.push({
+        description: desc.slice(0, 300),
+        sessionId,
+        timestamp: new Date().toISOString(),
+        turnIndex: matchingTurn?.turnIndex ?? turnIdx,
+      });
+    }
+  }
+
+  return { entities, relations, events };
+}
+
+/**
+ * LLM-backed fact extractor. Sends turn batches to the LLM, parses
+ * structured JSON. Falls back to regex extraction on any LLM failure.
+ */
+export function createLLMFactExtractor(opts: LLMExtractorOptions): LLMFactExtractor {
+  const batchSize = opts.batchSize ?? 10;
+
+  return {
+    async extractFacts(
+      turns: SessionTurn[],
+      sessionId: string,
+    ): Promise<{ entities: ExtractedEntity[]; relations: ExtractedRelation[]; events: ExtractedEvent[] }> {
+      const allEntities: ExtractedEntity[] = [];
+      const allRelations: ExtractedRelation[] = [];
+      const allEvents: ExtractedEvent[] = [];
+
+      for (let i = 0; i < turns.length; i += batchSize) {
+        const batch = turns.slice(i, i + batchSize);
+        try {
+          const prompt = buildPrompt(batch, sessionId);
+          const response = await opts.client.complete(prompt);
+          const parsed = parseLLMResponse(response, sessionId, batch);
+          allEntities.push(...parsed.entities);
+          allRelations.push(...parsed.relations);
+          allEvents.push(...parsed.events);
+        } catch {
+          const fallbackSession: LoadedSession = {
+            id: sessionId,
+            title: "",
+            createdAt: new Date().toISOString(),
+            filePath: "",
+            modifiedAt: new Date(),
+            turns: batch,
+          };
+          const fallback = await extractFacts(fallbackSession);
+          allEntities.push(...fallback.entities);
+          allRelations.push(...fallback.relations);
+          allEvents.push(...fallback.events);
+        }
+      }
+
+      return { entities: allEntities, relations: allRelations, events: allEvents };
+    },
+  };
+}

package/src/trainer/marker.ts

@@ -0,0 +1,30 @@
+import fs from "node:fs/promises";
+import path from "node:path";
+
+/**
+ * Read the last-trained marker timestamp.
+ * Returns null if marker file does not exist or is unreadable.
+ */
+export async function readMarker(memoryDir: string): Promise<Date | null> {
+  const markerPath = path.join(memoryDir, ".train_marker");
+  try {
+    const raw = (await fs.readFile(markerPath, "utf8")).trim();
+    if (!raw) return null;
+    const d = new Date(raw);
+    return isNaN(d.getTime()) ? null : d;
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * Write the last-trained marker timestamp (ISO string).
+ */
+export async function writeMarker(
+  memoryDir: string,
+  ts: Date,
+): Promise<void> {
+  const markerPath = path.join(memoryDir, ".train_marker");
+  await fs.mkdir(memoryDir, { recursive: true });
+  await fs.writeFile(markerPath, ts.toISOString() + "\n", "utf8");
+}

package/src/trainer/scheduler.ts

@@ -0,0 +1,104 @@
+import { trainBundle, type TrainBundleConfig, type TrainBundleResult } from "./index.js";
+
+export interface SchedulerConfig {
+  /** Interval string: "1h", "6h", "12h", "24h", or null to disable. */
+  interval: string | null;
+  /** Passed through to trainBundle. */
+  trainConfig?: Omit<TrainBundleConfig, "full" | "dryRun">;
+}
+
+export interface TrainScheduler {
+  /** Stop the scheduler. Safe to call multiple times. */
+  stop(): void;
+  /** Whether the scheduler is currently running. */
+  running(): boolean;
+}
+
+export interface SchedulerLog {
+  timestamp: string;
+  sessionsProcessed: number;
+  entityCount: number;
+  relationCount: number;
+  eventCount: number;
+  durationMs: number;
+  error?: string;
+}
+
+const INTERVAL_MAP: Record<string, number> = {
+  "1h": 3_600_000,
+  "6h": 21_600_000,
+  "12h": 43_200_000,
+  "24h": 86_400_000,
+};
+
+/** Parse interval string to milliseconds. Returns null if disabled/invalid. */
+export function parseInterval(interval: string | null | undefined): number | null {
+  if (!interval) return null;
+  const ms = INTERVAL_MAP[interval.trim().toLowerCase()];
+  return ms ?? null;
+}
+
+export type SchedulerLogger = (log: SchedulerLog) => void;
+
+/**
+ * Create an interval-based trainer that runs `trainBundle()` periodically.
+ * Runs one tick immediately on start, then repeats at the configured interval.
+ */
+export function createTrainScheduler(
+  config: SchedulerConfig,
+  logger?: SchedulerLogger,
+): TrainScheduler {
+  const intervalMs = parseInterval(config.interval);
+  if (intervalMs == null) {
+    return { stop() {}, running() { return false; } };
+  }
+
+  let timer: ReturnType<typeof setInterval> | null = null;
+  let stopped = false;
+  let tickInProgress = false;
+
+  async function tick(): Promise<void> {
+    if (stopped || tickInProgress) return;
+    tickInProgress = true;
+    const start = Date.now();
+    let result: TrainBundleResult | null = null;
+    let error: string | undefined;
+    try {
+      result = await trainBundle({
+        ...config.trainConfig,
+        full: false,
+        dryRun: false,
+      });
+    } catch (err) {
+      error = err instanceof Error ? err.message : String(err);
+    }
+    const durationMs = Date.now() - start;
+    logger?.({
+      timestamp: new Date().toISOString(),
+      sessionsProcessed: result?.sessionsProcessed ?? 0,
+      entityCount: result?.entityCount ?? 0,
+      relationCount: result?.relationCount ?? 0,
+      eventCount: result?.eventCount ?? 0,
+      durationMs,
+      error,
+    });
+    tickInProgress = false;
+  }
+
+  // Fire first tick asynchronously, then schedule repeating
+  void tick();
+  timer = setInterval(() => void tick(), intervalMs);
+
+  return {
+    stop() {
+      stopped = true;
+      if (timer != null) {
+        clearInterval(timer);
+        timer = null;
+      }
+    },
+    running() {
+      return !stopped && timer != null;
+    },
+  };
+}

package/src/trainer/sessionLoader.ts

@@ -0,0 +1,139 @@
+import { createHash } from "node:crypto";
+import fs from "node:fs/promises";
+import path from "node:path";
+
+export interface SessionTurn {
+  role: string;
+  content: string;
+  turnIndex: number;
+}
+
+export interface LoadedSession {
+  id: string;
+  title: string;
+  createdAt: string;
+  filePath: string;
+  modifiedAt: Date;
+  turns: SessionTurn[];
+}
+
+interface PiSessionMessage {
+  role?: string;
+  content?: unknown;
+}
+
+interface PiSessionFile {
+  id?: string;
+  title?: string;
+  created_at?: string;
+  messages?: PiSessionMessage[];
+}
+
+function messageText(content: unknown): string {
+  if (typeof content === "string") return content;
+  if (!Array.isArray(content)) return "";
+  const parts: string[] = [];
+  for (const block of content) {
+    if (typeof block === "string") {
+      parts.push(block);
+      continue;
+    }
+    if (block && typeof block === "object") {
+      const b = block as Record<string, unknown>;
+      if (typeof b.text === "string") parts.push(b.text);
+      else if (typeof b.content === "string") parts.push(b.content);
+    }
+  }
+  return parts.join("\n");
+}
+
+export interface SessionLoaderOptions {
+  sessionsDir: string;
+  modifiedAfter?: Date | null;
+}
+
+/**
+ * Scan session JSON files, parse Pi session format, optionally filter by
+ * modified-after timestamp for incremental training.
+ */
+export async function loadSessions(
+  opts: SessionLoaderOptions,
+): Promise<LoadedSession[]> {
+  const { sessionsDir, modifiedAfter } = opts;
+  if (!sessionsDir.trim()) return [];
+
+  let entries: string[];
+  try {
+    entries = await fs.readdir(sessionsDir);
+  } catch {
+    return [];
+  }
+
+  const sessions: LoadedSession[] = [];
+
+  for (const name of entries) {
+    if (!name.endsWith(".json")) continue;
+    const filePath = path.join(sessionsDir, name);
+
+    let st: Awaited<ReturnType<typeof fs.stat>>;
+    try {
+      st = await fs.stat(filePath);
+    } catch {
+      continue;
+    }
+    if (!st.isFile()) continue;
+
+    if (modifiedAfter && st.mtime <= modifiedAfter) {
+      continue;
+    }
+
+    let session: PiSessionFile;
+    try {
+      const raw = await fs.readFile(filePath, "utf8");
+      session = JSON.parse(raw) as PiSessionFile;
+    } catch {
+      continue;
+    }
+
+    if (!session.messages || session.messages.length === 0) continue;
+
+    const turns: SessionTurn[] = [];
+    for (let i = 0; i < session.messages.length; i++) {
+      const msg = session.messages[i]!;
+      const text = messageText(msg.content);
+      if (!text.trim()) continue;
+      turns.push({
+        role: msg.role ?? "unknown",
+        content: text,
+        turnIndex: i,
+      });
+    }
+
+    if (turns.length === 0) continue;
+
+    sessions.push({
+      id: session.id ?? path.basename(name, ".json"),
+      title: session.title ?? "",
+      createdAt: session.created_at ?? "",
+      filePath,
+      modifiedAt: st.mtime,
+      turns,
+    });
+  }
+
+  sessions.sort((a, b) => a.modifiedAt.getTime() - b.modifiedAt.getTime());
+  return deduplicateSessions(sessions);
+}
+
+function deduplicateSessions(sessions: LoadedSession[]): LoadedSession[] {
+  const seen = new Set<string>();
+  return sessions.filter((s) => {
+    const fingerprint = createHash("sha256")
+      .update(s.turns.map((t) => t.content).join("\n"))
+      .digest("hex")
+      .slice(0, 16);
+    if (seen.has(fingerprint)) return false;
+    seen.add(fingerprint);
+    return true;
+  });
+}