@dbx-tools/appkit-memory 0.0.0-stage → 0.9.131

package/src/databricks/_runtime.ts

@@ -0,0 +1,103 @@
+/**
+ * Request-scoped Databricks managed-memory routing.
+ *
+ * @module
+ */
+
+import { getExecutionContext, type WorkspaceClient } from "@databricks/appkit";
+import { createCache, getOrSet } from "@dbx-tools/cache/cache";
+
+import {
+  RecallMemoryInputSchema,
+  type RecallMemoryInput,
+  type RecallMemoryResult,
+  type RememberMemoryInput,
+  type RememberMemoryResult,
+} from "../_contracts.ts";
+import type { ResolvedDatabricksMemoryOptions } from "../config.ts";
+import { ManagedAgentMemoryClient } from "./client.ts";
+
+const ACCESS_CACHE_TTL_MS = 5 * 60_000;
+const accessCache = createCache({
+  maxEntries: 1_000,
+  ttl: ACCESS_CACHE_TTL_MS,
+  useClone: false,
+});
+
+/** Databricks backend that selects the active request credential for each call. */
+export class DatabricksMemoryRuntime {
+  constructor(
+    private readonly workspaceClient: WorkspaceClient,
+    readonly options: ResolvedDatabricksMemoryOptions,
+  ) {}
+
+  async available(actorId: string, signal?: AbortSignal): Promise<boolean> {
+    const request = this.requestClient(actorId);
+    return (
+      (await getOrSet(accessCache, request.cacheKey, () => request.client.available(signal))) ??
+      false
+    );
+  }
+
+  async remember(
+    actorId: string,
+    input: RememberMemoryInput,
+    signal?: AbortSignal,
+  ): Promise<RememberMemoryResult> {
+    const entry = await this.requestClient(actorId).client.saveMemory({
+      actorId,
+      ...input,
+      ...(signal ? { signal } : {}),
+    });
+    return {
+      backend: "databricks",
+      memory: {
+        path: entry.path,
+        content: entry.content ?? input.content,
+        ...(entry.description ? { description: entry.description } : {}),
+      },
+    };
+  }
+
+  async recall(
+    actorId: string,
+    input: RecallMemoryInput,
+    signal?: AbortSignal,
+  ): Promise<RecallMemoryResult> {
+    const parsed = RecallMemoryInputSchema.parse(input);
+    const results = await this.requestClient(actorId).client.searchMemory({
+      actorId,
+      query: parsed.query,
+      limit: parsed.limit,
+      ...(signal ? { signal } : {}),
+    });
+    return {
+      backend: "databricks",
+      memories: results.map(({ entry, score }) => ({
+        path: entry.path,
+        content: entry.content ?? "",
+        ...(entry.description ? { description: entry.description } : {}),
+        ...(score === undefined ? {} : { score }),
+      })),
+    };
+  }
+
+  private requestClient(actorId: string): {
+    client: ManagedAgentMemoryClient;
+    cacheKey: string;
+  } {
+    let workspaceClient = this.workspaceClient;
+    let credentialId = actorId;
+    try {
+      const execution = getExecutionContext();
+      workspaceClient = execution.client;
+      credentialId = "userId" in execution ? execution.userId : execution.serviceUserId;
+    } catch {
+      // Startup and tests have no active AppKit request context.
+    }
+    return {
+      client: new ManagedAgentMemoryClient(workspaceClient, this.options),
+      cacheKey: `${this.options.memoryStore}:${credentialId}:${actorId}`,
+    };
+  }
+}

package/src/databricks/client.ts

@@ -0,0 +1,251 @@
+/**
+ * Databricks Agent Memory REST client.
+ *
+ * @module
+ */
+
+import type { WorkspaceClient } from "@databricks/appkit";
+import { databricks } from "@dbx-tools/appkit";
+import { errorUtils } from "@dbx-tools/shared-core";
+import { z } from "zod";
+
+import type { ResolvedDatabricksMemoryOptions } from "../config.ts";
+
+export const ManagedMemoryEntrySchema = z
+  .object({
+    name: z.string().optional().describe("Managed memory entry resource name."),
+    actor_id: z.string().describe("Trusted application actor identifier."),
+    session_id: z.string().optional().describe("Associated session identifier."),
+    path: z.string().describe("Absolute memory path."),
+    content: z.string().optional().describe("Memory content."),
+    description: z.string().optional().describe("Memory retrieval description."),
+    source_type: z.string().optional().describe("Entry writer type."),
+    create_time: z.string().optional().describe("Entry creation timestamp."),
+    update_time: z.string().optional().describe("Entry update timestamp."),
+  })
+  .loose()
+  .describe("Databricks managed memory entry.");
+
+export type ManagedMemoryEntry = z.infer<typeof ManagedMemoryEntrySchema>;
+
+const SearchMemoryResponseSchema = z
+  .object({
+    managed_memory_entries: z
+      .array(ManagedMemoryEntrySchema)
+      .optional()
+      .describe("Compatibility list of matching entries."),
+    results: z
+      .array(
+        z
+          .object({
+            managed_memory_entry: ManagedMemoryEntrySchema.describe("Matching memory entry."),
+            score: z.number().optional().describe("BM25 relevance score."),
+          })
+          .loose(),
+      )
+      .optional()
+      .describe("Ranked memory search results."),
+  })
+  .loose()
+  .describe("Databricks managed memory search response.");
+
+const ListMemoryResponseSchema = z
+  .object({
+    managed_memory_entries: z
+      .array(ManagedMemoryEntrySchema)
+      .optional()
+      .describe("Managed memory entries in this page."),
+    next_page_token: z.string().optional().describe("Next page token."),
+  })
+  .loose()
+  .describe("Page of Databricks managed memory entries.");
+
+/** One managed-memory search result with its optional backend score. */
+export interface ManagedMemorySearchResult {
+  readonly entry: ManagedMemoryEntry;
+  readonly score?: number;
+}
+
+/** Thin REST client for actor-scoped Databricks Agent Memory operations. */
+export class ManagedAgentMemoryClient {
+  readonly memoryStore: string;
+
+  constructor(
+    private readonly client: WorkspaceClient,
+    options: ResolvedDatabricksMemoryOptions,
+  ) {
+    this.memoryStore = options.memoryStore;
+  }
+
+  async available(signal?: AbortSignal): Promise<boolean> {
+    try {
+      await this.client.apiClient.request(
+        {
+          path: `/api/2.0/agents/memory-stores/${encodeURIComponent(this.memoryStore)}`,
+          method: "GET",
+          headers: new Headers({ accept: "application/json" }),
+          raw: false,
+        },
+        signal ? databricks.toContext(signal) : undefined,
+      );
+      return true;
+    } catch (error) {
+      if (errorUtils.errorContext(error).hasStatusCode(403, 404)) return false;
+      throw error;
+    }
+  }
+
+  async searchMemory(input: {
+    actorId: string;
+    query: string;
+    limit: number;
+    signal?: AbortSignal;
+  }): Promise<ManagedMemorySearchResult[]> {
+    const response = SearchMemoryResponseSchema.parse(
+      await this.client.apiClient.request(
+        {
+          path: `/api/2.0/agents/memory-stores/${encodeURIComponent(this.memoryStore)}/entries:search`,
+          method: "POST",
+          headers: new Headers({
+            accept: "application/json",
+            "content-type": "application/json",
+          }),
+          raw: false,
+          payload: {
+            actor_id: input.actorId,
+            query: input.query,
+            page_size: input.limit,
+          },
+        },
+        input.signal ? databricks.toContext(input.signal) : undefined,
+      ),
+    );
+    if (response.results) {
+      return response.results.map((result) => ({
+        entry: result.managed_memory_entry,
+        ...(result.score === undefined ? {} : { score: result.score }),
+      }));
+    }
+    return (response.managed_memory_entries ?? []).map((entry) => ({ entry }));
+  }
+
+  async saveMemory(input: {
+    actorId: string;
+    path: string;
+    content: string;
+    description?: string;
+    signal?: AbortSignal;
+  }): Promise<ManagedMemoryEntry> {
+    const existing = (await this.listMemoryEntries(input))[0];
+    if (existing?.name) {
+      return this.updateMemory(existing as ManagedMemoryEntry & { name: string }, input);
+    }
+    try {
+      return await this.createMemory(input);
+    } catch (error) {
+      if (!errorUtils.errorContext(error).hasStatusCode(409)) throw error;
+      const raced = (await this.listMemoryEntries(input))[0];
+      if (!raced?.name) throw error;
+      return this.updateMemory(raced as ManagedMemoryEntry & { name: string }, input);
+    }
+  }
+
+  private async createMemory(input: {
+    actorId: string;
+    path: string;
+    content: string;
+    description?: string;
+    signal?: AbortSignal;
+  }): Promise<ManagedMemoryEntry> {
+    return ManagedMemoryEntrySchema.parse(
+      await this.client.apiClient.request(
+        {
+          path: `/api/2.0/agents/memory-stores/${encodeURIComponent(this.memoryStore)}/entries`,
+          method: "POST",
+          headers: new Headers({
+            accept: "application/json",
+            "content-type": "application/json",
+          }),
+          raw: false,
+          payload: {
+            actor_id: input.actorId,
+            path: input.path,
+            content: input.content,
+            ...(input.description ? { description: input.description } : {}),
+            source_type: "MANAGED_MEMORY_ENTRY_SOURCE_TYPE_AGENT",
+          },
+        },
+        input.signal ? databricks.toContext(input.signal) : undefined,
+      ),
+    );
+  }
+
+  private async listMemoryEntries(input: {
+    actorId: string;
+    path: string;
+    signal?: AbortSignal;
+  }): Promise<ManagedMemoryEntry[]> {
+    const response = ListMemoryResponseSchema.parse(
+      await this.client.apiClient.request(
+        {
+          path: `/api/2.0/agents/memory-stores/${encodeURIComponent(this.memoryStore)}/entries`,
+          method: "GET",
+          headers: new Headers({ accept: "application/json" }),
+          query: {
+            actor_id: input.actorId,
+            path: input.path,
+            page_size: 100,
+          },
+          raw: false,
+        },
+        input.signal ? databricks.toContext(input.signal) : undefined,
+      ),
+    );
+    return (response.managed_memory_entries ?? []).filter((entry) => entry.path === input.path);
+  }
+
+  private async updateMemory(
+    current: ManagedMemoryEntry & { name: string },
+    input: {
+      actorId: string;
+      path: string;
+      content: string;
+      description?: string;
+      signal?: AbortSignal;
+    },
+  ): Promise<ManagedMemoryEntry> {
+    const headers = new Headers({
+      accept: "application/json",
+      "content-type": "application/json",
+    });
+    await this.client.config.authenticate(headers);
+    const response = await fetch(
+      new URL(
+        `/api/2.0/agents/${current.name
+          .split("/")
+          .map((part) => encodeURIComponent(part))
+          .join("/")}?update_mask=content,description`,
+        (await this.client.config.getHost()).toString(),
+      ),
+      {
+        method: "PATCH",
+        headers,
+        body: JSON.stringify({
+          actor_id: input.actorId,
+          path: input.path,
+          content: input.content,
+          ...(input.description ? { description: input.description } : {}),
+        }),
+        ...(input.signal ? { signal: input.signal } : {}),
+      },
+    );
+    const body = await response.json().catch(() => undefined);
+    if (!response.ok) {
+      throw Object.assign(
+        new Error(`Databricks Agent Memory request failed (${response.status})`),
+        { statusCode: response.status, response: body },
+      );
+    }
+    return ManagedMemoryEntrySchema.parse(body);
+  }
+}

package/src/databricks/index.ts

@@ -0,0 +1,7 @@
+// GENERATED by projen watch - DO NOT EDIT.
+// Regenerated from the exporting modules in ./src.
+// Hand edits are overwritten on the next watch; this file is read-only.
+
+export * as client from "./client.ts";
+export { ManagedMemoryEntrySchema, ManagedAgentMemoryClient } from "./client.ts";
+export type { ManagedMemoryEntry, ManagedMemorySearchResult } from "./client.ts";

package/src/graphiti/_backend.ts

@@ -0,0 +1,282 @@
+/**
+ * Lazy Graphiti implementation of the backend-neutral memory contract.
+ *
+ * @module
+ */
+
+import { createHash, randomBytes } from "node:crypto";
+import { createServer } from "node:net";
+import { ConfigurationError } from "@databricks/appkit";
+import { openApiTools, type OpenApiTool } from "@dbx-tools/appkit/openapi-tool";
+import { profile as authProfile } from "@dbx-tools/auth";
+import { configUtils } from "@dbx-tools/core";
+import type { ChildProcessResult } from "@dbx-tools/core/exec";
+import { z } from "zod";
+
+import {
+  type GraphitiMemoryBackend,
+  type MemoryTurn,
+  type RecallMemoryInput,
+  type RecallMemoryResult,
+  type RememberMemoryInput,
+  type RememberMemoryResult,
+} from "../_contracts.ts";
+import {
+  graphitiOptionOverrides,
+  graphitiOptionsFromEnvironment,
+  resolveGraphitiOptions,
+  type GraphitiOptions,
+  type ResolvedGraphitiOptions,
+} from "./options.ts";
+import {
+  createGraphitiChildProcess,
+  graphitiHealthCheck,
+  graphitiHttpUrl,
+  graphitiRequestHeaders,
+} from "./runtime.ts";
+
+const GraphitiFactSchema = z
+  .object({
+    name: z.string().optional().describe("Graphiti fact name."),
+    fact: z.string().describe("Retrieved graph fact."),
+  })
+  .loose()
+  .describe("Graphiti fact search result.");
+
+const GraphitiFactSearchSchema = z
+  .object({
+    facts: z.array(GraphitiFactSchema).optional().describe("Matching graph facts."),
+    error: z.string().optional().describe("Graphiti operation error."),
+  })
+  .loose()
+  .describe("Graphiti fact search response.");
+
+const GraphitiWriteSchema = z
+  .object({
+    message: z.string().optional().describe("Graphiti write confirmation."),
+    error: z.string().optional().describe("Graphiti operation error."),
+  })
+  .loose()
+  .describe("Graphiti write response.");
+
+class GraphitiBackend implements GraphitiMemoryBackend {
+  constructor(
+    private readonly options: ResolvedGraphitiOptions,
+    private readonly child: ChildProcessResult,
+    private readonly writeTool: OpenApiTool,
+    private readonly recallTool: OpenApiTool,
+  ) {}
+
+  get running(): boolean {
+    return (
+      (this.child.exitCode === null || this.child.exitCode === undefined) &&
+      (this.child.signalCode === null || this.child.signalCode === undefined)
+    );
+  }
+
+  async remember(
+    actorId: string,
+    input: RememberMemoryInput,
+    signal?: AbortSignal,
+  ): Promise<RememberMemoryResult> {
+    await this.write(
+      actorId,
+      {
+        name: input.path,
+        episode_body: input.content,
+        source: "text",
+        source_description: input.description ? `${input.path}: ${input.description}` : input.path,
+      },
+      signal,
+    );
+    return {
+      backend: "graphiti",
+      memory: {
+        path: input.path,
+        content: input.content,
+        ...(input.description ? { description: input.description } : {}),
+      },
+    };
+  }
+
+  async recall(
+    actorId: string,
+    input: RecallMemoryInput,
+    signal?: AbortSignal,
+  ): Promise<RecallMemoryResult> {
+    const response = GraphitiFactSearchSchema.parse(
+      await this.request(
+        this.recallTool,
+        {
+          query: input.query,
+          max_facts: input.limit,
+          group_ids: [userScope(actorId)],
+        },
+        signal,
+      ),
+    );
+    if (response.error) throw new Error(response.error);
+    return {
+      backend: "graphiti",
+      memories: (response.facts ?? []).map((fact) => ({
+        content: fact.fact,
+        ...(fact.name ? { description: fact.name } : {}),
+      })),
+    };
+  }
+
+  async saveTurn(turn: MemoryTurn, signal?: AbortSignal): Promise<void> {
+    await this.write(
+      turn.actorId,
+      {
+        name: `Mastra turn ${turn.agentId}/${turn.threadId}`,
+        episode_body: `User: ${turn.user}\n\nAssistant: ${turn.assistant}`,
+        source: "message",
+        source_description: `Completed ${turn.agentId} conversation turn`,
+      },
+      signal,
+    );
+  }
+
+  shutdown(): Promise<void> {
+    return this.child.shutdown();
+  }
+
+  private async write(
+    actorId: string,
+    args: Record<string, unknown>,
+    signal?: AbortSignal,
+  ): Promise<void> {
+    const response = GraphitiWriteSchema.parse(
+      await this.request(this.writeTool, { ...args, group_id: userScope(actorId) }, signal),
+    );
+    if (response.error) throw new Error(response.error);
+  }
+
+  private async request(
+    tool: OpenApiTool,
+    args: Record<string, unknown>,
+    signal?: AbortSignal,
+  ): Promise<unknown> {
+    if (!this.running) throw new Error("Graphiti sidecar is not running");
+    const response = await fetch(tool.url, {
+      method: tool.method,
+      headers: {
+        ...graphitiRequestHeaders(this.options),
+        "content-type": "application/json",
+      },
+      body: JSON.stringify(args),
+      signal,
+    });
+    if (!response.ok) {
+      throw new Error(
+        `Graphiti operation ${tool.id} failed with HTTP ${response.status}: ${await response.text()}`,
+      );
+    }
+    return response.json();
+  }
+}
+
+/** Start a fully ready Graphiti backend for the first fallback operation. */
+export async function startGraphitiBackend(
+  options: GraphitiOptions,
+  signal?: AbortSignal,
+): Promise<GraphitiMemoryBackend> {
+  const configured = resolveGraphitiConfig(options);
+  if (configured.listen.scheme !== "tcp") {
+    throw new Error("AppKit memory requires a TCP Graphiti listener");
+  }
+  const [graphitiPort] = await distinctPorts(
+    configUtils.port(undefined, "DATABRICKS_APP_PORT", 8000, configUtils.ENV_ONLY),
+    configured.listen.port,
+  );
+  const resolved = resolveGraphitiOptions({
+    ...configured,
+    bearer: randomBytes(32).toString("base64url"),
+    listen: { ...configured.listen, port: graphitiPort },
+  });
+  let tools: OpenApiTool[] | undefined;
+  const { child } = await createGraphitiChildProcess({
+    ...resolved,
+    ...(signal ? { signal } : {}),
+    healthCheck: async ({ signal }) => {
+      try {
+        if (!(await graphitiHealthCheck(resolved, signal))) return false;
+        const loaded = await openApiTools(graphitiHttpUrl(resolved, "/openapi.json"), {
+          signal,
+          headers: graphitiRequestHeaders(resolved),
+        });
+        tools = [
+          requiredOperation(loaded, "add_memory_sync"),
+          requiredOperation(loaded, "search_memory_facts"),
+        ];
+        return true;
+      } catch (error) {
+        if (signal.aborted) throw error;
+        return false;
+      }
+    },
+  });
+  if (!tools) {
+    await child.shutdown();
+    throw new Error("Graphiti OpenAPI tools were not captured during readiness");
+  }
+  return new GraphitiBackend(resolved, child, tools[0]!, tools[1]!);
+}
+
+function resolveGraphitiConfig(options: GraphitiOptions): ResolvedGraphitiOptions {
+  const environment = graphitiOptionsFromEnvironment(process.env);
+  if (environment.profile && !authProfile.resolveProfile({ profile: environment.profile })) {
+    environment.profile = undefined;
+  }
+  const overrides = graphitiOptionOverrides(options);
+  return resolveGraphitiOptions({
+    ...environment,
+    ...overrides,
+    listen: overrides.listen ??
+      environment.listen ?? {
+        scheme: "tcp",
+        host: "127.0.0.1",
+        port: 0,
+      },
+  });
+}
+
+function requiredOperation(tools: readonly OpenApiTool[], name: string): OpenApiTool {
+  const tool = tools.find(({ id }) => id === name);
+  if (!tool) throw new Error(`Graphiti OpenAPI is missing operation: ${name}`);
+  return tool;
+}
+
+function userScope(userId: string): string {
+  return `user_${createHash("sha256").update(userId).digest("base64url")}`;
+}
+
+async function availablePort(): Promise<number> {
+  return new Promise((resolve, reject) => {
+    const server = createServer();
+    server.once("error", reject);
+    server.listen(0, "127.0.0.1", () => {
+      const address = server.address();
+      if (!address || typeof address === "string") {
+        server.close();
+        reject(new Error("Could not allocate a Graphiti loopback port"));
+        return;
+      }
+      server.close((error) => (error ? reject(error) : resolve(address.port)));
+    });
+  });
+}
+
+async function distinctPorts(appPort: number, graphitiPort: number): Promise<[number]> {
+  const ports = [appPort];
+  for (const configuredPort of [graphitiPort]) {
+    if (configuredPort && ports.includes(configuredPort)) {
+      throw new ConfigurationError("Graphiti sidecar ports must differ from DATABRICKS_APP_PORT");
+    }
+    let port = configuredPort || (await availablePort());
+    while (ports.includes(port)) port = await availablePort();
+    ports.push(port);
+  }
+  return ports.slice(1) as [number];
+}

package/src/graphiti/index.ts

@@ -0,0 +1,7 @@
+// GENERATED by projen watch - DO NOT EDIT.
+// Regenerated from the exporting modules in ./src.
+// Hand edits are overwritten on the next watch; this file is read-only.
+
+export * as options from "./options.ts";
+export * as runtime from "./runtime.ts";
+export type { GraphitiHealthCheckContext, GraphitiHealthCheck, GraphitiChildProcess, GraphitiRuntimeOptions, GraphitiPythonCommand } from "./runtime.ts";

package/src/graphiti/options.ts

@@ -0,0 +1,17 @@
+/**
+ * Node export of the shared Graphiti option contract.
+ *
+ * @module
+ */
+
+export {
+  GRAPHITI_DEFAULTS,
+  GraphitiCliOptionsSchema,
+  GraphitiOptionsSchema,
+  ResolvedGraphitiOptionsSchema,
+  graphitiOptionOverrides,
+  graphitiOptionsEnvironment,
+  graphitiOptionsFromEnvironment,
+  resolveGraphitiOptions,
+} from "@dbx-tools/shared-graphiti";
+export type { GraphitiOptions, ResolvedGraphitiOptions } from "@dbx-tools/shared-graphiti";