@workglow/knowledge-base 0.1.1 → 0.2.0

package/dist/bun.js

@@ -343,9 +343,7 @@ var ChunkRecordArraySchema = {
   description: "Array of chunk records"
 };
 // src/chunk/ChunkVectorStorageSchema.ts
-import {
-  TypedArraySchema
-} from "@workglow/util/schema";
+import { TypedArraySchema } from "@workglow/util/schema";
 var ChunkVectorStorageSchema = {
   type: "object",
   properties: {
@@ -564,6 +562,12 @@ class KnowledgeBase {
     this.tabularStorage.destroy();
     this.chunkStorage.destroy();
   }
+  async[Symbol.asyncDispose]() {
+    this.destroy();
+  }
+  [Symbol.dispose]() {
+    this.destroy();
+  }
   async getChunk(chunk_id) {
     return this.chunkStorage.get({ chunk_id });
   }
@@ -600,6 +604,55 @@ class KnowledgeBase {
     return doc.findChunksByNodeId(nodeId);
   }
 }
+// src/knowledge-base/SharedTableSchemas.ts
+import { TypedArraySchema as TypedArraySchema2 } from "@workglow/util/schema";
+var SHARED_DOCUMENT_TABLE = "shared_documents";
+var SHARED_CHUNK_TABLE = "shared_chunks";
+var SharedDocumentStorageSchema = {
+  type: "object",
+  properties: {
+    doc_id: {
+      type: "string",
+      "x-auto-generated": true,
+      title: "Document ID",
+      description: "Unique identifier for the document"
+    },
+    kb_id: {
+      type: "string",
+      title: "Knowledge Base ID",
+      description: "Owning knowledge base identifier"
+    },
+    data: {
+      type: "string",
+      title: "Document Data",
+      description: "JSON-serialized document"
+    },
+    metadata: {
+      type: "object",
+      title: "Metadata",
+      description: "Metadata of the document"
+    }
+  },
+  required: ["doc_id", "kb_id", "data"],
+  additionalProperties: true
+};
+var SharedChunkVectorStorageSchema = {
+  type: "object",
+  properties: {
+    chunk_id: { type: "string", "x-auto-generated": true },
+    kb_id: { type: "string" },
+    doc_id: { type: "string" },
+    vector: TypedArraySchema2(),
+    metadata: { type: "object", format: "metadata", additionalProperties: true }
+  },
+  required: ["chunk_id", "kb_id", "doc_id", "vector", "metadata"],
+  additionalProperties: false
+};
+var SharedDocumentPrimaryKey = ["kb_id", "doc_id"];
+var SharedChunkPrimaryKey = ["kb_id", "chunk_id"];
+var SharedDocumentIndexes = [["kb_id"]];
+var SharedChunkIndexes = [["kb_id"], ["kb_id", "doc_id"]];
+
 // src/knowledge-base/KnowledgeBaseSchema.ts
 var KnowledgeBaseRecordSchema = {
   type: "object",
@@ -633,6 +686,9 @@ function knowledgeBaseTableNames(kbId) {
     chunkTable: `kb_chunks_${safe}`
   };
 }
+function isSharedTableMode(record) {
+  return record.document_table === SHARED_DOCUMENT_TABLE && record.chunk_table === SHARED_CHUNK_TABLE;
+}
 // src/knowledge-base/KnowledgeBaseRepository.ts
 import { EventEmitter } from "@workglow/util";
 
@@ -697,16 +753,13 @@ class InMemoryKnowledgeBaseRepository extends KnowledgeBaseRepository {
 import {
   createServiceToken,
   globalServiceRegistry,
+  registerInputCompactor,
   registerInputResolver
 } from "@workglow/util";
 var KNOWLEDGE_BASES = createServiceToken("knowledge-base.registry");
 var KNOWLEDGE_BASE_REPOSITORY = createServiceToken("knowledge-base.repository");
-if (!globalServiceRegistry.has(KNOWLEDGE_BASES)) {
-  globalServiceRegistry.register(KNOWLEDGE_BASES, () => new Map, true);
-}
-if (!globalServiceRegistry.has(KNOWLEDGE_BASE_REPOSITORY)) {
-  globalServiceRegistry.register(KNOWLEDGE_BASE_REPOSITORY, () => new InMemoryKnowledgeBaseRepository, true);
-}
+globalServiceRegistry.registerIfAbsent(KNOWLEDGE_BASES, () => new Map, true);
+globalServiceRegistry.registerIfAbsent(KNOWLEDGE_BASE_REPOSITORY, () => new InMemoryKnowledgeBaseRepository, true);
 function getGlobalKnowledgeBases() {
   return globalServiceRegistry.get(KNOWLEDGE_BASES);
 }
@@ -716,23 +769,53 @@ function getGlobalKnowledgeBaseRepository() {
 function setGlobalKnowledgeBaseRepository(repository) {
   globalServiceRegistry.registerInstance(KNOWLEDGE_BASE_REPOSITORY, repository);
 }
-async function registerKnowledgeBase(id, kb) {
-  const kbs = getGlobalKnowledgeBases();
-  const now = new Date().toISOString();
-  const tableNames = knowledgeBaseTableNames(id);
-  const record = {
-    kb_id: id,
-    title: kb.title,
-    description: kb.description,
-    vector_dimensions: kb.getVectorDimensions(),
-    document_table: tableNames.documentTable,
-    chunk_table: tableNames.chunkTable,
-    created_at: now,
-    updated_at: now
+var pendingOps = new Map;
+function withIdLock(id, fn) {
+  const prev = pendingOps.get(id) ?? Promise.resolve();
+  const next = prev.then(fn, fn);
+  pendingOps.set(id, next);
+  const cleanup = () => {
+    if (pendingOps.get(id) === next) {
+      pendingOps.delete(id);
+    }
   };
+  next.finally(cleanup);
+  return next;
+}
+function registerKnowledgeBase(id, kb, options) {
+  return withIdLock(id, async () => {
+    const kbs = getGlobalKnowledgeBases();
+    const now = new Date().toISOString();
+    const useShared = options?.sharedTables === true;
+    const tableNames = useShared ? { documentTable: SHARED_DOCUMENT_TABLE, chunkTable: SHARED_CHUNK_TABLE } : knowledgeBaseTableNames(id);
+    const record = {
+      kb_id: id,
+      title: kb.title,
+      description: kb.description,
+      vector_dimensions: kb.getVectorDimensions(),
+      document_table: tableNames.documentTable,
+      chunk_table: tableNames.chunkTable,
+      created_at: now,
+      updated_at: now
+    };
+    const repo = getGlobalKnowledgeBaseRepository();
+    await repo.addKnowledgeBase(record);
+    kbs.set(id, kb);
+  });
+}
+function unregisterKnowledgeBase(id) {
+  return withIdLock(id, async () => {
+    const repo = getGlobalKnowledgeBaseRepository();
+    await repo.removeKnowledgeBase(id);
+    const kbs = getGlobalKnowledgeBases();
+    kbs.delete(id);
+  });
+}
+async function deregisterKnowledgeBase(id) {
   const repo = getGlobalKnowledgeBaseRepository();
-  await repo.addKnowledgeBase(record);
-  kbs.set(id, kb);
+  await repo.removeKnowledgeBase(id);
+  const kbs = getGlobalKnowledgeBases();
+  kbs.delete(id);
 }
 function getKnowledgeBase(id) {
   return getGlobalKnowledgeBases().get(id);
@@ -746,6 +829,14 @@ async function resolveKnowledgeBaseFromRegistry(id, format, registry) {
   return kb;
 }
 registerInputResolver("knowledge-base", resolveKnowledgeBaseFromRegistry);
+registerInputCompactor("knowledge-base", (value, _format, registry) => {
+  const kbs = registry.has(KNOWLEDGE_BASES) ? registry.get(KNOWLEDGE_BASES) : getGlobalKnowledgeBases();
+  for (const [id, kb] of kbs) {
+    if (kb === value)
+      return id;
+  }
+  return;
+});
 // src/knowledge-base/createKnowledgeBase.ts
 import { InMemoryTabularStorage as InMemoryTabularStorage2, InMemoryVectorStorage } from "@workglow/storage";
 
@@ -801,6 +892,207 @@ async function createKnowledgeBase(options) {
   }
   return kb;
 }
+// src/knowledge-base/ScopedTabularStorage.ts
+import { EventEmitter as EventEmitter2 } from "@workglow/util";
+
+class ScopedTabularStorage {
+  inner;
+  kbId;
+  events = new EventEmitter2;
+  constructor(inner, kbId) {
+    this.inner = inner;
+    this.kbId = kbId;
+  }
+  inject(value) {
+    return { ...value, kb_id: this.kbId };
+  }
+  strip(entity) {
+    if (!entity)
+      return entity;
+    const { kb_id: _, ...rest } = entity;
+    return rest;
+  }
+  stripArray(entities) {
+    if (!entities)
+      return;
+    return entities.map((e) => this.strip(e));
+  }
+  async put(value) {
+    const result = await this.inner.put(this.inject(value));
+    const stripped = this.strip(result);
+    this.events.emit("put", stripped);
+    return stripped;
+  }
+  async putBulk(values) {
+    const injected = values.map((v) => this.inject(v));
+    const results = await this.inner.putBulk(injected);
+    const stripped = results.map((r) => this.strip(r));
+    for (const entity of stripped) {
+      this.events.emit("put", entity);
+    }
+    return stripped;
+  }
+  async get(key) {
+    const result = await this.inner.get({ ...key, kb_id: this.kbId });
+    if (!result)
+      return;
+    const stripped = this.strip(result);
+    this.events.emit("get", key, stripped);
+    return stripped;
+  }
+  async delete(key) {
+    await this.inner.deleteSearch({ ...key, kb_id: this.kbId });
+    this.events.emit("delete", key);
+  }
+  async getAll(options) {
+    const results = await this.inner.query({ kb_id: this.kbId }, options);
+    return this.stripArray(results);
+  }
+  async deleteAll() {
+    await this.inner.deleteSearch({ kb_id: this.kbId });
+    this.events.emit("clearall");
+  }
+  async size() {
+    let count = 0;
+    const pageSize = 1000;
+    let offset = 0;
+    while (true) {
+      const page = await this.inner.query({ kb_id: this.kbId }, { offset, limit: pageSize });
+      if (!page || page.length === 0)
+        break;
+      count += page.length;
+      if (page.length < pageSize)
+        break;
+      offset += pageSize;
+    }
+    return count;
+  }
+  async query(criteria, options) {
+    const results = await this.inner.query({ ...criteria, kb_id: this.kbId }, options);
+    const stripped = this.stripArray(results);
+    this.events.emit("query", criteria, stripped);
+    return stripped;
+  }
+  async deleteSearch(criteria) {
+    await this.inner.deleteSearch({ ...criteria, kb_id: this.kbId });
+  }
+  async getBulk(offset, limit) {
+    const results = await this.inner.query({ kb_id: this.kbId }, { offset, limit });
+    return this.stripArray(results);
+  }
+  async* records(pageSize = 100) {
+    if (pageSize <= 0) {
+      throw new RangeError(`pageSize must be greater than 0, got ${pageSize}`);
+    }
+    let offset = 0;
+    while (true) {
+      const page = await this.getBulk(offset, pageSize);
+      if (!page || page.length === 0) {
+        break;
+      }
+      for (const entity of page) {
+        yield entity;
+      }
+      if (page.length < pageSize)
+        break;
+      offset += pageSize;
+    }
+  }
+  async* pages(pageSize = 100) {
+    if (pageSize <= 0) {
+      throw new RangeError(`pageSize must be greater than 0, got ${pageSize}`);
+    }
+    let offset = 0;
+    while (true) {
+      const page = await this.getBulk(offset, pageSize);
+      if (!page || page.length === 0) {
+        break;
+      }
+      yield page;
+      if (page.length < pageSize)
+        break;
+      offset += pageSize;
+    }
+  }
+  on(name, fn) {
+    this.events.on(name, fn);
+  }
+  off(name, fn) {
+    this.events.off(name, fn);
+  }
+  emit(name, ...args) {
+    this.events.emit(name, ...args);
+  }
+  once(name, fn) {
+    this.events.once(name, fn);
+  }
+  waitOn(name) {
+    return this.events.waitOn(name);
+  }
+  subscribeToChanges(callback, options) {
+    return this.inner.subscribeToChanges((change) => {
+      const newKbId = change.new?.kb_id;
+      const oldKbId = change.old?.kb_id;
+      if (newKbId !== undefined && newKbId !== this.kbId)
+        return;
+      if (oldKbId !== undefined && oldKbId !== this.kbId)
+        return;
+      if (newKbId === undefined && oldKbId === undefined)
+        return;
+      callback({
+        type: change.type,
+        ...change.old ? { old: this.strip(change.old) } : {},
+        ...change.new ? { new: this.strip(change.new) } : {}
+      });
+    }, options);
+  }
+  async setupDatabase() {}
+  destroy() {}
+  [Symbol.dispose]() {}
+  async[Symbol.asyncDispose]() {}
+}
+// src/knowledge-base/ScopedVectorStorage.ts
+class ScopedVectorStorage extends ScopedTabularStorage {
+  inner;
+  overFetchMultiplier;
+  constructor(inner, kbId, overFetchMultiplier = 3) {
+    super(inner, kbId);
+    this.inner = inner;
+    this.overFetchMultiplier = overFetchMultiplier;
+  }
+  getVectorDimensions() {
+    return this.inner.getVectorDimensions();
+  }
+  filterAndStrip(results, topK, overfetchLimit) {
+    const filtered = results.filter((r) => r.kb_id === this.kbId).slice(0, topK);
+    if (topK && overfetchLimit && results.length >= overfetchLimit && filtered.length < topK) {
+      console.warn(`ScopedVectorStorage: search returned ${filtered.length}/${topK} results after ` + `kb_id filtering. Consider increasing overFetchMultiplier (currently ${this.overFetchMultiplier}).`);
+    }
+    return filtered.map((r) => {
+      const { kb_id: _, ...rest } = r;
+      return rest;
+    });
+  }
+  async similaritySearch(query, options) {
+    const overfetchLimit = options?.topK ? options.topK * this.overFetchMultiplier : undefined;
+    const results = await this.inner.similaritySearch(query, {
+      ...options,
+      topK: overfetchLimit
+    });
+    return this.filterAndStrip(results, options?.topK, overfetchLimit);
+  }
+  async hybridSearch(query, options) {
+    if (typeof this.inner.hybridSearch !== "function") {
+      throw new Error("Hybrid search is not supported by the configured chunk storage backend. " + "Please use a vector storage implementation that provides `hybridSearch`.");
+    }
+    const overfetchLimit = options?.topK ? options.topK * this.overFetchMultiplier : undefined;
+    const results = await this.inner.hybridSearch(query, {
+      ...options,
+      topK: overfetchLimit
+    });
+    return this.filterAndStrip(results, options?.topK, overfetchLimit);
+  }
+}
 // src/util/DatasetSchema.ts
 function TypeTabularStorage(options = {}) {
   return {
@@ -1054,10 +1346,12 @@ class StructuralParser {
   }
 }
 export {
+  unregisterKnowledgeBase,
   traverseDepthFirst,
   setGlobalKnowledgeBaseRepository,
   registerKnowledgeBase,
   knowledgeBaseTableNames,
+  isSharedTableMode,
   hasChildren,
   getNodePath,
   getKnowledgeBase,
@@ -1066,14 +1360,25 @@ export {
   getDocumentRange,
   getChildren,
   estimateTokens,
+  deregisterKnowledgeBase,
   createKnowledgeBase,
   TypeTabularStorage,
   TypeKnowledgeBase,
   TopicNodeSchema,
   TokenBudgetSchema,
   StructuralParser,
+  SharedDocumentStorageSchema,
+  SharedDocumentPrimaryKey,
+  SharedDocumentIndexes,
+  SharedChunkVectorStorageSchema,
+  SharedChunkPrimaryKey,
+  SharedChunkIndexes,
   SentenceNodeSchema,
   SectionNodeSchema,
+  ScopedVectorStorage,
+  ScopedTabularStorage,
+  SHARED_DOCUMENT_TABLE,
+  SHARED_CHUNK_TABLE,
   ParagraphNodeSchema,
   NodeRangeSchema,
   NodeKind,
@@ -1099,4 +1404,4 @@ export {
   ChunkRecordArraySchema
 };
 
-//# debugId=E518FA75797167F364756E2164756E21
+//# debugId=EC2844C75568481264756E2164756E21