convex-agent-knowledge 0.0.1

package/src/component/queries.ts

@@ -0,0 +1,202 @@
+// @ts-nocheck
+import { v } from "convex/values";
+import { internalQuery, query } from "./_generated/server.js";
+import { memoryCardValidator, vectorTableForDimension } from "./validators.js";
+
+type QueryCtx = any;
+
+async function getActiveMemory(ctx: QueryCtx, memoryId: string) {
+  const id = ctx.db.normalizeId("memories", memoryId);
+  if (!id) {
+    return null;
+  }
+  const memory = await ctx.db.get(id);
+  if (!memory || memory.status !== "active") {
+    return null;
+  }
+  return memory;
+}
+
+async function buildMemoryCard(
+  ctx: QueryCtx,
+  memoryId: string,
+  score: number,
+  scores?: { semanticScore?: number; graphScore?: number },
+) {
+  const memory = await getActiveMemory(ctx, memoryId);
+  if (!memory) {
+    return null;
+  }
+  const [entities, relationships] = await Promise.all([
+    ctx.db
+      .query("entities")
+      .withIndex("by_memory", (q) => q.eq("memoryId", memory._id))
+      .collect(),
+    ctx.db
+      .query("relationships")
+      .withIndex("by_memory", (q) => q.eq("memoryId", memory._id))
+      .collect(),
+  ]);
+
+  return {
+    memoryId: memory._id,
+    namespace: memory.namespace,
+    ...(memory.key === undefined ? {} : { key: memory.key }),
+    ...(memory.agentId === undefined ? {} : { agentId: memory.agentId }),
+    text: memory.text,
+    score,
+    ...(scores?.semanticScore === undefined ? {} : { semanticScore: scores.semanticScore }),
+    ...(scores?.graphScore === undefined ? {} : { graphScore: scores.graphScore }),
+    importance: memory.importance,
+    ...(memory.source === undefined ? {} : { source: memory.source }),
+    ...(memory.metadata === undefined ? {} : { metadata: memory.metadata }),
+    entities: entities.map((entity) => ({
+      externalId: entity.externalId,
+      name: entity.name,
+      type: entity.type,
+      ...(entity.description === undefined ? {} : { description: entity.description }),
+      confidence: entity.confidence,
+    })),
+    relationships: relationships.map((relationship) => ({
+      fromEntityExternalId: relationship.fromEntityExternalId,
+      toEntityExternalId: relationship.toEntityExternalId,
+      type: relationship.type,
+      ...(relationship.description === undefined ? {} : { description: relationship.description }),
+      confidence: relationship.confidence,
+      weight: relationship.weight,
+    })),
+  };
+}
+
+export const getMemory = query({
+  args: {
+    memoryId: v.string(),
+  },
+  returns: v.nullable(memoryCardValidator),
+  handler: async (ctx, args) => {
+    return await buildMemoryCard(ctx, args.memoryId, 1);
+  },
+});
+
+export const listMemories = query({
+  args: {
+    namespace: v.string(),
+    limit: v.optional(v.number()),
+    cursor: v.optional(v.string()),
+  },
+  returns: v.object({
+    page: v.array(memoryCardValidator),
+    continueCursor: v.union(v.string(), v.null()),
+    isDone: v.boolean(),
+  }),
+  handler: async (ctx, args) => {
+    const page = await ctx.db
+      .query("memories")
+      .withIndex("by_namespace", (q) => q.eq("namespace", args.namespace))
+      .filter((q) => q.eq(q.field("status"), "active"))
+      .order("desc")
+      .paginate({
+        cursor: args.cursor ?? null,
+        numItems: args.limit ?? 25,
+      });
+
+    const cards = [];
+    for (const memory of page.page) {
+      const card = await buildMemoryCard(ctx, memory._id, memory.importance);
+      if (card) {
+        cards.push(card);
+      }
+    }
+    return { ...page, page: cards };
+  },
+});
+
+export const fetchMemoryCards = query({
+  args: {
+    matches: v.array(
+      v.object({
+        memoryId: v.string(),
+        score: v.number(),
+        semanticScore: v.optional(v.number()),
+        graphScore: v.optional(v.number()),
+      }),
+    ),
+  },
+  returns: v.array(memoryCardValidator),
+  handler: async (ctx, args) => {
+    const cards = [];
+    for (const match of args.matches) {
+      const card = await buildMemoryCard(ctx, match.memoryId, match.score, {
+        semanticScore: match.semanticScore,
+        graphScore: match.graphScore,
+      });
+      if (card) {
+        cards.push(card);
+      }
+    }
+    return cards;
+  },
+});
+
+export const fetchMemoryCardsByVectorMatches = internalQuery({
+  args: {
+    embeddingDimension: v.number(),
+    matches: v.array(
+      v.object({
+        vectorId: v.string(),
+        score: v.number(),
+      }),
+    ),
+  },
+  returns: v.array(memoryCardValidator),
+  handler: async (ctx, args) => {
+    const table = vectorTableForDimension(args.embeddingDimension);
+    const cards = [];
+    for (const match of args.matches) {
+      const vectorId = ctx.db.normalizeId(table, match.vectorId);
+      if (!vectorId) {
+        continue;
+      }
+      const vectorRow = await ctx.db.get(vectorId);
+      if (!vectorRow) {
+        continue;
+      }
+      const card = await buildMemoryCard(ctx, vectorRow.memoryId, match.score, {
+        semanticScore: match.score,
+      });
+      if (card) {
+        cards.push(card);
+      }
+    }
+    return cards;
+  },
+});
+
+export const listPendingGraphSyncJobs = query({
+  args: {
+    limit: v.number(),
+  },
+  returns: v.array(
+    v.object({
+      jobId: v.string(),
+      namespace: v.string(),
+      operation: v.string(),
+      attempts: v.number(),
+      payload: v.any(),
+    }),
+  ),
+  handler: async (ctx, args) => {
+    const jobs = await ctx.db
+      .query("graphSyncJobs")
+      .withIndex("by_status", (q) => q.eq("status", "pending"))
+      .order("asc")
+      .take(args.limit);
+    return jobs.map((job) => ({
+      jobId: job._id,
+      namespace: job.namespace,
+      operation: job.operation,
+      attempts: job.attempts,
+      payload: job.payload,
+    }));
+  },
+});

package/src/component/schema.ts

@@ -0,0 +1,161 @@
+import { defineSchema, defineTable } from "convex/server";
+import { v } from "convex/values";
+
+const sourceValidator = v.object({
+  type: v.string(),
+  id: v.optional(v.string()),
+  url: v.optional(v.string()),
+  title: v.optional(v.string()),
+});
+
+const memoryStatus = v.union(v.literal("active"), v.literal("deleted"), v.literal("pending"));
+
+const syncStatus = v.union(
+  v.literal("pending"),
+  v.literal("running"),
+  v.literal("succeeded"),
+  v.literal("failed"),
+);
+
+const vectorRowFields = {
+  namespace: v.string(),
+  memoryId: v.id("memories"),
+  chunkId: v.id("chunks"),
+  agentId: v.string(),
+  kind: v.string(),
+  embedding: v.array(v.float64()),
+};
+
+function vectorTable(dimensions: number) {
+  return defineTable(vectorRowFields)
+    .index("by_memory", ["memoryId"])
+    .index("by_chunk", ["chunkId"])
+    .vectorIndex("by_embedding", {
+      vectorField: "embedding",
+      dimensions,
+      filterFields: ["namespace", "agentId", "kind"],
+    });
+}
+
+export default defineSchema({
+  namespaces: defineTable({
+    namespace: v.string(),
+    metadata: v.optional(v.any()),
+    createdAt: v.number(),
+    updatedAt: v.number(),
+  }).index("by_namespace", ["namespace"]),
+
+  memories: defineTable({
+    namespace: v.string(),
+    key: v.optional(v.string()),
+    agentId: v.optional(v.string()),
+    text: v.string(),
+    contentHash: v.string(),
+    source: v.optional(sourceValidator),
+    metadata: v.optional(v.any()),
+    status: memoryStatus,
+    importance: v.number(),
+    observationScore: v.number(),
+    embeddingDimension: v.number(),
+    chunkCount: v.number(),
+    entityCount: v.number(),
+    relationshipCount: v.number(),
+    createdAt: v.number(),
+    updatedAt: v.number(),
+    deletedAt: v.optional(v.number()),
+  })
+    .index("by_namespace", ["namespace"])
+    .index("by_namespace_key", ["namespace", "key"])
+    .index("by_namespace_hash", ["namespace", "contentHash"])
+    .index("by_agent", ["namespace", "agentId"]),
+
+  chunks: defineTable({
+    namespace: v.string(),
+    memoryId: v.id("memories"),
+    agentId: v.optional(v.string()),
+    order: v.number(),
+    text: v.string(),
+    summary: v.optional(v.string()),
+    tokenCount: v.optional(v.number()),
+    metadata: v.optional(v.any()),
+    createdAt: v.number(),
+  })
+    .index("by_memory", ["memoryId", "order"])
+    .index("by_namespace", ["namespace"]),
+
+  entities: defineTable({
+    namespace: v.string(),
+    memoryId: v.id("memories"),
+    chunkId: v.optional(v.id("chunks")),
+    externalId: v.string(),
+    name: v.string(),
+    type: v.string(),
+    description: v.optional(v.string()),
+    aliases: v.optional(v.array(v.string())),
+    confidence: v.number(),
+    metadata: v.optional(v.any()),
+    createdAt: v.number(),
+    updatedAt: v.number(),
+  })
+    .index("by_memory", ["memoryId"])
+    .index("by_namespace_name", ["namespace", "name"])
+    .index("by_external_id", ["namespace", "externalId"]),
+
+  relationships: defineTable({
+    namespace: v.string(),
+    memoryId: v.id("memories"),
+    fromEntityExternalId: v.string(),
+    toEntityExternalId: v.string(),
+    type: v.string(),
+    description: v.optional(v.string()),
+    confidence: v.number(),
+    weight: v.number(),
+    metadata: v.optional(v.any()),
+    createdAt: v.number(),
+    updatedAt: v.number(),
+  })
+    .index("by_memory", ["memoryId"])
+    .index("by_from", ["namespace", "fromEntityExternalId"])
+    .index("by_to", ["namespace", "toEntityExternalId"]),
+
+  observations: defineTable({
+    namespace: v.string(),
+    memoryId: v.id("memories"),
+    query: v.string(),
+    outcome: v.union(v.literal("helpful"), v.literal("not_helpful"), v.literal("neutral")),
+    feedback: v.optional(v.string()),
+    metadata: v.optional(v.any()),
+    createdAt: v.number(),
+  })
+    .index("by_memory", ["memoryId"])
+    .index("by_namespace", ["namespace"]),
+
+  graphSyncJobs: defineTable({
+    namespace: v.string(),
+    memoryId: v.optional(v.id("memories")),
+    operation: v.union(
+      v.literal("upsert_memory"),
+      v.literal("delete_memory"),
+      v.literal("promote_memory"),
+    ),
+    status: syncStatus,
+    attempts: v.number(),
+    payload: v.any(),
+    lastError: v.optional(v.string()),
+    createdAt: v.number(),
+    updatedAt: v.number(),
+  })
+    .index("by_status", ["status", "createdAt"])
+    .index("by_memory", ["memoryId"]),
+
+  vectors_128: vectorTable(128),
+  vectors_256: vectorTable(256),
+  vectors_384: vectorTable(384),
+  vectors_512: vectorTable(512),
+  vectors_768: vectorTable(768),
+  vectors_1024: vectorTable(1024),
+  vectors_1536: vectorTable(1536),
+  vectors_2048: vectorTable(2048),
+  vectors_3072: vectorTable(3072),
+  vectors_4096: vectorTable(4096),
+});

package/src/component/validators.ts

@@ -0,0 +1,104 @@
+import { v } from "convex/values";
+
+export const supportedEmbeddingDimensions = [
+  128, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096,
+] as const;
+
+export type SupportedEmbeddingDimension = (typeof supportedEmbeddingDimensions)[number];
+
+export const sourceValidator = v.object({
+  type: v.string(),
+  id: v.optional(v.string()),
+  url: v.optional(v.string()),
+  title: v.optional(v.string()),
+});
+
+export const entityInputValidator = v.object({
+  externalId: v.string(),
+  name: v.string(),
+  type: v.string(),
+  description: v.optional(v.string()),
+  aliases: v.optional(v.array(v.string())),
+  confidence: v.optional(v.number()),
+  metadata: v.optional(v.any()),
+});
+
+export const relationshipInputValidator = v.object({
+  fromEntityExternalId: v.string(),
+  toEntityExternalId: v.string(),
+  type: v.string(),
+  description: v.optional(v.string()),
+  confidence: v.optional(v.number()),
+  weight: v.optional(v.number()),
+  metadata: v.optional(v.any()),
+});
+
+export const chunkInputValidator = v.object({
+  text: v.string(),
+  embedding: v.array(v.float64()),
+  summary: v.optional(v.string()),
+  tokenCount: v.optional(v.number()),
+  metadata: v.optional(v.any()),
+});
+
+export const memoryCardValidator = v.object({
+  memoryId: v.string(),
+  namespace: v.string(),
+  key: v.optional(v.string()),
+  agentId: v.optional(v.string()),
+  text: v.string(),
+  score: v.number(),
+  semanticScore: v.optional(v.number()),
+  graphScore: v.optional(v.number()),
+  importance: v.number(),
+  source: v.optional(sourceValidator),
+  metadata: v.optional(v.any()),
+  entities: v.array(
+    v.object({
+      externalId: v.string(),
+      name: v.string(),
+      type: v.string(),
+      description: v.optional(v.string()),
+      confidence: v.number(),
+    }),
+  ),
+  relationships: v.array(
+    v.object({
+      fromEntityExternalId: v.string(),
+      toEntityExternalId: v.string(),
+      type: v.string(),
+      description: v.optional(v.string()),
+      confidence: v.number(),
+      weight: v.number(),
+    }),
+  ),
+});
+
+export function vectorTableForDimension(dimension: number) {
+  switch (dimension) {
+    case 128:
+      return "vectors_128";
+    case 256:
+      return "vectors_256";
+    case 384:
+      return "vectors_384";
+    case 512:
+      return "vectors_512";
+    case 768:
+      return "vectors_768";
+    case 1024:
+      return "vectors_1024";
+    case 1536:
+      return "vectors_1536";
+    case 2048:
+      return "vectors_2048";
+    case 3072:
+      return "vectors_3072";
+    case 4096:
+      return "vectors_4096";
+    default:
+      throw new Error(
+        `Unsupported embedding dimension ${dimension}. Supported dimensions: ${supportedEmbeddingDimensions.join(", ")}`,
+      );
+  }
+}

package/src/node/index.ts

@@ -0,0 +1,2 @@
+export { createNeo4jGraphStore } from "./neo4j.js";
+export type { Neo4jGraphStoreOptions } from "./neo4j.js";

package/src/node/neo4j.ts

@@ -0,0 +1,210 @@
+import neo4j from "neo4j-driver";
+import type {
+  GraphExpandInput,
+  GraphMemoryScore,
+  GraphStore,
+  GraphSyncJob,
+  Neo4jConfig,
+} from "../client/types.js";
+
+export type Neo4jGraphStoreOptions = Neo4jConfig;
+
+export function createNeo4jGraphStore(options: Neo4jGraphStoreOptions): GraphStore {
+  return {
+    async syncJob(job) {
+      await withNeo4j(options, async (session) => {
+        if (job.operation === "upsert_memory") {
+          await upsertMemoryGraph(session, job.payload);
+        } else if (job.operation === "delete_memory") {
+          await deleteMemoryGraph(session, job.payload);
+        } else if (job.operation === "promote_memory") {
+          await promoteMemoryGraph(session, job.payload);
+        } else {
+          throw new Error(`Unsupported graph sync operation ${job.operation}`);
+        }
+      });
+    },
+    async expand(input) {
+      return await expandGraph(options, input);
+    },
+  };
+}
+
+async function withNeo4j<T>(config: Neo4jConfig, fn: (session: neo4j.Session) => Promise<T>) {
+  const driver = neo4j.driver(config.uri, neo4j.auth.basic(config.user, config.password));
+  const session = driver.session(config.database ? { database: config.database } : undefined);
+  try {
+    return await fn(session);
+  } finally {
+    await session.close();
+    await driver.close();
+  }
+}
+
+async function upsertMemoryGraph(session: neo4j.Session, payload: unknown) {
+  const graphPayload = payload as {
+    memory: Record<string, unknown> & {
+      id: string;
+      namespace: string;
+      importance?: number;
+    };
+    entities?: Array<Record<string, unknown>>;
+    relationships?: Array<Record<string, unknown>>;
+  };
+  await session.executeWrite((tx) =>
+    tx.run(
+      `
+      MERGE (m:Memory {id: $memory.id})
+      SET m.namespace = $memory.namespace,
+          m.key = $memory.key,
+          m.agentId = $memory.agentId,
+          m.text = $memory.text,
+          m.importance = $memory.importance,
+          m.updatedAt = timestamp()
+      WITH m
+      UNWIND $entities AS entity
+      MERGE (e:Entity {id: entity.externalId})
+      SET e.namespace = $memory.namespace,
+          e.name = entity.name,
+          e.type = entity.type,
+          e.description = entity.description,
+          e.confidence = coalesce(entity.confidence, 0.75),
+          e.updatedAt = timestamp()
+      MERGE (m)-[mention:MENTIONS]->(e)
+      SET mention.namespace = $memory.namespace,
+          mention.updatedAt = timestamp()
+      `,
+      {
+        memory: graphPayload.memory,
+        entities: graphPayload.entities ?? [],
+      },
+    ),
+  );
+
+  await session.executeWrite((tx) =>
+    tx.run(
+      `
+      UNWIND $relationships AS relationship
+      MATCH (from:Entity {id: relationship.fromEntityExternalId})
+      MATCH (to:Entity {id: relationship.toEntityExternalId})
+      MERGE (from)-[r:RELATED {namespace: $namespace, type: relationship.type}]->(to)
+      SET r.description = relationship.description,
+          r.confidence = coalesce(relationship.confidence, 0.75),
+          r.weight = coalesce(relationship.weight, 0.5),
+          r.memoryId = $memoryId,
+          r.updatedAt = timestamp()
+      `,
+      {
+        namespace: graphPayload.memory.namespace,
+        memoryId: graphPayload.memory.id,
+        relationships: graphPayload.relationships ?? [],
+      },
+    ),
+  );
+}
+
+async function deleteMemoryGraph(session: neo4j.Session, payload: unknown) {
+  const graphPayload = payload as { memoryId: string; namespace: string };
+  await session.executeWrite((tx) =>
+    tx.run(
+      `
+      MATCH (m:Memory {id: $memoryId, namespace: $namespace})
+      DETACH DELETE m
+      `,
+      graphPayload,
+    ),
+  );
+}
+
+async function promoteMemoryGraph(session: neo4j.Session, payload: unknown) {
+  const graphPayload = payload as {
+    memoryId: string;
+    namespace: string;
+    importance: number;
+    observationScore: number;
+  };
+  await session.executeWrite((tx) =>
+    tx.run(
+      `
+      MATCH (m:Memory {id: $memoryId, namespace: $namespace})
+      SET m.importance = $importance,
+          m.observationScore = $observationScore,
+          m.updatedAt = timestamp()
+      `,
+      graphPayload,
+    ),
+  );
+}
+
+async function expandGraph(
+  config: Neo4jConfig,
+  input: GraphExpandInput,
+): Promise<GraphMemoryScore[]> {
+  const seedMemoryIds = input.seedMemoryIds;
+  const entityHints = input.entityHints ?? [];
+  if (seedMemoryIds.length === 0 && entityHints.length === 0) {
+    return [];
+  }
+  const safeHops = Math.min(Math.max(Math.trunc(input.hops ?? 2), 1), 4);
+  const limit = Math.min(Math.max(input.limit ?? 32, 1), 128);
+  return await withNeo4j(config, async (session) => {
+    const result =
+      seedMemoryIds.length > 0
+        ? await session.executeRead((tx) =>
+            tx.run(
+              `
+              MATCH (seed:Memory {namespace: $namespace})
+              WHERE seed.id IN $seedMemoryIds
+              MATCH path = (seed)-[:MENTIONS]->(:Entity)-[:RELATED*1..${safeHops}]-(:Entity)<-[:MENTIONS]-(related:Memory {namespace: $namespace})
+              WHERE NOT related.id IN $seedMemoryIds
+              RETURN related.id AS memoryId, count(path) AS graphScore
+              ORDER BY graphScore DESC
+              LIMIT $limit
+              `,
+              {
+                namespace: input.namespace,
+                seedMemoryIds,
+                limit: neo4j.int(limit),
+              },
+            ),
+          )
+        : await session.executeRead((tx) =>
+            tx.run(
+              `
+              MATCH (seed:Entity {namespace: $namespace})
+              WHERE toLower(seed.name) IN $entityHints
+              MATCH path = (seed)-[:RELATED*0..${safeHops}]-(:Entity)<-[:MENTIONS]-(related:Memory {namespace: $namespace})
+              RETURN related.id AS memoryId, count(path) AS graphScore
+              ORDER BY graphScore DESC
+              LIMIT $limit
+              `,
+              {
+                namespace: input.namespace,
+                entityHints: entityHints.map((hint) => hint.toLowerCase()),
+                limit: neo4j.int(limit),
+              },
+            ),
+          );
+    return result.records.map((record) => ({
+      memoryId: record.get("memoryId") as string,
+      graphScore: toNumber(record.get("graphScore")),
+    }));
+  });
+}
+
+function toNumber(value: unknown) {
+  if (typeof value === "number") {
+    return value;
+  }
+  if (
+    value &&
+    typeof value === "object" &&
+    "toNumber" in value &&
+    typeof value.toNumber === "function"
+  ) {
+    return value.toNumber();
+  }
+  return Number(value);
+}
+
+export type { GraphStore, GraphSyncJob };

package/src/shared/ranking.test.ts

@@ -0,0 +1,30 @@
+import { describe, expect, it } from "vitest";
+import { clamp, fuseMemoryScores } from "./ranking.js";
+
+describe("clamp", () => {
+  it("limits values to the target range", () => {
+    expect(clamp(-1)).toBe(0);
+    expect(clamp(0.5)).toBe(0.5);
+    expect(clamp(2)).toBe(1);
+  });
+});
+
+describe("fuseMemoryScores", () => {
+  it("merges semantic and graph scores by memory id", () => {
+    const results = fuseMemoryScores(
+      [{ memoryId: "a", score: 0.8, semanticScore: 0.8, importance: 0.2 }],
+      [
+        { memoryId: "a", score: 2, graphScore: 2, importance: 0.2 },
+        { memoryId: "b", score: 3, graphScore: 3, importance: 0.1 },
+      ],
+      { limit: 2 },
+    );
+
+    expect(results).toHaveLength(2);
+    expect(results[0]!.memoryId).toBe("a");
+    expect(results.find((result) => result.memoryId === "a")).toMatchObject({
+      semanticScore: 0.8,
+      graphScore: 2,
+    });
+  });
+});