@andreprado/agentkit 0.1.0 → 0.1.1

package/src/runtime/knowledge/embeddings.ts

@@ -5,8 +5,15 @@ import { AgentKitError } from "../errors";
 
 export type EmbeddingProvider = {
   embed(input: string): Promise<number[]>;
+  /** Embeds inputs in as few provider requests as possible; results keep input order. */
+  embedMany(inputs: string[]): Promise<number[][]>;
 };
 
+// Each valid input is at most 8,192 tokens; 36 inputs cannot exceed the
+// 300,000-token request limit, even for long CSV rows or dense text.
+// ponytail: conservative batches; add token-aware packing if throughput needs it.
+const OPENAI_EMBEDDING_BATCH_SIZE = Math.floor(300_000 / 8_192);
+
 export type CreateEmbeddingProviderOptions = {
   config: ResolvedKnowledgeConfig["embedding"];
   env?: Record<string, string | undefined>;
@@ -23,13 +30,26 @@ export function createEmbeddingProvider(options: CreateEmbeddingProviderOptions)
       async embed(input) {
         return fakeEmbedding(input, options.config.dimensions ?? 32);
       },
+      async embedMany(inputs) {
+        return inputs.map((input) => fakeEmbedding(input, options.config.dimensions ?? 32));
+      },
     };
   }
 
   if (options.config.provider === "openai") {
     return {
       async embed(input) {
-        return openAiEmbedding(input, options);
+        const [embedding] = await openAiEmbeddings([input], options);
+        return embedding;
+      },
+      async embedMany(inputs) {
+        const embeddings: number[][] = [];
+
+        for (let start = 0; start < inputs.length; start += OPENAI_EMBEDDING_BATCH_SIZE) {
+          embeddings.push(...(await openAiEmbeddings(inputs.slice(start, start + OPENAI_EMBEDDING_BATCH_SIZE), options)));
+        }
+
+        return embeddings;
       },
     };
   }
@@ -74,7 +94,7 @@ export function cosineSimilarity(left: number[], right: number[]): number {
   return dot / (Math.sqrt(leftNorm) * Math.sqrt(rightNorm));
 }
 
-async function openAiEmbedding(input: string, options: CreateEmbeddingProviderOptions): Promise<number[]> {
+async function openAiEmbeddings(inputs: string[], options: CreateEmbeddingProviderOptions): Promise<number[][]> {
   const secretName = options.config.secret ?? "OPENAI_API_KEY";
   const apiKey = options.env?.[secretName] ?? process.env[secretName];
 
@@ -94,7 +114,7 @@ async function openAiEmbedding(input: string, options: CreateEmbeddingProviderOp
     },
     body: JSON.stringify({
       model: options.config.model ?? "text-embedding-3-small",
-      input,
+      input: inputs,
       encoding_format: "float",
       ...(options.config.dimensions ? { dimensions: options.config.dimensions } : {}),
     }),
@@ -109,13 +129,31 @@ async function openAiEmbedding(input: string, options: CreateEmbeddingProviderOp
   }
 
   const payload = await response.json();
-  const embedding = payload?.data?.[0]?.embedding;
+  const data: unknown[] = Array.isArray(payload?.data) ? payload.data : [];
+  const embeddings: number[][] = new Array(inputs.length);
+
+  for (const [position, item] of data.entries()) {
+    const record = item as { index?: unknown; embedding?: unknown };
+    const index = typeof record.index === "number" ? record.index : position;
+    const embedding = record.embedding;
+
+    if (
+      index >= 0 &&
+      index < inputs.length &&
+      Array.isArray(embedding) &&
+      embedding.every((value) => typeof value === "number")
+    ) {
+      embeddings[index] = embedding;
+    }
+  }
 
-  if (!Array.isArray(embedding) || embedding.some((value) => typeof value !== "number")) {
-    throw new AgentKitError("knowledge_embedding_failed", "OpenAI embedding response did not contain a numeric vector.");
+  for (let index = 0; index < inputs.length; index += 1) {
+    if (!embeddings[index]) {
+      throw new AgentKitError("knowledge_embedding_failed", "OpenAI embedding response did not contain a numeric vector.");
+    }
   }
 
-  return embedding;
+  return embeddings;
 }
 
 function normalizeVector(values: number[]): number[] {

package/src/runtime/knowledge/ingest.ts

@@ -151,7 +151,9 @@ async function ingestKnowledgeSources(
       }
     }
 
-    await checkpointLocalKnowledgeDatabase(db);
+    if (results.some((result) => result.status === "indexed")) {
+      await checkpointLocalKnowledgeDatabase(db);
+    }
   } finally {
     store.close();
   }
@@ -304,13 +306,28 @@ async function ingestKnowledgeSource(
     id: string;
   }> = [];
 
-  for (const chunk of chunks) {
-    const embedding = embeddingProvider ? await embeddingProvider.embed(chunk.content) : null;
+  const chunkHashes = chunks.map((chunk) => sha256(chunk.content));
+  const embeddingsByHash = embeddingProvider
+    ? await loadReusableChunkEmbeddings(db, sourceId, resolvedKnowledge)
+    : new Map<string, number[]>();
+
+  if (embeddingProvider) {
+    const missingHashes = [...new Set(chunkHashes.filter((hash) => !embeddingsByHash.has(hash)))];
+
+    if (missingHashes.length > 0) {
+      const contentByHash = new Map(chunks.map((chunk, index) => [chunkHashes[index], chunk.content]));
+      const embeddings = await embeddingProvider.embedMany(missingHashes.map((hash) => contentByHash.get(hash)!));
+      missingHashes.forEach((hash, index) => embeddingsByHash.set(hash, embeddings[index]));
+    }
+  }
+
+  for (const [index, chunk] of chunks.entries()) {
+    const contentSha256 = chunkHashes[index];
     embeddedChunks.push({
       chunk,
-      embedding,
-      contentSha256: sha256(chunk.content),
-      id: stableId("kchk", `${sourceId}:${chunk.ordinal}:${sha256(chunk.content)}`),
+      embedding: embeddingProvider ? embeddingsByHash.get(contentSha256) ?? null : null,
+      contentSha256,
+      id: stableId("kchk", `${sourceId}:${chunk.ordinal}:${contentSha256}`),
     });
   }
 
@@ -491,6 +508,52 @@ function capsuleRelativePath(root: string, absolutePath: string): string {
   return relativePath.split("\\").join("/");
 }
 
+// Chunks whose text is unchanged keep their stored vector, so editing one part
+// of a file only pays for embedding the chunks that actually changed.
+async function loadReusableChunkEmbeddings(
+  db: DatabaseRunner,
+  sourceId: string,
+  knowledge: ReturnType<typeof resolveKnowledgeConfig>,
+): Promise<Map<string, number[]>> {
+  const result = await db.query(
+    `
+    SELECT content_sha256, embedding_json
+    FROM agentkit_knowledge_chunks
+    WHERE source_id = ?
+      AND embedding_json IS NOT NULL
+      AND embedding_provider = ?
+      AND embedding_model = ?
+      AND (? IS NULL OR embedding_dimensions = ?)
+  `,
+    [
+      sourceId,
+      knowledge.embedding.provider,
+      knowledge.embedding.model,
+      knowledge.embedding.dimensions,
+      knowledge.embedding.dimensions,
+    ],
+  );
+  const embeddings = new Map<string, number[]>();
+
+  for (const row of result.rows) {
+    if (typeof row.content_sha256 !== "string" || typeof row.embedding_json !== "string") {
+      continue;
+    }
+
+    try {
+      const embedding = JSON.parse(row.embedding_json);
+
+      if (Array.isArray(embedding) && embedding.every((value) => typeof value === "number")) {
+        embeddings.set(row.content_sha256, embedding);
+      }
+    } catch {
+      // A corrupt stored vector is simply re-embedded.
+    }
+  }
+
+  return embeddings;
+}
+
 async function findExistingSource(db: DatabaseRunner, sourceKey: string): Promise<DatabaseRow | undefined> {
   const result = await db.query(
     "SELECT id, content_sha256, status FROM agentkit_knowledge_sources WHERE source_key = ?",

package/src/runtime/knowledge/vector.ts

@@ -2,7 +2,7 @@ import { access } from "node:fs/promises";
 import { basename, dirname, extname, join } from "node:path";
 import { pathToFileURL } from "node:url";
 
-import { createClient } from "@libsql/client";
+import { createClient, type InStatement } from "@libsql/client";
 
 import type { DatabaseRow } from "../../index";
 
@@ -48,32 +48,6 @@ export async function rebuildLocalKnowledgeVectorIndex(options: {
   const dimensions = options.dimensions;
 
   try {
-    await vectorClient.batch(
-      [
-        `DROP INDEX IF EXISTS ${KNOWLEDGE_VECTOR_INDEX}`,
-        `DROP INDEX IF EXISTS ${KNOWLEDGE_VECTOR_SOURCE_INDEX}`,
-        `DROP TABLE IF EXISTS ${KNOWLEDGE_VECTOR_TABLE}`,
-        `
-          CREATE TABLE ${KNOWLEDGE_VECTOR_TABLE} (
-            id INTEGER PRIMARY KEY,
-            chunk_id TEXT NOT NULL UNIQUE,
-            source_id TEXT NOT NULL,
-            source_key TEXT NOT NULL,
-            content TEXT NOT NULL,
-            title TEXT,
-            section TEXT,
-            locator TEXT,
-            embedding F32_BLOB(${dimensions}) NOT NULL,
-            embedding_hash TEXT NOT NULL,
-            created_at TEXT NOT NULL
-          )
-        `,
-        `CREATE INDEX ${KNOWLEDGE_VECTOR_SOURCE_INDEX} ON ${KNOWLEDGE_VECTOR_TABLE}(source_id)`,
-        `CREATE INDEX ${KNOWLEDGE_VECTOR_INDEX} ON ${KNOWLEDGE_VECTOR_TABLE}(libsql_vector_idx(embedding, 'metric=cosine'))`,
-      ],
-      "write",
-    );
-
     const rows = await mainClient.execute({
       sql: `
         SELECT
@@ -97,7 +71,31 @@ export async function rebuildLocalKnowledgeVectorIndex(options: {
       args: [dimensions],
     });
 
-    let chunksIndexed = 0;
+    // Drop, recreate, and refill in one write batch so a failed rebuild never
+    // leaves a half-populated index behind, and inserts share one transaction.
+    const schemaStatements: InStatement[] = [
+      `DROP INDEX IF EXISTS ${KNOWLEDGE_VECTOR_INDEX}`,
+      `DROP INDEX IF EXISTS ${KNOWLEDGE_VECTOR_SOURCE_INDEX}`,
+      `DROP TABLE IF EXISTS ${KNOWLEDGE_VECTOR_TABLE}`,
+      `
+        CREATE TABLE ${KNOWLEDGE_VECTOR_TABLE} (
+          id INTEGER PRIMARY KEY,
+          chunk_id TEXT NOT NULL UNIQUE,
+          source_id TEXT NOT NULL,
+          source_key TEXT NOT NULL,
+          content TEXT NOT NULL,
+          title TEXT,
+          section TEXT,
+          locator TEXT,
+          embedding F32_BLOB(${dimensions}) NOT NULL,
+          embedding_hash TEXT NOT NULL,
+          created_at TEXT NOT NULL
+        )
+      `,
+      `CREATE INDEX ${KNOWLEDGE_VECTOR_SOURCE_INDEX} ON ${KNOWLEDGE_VECTOR_TABLE}(source_id)`,
+      `CREATE INDEX ${KNOWLEDGE_VECTOR_INDEX} ON ${KNOWLEDGE_VECTOR_TABLE}(libsql_vector_idx(embedding, 'metric=cosine'))`,
+    ];
+    const inserts: InStatement[] = [];
 
     for (const row of rows.rows) {
       const chunkId = readString(row, "chunk_id");
@@ -112,7 +110,7 @@ export async function rebuildLocalKnowledgeVectorIndex(options: {
         continue;
       }
 
-      await vectorClient.execute({
+      inserts.push({
         sql: `
           INSERT INTO ${KNOWLEDGE_VECTOR_TABLE} (
             chunk_id,
@@ -141,9 +139,11 @@ export async function rebuildLocalKnowledgeVectorIndex(options: {
           createdAt,
         ],
       });
-      chunksIndexed += 1;
     }
 
+    await vectorClient.batch([...schemaStatements, ...inserts], "write");
+    const chunksIndexed = inserts.length;
+
     await vectorClient.execute("PRAGMA wal_checkpoint(TRUNCATE)").catch(() => undefined);
 
     return {