@axiom-lattice/gateway 1.0.11 → 1.0.13

package/src/schemas/index.ts

@@ -0,0 +1,194 @@
+import { FastifySchema } from "fastify";
+
+// Create Run Schemas
+export const createRunSchema: FastifySchema = {
+  description: "Create a new agent run",
+  tags: ["Runs"],
+  summary: "Create Agent Run",
+  body: {
+    type: "object",
+    properties: {
+      thread_id: { type: "string", description: "Thread ID" },
+      assistant_id: { type: "string", description: "Assistant ID" },
+      message: { type: "string", description: "Message data for the run" },
+      command: {
+        type: "object",
+        description: "Command data for the run",
+        nullable: true,
+      },
+      streaming: {
+        type: "boolean",
+        description: "Whether to stream the response",
+        nullable: true,
+      },
+      background: {
+        type: "boolean",
+        description: "Whether to run in background",
+        nullable: true,
+      },
+    },
+    required: ["thread_id", "assistant_id", "message"],
+  },
+  response: {
+    200: {},
+    400: {},
+  },
+};
+
+// Memory Schemas
+export const getAllMemoryItemsSchema: FastifySchema = {
+  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: {},
+  },
+};
+
+export const getAgentStateSchema: FastifySchema = {
+  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: {},
+  },
+};
+
+export const getMemoryItemSchema: FastifySchema = {
+  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: {},
+  },
+};
+
+export const setMemoryItemSchema: FastifySchema = {
+  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: {},
+  },
+};
+
+export const deleteMemoryItemSchema: FastifySchema = {
+  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: {},
+  },
+};
+
+export const clearMemorySchema: FastifySchema = {
+  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: {},
+  },
+};
+
+// Graph Schema
+export const getAgentGraphSchema: FastifySchema = {
+  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: {},
+  },
+};
+
+// Resume Stream Schema
+export const resumeStreamSchema: FastifySchema = {
+  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: {},
+  },
+};

package/src/services/agent_service.ts

@@ -10,7 +10,14 @@ import {
 import { Command, CommandParams } from "@langchain/langgraph";
 import { v4 } from "uuid";
 // 修改导入路径，使用 lattice_core 包
-import { getAgentClient, getAgentLattice } from "@axiom-lattice/core";
+import {
+  getAgentClient,
+  getAgentLattice,
+  InMemoryChunkBuffer,
+  registerChunkBuffer,
+  getChunkBuffer,
+  hasChunkBuffer,
+} from "@axiom-lattice/core";
 
 function isAIMessageChunk(msg: any): msg is AIMessageChunk {
   return msg && msg.constructor.name === "AIMessageChunk";
@@ -25,6 +32,20 @@ function isToolMessage(msg: any): msg is ToolMessage {
   return msg && msg.constructor.name === "ToolMessage";
 }
 
+/**
+ * Get or create the global ChunkBuffer instance
+ */
+function getOrCreateChunkBuffer(): InMemoryChunkBuffer {
+  if (!hasChunkBuffer("default")) {
+    const buffer = new InMemoryChunkBuffer({
+      ttl: 60 * 60 * 1000, // 1 hour TTL
+      cleanupInterval: 5 * 60 * 1000, // Clean every 5 minutes
+    });
+    registerChunkBuffer("default", buffer);
+  }
+  return getChunkBuffer("default") as InMemoryChunkBuffer;
+}
+
 export async function agent_invoke({
   input,
   thread_id,
@@ -40,7 +61,7 @@ export async function agent_invoke({
   run_id?: string;
   command?: CommandParams<any>;
 }) {
-  const runnable_agent = getAgentLattice(assistant_id)?.client;
+  const runnable_agent = getAgentClient(assistant_id);
   const { files, message, ...rest } = input;
   const humanMessage = new HumanMessage(message || "");
   humanMessage.additional_kwargs = { files: files };
@@ -57,15 +78,22 @@ export async function agent_invoke({
       configurable: {
         thread_id: 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((message: BaseMessage) => {
+    const { type, data } = message.toDict();
+    return {
+      ...data,
+      role: type,
+    };
+  });
+  return { messages: data };
 }
 
 export async function agent_stream({
@@ -92,6 +120,9 @@ export async function agent_stream({
     messages = [humanMessage];
   }
 
+  // Get ChunkBuffer instance
+  const chunkBuffer = getOrCreateChunkBuffer();
+
   try {
     if (!runnable_agent) {
       throw new Error(`Agent ${assistant_id} not found`);
@@ -115,15 +146,18 @@ export async function agent_stream({
         },
         streamMode: ["updates", "messages"],
         subgraphs: false,
+        recursionLimit: 200,
       }
     );
 
-    // 创建一个可迭代的 ReadableStream
+    // 创建一个可迭代的 ReadableStream with ChunkBuffer integration
     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);
@@ -138,31 +172,35 @@ export async function agent_stream({
             }
 
             if ((chunk?.[1] as any)?.__interrupt__) {
-              data = chunk?.[1]?.[0]?.toDict();
+              //data = chunk?.[1]?.[0]?.toDict();
               // 原有的中断消息处理
-              // data = {
-              //   messages: [
-              //     {
-              //       role: "ai",
-              //       content: (chunk?.[1] as any)?.__interrupt__[0].value,
-              //       type: "action",
-              //     },
-              //   ],
-              // };
+              data = {
+                type: "ai",
+                data: { content: (chunk?.[1] as any)?.__interrupt__[0].value },
+              };
             }
 
             if (data) {
               //console.log(data);
+              await chunkBuffer.addChunk(thread_id, data);
+
               yield data;
             }
           }
+
+          // Mark thread as completed when streaming finishes successfully
+          await chunkBuffer.completeThread(thread_id);
         } catch (error) {
           console.error("Stream error:", error);
+          // Mark thread as aborted on error
+          await chunkBuffer.abortThread(thread_id);
           throw error;
         }
       },
     };
   } catch (error) {
+    // Mark thread as aborted on initialization error
+    await chunkBuffer.abortThread(thread_id);
     throw error;
   }
 }
@@ -241,3 +279,79 @@ export async function draw_graph(assistant_id: string) {
   const image = await drawableGraph.drawMermaid();
   return image;
 }
+
+/**
+ * Get accumulated content for a thread from ChunkBuffer
+ */
+export async function get_accumulated_content(thread_id: string) {
+  const chunkBuffer = getOrCreateChunkBuffer();
+  return await chunkBuffer.getAccumulatedContent(thread_id);
+}
+
+/**
+ * Get thread status and metadata from ChunkBuffer
+ */
+export async function get_thread_status(thread_id: string) {
+  const chunkBuffer = getOrCreateChunkBuffer();
+  const threadBuffer = await chunkBuffer.getThreadBuffer(thread_id);
+  return {
+    exists: !!threadBuffer,
+    status: threadBuffer?.status,
+    chunkCount: threadBuffer?.chunks.length || 0,
+    createdAt: threadBuffer?.createdAt,
+    updatedAt: threadBuffer?.updatedAt,
+  };
+}
+
+/**
+ * Get all active streaming threads
+ */
+export async function get_active_threads() {
+  const chunkBuffer = getOrCreateChunkBuffer();
+  return await chunkBuffer.getActiveThreads();
+}
+
+/**
+ * Get ChunkBuffer statistics
+ */
+export async function get_buffer_stats() {
+  const chunkBuffer = getOrCreateChunkBuffer();
+  return chunkBuffer.getStats();
+}
+
+/**
+ * Clear specific thread buffer
+ */
+export async function clear_thread_buffer(thread_id: string) {
+  const chunkBuffer = getOrCreateChunkBuffer();
+  await chunkBuffer.clearThread(thread_id);
+}
+
+/**
+ * Resume streaming from a known position
+ * Creates an async iterator that yields new chunks as they arrive
+ *
+ * @param thread_id - Thread identifier
+ * @param message_id - Message identifier (usually run_id)
+ * @param known_content - Content already received (used to find resume position)
+ * @param poll_interval - Polling interval in milliseconds (default: 100ms)
+ */
+export async function* resume_stream({
+  thread_id,
+  message_id,
+  known_content,
+  poll_interval = 100,
+}: {
+  thread_id: string;
+  message_id: string;
+  known_content: string;
+  poll_interval?: number;
+}) {
+  const chunkBuffer = getOrCreateChunkBuffer();
+
+  return chunkBuffer.getNewChunksSinceContent(
+    thread_id,
+    message_id,
+    known_content
+  );
+}

package/src/swagger.ts

@@ -0,0 +1,77 @@
+import { FastifyInstance } from "fastify";
+import swagger from "@fastify/swagger";
+import swaggerUi from "@fastify/swagger-ui";
+// Example usage:
+// configureSwagger(app)
+// configureSwagger(app, { openapi: { info: { version: "2.0.0" } } })
+// configureSwagger(app, undefined, { routePrefix: "/docs" })
+
+// Default swagger configuration
+export const 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" as const,
+          scheme: "bearer" as const,
+          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" },
+    ],
+  },
+};
+
+// Default swagger UI configuration
+export const defaultSwaggerUiConfig = {
+  routePrefix: "/api-docs",
+  uiConfig: {
+    docExpansion: "full" as const,
+    deepLinking: false,
+  },
+  staticCSP: true,
+  transformStaticCSP: (header: string) => header,
+};
+
+// Configure Swagger with optional custom configuration
+export const configureSwagger = async (
+  app: FastifyInstance,
+  customSwaggerConfig?: Partial<typeof defaultSwaggerConfig>,
+  customSwaggerUiConfig?: Partial<typeof defaultSwaggerUiConfig>
+) => {
+  // Merge default config with custom config
+  const swaggerConfig = { ...defaultSwaggerConfig, ...customSwaggerConfig };
+  const swaggerUiConfig = {
+    ...defaultSwaggerUiConfig,
+    ...customSwaggerUiConfig,
+  };
+
+  await app.register(swagger, swaggerConfig);
+  await app.register(swaggerUi, swaggerUiConfig);
+};