@axiom-lattice/gateway 2.1.44 → 2.1.46

package/dist/index.mjs

@@ -7,279 +7,10 @@ import websocket from "@fastify/websocket";
 import staticPlugin from "@fastify/static";
 import path2 from "path";
 
-// src/services/agent_service.ts
-import {
-  filterMessages,
-  HumanMessage
-} from "@langchain/core/messages";
-import { Command } from "@langchain/langgraph";
-import { v4 } from "uuid";
-import {
-  getAgentClient,
-  agentLatticeManager,
-  InMemoryChunkBuffer,
-  registerChunkBuffer,
-  getChunkBuffer,
-  hasChunkBuffer
-} from "@axiom-lattice/core";
-async function checkAgentExists(tenant_id, assistant_id) {
-  try {
-    const agentLattice = agentLatticeManager.getAgentLatticeWithTenant(tenant_id, assistant_id);
-    return agentLattice !== void 0;
-  } catch {
-    return false;
-  }
-}
-function getOrCreateChunkBuffer() {
-  if (!hasChunkBuffer("default")) {
-    const buffer = new InMemoryChunkBuffer({
-      ttl: 60 * 60 * 1e3,
-      // 1 hour TTL
-      cleanupInterval: 5 * 60 * 1e3
-      // Clean every 5 minutes
-    });
-    registerChunkBuffer("default", buffer);
-  }
-  return getChunkBuffer("default");
-}
-async function agent_invoke({
-  input,
-  thread_id,
-  assistant_id,
-  tenant_id,
-  workspace_id,
-  project_id,
-  command,
-  run_id,
-  custom_run_config
-}) {
-  const agentLattice = agentLatticeManager.getAgentLatticeWithTenant(tenant_id, assistant_id);
-  const runnable_agent = agentLattice?.client;
-  const { message, ...rest } = input;
-  const humanMessage = new HumanMessage(message || "");
-  const messages = [humanMessage];
-  if (!runnable_agent) {
-    throw new Error(`Agent ${assistant_id} not found`);
-  }
-  const runConfig = {
-    ...agentLattice?.config?.runConfig || {},
-    tenantId: tenant_id,
-    workspaceId: workspace_id,
-    projectId: project_id,
-    ...custom_run_config || {},
-    assistant_id
-  };
-  const result = await runnable_agent.invoke(
-    command ? new Command(command) : { ...rest, messages, "x-tenant-id": tenant_id },
-    {
-      context: {
-        runConfig
-      },
-      configurable: {
-        thread_id,
-        run_id: run_id || v4(),
-        "x-tenant-id": tenant_id,
-        "x-workspace-id": workspace_id,
-        "x-project-id": project_id,
-        "x-request-id": run_id,
-        "x-thread-id": thread_id,
-        "x-assistant-id": assistant_id,
-        runConfig
-      },
-      recursionLimit: 200
-    }
-  );
-  const data = result.messages.map((message2) => {
-    const { type, data: data2 } = message2.toDict();
-    return {
-      ...data2,
-      role: type
-    };
-  });
-  return { messages: data };
-}
-async function agent_stream({
-  input,
-  thread_id,
-  command,
-  tenant_id,
-  workspace_id,
-  project_id,
-  assistant_id,
-  run_id,
-  custom_run_config
-}) {
-  const runnable_agent = await getAgentClient(tenant_id, assistant_id);
-  const agentLattice = agentLatticeManager.getAgentLatticeWithTenant(tenant_id, assistant_id);
-  const { message, ...rest } = input;
-  let messages = [];
-  if (!command) {
-    const humanMessage = new HumanMessage(message);
-    messages = [humanMessage];
-  }
-  const chunkBuffer = getOrCreateChunkBuffer();
-  const runConfig = {
-    ...agentLattice?.config?.runConfig || {},
-    tenantId: tenant_id,
-    workspaceId: workspace_id,
-    projectId: project_id,
-    ...custom_run_config || {},
-    assistant_id
-  };
-  try {
-    if (!runnable_agent) {
-      throw new Error(`Agent ${assistant_id} not found`);
-    }
-    const agentStream = await runnable_agent.stream(
-      command ? new Command(command) : {
-        ...rest,
-        messages,
-        "x-tenant-id": tenant_id
-      },
-      {
-        context: {
-          runConfig
-        },
-        configurable: {
-          thread_id,
-          run_id: run_id || v4(),
-          "x-tenant-id": tenant_id,
-          "x-workspace-id": workspace_id,
-          "x-project-id": project_id,
-          "x-request-id": run_id,
-          "x-thread-id": thread_id,
-          "x-assistant-id": assistant_id,
-          runConfig
-          // Inject runConfig for tools to access
-        },
-        streamMode: ["updates", "messages"],
-        subgraphs: false,
-        recursionLimit: 200
-      }
-    );
-    return {
-      [Symbol.asyncIterator]: async function* () {
-        try {
-          for await (const chunk of agentStream) {
-            let data;
-            let chunkContent = "";
-            if (chunk[0] === "updates") {
-              const update = chunk[1];
-              const values = Object.values(update);
-              const messages2 = values[0]?.messages;
-              if (messages2?.[0]?.tool_call_id) {
-                data = messages2[0].toDict();
-              }
-            } else if (chunk[0] === "messages") {
-              const messages2 = chunk[1];
-              data = messages2?.[0]?.toDict();
-            }
-            if (chunk?.[1]?.__interrupt__) {
-              data = {
-                type: "interrupt",
-                id: chunk?.[1]?.__interrupt__[0].id,
-                data: { content: chunk?.[1]?.__interrupt__[0].value }
-              };
-            }
-            if (data) {
-              if (data.type !== "interrupt") {
-                await chunkBuffer.addChunk(thread_id, data);
-              }
-              yield data;
-            }
-          }
-          await chunkBuffer.completeThread(thread_id);
-        } catch (error) {
-          console.error("Stream error:", error);
-          await chunkBuffer.abortThread(thread_id);
-          throw error;
-        }
-      }
-    };
-  } catch (error) {
-    await chunkBuffer.abortThread(thread_id);
-    throw error;
-  }
-}
-async function agent_state({
-  assistant_id,
-  thread_id,
-  tenant_id
-}) {
-  const runnable_agent = await getAgentClient(tenant_id, assistant_id);
-  if (!runnable_agent) {
-    throw new Error(`Agent ${assistant_id} not found`);
-  }
-  const state = await runnable_agent.getState({
-    configurable: { thread_id, subgraphs: false }
-  });
-  return state;
-}
-async function agent_messages({
-  thread_id,
-  tenant_id,
-  assistant_id
-}) {
-  const runnable_agent = await getAgentClient(tenant_id, assistant_id);
-  if (!runnable_agent) {
-    throw new Error(`Agent ${assistant_id} not found`);
-  }
-  const state = await runnable_agent.getState({
-    configurable: { thread_id, subgraphs: false }
-  });
-  const messages = state.values.messages || [];
-  const filteredMessages = filterMessages(messages, {
-    includeTypes: ["ai", "human", "tool"]
-    //["human", "ai", "tool"],
-  });
-  let messagesArray = filteredMessages.map((message) => ({
-    id: message.id,
-    role: message.getType(),
-    content: message.content,
-    ...message.lc_kwargs
-  }));
-  const new_messages = messagesArray;
-  return new_messages;
-}
-async function draw_graph(assistant_id, tenant_id) {
-  const runnable_agent = await getAgentClient(tenant_id, assistant_id);
-  if (!runnable_agent) {
-    throw new Error(`Agent ${assistant_id} not found`);
-  }
-  const drawableGraph = await runnable_agent.getGraphAsync();
-  const image = await drawableGraph.drawMermaid();
-  return image;
-}
-async function resume_stream({
-  thread_id,
-  message_id,
-  known_content,
-  poll_interval = 100
-}) {
-  const chunkBuffer = getOrCreateChunkBuffer();
-  const stream = await chunkBuffer.getNewChunksSinceContentIterator(
-    thread_id,
-    message_id,
-    known_content
-  );
-  return {
-    [Symbol.asyncIterator]: async function* () {
-      try {
-        for await (const chunk of stream) {
-          yield chunk;
-        }
-      } catch (error) {
-        console.error("Resume stream error:", error);
-        throw error;
-      }
-    }
-  };
-}
-
 // src/controllers/assistant.ts
-import { getStoreLattice } from "@axiom-lattice/core";
+import { agentInstanceManager, getStoreLattice } from "@axiom-lattice/core";
 import { randomUUID } from "crypto";
-import { agentLatticeManager as agentLatticeManager2, eventBus } from "@axiom-lattice/core";
+import { agentLatticeManager, eventBus } from "@axiom-lattice/core";
 function getTenantId(request) {
   const userTenantId = request.user?.tenantId;
   if (userTenantId) {
@@ -304,7 +35,7 @@ function convertAgentConfigToAssistant(config) {
 }
 async function getAssistantList(request, reply) {
   const tenantId = getTenantId(request);
-  const agentConfigs = await agentLatticeManager2.getAllAgentConfigsByTenant(tenantId);
+  const agentConfigs = await agentLatticeManager.getAllAgentConfigsByTenant(tenantId);
   const codeConfiguredAssistants = agentConfigs.map(
     convertAgentConfigToAssistant
   );
@@ -335,7 +66,7 @@ async function getAssistant(request, reply) {
   const assistantStore = storeLattice.store;
   let assistant = await assistantStore.getAssistantById(tenantId, id);
   if (!assistant) {
-    const agentConfig = await agentLatticeManager2.getAgentConfigWithTenant(tenantId, id);
+    const agentConfig = await agentLatticeManager.getAgentConfigWithTenant(tenantId, id);
     if (agentConfig) {
       assistant = convertAgentConfigToAssistant(agentConfig);
     }
@@ -412,7 +143,7 @@ async function deleteAssistant(request, reply) {
   const { id } = request.params;
   const storeLattice = getStoreLattice("default", "assistant");
   const assistantStore = storeLattice.store;
-  const agentConfig = await agentLatticeManager2.getAgentConfigWithTenant(tenantId, id);
+  const agentConfig = await agentLatticeManager.getAgentConfigWithTenant(tenantId, id);
   const isCodeConfigured = !!agentConfig;
   if (isCodeConfigured) {
     const exists2 = await assistantStore.hasAssistant(tenantId, id);
@@ -453,7 +184,8 @@ var getAgentGraph = async (request, reply) => {
   try {
     const { assistantId } = request.params;
     const tenant_id = getTenantId(request);
-    const imageData = await draw_graph(assistantId, tenant_id);
+    const agent = agentInstanceManager.getAgent({ assistant_id: assistantId, tenant_id, thread_id: "" });
+    const imageData = await agent.get_draw_graph();
     reply.header("Content-Type", "application/json").send({
       image: imageData
     });
@@ -466,7 +198,11 @@ var getAgentGraph = async (request, reply) => {
 };
 
 // src/controllers/run.ts
-import { v4 as v42 } from "uuid";
+import { v4 } from "uuid";
+import {
+  Agent,
+  agentInstanceManager as agentInstanceManager2
+} from "@axiom-lattice/core";
 var createRun = async (request, reply) => {
   try {
     const {
@@ -481,7 +217,7 @@ var createRun = async (request, reply) => {
     const tenant_id = request.headers["x-tenant-id"];
     const workspace_id = request.headers["x-workspace-id"];
     const project_id = request.headers["x-project-id"];
-    const x_request_id = request.headers["x-request-id"] || v42();
+    const x_request_id = request.headers["x-request-id"] || v4();
     if (!assistant_id) {
       reply.status(400).send({
         success: false,
@@ -489,26 +225,15 @@ var createRun = async (request, reply) => {
       });
       return;
     }
+    const agent = agentInstanceManager2.getAgent({
+      assistant_id,
+      thread_id,
+      tenant_id,
+      workspace_id,
+      project_id,
+      custom_run_config
+    });
     if (streaming) {
-      const agentExists = await checkAgentExists(tenant_id, assistant_id);
-      if (!agentExists) {
-        reply.status(404).send({
-          success: false,
-          error: `Agent ${assistant_id} not found for tenant ${tenant_id}`
-        });
-        return;
-      }
-      const stream = await agent_stream({
-        assistant_id,
-        input,
-        thread_id,
-        command,
-        tenant_id,
-        workspace_id,
-        project_id,
-        run_id: x_request_id,
-        custom_run_config
-      });
       reply.hijack();
       reply.raw.writeHead(200, {
         "Content-Type": "text/event-stream",
@@ -517,9 +242,14 @@ var createRun = async (request, reply) => {
         "Access-Control-Allow-Origin": "*"
       });
       try {
-        let chunkCount = 0;
+        const result = await agent.addMessage({
+          input,
+          command,
+          custom_run_config
+        });
+        const stream = agent.chunkStream(result.messageId);
         for await (const chunk of stream) {
-          chunkCount++;
+          console.log(input.message, chunk.data.content);
           const success = reply.raw.write(`data: ${JSON.stringify(chunk)}
 
 `);
@@ -531,7 +261,7 @@ var createRun = async (request, reply) => {
         const errorEvent = {
           type: "error",
           data: {
-            id: v42(),
+            id: v4(),
             content: error.message || "Stream processing error"
           }
         };
@@ -542,18 +272,16 @@ var createRun = async (request, reply) => {
         reply.raw.end();
       }
     } else {
-      const result = await agent_invoke({
-        assistant_id,
-        input,
+      const { message: msg, ...restInputNonStream } = input;
+      const result = await agent.invoke({
+        input: { message: msg, ...restInputNonStream },
         command,
-        thread_id,
-        tenant_id,
-        workspace_id,
-        project_id,
-        run_id: x_request_id,
         custom_run_config
       });
-      reply.status(200).send(result);
+      reply.status(200).send({
+        success: true,
+        ...result
+      });
     }
   } catch (error) {
     reply.status(500).send({
@@ -580,12 +308,12 @@ var resumeStream = async (request, reply) => {
       "Access-Control-Allow-Origin": "*"
     });
     try {
-      const stream = await resume_stream({
+      const agent = new Agent({
+        assistant_id: "",
         thread_id,
-        message_id,
-        known_content,
-        poll_interval: poll_interval || 100
+        tenant_id: ""
       });
+      const stream = agent.chunkStream(message_id, known_content);
       for await (const chunk of stream) {
         reply.raw.write(`data: ${JSON.stringify(chunk)}
 
@@ -595,7 +323,7 @@ var resumeStream = async (request, reply) => {
       const errorEvent = {
         type: "error",
         data: {
-          id: v42(),
+          id: v4(),
           content: error.message || "Resume stream processing error"
         }
       };
@@ -614,6 +342,7 @@ var resumeStream = async (request, reply) => {
 };
 
 // src/controllers/memory.ts
+import { agentInstanceManager as agentInstanceManager3 } from "@axiom-lattice/core";
 var setMemoryItem = async (request, reply) => {
   try {
     const { assistantId, key } = request.params;
@@ -678,11 +407,8 @@ var getAllMemoryItems = async (request, reply) => {
       });
       return;
     }
-    const result = await agent_messages({
-      assistant_id: assistantId,
-      thread_id,
-      tenant_id
-    });
+    const agent = agentInstanceManager3.getAgent({ assistant_id: assistantId, tenant_id, thread_id });
+    const result = await agent.getCurrentMessages();
     if (!result) {
       reply.status(500).send(result);
       return;
@@ -713,16 +439,18 @@ var getAgentState = async (request, reply) => {
       return;
     }
     const tenant_id = request.headers["x-tenant-id"];
-    const result = await agent_state({
-      assistant_id: assistantId,
-      thread_id,
-      tenant_id
-    });
+    const agent = agentInstanceManager3.getAgent({ assistant_id: assistantId, tenant_id, thread_id });
+    const result = await agent.getCurrentState();
+    const pendingMessages = await agent.getPendingMessages();
     if (!result) {
       reply.status(500).send(result);
       return;
     }
-    reply.send(result);
+    const mergedResult = {
+      ...result,
+      pendingMessages
+    };
+    reply.send(mergedResult);
   } catch (error) {
     reply.status(500).send({
       success: false,
@@ -1329,10 +1057,12 @@ async function getModels(request, reply) {
     const allLattices = modelLatticeManager.getAllLattices();
     const models = allLattices.map((lattice) => {
       const config = lattice.client.config || {};
+      const displayName = config.displayName || lattice.key.split("-").map((word) => word.charAt(0).toUpperCase() + word.slice(1)).join(" ");
       return {
         key: lattice.key,
         model: config.model || "",
         provider: config.provider || "openai",
+        displayName,
         streaming: config.streaming || false,
         apiKey: config.apiKey || "",
         baseURL: config.baseURL || "",
@@ -1382,6 +1112,7 @@ async function updateModels(request, reply) {
         const llmConfig = {
           provider: modelConfig.provider,
           model: modelConfig.model,
+          displayName: modelConfig.displayName,
           streaming: modelConfig.streaming ?? false,
           apiKey: modelConfig.apiKey,
           baseURL: modelConfig.baseURL,
@@ -5182,7 +4913,7 @@ var configureSwagger = async (app2, customSwaggerConfig, customSwaggerUiConfig)
 };
 
 // src/services/agent_task_consumer.ts
-import { eventBus as eventBus2, AGENT_TASK_EVENT } from "@axiom-lattice/core";
+import { eventBus as eventBus2, AGENT_TASK_EVENT, agentInstanceManager as agentInstanceManager4, QueueMode } from "@axiom-lattice/core";
 var handleAgentTask = async (taskRequest, retryCount = 0) => {
   const {
     assistant_id,
@@ -5199,82 +4930,18 @@ var handleAgentTask = async (taskRequest, retryCount = 0) => {
     );
     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,
-        custom_run_config: runConfig
-      }),
-      headers: {
-        "Content-Type": "application/json",
-        "x-tenant-id": tenant_id,
-        "x-workspace-id": runConfig?.workspaceId,
-        "x-project-id": runConfig?.projectId
-      }
-    }).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}`);
-    }
-    const contentType = response.headers.get("content-type");
-    if (contentType?.includes("text/event-stream")) {
-      const reader = response.body?.getReader();
-      const decoder = new TextDecoder();
-      if (!reader) {
-        throw new Error("Response body is not readable");
-      }
-      let buffer = "";
-      let streamEnded = false;
-      try {
-        while (true) {
-          const { done, value } = await reader.read();
-          if (done) {
-            streamEnded = true;
-            console.log(
-              `SSE\u6D41\u5DF2\u7ED3\u675F [assistant_id: ${assistant_id}, thread_id: ${thread_id}]`
-            );
-            break;
-          }
-        }
-      } catch (streamError) {
-        console.error("Error reading SSE stream:", streamError);
-        throw streamError;
-      } finally {
-        reader.releaseLock();
-      }
-      if (callback_event) {
-        const state = await agent_state({ assistant_id, thread_id, tenant_id });
-        eventBus2.publish(callback_event, {
-          success: true,
-          state,
-          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;
-    } else {
-      await response.text();
-      if (callback_event) {
-        const state = await agent_state({ assistant_id, thread_id, tenant_id });
+    const agent = agentInstanceManager4.getAgent({ assistant_id, thread_id, tenant_id, workspace_id: runConfig?.workspaceId, project_id: runConfig?.projectId, custom_run_config: runConfig });
+    await agent.addMessage({ input, command }, QueueMode.STEER);
+    if (callback_event) {
+      agent.subscribeOnce("message:completed", (evt) => {
         eventBus2.publish(callback_event, {
           success: true,
-          state,
+          state: evt.state,
           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;
+      });
     }
+    return true;
   } catch (error) {
     console.error(
       `Agent\u4EFB\u52A1\u6267\u884C\u5931\u8D25: ${assistant_id}, \u7EBF\u7A0B: ${thread_id}`,
@@ -5536,7 +5203,6 @@ app.addHook("onResponse", (request, reply, done) => {
     "x-tenant-id": getHeaderValue(request.headers["x-tenant-id"]),
     "x-request-id": getHeaderValue(request.headers["x-request-id"])
   };
-  loggerLattice.info(`${request.method} ${request.url} - ${reply.statusCode}`);
   done();
 });
 app.register(cors, {
@@ -5614,6 +5280,11 @@ var start = async (config) => {
     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: ${target_port}`);
+    try {
+      logger.info("AgentLifecycleManager initialized");
+    } catch (error) {
+      logger.warn("Failed to initialize AgentLifecycleManager", { error });
+    }
     const queueServiceConfig = config?.queueServiceConfig;
     if (queueServiceConfig) {
       setQueueServiceType(queueServiceConfig.type);