@axiom-lattice/core 2.1.3 → 2.1.4

package/dist/index.mjs

@@ -1478,6 +1478,407 @@ import {
 } from "langchain";
 import { Command as Command2, getCurrentTaskInput as getCurrentTaskInput2 } from "@langchain/langgraph";
 import { HumanMessage } from "@langchain/core/messages";
+
+// src/agent_worker/agent_worker_graph.ts
+import {
+  StateGraph,
+  END as END2,
+  interrupt as interrupt2,
+  Annotation
+} from "@langchain/langgraph";
+
+// src/constants/agent_task_types.ts
+var AGENT_TASK_EVENT = "agent:execute";
+
+// src/services/event_bus.ts
+import { EventEmitter } from "events";
+
+// src/queue_lattice/QueueLatticeManager.ts
+import {
+  QueueType
+} from "@axiom-lattice/protocols";
+
+// src/queue_lattice/MemoryQueueClient.ts
+var queues = /* @__PURE__ */ new Map();
+var ensureQueue = (queueName) => {
+  if (!queues.has(queueName)) {
+    queues.set(queueName, []);
+  }
+};
+var MemoryQueueClient = class {
+  constructor(queueName = "tasks") {
+    this.queueName = queueName;
+    ensureQueue(this.queueName);
+  }
+  /**
+   * Enqueue message (equivalent to lPush - push to left/beginning)
+   */
+  async push(item) {
+    try {
+      ensureQueue(this.queueName);
+      const queue = queues.get(this.queueName);
+      queue.unshift(item);
+      const result = queue.length;
+      console.log("lPush (memory)", result);
+      return { data: result, error: null };
+    } catch (error) {
+      console.error(error);
+      return { data: null, error };
+    }
+  }
+  /**
+   * Dequeue message (equivalent to rPop - pop from right/end)
+   */
+  async pop() {
+    try {
+      ensureQueue(this.queueName);
+      const queue = queues.get(this.queueName);
+      const message = queue.pop();
+      if (message) {
+        return { data: message, error: null };
+      }
+      return { data: null, error: null };
+    } catch (error) {
+      console.error(error);
+      return { data: null, error };
+    }
+  }
+  /**
+   * Create queue (initialize in memory)
+   */
+  async createQueue() {
+    try {
+      ensureQueue(this.queueName);
+      const queue = queues.get(this.queueName);
+      const exists = queue !== void 0;
+      return { success: true, queue_name: this.queueName };
+    } catch (error) {
+      console.error(error);
+      return { success: false, error };
+    }
+  }
+};
+
+// src/queue_lattice/QueueLatticeManager.ts
+var QueueLatticeManager = class _QueueLatticeManager extends BaseLatticeManager {
+  /**
+   * Get QueueLatticeManager singleton instance
+   */
+  static getInstance() {
+    if (!_QueueLatticeManager._instance) {
+      _QueueLatticeManager._instance = new _QueueLatticeManager();
+    }
+    return _QueueLatticeManager._instance;
+  }
+  /**
+   * Get Lattice type prefix
+   */
+  getLatticeType() {
+    return "queues";
+  }
+  /**
+   * Register queue Lattice
+   * @param key Lattice key name
+   * @param config Queue configuration
+   * @param client Optional queue client. If not provided, will create based on config type.
+   *               For REDIS type, client must be provided (use @axiom-lattice/queue-redis).
+   */
+  registerLattice(key, config, client) {
+    let queueClient;
+    if (client) {
+      queueClient = client;
+    } else {
+      const queueName = config.queueName || process.env.QUEUE_NAME || "tasks";
+      if (config.type === QueueType.REDIS) {
+        throw new Error(
+          `Redis queue client must be provided. Please install @axiom-lattice/queue-redis and pass the client to registerLattice. Example: registerLattice(key, config, new RedisQueueClient(queueName, options))`
+        );
+      } else {
+        queueClient = new MemoryQueueClient(queueName);
+      }
+    }
+    const queueLattice = {
+      key,
+      config,
+      client: queueClient,
+      push: async (item) => {
+        return queueClient.push(item);
+      },
+      pop: async () => {
+        return queueClient.pop();
+      },
+      createQueue: queueClient.createQueue ? async () => {
+        return queueClient.createQueue();
+      } : void 0
+    };
+    this.register(key, queueLattice);
+  }
+  /**
+   * Get QueueLattice
+   * @param key Lattice key name
+   */
+  getQueueLattice(key) {
+    const queueLattice = this.get(key);
+    if (!queueLattice) {
+      throw new Error(`QueueLattice ${key} not found`);
+    }
+    return queueLattice;
+  }
+  /**
+   * Get all Lattices
+   */
+  getAllLattices() {
+    return this.getAll();
+  }
+  /**
+   * Check if Lattice exists
+   * @param key Lattice key name
+   */
+  hasLattice(key) {
+    return this.has(key);
+  }
+  /**
+   * Remove Lattice
+   * @param key Lattice key name
+   */
+  removeLattice(key) {
+    return this.remove(key);
+  }
+  /**
+   * Clear all Lattices
+   */
+  clearLattices() {
+    this.clear();
+  }
+  /**
+   * Get Lattice count
+   */
+  getLatticeCount() {
+    return this.count();
+  }
+  /**
+   * Get Lattice key list
+   */
+  getLatticeKeys() {
+    return this.keys();
+  }
+};
+var queueLatticeManager = QueueLatticeManager.getInstance();
+var registerQueueLattice = (key, config, client) => queueLatticeManager.registerLattice(key, config, client);
+var getQueueLattice = (key) => queueLatticeManager.getQueueLattice(key);
+
+// src/services/event_bus.ts
+var EventBus = class {
+  constructor() {
+    this.defaultQueueKey = "default";
+    this.emitter = new EventEmitter();
+    this.emitter.setMaxListeners(100);
+  }
+  /**
+   * Set the default queue key for queue operations
+   * @param queueKey Queue key name
+   */
+  setDefaultQueueKey(queueKey) {
+    this.defaultQueueKey = queueKey;
+  }
+  /**
+   * Publish event
+   * @param eventName Event name
+   * @param data Event data
+   * @param useQueue Whether to use queue for publishing
+   */
+  publish(eventName, data, useQueue = false) {
+    if (useQueue) {
+      if (queueLatticeManager.hasLattice(this.defaultQueueKey)) {
+        const queue = queueLatticeManager.getQueueLattice(this.defaultQueueKey);
+        queue.push(data).catch((error) => {
+          console.error("Failed to push event to queue:", error);
+        });
+      } else {
+        console.warn(
+          `Queue "${this.defaultQueueKey}" not found. Event will not be queued.`
+        );
+      }
+    } else {
+      this.emitter.emit(eventName, data);
+    }
+  }
+  /**
+   * Subscribe to event
+   * @param eventName Event name
+   * @param callback Callback function
+   */
+  subscribe(eventName, callback) {
+    this.emitter.on(eventName, callback);
+  }
+  /**
+   * Unsubscribe from event
+   * @param eventName Event name
+   * @param callback Callback function
+   */
+  unsubscribe(eventName, callback) {
+    this.emitter.off(eventName, callback);
+  }
+  /**
+   * Subscribe to event once
+   * @param eventName Event name
+   * @param callback Callback function
+   */
+  subscribeOnce(eventName, callback) {
+    this.emitter.once(eventName, callback);
+  }
+};
+var eventBus = new EventBus();
+var event_bus_default = eventBus;
+
+// src/services/AgentManager.ts
+var AgentManager = class _AgentManager {
+  constructor() {
+    this.agents = [];
+  }
+  static getInstance() {
+    if (!_AgentManager.instance) {
+      _AgentManager.instance = new _AgentManager();
+    }
+    return _AgentManager.instance;
+  }
+  callAgentInQueue(queue, return_agent_state) {
+    return new Promise((resolve, reject) => {
+      const callback_event = `${queue.assistant_id}::${queue.thread_id}`;
+      if (return_agent_state) {
+        event_bus_default.subscribeOnce(callback_event, (data) => {
+          if (data.success) {
+            console.log("AgentManager callAgentInQueue success", data);
+            resolve(data.state);
+          } else {
+            console.log("AgentManager callAgentInQueue error", data);
+            reject(data.error);
+          }
+        });
+      }
+      try {
+        event_bus_default.publish(
+          AGENT_TASK_EVENT,
+          {
+            ...queue,
+            callback_event
+          },
+          true
+        );
+        !return_agent_state && resolve({ callback_event_id: callback_event, success: true });
+      } catch (error) {
+        !return_agent_state && reject({
+          callback_event_id: callback_event,
+          success: false,
+          error
+        });
+      }
+    });
+  }
+};
+
+// src/agent_worker/agent_worker_graph.ts
+var AgentWorkerState = Annotation.Root({
+  // Input parameters
+  assistant_id: Annotation({
+    reducer: (x, y) => y ?? x,
+    default: () => ""
+  }),
+  thread_id: Annotation({
+    reducer: (x, y) => y ?? x,
+    default: () => ""
+  }),
+  input: Annotation({
+    reducer: (x, y) => y ?? x,
+    default: () => ({})
+  }),
+  command2: Annotation({
+    reducer: (x, y) => y ?? x,
+    default: () => void 0
+  }),
+  // Execution result
+  result: Annotation({
+    reducer: (x, y) => y ?? x,
+    default: () => void 0
+  }),
+  // Interrupt data
+  interruptValue: Annotation({
+    reducer: (x, y) => y ?? x,
+    default: () => void 0
+  }),
+  // Final state to return
+  finalState: Annotation({
+    reducer: (x, y) => y ?? x,
+    default: () => void 0
+  })
+});
+async function executeNode(state) {
+  const { assistant_id, thread_id, input, command2 } = state;
+  const callParams = {
+    assistant_id,
+    thread_id,
+    input,
+    command: command2
+  };
+  const result = await AgentManager.getInstance().callAgentInQueue(
+    callParams,
+    true
+  );
+  return {
+    result
+  };
+}
+async function interruptNode(state) {
+  const { result } = state;
+  const interruptValue = result?.tasks?.[0]?.interrupts?.[0]?.value;
+  if (!interruptValue) {
+    throw new Error(
+      "Interrupt node called but no interrupt value found in result"
+    );
+  }
+  const response = await interrupt2(interruptValue);
+  return {
+    command2: {
+      resume: response
+    }
+  };
+}
+function shouldInterrupt(state) {
+  const { result } = state;
+  if (result?.tasks && Array.isArray(result.tasks) && result.tasks.length > 0) {
+    if (result.tasks[0]?.interrupts && result.tasks[0].interrupts.length > 0) {
+      return "interrupt";
+    }
+  }
+  return "end";
+}
+function endNode(state) {
+  const finalState = state.result;
+  return {
+    finalState,
+    // Clear intermediate states
+    result: void 0,
+    interruptValue: void 0,
+    command2: void 0
+  };
+}
+function createAgentWorkerGraph() {
+  const workflow = new StateGraph(AgentWorkerState);
+  workflow.addNode("execute", executeNode);
+  workflow.addNode("interrupt", interruptNode);
+  workflow.addNode("end", endNode);
+  workflow.setEntryPoint("execute");
+  workflow.addConditionalEdges("execute", shouldInterrupt, {
+    interrupt: "interrupt",
+    end: "end"
+  });
+  workflow.addEdge("interrupt", "execute");
+  workflow.addEdge("end", END2);
+  return workflow.compile();
+}
+var agentWorkerGraph = createAgentWorkerGraph();
+
+// src/deep_agent_new/middleware/subagents.ts
 var DEFAULT_SUBAGENT_PROMPT = "In order to complete the objective that the user asks of you, you have access to a number of standard tools.";
 var EXCLUDED_STATE_KEYS = ["messages", "todos", "jumpTo"];
 var DEFAULT_GENERAL_PURPOSE_DESCRIPTION = "General-purpose agent for researching complex questions, searching for files and content, and executing multi-step tasks. When you are searching for a keyword or file and are not confident that you will find the right match in the first few tries use this agent to perform the search for you. This agent has access to all tools as the main agent.";
@@ -1730,7 +2131,15 @@ function createTaskTool(options) {
       const currentState = getCurrentTaskInput2();
       const subagentState = filterStateForSubagent(currentState);
       subagentState.messages = [new HumanMessage({ content: description })];
-      const result = await subagent.invoke(subagentState, config);
+      const subagent_thread_id = config.configurable?.thread_id + "_" + subagent_type + "_" + config.toolCall.id;
+      console.log(subagent_thread_id);
+      const workerResult = await agentWorkerGraph.invoke({
+        assistant_id: subagent_type,
+        thread_id: subagent_thread_id,
+        input: { ...subagentState, message: description }
+      });
+      const result = workerResult.finalState?.values;
+      console.log("workerResult", workerResult);
       if (!config.toolCall?.id) {
         throw new Error("Tool call ID is required for subagent invocation");
       }
@@ -2840,20 +3249,26 @@ var hasChunkBuffer = (key) => ChunkBufferLatticeManager.getInstance().has(key);
 // src/index.ts
 import * as Protocols from "@axiom-lattice/protocols";
 export {
+  AGENT_TASK_EVENT,
   AgentConfig,
   AgentLatticeManager,
+  AgentManager,
   AgentType,
   ChunkBuffer,
   ChunkBufferLatticeManager,
   GraphBuildOptions,
   InMemoryChunkBuffer,
   MemoryLatticeManager,
+  MemoryQueueClient,
   MemoryType,
   ModelLatticeManager,
   Protocols,
+  QueueLatticeManager,
   ThreadStatus,
   ToolLatticeManager,
   agentLatticeManager,
+  eventBus,
+  event_bus_default as eventBusDefault,
   getAgentClient,
   getAgentConfig,
   getAgentLattice,
@@ -2862,16 +3277,19 @@ export {
   getCheckpointSaver,
   getChunkBuffer,
   getModelLattice,
+  getQueueLattice,
   getToolClient,
   getToolDefinition,
   getToolLattice,
   hasChunkBuffer,
   modelLatticeManager,
+  queueLatticeManager,
   registerAgentLattice,
   registerAgentLattices,
   registerCheckpointSaver,
   registerChunkBuffer,
   registerModelLattice,
+  registerQueueLattice,
   registerToolLattice,
   toolLatticeManager,
   validateAgentInput,