@axiom-lattice/gateway 2.1.3 → 2.1.5

package/dist/index.js

@@ -534,6 +534,305 @@ var getAgentGraph = async (request, reply) => {
   }
 };
 
+// src/services/agent_task_types.ts
+var AGENT_TASK_EVENT = "agent:execute";
+
+// src/services/event_bus.ts
+var import_events = require("events");
+
+// src/services/queue_service_memory.ts
+var queue_service_memory_exports = {};
+__export(queue_service_memory_exports, {
+  popAgentTaskFromQueue: () => popAgentTaskFromQueue,
+  pushAgentTaskToQueue: () => pushAgentTaskToQueue
+});
+var queue_name = process.env.QUEUE_NAME || "tasks";
+var queues = /* @__PURE__ */ new Map();
+var ensureQueue = (queueName) => {
+  if (!queues.has(queueName)) {
+    queues.set(queueName, []);
+  }
+};
+var sendToQueue = async (queueName, message) => {
+  try {
+    ensureQueue(queueName);
+    const queue = queues.get(queueName);
+    queue.unshift(message);
+    const result = queue.length;
+    console.log("lPush (memory)", result);
+    return { data: result, error: null };
+  } catch (error) {
+    console.error(error);
+    return { data: null, error };
+  }
+};
+var popFromQueue = async (queueName) => {
+  try {
+    ensureQueue(queueName);
+    const queue = queues.get(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 };
+  }
+};
+var createTenantQueue = async () => {
+  try {
+    ensureQueue(queue_name);
+    const queue = queues.get(queue_name);
+    const exists = queue !== void 0;
+    return { success: true, queue_name };
+  } catch (error) {
+    console.error(error);
+    return { success: false, error };
+  }
+};
+var pushAgentTaskToQueue = async (agentTask) => {
+  const tenantId = agentTask["x-tenant-id"];
+  try {
+    await createTenantQueue();
+    const result = await sendToQueue(`${queue_name}`, agentTask);
+    return result;
+  } catch (error) {
+    console.error(error);
+    return null;
+  }
+};
+var popAgentTaskFromQueue = async () => {
+  const result = await popFromQueue(`${queue_name}`);
+  return result;
+};
+
+// src/services/queue_service_redis.ts
+var queue_service_redis_exports = {};
+__export(queue_service_redis_exports, {
+  popAgentTaskFromQueue: () => popAgentTaskFromQueue2,
+  pushAgentTaskToQueue: () => pushAgentTaskToQueue2
+});
+var import_redis = require("redis");
+var queue_name2 = process.env.QUEUE_NAME || "tasks";
+var redisClient = null;
+var isConnecting = false;
+var connectionPromise = null;
+var getRedisClient = async () => {
+  if (redisClient?.isOpen) {
+    return redisClient;
+  }
+  if (isConnecting && connectionPromise) {
+    await connectionPromise;
+    if (redisClient?.isOpen) {
+      return redisClient;
+    }
+  }
+  isConnecting = true;
+  connectionPromise = (async () => {
+    try {
+      const redisOptions = {
+        url: process.env.REDIS_URL || "redis://localhost:6379",
+        password: process.env.REDIS_PASSWORD
+      };
+      redisClient = (0, import_redis.createClient)(redisOptions);
+      redisClient.on(
+        "error",
+        (err) => console.error("Redis Client Error", err)
+      );
+      await redisClient.connect();
+      console.log("Redis connection successful");
+      isConnecting = false;
+    } catch (error) {
+      console.error("Redis connection failed:", error);
+      isConnecting = false;
+      throw error;
+    }
+  })();
+  await connectionPromise;
+  return redisClient;
+};
+var sendToQueue2 = async (queueName, message) => {
+  try {
+    const client = await getRedisClient();
+    const result = await client.lPush(queueName, JSON.stringify(message));
+    console.log("lPush", result);
+    return { data: result, error: null };
+  } catch (error) {
+    console.error(error);
+    return { data: null, error };
+  }
+};
+var popFromQueue2 = async (queueName) => {
+  try {
+    const client = await getRedisClient();
+    const message = await client.rPop(queueName);
+    if (message) {
+      return { data: JSON.parse(message), error: null };
+    }
+    return { data: null, error: null };
+  } catch (error) {
+    console.error(error);
+    return { data: null, error };
+  }
+};
+var createTenantQueue2 = async () => {
+  try {
+    const client = await getRedisClient();
+    const exists = await client.exists(queue_name2);
+    return { success: true, queue_name: queue_name2 };
+  } catch (error) {
+    console.error(error);
+    return { success: false, error };
+  }
+};
+var pushAgentTaskToQueue2 = async (agentTask) => {
+  const tenantId = agentTask["x-tenant-id"];
+  try {
+    await createTenantQueue2();
+    const result = await sendToQueue2(`${queue_name2}`, agentTask);
+    return result;
+  } catch (error) {
+    console.error(error);
+    return null;
+  }
+};
+var popAgentTaskFromQueue2 = async () => {
+  const result = await popFromQueue2(`${queue_name2}`);
+  return result;
+};
+
+// src/services/queue_service.ts
+var queueServiceType = process.env.QUEUE_SERVICE_TYPE || "memory";
+var setQueueServiceType = (type) => {
+  queueServiceType = type;
+  console.log(`Queue service type set to: ${type}`);
+};
+var getQueueService = () => {
+  return queueServiceType === "redis" ? queue_service_redis_exports : queue_service_memory_exports;
+};
+var pushAgentTaskToQueue3 = async (agentTask) => {
+  const service = getQueueService();
+  return service.pushAgentTaskToQueue(agentTask);
+};
+var popAgentTaskFromQueue3 = async () => {
+  const service = getQueueService();
+  return service.popAgentTaskFromQueue();
+};
+
+// src/services/event_bus.ts
+var EventBus = class {
+  constructor() {
+    this.emitter = new import_events.EventEmitter();
+    this.emitter.setMaxListeners(100);
+  }
+  /**
+   * 发布事件
+   * @param eventName 事件名称
+   * @param data 事件数据
+   */
+  publish(eventName, data, useQueue = false) {
+    if (useQueue) {
+      pushAgentTaskToQueue3(data);
+    } else {
+      this.emitter.emit(eventName, data);
+    }
+  }
+  /**
+   * 订阅事件
+   * @param eventName 事件名称
+   * @param callback 回调函数
+   */
+  subscribe(eventName, callback) {
+    this.emitter.on(eventName, callback);
+  }
+  /**
+   * 取消订阅事件
+   * @param eventName 事件名称
+   * @param callback 回调函数
+   */
+  unsubscribe(eventName, callback) {
+    this.emitter.off(eventName, callback);
+  }
+  /**
+   * 只订阅一次事件
+   * @param eventName 事件名称
+   * @param callback 回调函数
+   */
+  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 new Promise((resolve, reject) => {
+      try {
+        event_bus_default.publish(
+          AGENT_TASK_EVENT,
+          {
+            ...queue
+          },
+          true
+        );
+        resolve(true);
+      } catch (error) {
+        reject(error);
+      }
+    });
+  }
+};
+
+// src/controllers/agent_task.ts
+var triggerAgentTask = async (request, reply) => {
+  try {
+    const { assistant_id, thread_id, input, command } = request.body;
+    const tenant_id = request.headers["x-tenant-id"];
+    if (!assistant_id) {
+      reply.status(400).send({
+        success: false,
+        error: "assistant_id is required"
+      });
+      return;
+    }
+    if (!thread_id) {
+      reply.status(400).send({
+        success: false,
+        error: "thread_id is required"
+      });
+      return;
+    }
+    const agentManager = AgentManager.getInstance();
+    const result = await agentManager.callAgentInQueue({
+      assistant_id,
+      thread_id,
+      input,
+      command,
+      "x-tenant-id": tenant_id
+    });
+    reply.status(200).send({
+      success: true
+    });
+  } catch (error) {
+    reply.status(500).send({
+      success: false,
+      error: `Failed to trigger agent task: ${error.message}`
+    });
+  }
+};
+
 // src/schemas/index.ts
 var getAllMemoryItemsSchema = {
   description: "Get all memory items for an assistant thread",
@@ -651,6 +950,57 @@ var getAgentGraphSchema = {
     200: {}
   }
 };
+var triggerAgentTaskSchema = {
+  description: "Trigger an agent task",
+  tags: ["Agent Tasks"],
+  summary: "Trigger Agent Task",
+  body: {
+    type: "object",
+    properties: {
+      assistant_id: {
+        type: "string",
+        description: "Assistant ID"
+      },
+      thread_id: {
+        type: "string",
+        description: "Thread ID"
+      },
+      input: {
+        type: "object",
+        description: "Task input data"
+      },
+      command: {
+        type: "object",
+        description: "Command data for the task",
+        nullable: true
+      }
+    },
+    required: ["assistant_id", "thread_id"]
+  },
+  response: {
+    200: {
+      type: "object",
+      properties: {
+        success: { type: "boolean" },
+        result: { type: "object" }
+      }
+    },
+    400: {
+      type: "object",
+      properties: {
+        success: { type: "boolean" },
+        error: { type: "string" }
+      }
+    },
+    500: {
+      type: "object",
+      properties: {
+        success: { type: "boolean" },
+        error: { type: "string" }
+      }
+    }
+  }
+};
 
 // src/routes/index.ts
 var registerLatticeRoutes = (app2) => {
@@ -691,6 +1041,11 @@ var registerLatticeRoutes = (app2) => {
     { schema: getAgentGraphSchema },
     getAgentGraph
   );
+  app2.post(
+    "/api/agent-tasks/trigger",
+    { schema: triggerAgentTaskSchema },
+    triggerAgentTask
+  );
 };
 
 // src/logger/Logger.ts
@@ -893,6 +1248,244 @@ var configureSwagger = async (app2, customSwaggerConfig, customSwaggerUiConfig)
   await app2.register(import_swagger_ui.default, swaggerUiConfig);
 };
 
+// src/services/agent_task_consumer.ts
+var handleAgentTask = async (taskRequest, retryCount = 0) => {
+  const {
+    assistant_id,
+    input = {},
+    thread_id,
+    "x-tenant-id": tenant_id,
+    command,
+    callback_event
+  } = taskRequest;
+  try {
+    console.log(
+      `\u5F00\u59CB\u5904\u7406\u4EFB\u52A1 [assistant_id: ${assistant_id}, thread_id: ${thread_id}]`
+    );
+    const apiUrl = AgentTaskConsumer.agent_run_endpoint;
+    console.log(`apiUrl: ${apiUrl}`);
+    const response = await fetch(apiUrl, {
+      method: "POST",
+      body: JSON.stringify({
+        assistant_id,
+        streaming: true,
+        ...input,
+        thread_id,
+        command
+      }),
+      headers: {
+        "Content-Type": "application/json",
+        "x-tenant-id": tenant_id
+      }
+    }).catch((err) => {
+      console.error(`fetch\u8BF7\u6C42\u5931\u8D25: ${err.message || String(err)}`);
+      throw new Error(`fetch\u5931\u8D25: ${err.message || String(err)}`);
+    });
+    if (!response.ok) {
+      throw new Error(`API\u8BF7\u6C42\u5931\u8D25: ${response.status} ${response.statusText}`);
+    }
+    if (callback_event) {
+      event_bus_default.publish(callback_event, {
+        success: true,
+        config: { assistant_id, thread_id, tenant_id }
+      });
+    }
+    console.log(
+      `\u4EFB\u52A1\u5904\u7406\u6210\u529F [assistant_id: ${assistant_id}, thread_id: ${thread_id}]`
+    );
+    return true;
+  } catch (error) {
+    console.error(
+      `Agent\u4EFB\u52A1\u6267\u884C\u5931\u8D25: ${assistant_id}, \u7EBF\u7A0B: ${thread_id}`,
+      error
+    );
+    const maxRetries = 0;
+    if (retryCount < maxRetries) {
+      const nextRetryCount = retryCount + 1;
+      const delayMs = Math.pow(2, nextRetryCount) * 1e3;
+      console.log(
+        `\u5C06\u5728 ${delayMs}ms \u540E\u91CD\u8BD5\u4EFB\u52A1 (${nextRetryCount}/${maxRetries})`
+      );
+      await new Promise((resolve) => setTimeout(resolve, delayMs));
+      return handleAgentTask(taskRequest, nextRetryCount);
+    }
+    if (callback_event) {
+      event_bus_default.publish(callback_event, {
+        success: false,
+        error: error instanceof Error ? error.message : String(error),
+        config: { assistant_id, thread_id, tenant_id }
+      });
+    }
+    console.error(
+      `\u4EFB\u52A1\u5904\u7406\u5931\u8D25\uFF0C\u5DF2\u8FBE\u5230\u6700\u5927\u91CD\u8BD5\u6B21\u6570 [assistant_id: ${assistant_id}, thread_id: ${thread_id}]`
+    );
+    return false;
+  }
+};
+var _AgentTaskConsumer = class _AgentTaskConsumer {
+  constructor(gatewayPort, pollingIntervalMs) {
+    this.isPolling = false;
+    this.pollingInterval = null;
+    this.pollingIntervalMs = 5e3;
+    // 默认5秒轮询一次
+    this.maxConcurrentTasks = 15;
+    // 最大并发任务数
+    this.activeTasks = 0;
+    // 当前活跃的任务数
+    this.processing = false;
+    // 是否正在处理任务批次
+    this.immediateProcessingEnabled = true;
+    // 是否启用即时处理模式
+    this.gatewayPort = 4001;
+    this.gatewayPort = gatewayPort;
+    _AgentTaskConsumer.agent_run_endpoint = `http://localhost:${this.gatewayPort}/api/runs`;
+    if (pollingIntervalMs) {
+      this.pollingIntervalMs = pollingIntervalMs;
+    }
+    this.initialize();
+  }
+  /**
+   * 初始化事件监听和队列轮询
+   */
+  initialize() {
+    event_bus_default.subscribe(AGENT_TASK_EVENT, this.trigger_agent_task.bind(this));
+    this.startPollingQueue();
+    console.log("Agent\u4EFB\u52A1\u6D88\u8D39\u8005\u5DF2\u542F\u52A8\u5E76\u76D1\u542C\u4EFB\u52A1\u4E8B\u4EF6\u548C\u961F\u5217");
+  }
+  /**
+   * 启动队列轮询
+   */
+  startPollingQueue() {
+    if (this.isPolling) {
+      return;
+    }
+    this.isPolling = true;
+    this.pollingInterval = setInterval(async () => {
+      try {
+        if (this.processing) {
+          console.log("\u961F\u5217\u5904\u7406\u4E2D\uFF0C\u8DF3\u8FC7\u672C\u6B21\u8F6E\u8BE2");
+          return;
+        }
+        await this.consumeFromQueue();
+      } catch (error) {
+        console.error("\u961F\u5217\u8F6E\u8BE2\u51FA\u9519:", error);
+      }
+    }, this.pollingIntervalMs);
+    console.log(
+      `\u5F00\u59CB\u8F6E\u8BE2\u961F\u5217\uFF0C\u95F4\u9694: ${this.pollingIntervalMs}ms\uFF0C\u6700\u5927\u5E76\u53D1\u4EFB\u52A1\u6570: ${this.maxConcurrentTasks}`
+    );
+  }
+  /**
+   * 停止队列轮询
+   */
+  stopPollingQueue() {
+    if (this.pollingInterval) {
+      clearInterval(this.pollingInterval);
+      this.pollingInterval = null;
+      this.isPolling = false;
+      console.log("\u5DF2\u505C\u6B62\u961F\u5217\u8F6E\u8BE2");
+    }
+  }
+  /**
+   * 处理单个任务并在完成后立即尝试处理下一个
+   */
+  async processNextTask() {
+    try {
+      const queueResult = await popAgentTaskFromQueue3();
+      if (queueResult && queueResult.data) {
+        const taskItem = queueResult.data;
+        if (taskItem && typeof taskItem === "object") {
+          const taskRequest = taskItem;
+          console.log(
+            `\u4ECE\u961F\u5217\u4E2D\u83B7\u53D6\u5230\u4EFB\u52A1 [assistant: ${taskRequest.assistant_id}, thread: ${taskRequest.thread_id}]`
+          );
+          this.activeTasks++;
+          handleAgentTask(taskRequest).then((success) => {
+            if (!success) {
+              console.error(`\u4EFB\u52A1 \u5904\u7406\u5931\u8D25`);
+            } else {
+              console.log(`\u4EFB\u52A1 \u5904\u7406\u6210\u529F`);
+            }
+          }).catch((error) => {
+            console.error(`\u4EFB\u52A1 \u5904\u7406\u65F6\u51FA\u9519:`, error);
+          }).finally(() => {
+            this.activeTasks--;
+            if (this.immediateProcessingEnabled && !this.processing) {
+              this.checkQueueForTasks();
+            }
+          });
+          return true;
+        } else {
+          console.log("\u961F\u5217\u4EFB\u52A1\u683C\u5F0F\u65E0\u6548:", taskItem);
+          return false;
+        }
+      } else {
+        return false;
+      }
+    } catch (error) {
+      console.error("\u5904\u7406\u4EFB\u52A1\u5931\u8D25:", error);
+      return false;
+    }
+  }
+  /**
+   * 检查队列中是否有任务并处理
+   */
+  async checkQueueForTasks() {
+    if (this.processing || this.activeTasks >= this.maxConcurrentTasks) {
+      return;
+    }
+    this.processing = true;
+    try {
+      while (this.activeTasks < this.maxConcurrentTasks) {
+        const taskProcessed = await this.processNextTask();
+        if (!taskProcessed) {
+          break;
+        }
+      }
+    } catch (error) {
+      console.error("\u68C0\u67E5\u961F\u5217\u4EFB\u52A1\u5931\u8D25:", error);
+    } finally {
+      this.processing = false;
+    }
+  }
+  /**
+   * 从队列中消费任务
+   */
+  async consumeFromQueue() {
+    if (this.processing) {
+      return;
+    }
+    await this.checkQueueForTasks();
+  }
+  /**
+   * 处理通过事件触发的任务
+   */
+  trigger_agent_task(taskRequest) {
+    console.log(
+      `\u901A\u8FC7\u4E8B\u4EF6\u89E6\u53D1\u4EFB\u52A1: [assistant: ${taskRequest.assistant_id}, thread: ${taskRequest.thread_id}]`
+    );
+    handleAgentTask(taskRequest).catch((error) => {
+      console.error("\u5904\u7406Agent\u4EFB\u52A1\u65F6\u53D1\u751F\u672A\u6355\u83B7\u7684\u9519\u8BEF:", error);
+      if (taskRequest.callback_event) {
+        event_bus_default.publish(taskRequest.callback_event, {
+          success: false,
+          error: error instanceof Error ? error.message : String(error),
+          config: {
+            assistant_id: taskRequest.assistant_id,
+            thread_id: taskRequest.thread_id,
+            tenant_id: taskRequest["x-tenant-id"]
+          }
+        });
+      }
+    });
+    if (this.immediateProcessingEnabled && this.activeTasks < this.maxConcurrentTasks) {
+      setImmediate(() => this.checkQueueForTasks());
+    }
+  }
+};
+_AgentTaskConsumer.agent_run_endpoint = "http://localhost:4001/api/runs";
+var AgentTaskConsumer = _AgentTaskConsumer;
+
 // src/index.ts
 process.on("unhandledRejection", (reason, promise) => {
   console.error("\u672A\u5904\u7406\u7684Promise\u62D2\u7EDD:", reason);
@@ -953,11 +1546,19 @@ app.setErrorHandler((error, request, reply) => {
   });
 });
 app.decorate("logger", logger);
-var start = async ({ port }) => {
+var start = async (config) => {
   try {
-    const target_port = port || Number(process.env.PORT) || 4001;
+    const target_port = config?.port || Number(process.env.PORT) || 4001;
     await app.listen({ port: target_port, host: "0.0.0.0" });
-    logger.info(`Lattice Gateway is running on port: ${port}`);
+    logger.info(`Lattice Gateway is running on port: ${target_port}`);
+    const queueServiceConfig = config?.queueServiceConfig;
+    if (queueServiceConfig) {
+      setQueueServiceType(queueServiceConfig.type);
+      if (queueServiceConfig.defaultStartPollingQueue) {
+        const agentTaskConsumer = new AgentTaskConsumer(target_port);
+        agentTaskConsumer.startPollingQueue();
+      }
+    }
   } catch (err) {
     logger.error("Server start failed", { error: err });
     process.exit(1);
@@ -967,7 +1568,8 @@ var LatticeGateway = {
   startAsHttpEndpoint: start,
   configureSwagger,
   registerLatticeRoutes,
-  app
+  app,
+  AgentTaskConsumer
 };
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {