goblin-agents 0.1.2 → 0.1.4

package/lib/llm/aiAgent.js

@@ -11,6 +11,7 @@ const {
   number,
   record,
 } = require('xcraft-core-stones');
+const {getMutex} = require('xcraft-core-utils/lib/locks.js');
 const LLMProvider = require('./providers.js');
 const {OCR_SYSTEM_PROMPT, getPDFImages} = require('./utils.js');
 
@@ -86,6 +87,17 @@ class OptionsShape {
   stop = option(array(string));
 }
 
+function vectorToSqlLiteral(vec) {
+  const array = new Float32Array(vec);
+  const view = new Uint8Array(array.buffer, array.byteOffset, array.byteLength);
+  let hex = '';
+  for (let i = 0; i < view.length; i++) {
+    const byte = view[i];
+    hex += byte < 16 ? '0' + byte.toString(16) : byte.toString(16);
+  }
+  return hex;
+}
+
 class AiAgentShape {
   id = id('aiAgent');
   version = number;
@@ -240,6 +252,7 @@ class AiAgent extends Elf {
 
   _provider;
   _conf;
+  _batchLock = getMutex;
 
   /**
    * @param {string} id
@@ -276,18 +289,19 @@ class AiAgent extends Elf {
   async _loadProvider() {
     const {state} = this;
     const {provider, host, headers} = state;
+    const headersStr = JSON.stringify(headers);
     if (
       this._provider &&
       this._conf.provider === provider &&
       this._conf.host === host &&
-      this._conf.headers === headers
+      this._conf.headersStr === headersStr
     ) {
       return;
     }
 
     const Instance = LLMProvider.get(this.state.provider);
     this._provider = new Instance(this.state.host, this.state.headers);
-    this._conf = {provider, host, headers};
+    this._conf = {provider, host, headersStr};
   }
 
   /**
@@ -378,21 +392,27 @@ class AiAgent extends Elf {
       return null;
     }
     await this._loadProvider();
-    return await this._provider.embed({
+    const vectors = await this._provider.embed({
       model: this.state.model,
       input: rawText,
     });
+    return vectorToSqlLiteral(vectors);
   }
 
   async embedInBatch(rawTexts = []) {
     if (rawTexts.length === 0) {
       return [];
     }
+
+    await this._batchLock.lock(this.id);
+    this.quest.defer(() => this._batchLock.unlock(this.id));
+
     await this._loadProvider();
-    return await this._provider.embedInBatch({
+    const res = await this._provider.embedInBatch({
       model: this.state.model,
       input: rawTexts,
     });
+    return res.map((v) => vectorToSqlLiteral(v));
   }
 
   async set(contextId, messages) {

package/lib/llm/utils.js

@@ -276,7 +276,60 @@ const buildPromptFromArticles = (articles, question, maxTokens) => {
   return context.trim();
 };
 
+/**
+ * Split an array into smaller chunks of a given size.
+ *
+ * @param {Array} array - The array to be chunked.
+ * @param {number} size - The size of each chunk.
+ * @returns {Array[]} An array of chunked subarrays.
+ */
+function chunkArray(array, size) {
+  const chunks = [];
+  for (let startIndex = 0; startIndex < array.length; startIndex += size) {
+    chunks.push(array.slice(startIndex, startIndex + size));
+  }
+  return chunks;
+}
+
+/**
+ * Perform batched embedding of content chunks, associating each embedding
+ * with its original contentId and key.
+ *
+ * @param {object} agent - The AI agent instance with an embedInBatch method.
+ * @param {Array<{ contentId: string, key: string, chunk: string }>} allChunks - Array of content to embed.
+ * @param {number} [batchSize=500] - Maximum number of chunks per embedding request.
+ * @returns {Promise<object>} A map of embeddings structured as:
+ *    { [contentId]: { [key]: { chunk, embedding } } }
+ */
+async function embedChunksInBatch(agent, allChunks, batchSize = 500) {
+  const embeddingsByContent = {};
+
+  for (let start = 0; start < allChunks.length; start += batchSize) {
+    const currentBatch = allChunks.slice(start, start + batchSize);
+    const inputTexts = currentBatch.map(({chunk}) => chunk);
+    const embeddedVectors = await agent.embedInBatch(inputTexts);
+
+    for (let i = 0; i < embeddedVectors.length; i++) {
+      const {contentId, key, chunk} = currentBatch[i];
+      const embedding = embeddedVectors[i];
+
+      if (!embeddingsByContent[contentId]) {
+        embeddingsByContent[contentId] = {};
+      }
+
+      embeddingsByContent[contentId][key] = {
+        chunk,
+        embedding,
+      };
+    }
+  }
+
+  return embeddingsByContent;
+}
+
 module.exports = {
+  chunkArray,
+  embedChunksInBatch,
   extractTitlesAndContent,
   html2chunks,
   splitLongText,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "goblin-agents",
-  "version": "0.1.2",
+  "version": "0.1.4",
   "description": "AI Agents framework",
   "author": "Samuel Loup",
   "license": "MIT",