@axiom-lattice/gateway 1.0.12 → 1.0.13

package/dist/index.mjs

@@ -10,7 +10,25 @@ import {
 } from "@langchain/core/messages";
 import { Command } from "@langchain/langgraph";
 import { v4 } from "uuid";
-import { getAgentClient } from "@axiom-lattice/core";
+import {
+  getAgentClient,
+  InMemoryChunkBuffer,
+  registerChunkBuffer,
+  getChunkBuffer,
+  hasChunkBuffer
+} from "@axiom-lattice/core";
+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,
@@ -33,14 +51,21 @@ async function agent_invoke({
       configurable: {
         thread_id,
         run_id: run_id || v4(),
-        recursionLimit: 200,
         "x-tenant-id": tenant_id,
         "x-request-id": run_id,
         "x-thread-id": thread_id
-      }
+      },
+      recursionLimit: 200
     }
   );
-  return result;
+  const data = result.messages.map((message2) => {
+    const { type, data: data2 } = message2.toDict();
+    return {
+      ...data2,
+      role: type
+    };
+  });
+  return { messages: data };
 }
 async function agent_stream({
   input,
@@ -58,6 +83,7 @@ async function agent_stream({
     humanMessage.additional_kwargs = { files };
     messages = [humanMessage];
   }
+  const chunkBuffer = getOrCreateChunkBuffer();
   try {
     if (!runnable_agent) {
       throw new Error(`Agent ${assistant_id} not found`);
@@ -77,7 +103,8 @@ async function agent_stream({
           "x-thread-id": thread_id
         },
         streamMode: ["updates", "messages"],
-        subgraphs: false
+        subgraphs: false,
+        recursionLimit: 200
       }
     );
     return {
@@ -85,6 +112,7 @@ async function agent_stream({
         try {
           for await (const chunk of agentStream) {
             let data;
+            let chunkContent = "";
             if (chunk[0] === "updates") {
               const update = chunk[1];
               const values = Object.values(update);
@@ -97,19 +125,26 @@ async function agent_stream({
               data = messages2?.[0]?.toDict();
             }
             if (chunk?.[1]?.__interrupt__) {
-              data = chunk?.[1]?.[0]?.toDict();
+              data = {
+                type: "ai",
+                data: { content: chunk?.[1]?.__interrupt__[0].value }
+              };
             }
             if (data) {
+              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;
   }
 }
@@ -171,6 +206,19 @@ async function draw_graph(assistant_id) {
   const image = await drawableGraph.drawMermaid();
   return image;
 }
+async function* resume_stream({
+  thread_id,
+  message_id,
+  known_content,
+  poll_interval = 100
+}) {
+  const chunkBuffer = getOrCreateChunkBuffer();
+  return chunkBuffer.getNewChunksSinceContent(
+    thread_id,
+    message_id,
+    known_content
+  );
+}
 
 // src/controllers/run.ts
 import { v4 as v42 } from "uuid";
@@ -215,6 +263,7 @@ var createRun = async (request, reply) => {
 `);
         }
       } catch (error) {
+        console.error("Stream processing error:", error);
       } finally {
         reply.raw.end();
         return reply.hijack();
@@ -237,6 +286,47 @@ var createRun = async (request, reply) => {
     });
   }
 };
+var resumeStream = async (request, reply) => {
+  try {
+    const { thread_id, message_id, known_content, poll_interval } = request.body;
+    if (!thread_id || !message_id || known_content === void 0) {
+      reply.status(400).send({
+        success: false,
+        error: "thread_id, message_id, and known_content are required"
+      });
+      return;
+    }
+    const stream = await resume_stream({
+      thread_id,
+      message_id,
+      known_content,
+      poll_interval: poll_interval || 100
+    });
+    reply.raw.writeHead(200, {
+      "Content-Type": "text/event-stream",
+      "Cache-Control": "no-cache",
+      Connection: "keep-alive",
+      "Access-Control-Allow-Origin": "*"
+    });
+    try {
+      for await (const chunk of stream) {
+        reply.raw.write(`data: ${JSON.stringify(chunk)}
+
+`);
+      }
+    } catch (error) {
+      console.error("Resume stream processing error:", error);
+    } finally {
+      reply.raw.end();
+      return reply.hijack();
+    }
+  } catch (error) {
+    reply.status(500).send({
+      success: false,
+      error: `Error resuming stream: ${error.message}`
+    });
+  }
+};
 
 // src/controllers/memory.ts
 var setMemoryItem = async (request, reply) => {
@@ -413,31 +503,198 @@ var getAgentGraph = async (request, reply) => {
   }
 };
 
+// src/schemas/index.ts
+var getAllMemoryItemsSchema = {
+  description: "Get all memory items for an assistant thread",
+  tags: ["Memory"],
+  summary: "Get All Memory Items",
+  params: {
+    type: "object",
+    properties: {
+      assistantId: { type: "string", description: "Assistant ID" },
+      thread_id: { type: "string", description: "Thread ID" }
+    },
+    required: ["assistantId", "thread_id"]
+  },
+  response: {
+    200: {},
+    400: {},
+    500: {}
+  }
+};
+var getAgentStateSchema = {
+  description: "Get agent state for an assistant thread",
+  tags: ["Memory"],
+  summary: "Get Agent State",
+  params: {
+    type: "object",
+    properties: {
+      assistantId: { type: "string", description: "Assistant ID" },
+      thread_id: { type: "string", description: "Thread ID" }
+    },
+    required: ["assistantId", "thread_id"]
+  },
+  response: {
+    200: {}
+  }
+};
+var getMemoryItemSchema = {
+  description: "Get a specific memory item by key",
+  tags: ["Memory"],
+  summary: "Get Memory Item",
+  params: {
+    type: "object",
+    properties: {
+      assistantId: { type: "string", description: "Assistant ID" },
+      key: { type: "string", description: "Memory key" }
+    },
+    required: ["assistantId", "key"]
+  },
+  response: {
+    200: {}
+  }
+};
+var setMemoryItemSchema = {
+  description: "Set or update a memory item",
+  tags: ["Memory"],
+  summary: "Set Memory Item",
+  params: {
+    type: "object",
+    properties: {
+      assistantId: { type: "string", description: "Assistant ID" },
+      key: { type: "string", description: "Memory key" }
+    },
+    required: ["assistantId", "key"]
+  },
+  body: {
+    type: "object",
+    description: "Memory item data"
+  },
+  response: {
+    200: {}
+  }
+};
+var deleteMemoryItemSchema = {
+  description: "Delete a specific memory item",
+  tags: ["Memory"],
+  summary: "Delete Memory Item",
+  params: {
+    type: "object",
+    properties: {
+      assistantId: { type: "string", description: "Assistant ID" },
+      key: { type: "string", description: "Memory key" }
+    },
+    required: ["assistantId", "key"]
+  },
+  response: {
+    200: {}
+  }
+};
+var clearMemorySchema = {
+  description: "Clear all memory items for an assistant",
+  tags: ["Memory"],
+  summary: "Clear Memory",
+  params: {
+    type: "object",
+    properties: {
+      assistantId: { type: "string", description: "Assistant ID" }
+    },
+    required: ["assistantId"]
+  },
+  response: {
+    200: {}
+  }
+};
+var getAgentGraphSchema = {
+  description: "Get agent graph visualization",
+  tags: ["Graph"],
+  summary: "Get Agent Graph",
+  params: {
+    type: "object",
+    properties: {
+      assistantId: { type: "string", description: "Assistant ID" }
+    },
+    required: ["assistantId"]
+  },
+  response: {
+    200: {}
+  }
+};
+var resumeStreamSchema = {
+  description: "Resume streaming from a known position",
+  tags: ["Streaming"],
+  summary: "Resume Stream",
+  body: {
+    type: "object",
+    properties: {
+      thread_id: { type: "string", description: "Thread ID" },
+      message_id: {
+        type: "string",
+        description: "Message ID (usually run_id)"
+      },
+      known_content: {
+        type: "string",
+        description: "Content already received"
+      },
+      poll_interval: {
+        type: "number",
+        description: "Polling interval in milliseconds",
+        nullable: true,
+        default: 100
+      }
+    },
+    required: ["thread_id", "message_id"]
+  },
+  response: {
+    200: {},
+    400: {},
+    500: {}
+  }
+};
+
 // src/routes/index.ts
-var registerRoutes = (app2) => {
+var registerLatticeRoutes = (app2) => {
   app2.post("/api/runs", createRun);
+  app2.post(
+    "/api/resume_stream",
+    { schema: resumeStreamSchema },
+    resumeStream
+  );
   app2.get(
     "/api/assistants/:assistantId/:thread_id/memory",
+    { schema: getAllMemoryItemsSchema },
     getAllMemoryItems
   );
   app2.get(
     "/api/assistants/:assistantId/:thread_id/state",
+    { schema: getAgentStateSchema },
     getAgentState
   );
   app2.get(
     "/api/assistants/:assistantId/memory/:key",
+    { schema: getMemoryItemSchema },
     getMemoryItem
   );
   app2.put(
     "/api/assistants/:assistantId/memory/:key",
+    { schema: setMemoryItemSchema },
     setMemoryItem
   );
   app2.delete(
     "/api/assistants/:assistantId/memory/:key",
+    { schema: deleteMemoryItemSchema },
     deleteMemoryItem
   );
-  app2.delete("/api/assistants/:assistantId/memory", clearMemory);
-  app2.get("/api/assistants/:assistantId/graph", getAgentGraph);
+  app2.delete(
+    "/api/assistants/:assistantId/memory",
+    { schema: clearMemorySchema },
+    clearMemory
+  );
+  app2.get(
+    "/api/assistants/:assistantId/graph",
+    { schema: getAgentGraphSchema },
+    getAgentGraph
+  );
 };
 
 // src/logger/Logger.ts
@@ -578,6 +835,68 @@ var Logger = class _Logger {
   }
 };
 
+// src/swagger.ts
+import swagger from "@fastify/swagger";
+import swaggerUi from "@fastify/swagger-ui";
+var defaultSwaggerConfig = {
+  openapi: {
+    openapi: "3.0.0",
+    info: {
+      title: "Axiom Lattice Gateway API",
+      description: "API Gateway for LangGraph agent-based applications",
+      version: "1.0.0",
+      contact: {
+        name: "Axiom Lattice Team",
+        email: "support@axiom-lattice.com"
+      }
+    },
+    servers: [
+      {
+        url: "http://localhost:4001",
+        description: "Development environment"
+      }
+    ],
+    components: {
+      securitySchemes: {
+        bearerAuth: {
+          type: "http",
+          scheme: "bearer",
+          bearerFormat: "JWT"
+        }
+      }
+    },
+    security: [
+      {
+        bearerAuth: []
+      }
+    ],
+    tags: [
+      { name: "Runs", description: "Agent run management" },
+      { name: "Memory", description: "Agent memory management" },
+      { name: "Graph", description: "Agent graph visualization" },
+      { name: "Health", description: "System health checks" }
+    ]
+  }
+};
+var defaultSwaggerUiConfig = {
+  routePrefix: "/api-docs",
+  uiConfig: {
+    docExpansion: "full",
+    deepLinking: false
+  },
+  staticCSP: true,
+  transformStaticCSP: (header) => header
+};
+var configureSwagger = async (app2, customSwaggerConfig, customSwaggerUiConfig) => {
+  const swaggerConfig = { ...defaultSwaggerConfig, ...customSwaggerConfig };
+  const swaggerUiConfig = {
+    ...defaultSwaggerUiConfig,
+    ...customSwaggerUiConfig
+  };
+  await app2.register(swagger, swaggerConfig);
+  await app2.register(swaggerUi, swaggerUiConfig);
+};
+
 // src/index.ts
 process.on("unhandledRejection", (reason, promise) => {
   console.error("\u672A\u5904\u7406\u7684Promise\u62D2\u7EDD:", reason);
@@ -638,7 +957,6 @@ app.setErrorHandler((error, request, reply) => {
   });
 });
 app.decorate("logger", logger);
-registerRoutes(app);
 var start = async ({ port }) => {
   try {
     const target_port = port || Number(process.env.PORT) || 4001;
@@ -651,6 +969,8 @@ var start = async ({ port }) => {
 };
 var LatticeGateway = {
   startAsHttpEndpoint: start,
+  configureSwagger,
+  registerLatticeRoutes,
   app
 };
 export {