@andreprado/agentkit 0.1.0-alpha.5 → 0.1.0-alpha.6

package/src/runtime/knowledge/retrieve.ts

@@ -0,0 +1,283 @@
+import type { DatabaseRunner, DatabaseRow } from "../../index";
+import { createLocalDatabaseRunner, openPreparedCapsuleStore } from "../database";
+import { loadCapsuleEnv } from "../env";
+import { AgentKitError } from "../errors";
+import { loadAgentCapsule, type LoadedAgentCapsule } from "../config";
+import { resolveKnowledgeConfig } from "./config";
+import { cosineSimilarity, createEmbeddingProvider } from "./embeddings";
+import { searchLocalKnowledgeVectorRows } from "./vector";
+
+export type KnowledgeSearchResult = {
+  chunkId: string;
+  content: string;
+  score: number;
+  source: {
+    id: string;
+    path: string;
+    title: string | null;
+    section: string | null;
+    locator: string | null;
+  };
+  match: {
+    lexical: number;
+    semantic: number | null;
+  };
+};
+
+export async function searchKnowledgeFromCwd(
+  cwd: string,
+  query: string,
+  options: { topK?: number } = {},
+): Promise<{ databasePath: string; results: KnowledgeSearchResult[] }> {
+  const capsule = await loadAgentCapsule(cwd);
+  const { store, databasePath } = await openPreparedCapsuleStore(capsule);
+
+  try {
+    const db = createLocalDatabaseRunner(store);
+    return {
+      databasePath,
+      results: await searchKnowledge(capsule, db, query, options),
+    };
+  } finally {
+    store.close();
+  }
+}
+
+export async function searchKnowledge(
+  capsule: LoadedAgentCapsule,
+  db: DatabaseRunner,
+  query: string,
+  options: { topK?: number } = {},
+): Promise<KnowledgeSearchResult[]> {
+  const normalizedQuery = query.trim();
+
+  if (!normalizedQuery) {
+    throw new AgentKitError("knowledge_query_invalid", "Knowledge search query must be a non-empty string.");
+  }
+
+  const knowledge = resolveKnowledgeConfig(capsule.config);
+  const topK = options.topK ?? knowledge.retrieval.topK;
+  const env = await loadCapsuleEnv(capsule.root);
+  const embeddingProvider = createEmbeddingProvider({
+    config: knowledge.embedding,
+    env,
+  });
+  const lexical = knowledge.retrieval.hybrid || !embeddingProvider
+    ? await searchLexical(db, normalizedQuery, Math.max(topK * 4, topK))
+    : [];
+  const semantic = embeddingProvider
+    ? await searchSemantic(capsule, db, normalizedQuery, topK * 4, embeddingProvider)
+    : [];
+
+  return combineResults(lexical, semantic, topK);
+}
+
+async function searchLexical(db: DatabaseRunner, query: string, limit: number): Promise<KnowledgeSearchResult[]> {
+  const tokens = queryTokens(query);
+
+  if (tokens.length === 0) {
+    return [];
+  }
+
+  try {
+    const result = await db.query(
+      `
+      SELECT
+        c.id AS chunk_id,
+        c.source_id AS source_id,
+        c.content AS content,
+        c.title AS title,
+        c.section AS section,
+        c.locator AS locator,
+        s.source_key AS source_key,
+        rank
+      FROM agentkit_knowledge_fts
+      JOIN agentkit_knowledge_chunks c ON c.id = agentkit_knowledge_fts.chunk_id
+      JOIN agentkit_knowledge_sources s ON s.id = c.source_id
+      WHERE agentkit_knowledge_fts MATCH ?
+      ORDER BY rank ASC
+      LIMIT ?
+    `,
+      [tokens.map((token) => `"${token.replace(/"/g, '""')}"`).join(" OR "), limit],
+    );
+
+    return result.rows.map((row, index) => knowledgeSearchResultFromRow(row, lexicalScore(index), null));
+  } catch {
+    return searchLexicalFallback(db, query, limit);
+  }
+}
+
+async function searchLexicalFallback(db: DatabaseRunner, query: string, limit: number): Promise<KnowledgeSearchResult[]> {
+  const like = `%${query.toLowerCase()}%`;
+  const result = await db.query(
+    `
+    SELECT
+      c.id AS chunk_id,
+      c.source_id AS source_id,
+      c.content AS content,
+      c.title AS title,
+      c.section AS section,
+      c.locator AS locator,
+      s.source_key AS source_key
+    FROM agentkit_knowledge_chunks c
+    JOIN agentkit_knowledge_sources s ON s.id = c.source_id
+    WHERE lower(c.content) LIKE ?
+      OR lower(COALESCE(c.title, '')) LIKE ?
+      OR lower(COALESCE(c.section, '')) LIKE ?
+    ORDER BY c.created_at DESC, c.ordinal ASC
+    LIMIT ?
+  `,
+    [like, like, like, limit],
+  );
+
+  return result.rows.map((row, index) => knowledgeSearchResultFromRow(row, lexicalScore(index), null));
+}
+
+async function searchSemantic(
+  capsule: LoadedAgentCapsule,
+  db: DatabaseRunner,
+  query: string,
+  limit: number,
+  embeddingProvider: NonNullable<ReturnType<typeof createEmbeddingProvider>>,
+): Promise<KnowledgeSearchResult[]> {
+  const queryEmbedding = await embeddingProvider.embed(query);
+
+  if (capsule.storagePath) {
+    try {
+      const vectorRows = await searchLocalKnowledgeVectorRows({
+        databasePath: capsule.storagePath,
+        queryEmbedding,
+        limit,
+      });
+
+      return vectorRows
+        .map((row) => {
+          const distance = Number(row.vector_distance);
+          const semantic = Number.isFinite(distance) ? 1 - distance : 0;
+          return knowledgeSearchResultFromRow(row, 0, semantic);
+        })
+        .filter((result) => result.match.semantic !== null && result.match.semantic > 0)
+        .slice(0, limit);
+    } catch {
+      // Older local databases and non-libSQL environments use the JSON fallback below.
+    }
+  }
+
+  const result = await db.query(
+    `
+    SELECT
+      c.id AS chunk_id,
+      c.source_id AS source_id,
+      c.content AS content,
+      c.title AS title,
+      c.section AS section,
+      c.locator AS locator,
+      c.embedding_json AS embedding_json,
+      s.source_key AS source_key
+    FROM agentkit_knowledge_chunks c
+    JOIN agentkit_knowledge_sources s ON s.id = c.source_id
+    WHERE c.embedding_json IS NOT NULL
+    ORDER BY c.created_at DESC, c.ordinal ASC
+  `,
+  );
+
+  return result.rows
+    .map((row) => {
+      const embedding = parseEmbedding(row.embedding_json);
+      const semantic = embedding ? cosineSimilarity(queryEmbedding, embedding) : 0;
+      return knowledgeSearchResultFromRow(row, 0, semantic);
+    })
+    .filter((result) => result.match.semantic !== null && result.match.semantic > 0)
+    .sort((left, right) => (right.match.semantic ?? 0) - (left.match.semantic ?? 0))
+    .slice(0, limit);
+}
+
+function combineResults(
+  lexical: KnowledgeSearchResult[],
+  semantic: KnowledgeSearchResult[],
+  topK: number,
+): KnowledgeSearchResult[] {
+  const byChunk = new Map<string, KnowledgeSearchResult>();
+
+  for (const result of lexical) {
+    byChunk.set(result.chunkId, result);
+  }
+
+  for (const result of semantic) {
+    const existing = byChunk.get(result.chunkId);
+
+    if (!existing) {
+      byChunk.set(result.chunkId, {
+        ...result,
+        score: combinedScore(0, result.match.semantic),
+      });
+      continue;
+    }
+
+    byChunk.set(result.chunkId, {
+      ...existing,
+      score: combinedScore(existing.match.lexical, result.match.semantic),
+      match: {
+        lexical: existing.match.lexical,
+        semantic: result.match.semantic,
+      },
+    });
+  }
+
+  return Array.from(byChunk.values())
+    .map((result) => ({
+      ...result,
+      score: result.score || combinedScore(result.match.lexical, result.match.semantic),
+    }))
+    .sort((left, right) => right.score - left.score || left.source.path.localeCompare(right.source.path))
+    .slice(0, topK);
+}
+
+function knowledgeSearchResultFromRow(
+  row: DatabaseRow,
+  lexical: number,
+  semantic: number | null,
+): KnowledgeSearchResult {
+  return {
+    chunkId: String(row.chunk_id),
+    content: String(row.content),
+    score: combinedScore(lexical, semantic),
+    source: {
+      id: String(row.source_id),
+      path: String(row.source_key),
+      title: typeof row.title === "string" ? row.title : null,
+      section: typeof row.section === "string" ? row.section : null,
+      locator: typeof row.locator === "string" ? row.locator : null,
+    },
+    match: {
+      lexical,
+      semantic,
+    },
+  };
+}
+
+function combinedScore(lexical: number, semantic: number | null): number {
+  const semanticScore = semantic === null ? 0 : Math.max(0, (semantic + 1) / 2);
+  return Number((lexical * 0.6 + semanticScore * 0.4).toFixed(6));
+}
+
+function lexicalScore(index: number): number {
+  return Number((1 / (index + 1)).toFixed(6));
+}
+
+function queryTokens(query: string): string[] {
+  return Array.from(new Set(query.toLowerCase().match(/[\p{L}\p{N}_-]+/gu) ?? [])).slice(0, 12);
+}
+
+function parseEmbedding(value: unknown): number[] | null {
+  if (typeof value !== "string") {
+    return null;
+  }
+
+  try {
+    const parsed = JSON.parse(value);
+    return Array.isArray(parsed) && parsed.every((item) => typeof item === "number") ? parsed : null;
+  } catch {
+    return null;
+  }
+}

package/src/runtime/knowledge/schema.ts

@@ -0,0 +1,56 @@
+export const KNOWLEDGE_MIGRATION_ID = "0006_create_knowledge_tables";
+
+export const KNOWLEDGE_SCHEMA_SQL = `
+  CREATE TABLE IF NOT EXISTS agentkit_knowledge_sources (
+    id TEXT PRIMARY KEY,
+    source_key TEXT NOT NULL UNIQUE,
+    kind TEXT NOT NULL,
+    path TEXT,
+    url TEXT,
+    title TEXT,
+    content_sha256 TEXT,
+    status TEXT NOT NULL DEFAULT 'pending',
+    chunk_count INTEGER NOT NULL DEFAULT 0,
+    metadata_json TEXT NOT NULL DEFAULT '{}',
+    created_at TEXT NOT NULL,
+    updated_at TEXT NOT NULL,
+    indexed_at TEXT,
+    error_message TEXT
+  );
+
+  CREATE INDEX IF NOT EXISTS agentkit_knowledge_sources_status_idx
+    ON agentkit_knowledge_sources(status, updated_at);
+
+  CREATE TABLE IF NOT EXISTS agentkit_knowledge_chunks (
+    id TEXT PRIMARY KEY,
+    source_id TEXT NOT NULL REFERENCES agentkit_knowledge_sources(id) ON DELETE CASCADE,
+    ordinal INTEGER NOT NULL,
+    content TEXT NOT NULL,
+    token_count INTEGER NOT NULL,
+    title TEXT,
+    section TEXT,
+    locator TEXT,
+    content_sha256 TEXT NOT NULL,
+    embedding_json TEXT,
+    embedding_hash TEXT,
+    embedding_provider TEXT,
+    embedding_model TEXT,
+    embedding_dimensions INTEGER,
+    metadata_json TEXT NOT NULL DEFAULT '{}',
+    created_at TEXT NOT NULL,
+    UNIQUE(source_id, ordinal)
+  );
+
+  CREATE INDEX IF NOT EXISTS agentkit_knowledge_chunks_source_ordinal_idx
+    ON agentkit_knowledge_chunks(source_id, ordinal);
+
+  CREATE VIRTUAL TABLE IF NOT EXISTS agentkit_knowledge_fts
+    USING fts5(
+      chunk_id UNINDEXED,
+      source_id UNINDEXED,
+      title,
+      section,
+      content,
+      tokenize='unicode61'
+    );
+`;

package/src/runtime/knowledge/tool.ts

@@ -0,0 +1,64 @@
+import type { AgentTool } from "../../index";
+import type { LoadedAgentCapsule } from "../config";
+import { resolveKnowledgeConfig } from "./config";
+import { searchKnowledge } from "./retrieve";
+
+export const KNOWLEDGE_SEARCH_TOOL_NAME = "agentkit_search_knowledge";
+
+export function createKnowledgeSearchTool(capsule: LoadedAgentCapsule): AgentTool | null {
+  const knowledge = resolveKnowledgeConfig(capsule.config);
+
+  if (!knowledge.enabled) {
+    return null;
+  }
+
+  return {
+    name: KNOWLEDGE_SEARCH_TOOL_NAME,
+    description: "Search the Agent Capsule's configured Knowledge sources for business-specific facts.",
+    visibility: "internal",
+    inputSchema: {
+      type: "object",
+      properties: {
+        query: {
+          type: "string",
+        },
+        topK: {
+          type: "integer",
+        },
+      },
+      required: ["query"],
+      additionalProperties: false,
+    },
+    outputSchema: {
+      type: "object",
+      properties: {
+        results: {
+          type: "array",
+          items: {
+            type: "object",
+            properties: {
+              content: { type: "string" },
+              score: { type: "number" },
+              source: { type: "object" },
+            },
+            required: ["content", "score", "source"],
+          },
+        },
+      },
+      required: ["results"],
+      additionalProperties: false,
+    },
+    async execute(input, ctx) {
+      const typed = input as { query: string; topK?: number };
+      const results = await searchKnowledge(capsule, ctx.db, typed.query, { topK: typed.topK });
+
+      return {
+        results: results.map((result) => ({
+          content: result.content,
+          score: result.score,
+          source: result.source,
+        })),
+      };
+    },
+  };
+}

package/src/runtime/knowledge/vector.ts

@@ -0,0 +1,258 @@
+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 type { DatabaseRow } from "../../index";
+
+export const KNOWLEDGE_VECTOR_TABLE = "agentkit_knowledge_vectors";
+export const KNOWLEDGE_VECTOR_INDEX = "agentkit_knowledge_vectors_embedding_idx";
+const KNOWLEDGE_VECTOR_SOURCE_INDEX = "agentkit_knowledge_vectors_source_idx";
+
+export type LocalKnowledgeVectorStatus = {
+  enabled: boolean;
+  provider: "libsql" | "none" | "json_fallback";
+  chunksIndexed: number;
+  message: string;
+};
+
+export function localKnowledgeVectorDatabasePath(databasePath: string): string {
+  const extension = extname(databasePath);
+  const stem = extension ? basename(databasePath, extension) : basename(databasePath);
+  return join(dirname(databasePath), `${stem}.vectors.db`);
+}
+
+export async function localKnowledgeVectorDatabaseExists(databasePath: string): Promise<boolean> {
+  return access(localKnowledgeVectorDatabasePath(databasePath)).then(
+    () => true,
+    () => false,
+  );
+}
+
+export async function rebuildLocalKnowledgeVectorIndex(options: {
+  databasePath: string;
+  dimensions: number | null;
+}): Promise<LocalKnowledgeVectorStatus> {
+  if (!isPositiveInteger(options.dimensions)) {
+    return {
+      enabled: false,
+      provider: "none",
+      chunksIndexed: 0,
+      message: "Knowledge embeddings are disabled, so no local vector index was created.",
+    };
+  }
+
+  const mainClient = createClient({ url: pathToFileURL(options.databasePath).href });
+  const vectorClient = createClient({ url: pathToFileURL(localKnowledgeVectorDatabasePath(options.databasePath)).href });
+  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
+          c.id AS chunk_id,
+          c.source_id,
+          c.content,
+          c.title,
+          c.section,
+          c.locator,
+          c.embedding_json,
+          c.embedding_hash,
+          c.created_at,
+          s.source_key
+        FROM agentkit_knowledge_chunks c
+        JOIN agentkit_knowledge_sources s ON s.id = c.source_id
+        WHERE c.embedding_json IS NOT NULL
+          AND c.embedding_hash IS NOT NULL
+          AND c.embedding_dimensions = ?
+        ORDER BY c.created_at DESC, c.ordinal ASC
+      `,
+      args: [dimensions],
+    });
+
+    let chunksIndexed = 0;
+
+    for (const row of rows.rows) {
+      const chunkId = readString(row, "chunk_id");
+      const sourceId = readString(row, "source_id");
+      const sourceKey = readString(row, "source_key");
+      const content = readString(row, "content");
+      const embeddingJson = readString(row, "embedding_json");
+      const embeddingHash = readString(row, "embedding_hash");
+      const createdAt = readString(row, "created_at");
+
+      if (!chunkId || !sourceId || !sourceKey || !content || !embeddingJson || !embeddingHash || !createdAt) {
+        continue;
+      }
+
+      await vectorClient.execute({
+        sql: `
+          INSERT INTO ${KNOWLEDGE_VECTOR_TABLE} (
+            chunk_id,
+            source_id,
+            source_key,
+            content,
+            title,
+            section,
+            locator,
+            embedding,
+            embedding_hash,
+            created_at
+          )
+          VALUES (?, ?, ?, ?, ?, ?, ?, vector32(?), ?, ?)
+        `,
+        args: [
+          chunkId,
+          sourceId,
+          sourceKey,
+          content,
+          readNullableString(row, "title"),
+          readNullableString(row, "section"),
+          readNullableString(row, "locator"),
+          embeddingJson,
+          embeddingHash,
+          createdAt,
+        ],
+      });
+      chunksIndexed += 1;
+    }
+
+    await vectorClient.execute("PRAGMA wal_checkpoint(TRUNCATE)").catch(() => undefined);
+
+    return {
+      enabled: true,
+      provider: "libsql",
+      chunksIndexed,
+      message: `Local libSQL vector database rebuilt with ${chunksIndexed} chunk${chunksIndexed === 1 ? "" : "s"}.`,
+    };
+  } catch (error) {
+    return {
+      enabled: false,
+      provider: "json_fallback",
+      chunksIndexed: 0,
+      message: `Local libSQL vector index unavailable; using JSON embedding fallback. ${formatVectorError(error)}`,
+    };
+  } finally {
+    mainClient.close();
+    vectorClient.close();
+  }
+}
+
+export async function searchLocalKnowledgeVectorRows(options: {
+  databasePath: string;
+  queryEmbedding: number[];
+  limit: number;
+}): Promise<DatabaseRow[]> {
+  const vectorDatabasePath = localKnowledgeVectorDatabasePath(options.databasePath);
+
+  await access(vectorDatabasePath);
+
+  const client = createClient({ url: pathToFileURL(vectorDatabasePath).href });
+  const vector = JSON.stringify(options.queryEmbedding);
+
+  try {
+    const result = await client.execute({
+      sql: `
+        SELECT
+          c.id AS chunk_id,
+          c.source_id AS source_id,
+          c.content AS content,
+          c.title AS title,
+          c.section AS section,
+          c.locator AS locator,
+          c.source_key AS source_key,
+          vector_distance_cos(c.embedding, vector32(?)) AS vector_distance
+        FROM vector_top_k('${KNOWLEDGE_VECTOR_INDEX}', vector32(?), ?) nearest
+        JOIN ${KNOWLEDGE_VECTOR_TABLE} c ON c.id = nearest.id
+        ORDER BY vector_distance ASC
+      `,
+      args: [vector, vector, options.limit],
+    });
+
+    return result.rows.map(normalizeLibsqlRow);
+  } finally {
+    client.close();
+  }
+}
+
+function normalizeLibsqlRow(row: unknown): DatabaseRow {
+  const result: DatabaseRow = {};
+
+  if (!row || typeof row !== "object") {
+    return result;
+  }
+
+  for (const [key, value] of Object.entries(row)) {
+    if (key === "length" || /^\d+$/.test(key)) {
+      continue;
+    }
+
+    if (typeof value === "bigint") {
+      result[key] = Number.isSafeInteger(Number(value)) ? Number(value) : String(value);
+      continue;
+    }
+
+    if (value === null || typeof value === "string" || typeof value === "number" || typeof value === "boolean") {
+      result[key] = value;
+      continue;
+    }
+
+    result[key] = String(value);
+  }
+
+  return result;
+}
+
+function isPositiveInteger(value: number | null): value is number {
+  return typeof value === "number" && Number.isInteger(value) && value > 0;
+}
+
+function readString(row: unknown, key: string): string | null {
+  if (!row || typeof row !== "object" || !(key in row)) {
+    return null;
+  }
+
+  const value = (row as Record<string, unknown>)[key];
+  return typeof value === "string" && value.length > 0 ? value : null;
+}
+
+function readNullableString(row: unknown, key: string): string | null {
+  if (!row || typeof row !== "object" || !(key in row)) {
+    return null;
+  }
+
+  const value = (row as Record<string, unknown>)[key];
+  return typeof value === "string" ? value : null;
+}
+
+function formatVectorError(error: unknown): string {
+  const message = error instanceof Error ? error.message : String(error);
+  return message.length <= 240 ? message : `${message.slice(0, 237)}...`;
+}