@chendpoc/pi-memory 0.1.0 → 0.1.11

package/src/local/graphQuery.ts

@@ -0,0 +1,228 @@
+import fs from "node:fs";
+import path from "node:path";
+
+import type {
+  ErrorClass,
+  MemoryBlock,
+  MemoryCandidateGroup,
+  QueryCandidate,
+  QueryIntent,
+  ResponseEnvelope,
+} from "../types.js";
+
+interface BundleEntity {
+  entity_id: string;
+  label: string;
+  type: string;
+  aliases: string[];
+  mention_count: number;
+  distinct_session_count: number;
+}
+
+interface BundleEdge {
+  head_entity_id: string;
+  relation: string;
+  tail_entity_id: string;
+  supporting_event_ids: string[];
+  evidence: string;
+}
+
+interface BundleEvent {
+  event_id: string;
+  description: string;
+  session_id: string;
+  timestamp: string;
+}
+
+interface GraphData {
+  entities: BundleEntity[];
+  edges: BundleEdge[];
+  events: BundleEvent[];
+}
+
+export class LocalGraphQuerier {
+  private entities: BundleEntity[] = [];
+  private edges: BundleEdge[] = [];
+  private events: BundleEvent[] = [];
+  private entityById = new Map<string, BundleEntity>();
+  private entityByLabel = new Map<string, BundleEntity>();
+  private loaded = false;
+
+  constructor(private readonly bundleRoot: string) {}
+
+  load(): boolean {
+    const graphPath = path.join(this.bundleRoot, "current", "graph.json");
+    try {
+      const raw = fs.readFileSync(graphPath, "utf8");
+      const data = JSON.parse(raw) as GraphData;
+      this.entities = data.entities ?? [];
+      this.edges = data.edges ?? [];
+      this.events = data.events ?? [];
+
+      this.entityById.clear();
+      this.entityByLabel.clear();
+      for (const e of this.entities) {
+        this.entityById.set(e.entity_id, e);
+        this.entityByLabel.set(e.label.toLowerCase(), e);
+        for (const alias of e.aliases) {
+          this.entityByLabel.set(alias.toLowerCase(), e);
+        }
+      }
+      this.loaded = true;
+      return true;
+    } catch {
+      return false;
+    }
+  }
+
+  isLoaded(): boolean {
+    return this.loaded;
+  }
+
+  private findEntity(mention: string): BundleEntity | null {
+    const key = mention.toLowerCase().trim();
+    return this.entityByLabel.get(key) ?? null;
+  }
+
+  private findEntities(mentions: string[]): BundleEntity[] {
+    const found: BundleEntity[] = [];
+    for (const m of mentions) {
+      const e = this.findEntity(m);
+      if (e) found.push(e);
+    }
+    return found;
+  }
+
+  query(intent: QueryIntent): {
+    env: ResponseEnvelope | null;
+    errorClass: ErrorClass;
+  } {
+    if (!this.loaded) {
+      return { env: null, errorClass: "unavailable" };
+    }
+
+    const anchors = this.findEntities(intent.anchor_mentions);
+    if (anchors.length === 0) {
+      return {
+        env: this.makeEnvelope(intent, [], [], "no matching entities found"),
+        errorClass: "ok",
+      };
+    }
+
+    const limit = intent.result_limit && intent.result_limit > 0 ? intent.result_limit : 10;
+    const candidates: QueryCandidate[] = [];
+    const groups: MemoryCandidateGroup[] = [];
+
+    for (const anchor of anchors) {
+      const relatedEdges = this.edges.filter((e) => {
+        const matchesHead = e.head_entity_id === anchor.entity_id;
+        const matchesTail = e.tail_entity_id === anchor.entity_id;
+        if (!matchesHead && !matchesTail) return false;
+
+        if (intent.relation_constraints?.length) {
+          return intent.relation_constraints.some((rc) => {
+            const clean = rc.replace(/\^-1$/, "").replace(/^\^/, "");
+            return e.relation === clean;
+          });
+        }
+        return true;
+      });
+
+      for (const edge of relatedEdges) {
+        if (candidates.length >= limit) break;
+        const isHead = edge.head_entity_id === anchor.entity_id;
+        const otherId = isHead ? edge.tail_entity_id : edge.head_entity_id;
+        const other = this.entityById.get(otherId);
+        if (!other) continue;
+
+        if (intent.candidate_type) {
+          if (other.type.toLowerCase() !== intent.candidate_type.toLowerCase()) continue;
+        }
+
+        const score = other.mention_count + (other.distinct_session_count * 2);
+
+        candidates.push({
+          value: other.label,
+          score,
+          evidence: edge.evidence.slice(0, 200),
+          supporting_event_ids: edge.supporting_event_ids,
+          entity_id: other.entity_id,
+          scope: `via_${anchor.label}`,
+          support_count: other.mention_count,
+          distinct_session_count: other.distinct_session_count,
+          observed_path: [{
+            from_entity_id: anchor.entity_id,
+            from_label: anchor.label,
+            relation: edge.relation,
+            direction: isHead ? "forward" : "inverse",
+            to_entity_id: other.entity_id,
+            to_label: other.label,
+            supporting_event_ids: edge.supporting_event_ids,
+          }],
+          path_collision_count: 0,
+        });
+
+        groups.push({
+          value: other.label,
+          score,
+          evidence: edge.evidence.slice(0, 200),
+          support_count: other.mention_count,
+          supporting_event_ids: edge.supporting_event_ids,
+          entity_ids: [other.entity_id],
+          scopes: [`via_${anchor.label}`],
+          via_relations: [edge.relation],
+          via_anchor_entity_ids: [anchor.entity_id],
+          observed_path: [{
+            from_entity_id: anchor.entity_id,
+            from_label: anchor.label,
+            relation: edge.relation,
+            direction: isHead ? "forward" : "inverse",
+            to_entity_id: other.entity_id,
+            to_label: other.label,
+            supporting_event_ids: edge.supporting_event_ids,
+          }],
+          path_collision_count: 0,
+        });
+      }
+    }
+
+    candidates.sort((a, b) => b.score - a.score);
+    groups.sort((a, b) => b.score - a.score);
+
+    return {
+      env: this.makeEnvelope(intent, candidates.slice(0, limit), groups.slice(0, limit)),
+      errorClass: "ok",
+    };
+  }
+
+  private makeEnvelope(
+    intent: QueryIntent,
+    candidates: QueryCandidate[],
+    groups: MemoryCandidateGroup[],
+    noDataReason?: string,
+  ): ResponseEnvelope {
+    const manifestPath = path.join(this.bundleRoot, "current", "manifest.json");
+    let bundleVersion: string | undefined;
+    try {
+      const m = JSON.parse(fs.readFileSync(manifestPath, "utf8")) as { bundle_version?: string };
+      bundleVersion = m.bundle_version;
+    } catch { /* ok */ }
+
+    const memoryBlock: MemoryBlock = {
+      groups,
+      notes: [],
+      ...(noDataReason ? { no_data_reason: noDataReason } : {}),
+    };
+
+    return {
+      protocol_version: 1,
+      bundle_version: bundleVersion,
+      request_id: `local-${Date.now()}`,
+      candidates,
+      memory_block: memoryBlock,
+      warnings: [],
+      reason: candidates.length > 0 ? "ok" : "no_data",
+      latency_ms: 0,
+    };
+  }
+}

package/src/paths.ts

@@ -24,5 +24,5 @@ export function defaultSocketPath(): string {
 }
 
 export function defaultSessionsDir(): string {
-  return path.join(defaultPiHome(), "sessions");
+  return path.join(defaultPiHome(), "agent", "sessions");
 }

package/src/pi-extension.ts

@@ -5,13 +5,29 @@ import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-a
 import { truncateHead } from "@earendil-works/pi-coding-agent";
 import { Type } from "typebox";
 
+import {
+  createOllamaLLMClient,
+  createOllamaMemoryHelper,
+  ollamaHealthCheck,
+  type OllamaConfig,
+} from "./adapters/ollamaClient.js";
+import {
+  createOpenAICompatLLMClient,
+  createOpenAICompatMemoryHelper,
+  openaiCompatHealthCheck,
+  type OpenAICompatConfig,
+} from "./adapters/openaiCompatClient.js";
 import {
   createPiLLMClient,
   DEFAULT_HELPER_MODEL,
   DEFAULT_HELPER_PROVIDER,
   resolveMemoryHelperLLM,
 } from "./adapters/piComplete.js";
-import { defaultMemoryConfig, type MemoryConfig } from "./config.js";
+import type { MemoryConfig } from "./config.js";
+import {
+  loadMemorySettings,
+  resolveHelperModelSpec,
+} from "./settings.js";
 import { createFallbackQuery } from "./fallback/index.js";
 import type { RerankOptions } from "./fallback/llmRerank.js";
 import type { MemoryHelperLLM } from "./preflight/detectIntents.js";
@@ -57,6 +73,9 @@ const RECALL_MAX_BYTES = 32_000;
 
 let sharedService: MemoryService | null = null;
 let sessionCfg: MemoryConfig | null = null;
+let settingsHelperModel: string | undefined;
+let settingsOllama: OllamaConfig | null = null;
+let settingsVllm: OpenAICompatConfig | null = null;
 let sharedHelper: MemoryHelperLLM | null = null;
 let sharedLLMClient: LLMClient | null = null;
 let preflightCache: { userText: string; privateContext: string } | null = null;
@@ -65,9 +84,37 @@ export function getSharedMemoryService(): MemoryService | null {
   return sharedService;
 }
 
+function getHelperModelSpec(pi: ExtensionAPI): string | undefined {
+  return resolveHelperModelSpec(pi.getFlag("memory-helper-model"), settingsHelperModel);
+}
+
+function isOllamaModel(spec: string | undefined): boolean {
+  return !!spec?.startsWith("ollama/");
+}
+
 async function refreshMemoryHelper(ctx: ExtensionContext, pi: ExtensionAPI): Promise<void> {
-  sharedHelper = await resolveMemoryHelperLLM(ctx, pi.getFlag("memory-helper-model"));
-  sharedLLMClient = createPiLLMClient(ctx, pi.getFlag("memory-helper-model"));
+  const helperModel = getHelperModelSpec(pi);
+
+  if (settingsVllm) {
+    const vllmOk = await openaiCompatHealthCheck(settingsVllm.baseUrl);
+    if (vllmOk) {
+      sharedHelper = createOpenAICompatMemoryHelper(settingsVllm);
+      sharedLLMClient = createOpenAICompatLLMClient(settingsVllm);
+      return;
+    }
+  }
+
+  if ((isOllamaModel(helperModel) || settingsOllama) && settingsOllama) {
+    const ollamaOk = await ollamaHealthCheck(settingsOllama.baseUrl);
+    if (ollamaOk) {
+      sharedHelper = createOllamaMemoryHelper(settingsOllama);
+      sharedLLMClient = createOllamaLLMClient(settingsOllama);
+      return;
+    }
+  }
+
+  sharedHelper = await resolveMemoryHelperLLM(ctx, helperModel);
+  sharedLLMClient = createPiLLMClient(ctx, helperModel);
 }
 
 function getRerankOpts(): RerankOptions | null {
@@ -100,7 +147,13 @@ function formatMemoryStatus(service: MemoryService): string {
   const snap = service.getStatus();
   const lines = [
     `status: ${snap.status}`,
+    snap.mode ? `mode: ${snap.mode}` : null,
     snap.reason ? `reason: ${snap.reason}` : null,
+    sharedHelper
+      ? `helper: ${settingsVllm ? `vllm/${settingsVllm.model}` : settingsOllama ? `ollama/${settingsOllama.model}` : (settingsHelperModel ?? "pi-ai")}`
+      : "helper: none (regex only)",
+    settingsVllm ? `vllm: ${settingsVllm.baseUrl} (${settingsVllm.model})` : null,
+    settingsOllama ? `ollama: ${settingsOllama.baseUrl} (${settingsOllama.model})` : null,
     snap.health ? `health: ${JSON.stringify(snap.health)}` : null,
   ].filter(Boolean);
   return lines.join("\n");
@@ -113,7 +166,11 @@ export default function piMemoryExtension(pi: ExtensionAPI): void {
   });
 
   pi.on("session_start", async (_event, ctx) => {
-    const cfg = defaultMemoryConfig();
+    const loaded = loadMemorySettings();
+    const cfg = loaded.config;
+    settingsHelperModel = loaded.helperModel;
+    settingsOllama = loaded.ollama;
+    settingsVllm = loaded.vllm;
     sessionCfg = cfg;
     preflightCache = null;
     sharedHelper = null;
@@ -123,18 +180,15 @@ export default function piMemoryExtension(pi: ExtensionAPI): void {
     const service = new MemoryService(cfg);
     sharedService = service;
 
+    await service.start();
+    service.startSessionIndex();
+    if (cfg.trainer.auto_interval) {
+      service.startAutoTrainer();
+    }
     try {
-      await service.start();
-      service.startSessionIndex();
-      if (cfg.trainer.auto_interval) {
-        service.startAutoTrainer();
-      }
       await refreshMemoryHelper(ctx, pi);
-    } catch (err) {
-      const message = err instanceof Error ? err.message : String(err);
-      if (ctx.hasUI) {
-        ctx.ui.notify(`pi-memory: sidecar start failed (${message}) — fallback mode active`, "warning");
-      }
+    } catch {
+      /* helper unavailable — regex-only preflight */
     }
   });
 
@@ -144,6 +198,9 @@ export default function piMemoryExtension(pi: ExtensionAPI): void {
       sharedService = null;
     }
     sessionCfg = null;
+    settingsHelperModel = undefined;
+    settingsOllama = null;
+    settingsVllm = null;
     sharedHelper = null;
     sharedLLMClient = null;
     preflightCache = null;
@@ -157,9 +214,14 @@ export default function piMemoryExtension(pi: ExtensionAPI): void {
     preflightCache = null;
   });
 
+  const initialSettings = loadMemorySettings();
+  settingsHelperModel = initialSettings.helperModel;
+  settingsOllama = initialSettings.ollama;
+  settingsVllm = initialSettings.vllm;
+
   const fallback = createFallbackQuery({
-    sessionsDir: defaultMemoryConfig().sessionsDir,
-    memoryMdPaths: defaultMemoryConfig().memoryMdPaths,
+    sessionsDir: initialSettings.config.sessionsDir,
+    memoryMdPaths: initialSettings.config.memoryMdPaths,
   });
 
   pi.registerCommand("memory", {
@@ -208,7 +270,7 @@ export default function piMemoryExtension(pi: ExtensionAPI): void {
     },
   });
 
-  const memoryMdPath = defaultMemoryConfig().memoryMdPaths[0];
+  const memoryMdPath = initialSettings.config.memoryMdPaths[0];
   if (memoryMdPath) {
     const appendTool = createMemoryAppendTool(memoryMdPath);
     pi.registerTool({

package/src/service.ts

@@ -1,5 +1,7 @@
+import fs from "node:fs";
 import path from "node:path";
 import type { MemoryConfig } from "./config.js";
+import { LocalGraphQuerier } from "./local/graphQuery.js";
 import { currentBundleReadable } from "./sidecar/bundle.js";
 import { SidecarClient } from "./sidecar/client.js";
 import { SidecarProcess } from "./sidecar/process.js";
@@ -16,6 +18,7 @@ import type {
 export interface MemoryServiceStatus {
   status: ServiceStatus;
   reason?: string;
+  mode?: "sidecar" | "local_graph";
   health?: HealthPayload | null;
 }
 
@@ -26,14 +29,19 @@ export interface QueryBatchResult {
 }
 
 /**
- * Mode B local memory: spawn tlm sidecar, query via Unix socket.
- * No Cloud puller — bundle must exist under bundleRoot/current.
+ * Local memory service with two query backends:
+ * 1. tlm sidecar (Unix socket) — when tlm binary is available
+ * 2. LocalGraphQuerier — direct graph.json query when tlm is missing
+ *
+ * Both require a bundle at bundleRoot/current.
  */
 export class MemoryService {
   private serviceStatus: ServiceStatus = "disabled";
   private reason = "";
+  private mode: "sidecar" | "local_graph" | null = null;
   private process: SidecarProcess | null = null;
   private client: SidecarClient | null = null;
+  private localQuerier: LocalGraphQuerier | null = null;
   private abort: AbortController | null = null;
   private scheduler: TrainScheduler | null = null;
   private sessionIndex: SessionIndex | null = null;
@@ -52,6 +60,7 @@ export class MemoryService {
     return {
       status: this.serviceStatus,
       reason: this.reason || undefined,
+      mode: this.mode ?? undefined,
     };
   }
 
@@ -79,36 +88,50 @@ export class MemoryService {
 
     this.serviceStatus = "initializing";
     this.abort = new AbortController();
-    this.process = new SidecarProcess(this.cfg);
 
+    if (await this.trySidecar()) return;
+
+    if (this.tryLocalGraph()) return;
+
+    this.serviceStatus = "unavailable";
+    this.reason = "no_query_backend";
+  }
+
+  private async trySidecar(): Promise<boolean> {
+    this.process = new SidecarProcess(this.cfg);
     try {
       await this.process.resolveBinary();
     } catch {
-      this.serviceStatus = "unavailable";
-      this.reason = "tlm_binary_missing";
-      return;
+      this.process = null;
+      return false;
     }
 
     try {
       await this.process.spawn();
-      await this.process.waitReady(this.abort.signal);
+      await this.process.waitReady(this.abort!.signal);
       this.client = this.process.getClient();
       this.serviceStatus = "ready";
+      this.mode = "sidecar";
       this.reason = "";
-    } catch (err) {
-      this.serviceStatus = "unavailable";
-      this.reason =
-        err instanceof Error ? err.message : "sidecar_startup_failed";
+      return true;
+    } catch {
       await this.process.stop();
       this.process = null;
       this.client = null;
+      return false;
     }
   }
 
-  /**
-   * Start interval-based auto-training. If already running, stops and restarts.
-   * Uses config.trainer.auto_interval.
-   */
+  private tryLocalGraph(): boolean {
+    const querier = new LocalGraphQuerier(this.cfg.bundleRoot);
+    if (!querier.load()) return false;
+    this.localQuerier = querier;
+    this.serviceStatus = "ready";
+    this.mode = "local_graph";
+    this.reason = "";
+    return true;
+  }
+
   startAutoTrainer(logger?: (log: SchedulerLog) => void): void {
     this.scheduler?.stop();
     this.scheduler = createTrainScheduler(
@@ -123,12 +146,14 @@ export class MemoryService {
     );
   }
 
-  /**
-   * Trigger incremental session index build in the background (non-blocking).
-   * Opens (or creates) the SQLite FTS5 DB at ~/.pi/memory/sessions.db.
-   */
   startSessionIndex(): void {
-    const dbPath = path.join(this.cfg.bundleRoot, "sessions.db");
+    const dbDir = this.cfg.bundleRoot;
+    try {
+      fs.mkdirSync(dbDir, { recursive: true });
+    } catch {
+      return;
+    }
+    const dbPath = path.join(dbDir, "sessions.db");
     const idx = openSessionIndex(dbPath);
     if (!idx) return;
     this.sessionIndex = idx;
@@ -148,6 +173,8 @@ export class MemoryService {
     await this.process?.stop();
     this.process = null;
     this.client = null;
+    this.localQuerier = null;
+    this.mode = null;
     if (this.cfg.provider === "disabled") {
       this.serviceStatus = "disabled";
     } else {
@@ -181,22 +208,42 @@ export class MemoryService {
     errorClass: ErrorClass;
     transportError?: Error;
   }> {
-    if (this.serviceStatus !== "ready" || !this.client) {
+    if (this.serviceStatus !== "ready") {
       return { env: null, errorClass: "unavailable" };
     }
-    const timeout = AbortSignal.timeout(this.cfg.queryTimeoutMs);
-    const combined = signal
-      ? AbortSignal.any([signal, timeout])
-      : timeout;
-    return this.client.query(intent, combined);
+
+    if (this.client && this.mode === "sidecar") {
+      const timeout = AbortSignal.timeout(this.cfg.queryTimeoutMs);
+      const combined = signal
+        ? AbortSignal.any([signal, timeout])
+        : timeout;
+      return this.client.query(intent, combined);
+    }
+
+    if (this.localQuerier && this.mode === "local_graph") {
+      return this.localQuerier.query(intent);
+    }
+
+    return { env: null, errorClass: "unavailable" };
   }
 
   async health(): Promise<HealthPayload | null> {
-    if (!this.client) return null;
-    try {
-      return await this.client.health();
-    } catch {
-      return null;
+    if (this.client) {
+      try {
+        return await this.client.health();
+      } catch {
+        return null;
+      }
+    }
+    if (this.localQuerier) {
+      return {
+        ready: true,
+        compatibility: "local_graph",
+        protocol_version: 1,
+        uptime_secs: 0,
+        status_message: `local graph query (${this.localQuerier.isLoaded() ? "loaded" : "not loaded"})`,
+      };
     }
+    return null;
   }
 }