@agentmemory/agentmemory 0.7.0

package/src/functions/graph.ts

@@ -0,0 +1,275 @@
+import type { ISdk } from "iii-sdk";
+import { getContext } from "iii-sdk";
+import type {
+  GraphNode,
+  GraphEdge,
+  GraphQueryResult,
+  CompressedObservation,
+  MemoryProvider,
+} from "../types.js";
+import { KV, generateId } from "../state/schema.js";
+import type { StateKV } from "../state/kv.js";
+import {
+  GRAPH_EXTRACTION_SYSTEM,
+  buildGraphExtractionPrompt,
+} from "../prompts/graph-extraction.js";
+import { recordAudit } from "./audit.js";
+
+function parseGraphXml(
+  xml: string,
+  observationIds: string[],
+): {
+  nodes: GraphNode[];
+  edges: GraphEdge[];
+} {
+  const nodes: GraphNode[] = [];
+  const edges: GraphEdge[] = [];
+  const now = new Date().toISOString();
+
+  const entityRegex =
+    /<entity\s+type="([^"]+)"\s+name="([^"]+)"[^>]*>([\s\S]*?)<\/entity>/g;
+  let match;
+  while ((match = entityRegex.exec(xml)) !== null) {
+    const type = match[1] as GraphNode["type"];
+    const name = match[2];
+    const propsBlock = match[3];
+    const properties: Record<string, string> = {};
+
+    const propRegex = /<property\s+key="([^"]+)">([^<]*)<\/property>/g;
+    let propMatch;
+    while ((propMatch = propRegex.exec(propsBlock)) !== null) {
+      properties[propMatch[1]] = propMatch[2];
+    }
+
+    nodes.push({
+      id: generateId("gn"),
+      type,
+      name,
+      properties,
+      sourceObservationIds: observationIds,
+      createdAt: now,
+    });
+  }
+
+  const relRegex =
+    /<relationship\s+type="([^"]+)"\s+source="([^"]+)"\s+target="([^"]+)"\s+weight="([^"]+)"\s*\/>/g;
+  while ((match = relRegex.exec(xml)) !== null) {
+    const type = match[1] as GraphEdge["type"];
+    const sourceName = match[2];
+    const targetName = match[3];
+    const parsedWeight = parseFloat(match[4]);
+    const weight = Number.isNaN(parsedWeight) ? 0.5 : parsedWeight;
+
+    const sourceNode = nodes.find((n) => n.name === sourceName);
+    const targetNode = nodes.find((n) => n.name === targetName);
+
+    if (sourceNode && targetNode) {
+      edges.push({
+        id: generateId("ge"),
+        type,
+        sourceNodeId: sourceNode.id,
+        targetNodeId: targetNode.id,
+        weight: Math.max(0, Math.min(1, weight)),
+        sourceObservationIds: observationIds,
+        createdAt: now,
+      });
+    }
+  }
+
+  return { nodes, edges };
+}
+
+export function registerGraphFunction(
+  sdk: ISdk,
+  kv: StateKV,
+  provider: MemoryProvider,
+): void {
+  sdk.registerFunction(
+    { id: "mem::graph-extract" },
+    async (data: { observations: CompressedObservation[] }) => {
+      const ctx = getContext();
+      if (!data.observations || data.observations.length === 0) {
+        return { success: false, error: "No observations provided" };
+      }
+
+      const prompt = buildGraphExtractionPrompt(
+        data.observations.map((o) => ({
+          title: o.title,
+          narrative: o.narrative,
+          concepts: o.concepts,
+          files: o.files,
+          type: o.type,
+        })),
+      );
+
+      try {
+        const response = await provider.compress(
+          GRAPH_EXTRACTION_SYSTEM,
+          prompt,
+        );
+
+        const obsIds = data.observations.map((o) => o.id);
+        const { nodes, edges } = parseGraphXml(response, obsIds);
+
+        const existingNodes = await kv.list<GraphNode>(KV.graphNodes);
+        const existingEdges = await kv.list<GraphEdge>(KV.graphEdges);
+
+        for (const node of nodes) {
+          const existing = existingNodes.find(
+            (n) => n.name === node.name && n.type === node.type,
+          );
+          if (existing) {
+            const merged = {
+              ...existing,
+              sourceObservationIds: [
+                ...new Set([...existing.sourceObservationIds, ...obsIds]),
+              ],
+              properties: { ...existing.properties, ...node.properties },
+            };
+            await kv.set(KV.graphNodes, existing.id, merged);
+            const idx = existingNodes.findIndex((n) => n.id === existing.id);
+            if (idx !== -1) existingNodes[idx] = merged;
+          } else {
+            await kv.set(KV.graphNodes, node.id, node);
+            existingNodes.push(node);
+          }
+        }
+
+        for (const edge of edges) {
+          const edgeKey = `${edge.sourceNodeId}|${edge.targetNodeId}|${edge.type}`;
+          const existingEdge = existingEdges.find(
+            (e) => `${e.sourceNodeId}|${e.targetNodeId}|${e.type}` === edgeKey,
+          );
+          if (existingEdge) {
+            existingEdge.sourceObservationIds = [
+              ...new Set([...existingEdge.sourceObservationIds, ...obsIds]),
+            ];
+            await kv.set(KV.graphEdges, existingEdge.id, existingEdge);
+          } else {
+            await kv.set(KV.graphEdges, edge.id, edge);
+            existingEdges.push(edge);
+          }
+        }
+
+        await recordAudit(kv, "observe", "mem::graph-extract", obsIds, {
+          nodesExtracted: nodes.length,
+          edgesExtracted: edges.length,
+        });
+
+        ctx.logger.info("Graph extraction complete", {
+          nodes: nodes.length,
+          edges: edges.length,
+        });
+        return {
+          success: true,
+          nodesAdded: nodes.length,
+          edgesAdded: edges.length,
+        };
+      } catch (err) {
+        const msg = err instanceof Error ? err.message : String(err);
+        ctx.logger.error("Graph extraction failed", { error: msg });
+        return { success: false, error: msg };
+      }
+    },
+  );
+
+  sdk.registerFunction(
+    { id: "mem::graph-query" },
+    async (data: {
+      startNodeId?: string;
+      nodeType?: string;
+      maxDepth?: number;
+      query?: string;
+    }): Promise<GraphQueryResult> => {
+      const allNodes = (await kv.list<GraphNode>(KV.graphNodes)).filter((n) => !n.stale);
+      const allEdges = (await kv.list<GraphEdge>(KV.graphEdges)).filter((e) => !e.stale);
+      const maxDepth = Math.min(data.maxDepth || 3, 5);
+
+      if (data.query) {
+        const lower = data.query.toLowerCase();
+        const matchingNodes = allNodes.filter(
+          (n) =>
+            n.name.toLowerCase().includes(lower) ||
+            Object.values(n.properties).some(
+              (v) => typeof v === "string" && v.toLowerCase().includes(lower),
+            ),
+        );
+        const nodeIds = new Set(matchingNodes.map((n) => n.id));
+        const relatedEdges = allEdges.filter(
+          (e) => nodeIds.has(e.sourceNodeId) || nodeIds.has(e.targetNodeId),
+        );
+        return { nodes: matchingNodes, edges: relatedEdges, depth: 0 };
+      }
+
+      if (data.startNodeId) {
+        const visited = new Set<string>();
+        const visitedEdges = new Set<string>();
+        const resultNodes: GraphNode[] = [];
+        const resultEdges: GraphEdge[] = [];
+        const queue: Array<{ nodeId: string; depth: number }> = [
+          { nodeId: data.startNodeId, depth: 0 },
+        ];
+
+        while (queue.length > 0) {
+          const { nodeId, depth } = queue.shift()!;
+          if (visited.has(nodeId) || depth > maxDepth) continue;
+          visited.add(nodeId);
+
+          const node = allNodes.find((n) => n.id === nodeId);
+          if (node) {
+            if (!data.nodeType || node.type === data.nodeType) {
+              resultNodes.push(node);
+            }
+          }
+
+          const neighborEdges = allEdges.filter(
+            (e) => e.sourceNodeId === nodeId || e.targetNodeId === nodeId,
+          );
+          for (const edge of neighborEdges) {
+            if (!visitedEdges.has(edge.id)) {
+              visitedEdges.add(edge.id);
+              resultEdges.push(edge);
+            }
+            const nextId =
+              edge.sourceNodeId === nodeId
+                ? edge.targetNodeId
+                : edge.sourceNodeId;
+            if (!visited.has(nextId)) {
+              queue.push({ nodeId: nextId, depth: depth + 1 });
+            }
+          }
+        }
+
+        return { nodes: resultNodes, edges: resultEdges, depth: maxDepth };
+      }
+
+      let filtered = allNodes;
+      if (data.nodeType) {
+        filtered = allNodes.filter((n) => n.type === data.nodeType);
+      }
+      return { nodes: filtered, edges: allEdges, depth: 0 };
+    },
+  );
+
+  sdk.registerFunction({ id: "mem::graph-stats" }, async () => {
+    const nodes = await kv.list<GraphNode>(KV.graphNodes);
+    const edges = await kv.list<GraphEdge>(KV.graphEdges);
+
+    const nodesByType: Record<string, number> = {};
+    for (const n of nodes) {
+      nodesByType[n.type] = (nodesByType[n.type] || 0) + 1;
+    }
+
+    const edgesByType: Record<string, number> = {};
+    for (const e of edges) {
+      edgesByType[e.type] = (edgesByType[e.type] || 0) + 1;
+    }
+
+    return {
+      totalNodes: nodes.length,
+      totalEdges: edges.length,
+      nodesByType,
+      edgesByType,
+    };
+  });
+}

package/src/functions/leases.ts

@@ -0,0 +1,216 @@
+import type { ISdk } from "iii-sdk";
+import type { StateKV } from "../state/kv.js";
+import { KV, generateId } from "../state/schema.js";
+import { withKeyedLock } from "../state/keyed-mutex.js";
+import type { Action, Lease } from "../types.js";
+
+const DEFAULT_LEASE_TTL_MS = 10 * 60 * 1000;
+const MAX_LEASE_TTL_MS = 60 * 60 * 1000;
+
+export function registerLeasesFunction(sdk: ISdk, kv: StateKV): void {
+  sdk.registerFunction(
+    { id: "mem::lease-acquire" },
+    async (data: { actionId: string; agentId: string; ttlMs?: number }) => {
+      if (!data.actionId || !data.agentId) {
+        return { success: false, error: "actionId and agentId are required" };
+      }
+
+      const rawTtl = typeof data.ttlMs === "number" && Number.isFinite(data.ttlMs) && data.ttlMs > 0
+        ? data.ttlMs
+        : DEFAULT_LEASE_TTL_MS;
+      const ttl = Math.min(rawTtl, MAX_LEASE_TTL_MS);
+
+      return withKeyedLock(`mem:action:${data.actionId}`, async () => {
+        const action = await kv.get<Action>(KV.actions, data.actionId);
+        if (!action) {
+          return { success: false, error: "action not found" };
+        }
+        if (action.status === "done" || action.status === "cancelled") {
+          return { success: false, error: "action already completed" };
+        }
+        if (action.status === "blocked") {
+          return { success: false, error: "action is blocked" };
+        }
+
+        const existingLeases = await kv.list<Lease>(KV.leases);
+        const activeLease = existingLeases.find(
+          (l) =>
+            l.actionId === data.actionId &&
+            l.status === "active" &&
+            new Date(l.expiresAt).getTime() > Date.now(),
+        );
+
+        if (activeLease) {
+          if (activeLease.agentId === data.agentId) {
+            return {
+              success: true,
+              lease: activeLease,
+              renewed: false,
+              message: "Already holding this lease",
+            };
+          }
+          return {
+            success: false,
+            error: "action already leased",
+            heldBy: activeLease.agentId,
+            expiresAt: activeLease.expiresAt,
+          };
+        }
+
+        const now = new Date();
+        const lease: Lease = {
+          id: generateId("lse"),
+          actionId: data.actionId,
+          agentId: data.agentId,
+          acquiredAt: now.toISOString(),
+          expiresAt: new Date(now.getTime() + ttl).toISOString(),
+          status: "active",
+        };
+
+        await kv.set(KV.leases, lease.id, lease);
+
+        action.status = "active";
+        action.assignedTo = data.agentId;
+        action.updatedAt = now.toISOString();
+        await kv.set(KV.actions, action.id, action);
+
+        return { success: true, lease, renewed: false };
+      });
+    },
+  );
+
+  sdk.registerFunction(
+    { id: "mem::lease-release" },
+    async (data: { actionId: string; agentId: string; result?: string }) => {
+      if (!data.actionId || !data.agentId) {
+        return { success: false, error: "actionId and agentId are required" };
+      }
+
+      return withKeyedLock(`mem:action:${data.actionId}`, async () => {
+        const leases = await kv.list<Lease>(KV.leases);
+        const activeLease = leases.find(
+          (l) =>
+            l.actionId === data.actionId &&
+            l.agentId === data.agentId &&
+            l.status === "active" &&
+            new Date(l.expiresAt).getTime() > Date.now(),
+        );
+
+        if (!activeLease) {
+          return { success: false, error: "no active lease found for this agent" };
+        }
+
+        activeLease.status = "released";
+        await kv.set(KV.leases, activeLease.id, activeLease);
+
+        const action = await kv.get<Action>(KV.actions, data.actionId);
+        if (action && action.status === "active" && action.assignedTo === data.agentId) {
+          if (data.result) {
+            action.status = "done";
+            action.result = data.result;
+          } else {
+            action.status = "pending";
+          }
+          action.assignedTo = undefined;
+          action.updatedAt = new Date().toISOString();
+          await kv.set(KV.actions, action.id, action);
+        }
+
+        return { success: true, released: true };
+      });
+    },
+  );
+
+  sdk.registerFunction(
+    { id: "mem::lease-renew" },
+    async (data: { actionId: string; agentId: string; ttlMs?: number }) => {
+      if (!data.actionId || !data.agentId) {
+        return { success: false, error: "actionId and agentId are required" };
+      }
+
+      const rawTtl = typeof data.ttlMs === "number" && Number.isFinite(data.ttlMs) && data.ttlMs > 0
+        ? data.ttlMs
+        : DEFAULT_LEASE_TTL_MS;
+      const ttl = Math.min(rawTtl, MAX_LEASE_TTL_MS);
+
+      return withKeyedLock(`mem:action:${data.actionId}`, async () => {
+        const leases = await kv.list<Lease>(KV.leases);
+        const activeLease = leases.find(
+          (l) =>
+            l.actionId === data.actionId &&
+            l.agentId === data.agentId &&
+            l.status === "active" &&
+            new Date(l.expiresAt).getTime() > Date.now(),
+        );
+
+        if (!activeLease) {
+          return { success: false, error: "no active (non-expired) lease to renew" };
+        }
+
+        const now = new Date();
+        const base = Math.max(now.getTime(), new Date(activeLease.expiresAt).getTime());
+        activeLease.expiresAt = new Date(base + ttl).toISOString();
+        activeLease.renewedAt = now.toISOString();
+        await kv.set(KV.leases, activeLease.id, activeLease);
+
+        return { success: true, lease: activeLease };
+      });
+    },
+  );
+
+  sdk.registerFunction(
+    { id: "mem::lease-cleanup" },
+    async () => {
+      const leases = await kv.list<Lease>(KV.leases);
+      const now = Date.now();
+      let expired = 0;
+
+      for (const lease of leases) {
+        if (
+          lease.status === "active" &&
+          new Date(lease.expiresAt).getTime() <= now
+        ) {
+          const didExpire = await withKeyedLock(
+            `mem:action:${lease.actionId}`,
+            async () => {
+              const currentLease = await kv.get<Lease>(KV.leases, lease.id);
+              if (
+                !currentLease ||
+                currentLease.status !== "active" ||
+                new Date(currentLease.expiresAt).getTime() > Date.now()
+              ) {
+                return false;
+              }
+              currentLease.status = "expired";
+              await kv.set(KV.leases, currentLease.id, currentLease);
+
+              const action = await kv.get<Action>(KV.actions, currentLease.actionId);
+              const otherActiveLease = (await kv.list<Lease>(KV.leases)).some(
+                (l) =>
+                  l.id !== currentLease.id &&
+                  l.actionId === currentLease.actionId &&
+                  l.status === "active" &&
+                  new Date(l.expiresAt).getTime() > Date.now(),
+              );
+              if (
+                action &&
+                !otherActiveLease &&
+                action.status === "active" &&
+                action.assignedTo === currentLease.agentId
+              ) {
+                action.status = "pending";
+                action.assignedTo = undefined;
+                action.updatedAt = new Date().toISOString();
+                await kv.set(KV.actions, action.id, action);
+              }
+              return true;
+            },
+          );
+          if (didExpire) expired++;
+        }
+      }
+
+      return { success: true, expired };
+    },
+  );
+}