@axiom-lattice/gateway 2.1.44 → 2.1.46

package/src/services/agent_task_consumer.ts

@@ -1,6 +1,5 @@
-import { eventBus, AGENT_TASK_EVENT } from "@axiom-lattice/core";
+import { eventBus, AGENT_TASK_EVENT, agentInstanceManager, QueueMode } from "@axiom-lattice/core";
 import { popAgentTaskFromQueue } from "./queue_service";
-import { agent_state } from "./agent_service";
 
 // 任务请求结构
 export interface AgentTaskRequest {
@@ -10,7 +9,7 @@ export interface AgentTaskRequest {
   "x-tenant-id": string;
   command?: any;
   callback_event?: string; // 可选的回调事件名称
-  runConfig?: Record<string, unknown>; // RunConfig for subagent execution
+  runConfig?: Record<string, any>; // RunConfig for subagent execution
 }
 
 /**
@@ -32,6 +31,7 @@ const handleAgentTask = async (
     runConfig,
   } = taskRequest;
 
+
   try {
     console.log(
       `开始处理任务 [assistant_id: ${assistant_id}, thread_id: ${thread_id}]`
@@ -41,117 +41,22 @@ const handleAgentTask = async (
     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 as string,
-        "x-project-id": runConfig?.projectId as string,
-      },
-    }).catch((err) => {
-      console.error(`fetch请求失败: ${err.message || String(err)}`);
-      throw new Error(`fetch失败: ${err.message || String(err)}`);
-    });
-
-    if (!response.ok) {
-      throw new Error(`API请求失败: ${response.status} ${response.statusText}`);
-    }
-
-    // Check if response is SSE stream
-    const contentType = response.headers.get("content-type");
-    if (contentType?.includes("text/event-stream")) {
-      // Handle SSE 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) {
-            // Stream has ended
-            streamEnded = true;
-            console.log(
-              `SSE流已结束 [assistant_id: ${assistant_id}, thread_id: ${thread_id}]`
-            );
-            break;
-          }
-
-          // // Decode the chunk and process SSE events
-          // buffer += decoder.decode(value, { stream: true });
-          // const lines = buffer.split("\n");
-          // buffer = lines.pop() || ""; // Keep incomplete line in buffer
-
-          // for (const line of lines) {
-          //   if (line.startsWith("data: ")) {
-          //     try {
-          //       const data = JSON.parse(line.slice(6));
-          //       // Process SSE event data here if needed
-          //       // You can emit events or handle the data as required
-          //       console.log("SSE event received:", data);
-          //     } catch (e) {
-          //       // Ignore invalid JSON
-          //       console.warn("Failed to parse SSE data:", line);
-          //     }
-          //   }
-          // }
-        }
-      } catch (streamError) {
-        console.error("Error reading SSE stream:", streamError);
-        throw streamError;
-      } finally {
-        reader.releaseLock();
-      }
-
-      // Stream has ended successfully
-      if (callback_event) {
-        const state = await agent_state({ assistant_id, thread_id, tenant_id });
+    const agent = agentInstanceManager.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) => {
         eventBus.publish(callback_event, {
           success: true,
-          state: state,
+          state: evt.state,
           config: { assistant_id, thread_id, tenant_id },
         });
-      }
+      })
+    }
+
+    return true;
 
-      console.log(
-        `任务处理成功 [assistant_id: ${assistant_id}, thread_id: ${thread_id}]`
-      );
-      return true;
-    } else {
-      // Non-streaming response
-      await response.text(); // Consume the response
 
-      if (callback_event) {
-        const state = await agent_state({ assistant_id, thread_id, tenant_id });
-        eventBus.publish(callback_event, {
-          success: true,
-          state,
-          config: { assistant_id, thread_id, tenant_id },
-        });
-      }
 
-      console.log(
-        `任务处理成功 [assistant_id: ${assistant_id}, thread_id: ${thread_id}]`
-      );
-      return true;
-    }
   } catch (error) {
     console.error(
       `Agent任务执行失败: ${assistant_id}, 线程: ${thread_id}`,

package/src/__tests__/agent_service.test.ts

@@ -1,238 +0,0 @@
-/**
- * Agent Service Tests
- * 
- * Tests for agent_service.ts to ensure proper async/await usage
- * and error handling when working with agent clients.
- */
-
-import { describe, it, expect, beforeEach, jest } from '@jest/globals';
-import { agent_state, agent_messages, draw_graph, agent_invoke, agent_stream } from '../services/agent_service';
-import { getAgentClient, agentLatticeManager } from '@axiom-lattice/core';
-
-// Mock the dependencies
-jest.mock('@axiom-lattice/core', () => ({
-  getAgentClient: jest.fn(),
-  agentLatticeManager: {
-    getAgentLatticeWithTenant: jest.fn(),
-  },
-  InMemoryChunkBuffer: jest.fn().mockImplementation(() => ({
-    addChunk: jest.fn(),
-    completeThread: jest.fn(),
-    abortThread: jest.fn(),
-  })),
-  registerChunkBuffer: jest.fn(),
-  getChunkBuffer: jest.fn(),
-  hasChunkBuffer: jest.fn().mockReturnValue(false),
-}));
-
-describe('Agent Service - getAgentClient async handling', () => {
-  const mockTenantId = 'test-tenant';
-  const mockAssistantId = 'test-assistant';
-  const mockThreadId = 'test-thread';
-  
-  beforeEach(() => {
-    jest.clearAllMocks();
-  });
-
-  describe('agent_state', () => {
-    it('should properly await getAgentClient before calling getState', async () => {
-      const mockAgent = {
-        getState: jest.fn().mockResolvedValue({
-          values: { messages: [] },
-          next: [],
-        }),
-      };
-      
-      // Mock getAgentClient to return a promise that resolves to mockAgent
-      (getAgentClient as jest.Mock).mockResolvedValue(mockAgent);
-
-      const result = await agent_state({
-        tenant_id: mockTenantId,
-        assistant_id: mockAssistantId,
-        thread_id: mockThreadId,
-      });
-
-      // Verify getAgentClient was called with correct arguments
-      expect(getAgentClient).toHaveBeenCalledWith(mockTenantId, mockAssistantId);
-      
-      // Verify getState was called on the resolved agent (not on a promise)
-      expect(mockAgent.getState).toHaveBeenCalledWith({
-        configurable: { thread_id: mockThreadId, subgraphs: false },
-      });
-      
-      expect(result).toEqual({
-        values: { messages: [] },
-        next: [],
-      });
-    });
-
-    it('should throw error when agent is not found', async () => {
-      (getAgentClient as jest.Mock).mockResolvedValue(null);
-
-      await expect(agent_state({
-        tenant_id: mockTenantId,
-        assistant_id: mockAssistantId,
-        thread_id: mockThreadId,
-      })).rejects.toThrow(`Agent ${mockAssistantId} not found`);
-    });
-
-    it('should throw error when getAgentClient returns undefined', async () => {
-      (getAgentClient as jest.Mock).mockResolvedValue(undefined);
-
-      await expect(agent_state({
-        tenant_id: mockTenantId,
-        assistant_id: mockAssistantId,
-        thread_id: mockThreadId,
-      })).rejects.toThrow(`Agent ${mockAssistantId} not found`);
-    });
-  });
-
-  describe('agent_messages', () => {
-    it('should properly await getAgentClient before calling getState', async () => {
-      const mockMessages = [
-        { id: '1', getType: () => 'human', content: 'Hello', lc_kwargs: {} },
-        { id: '2', getType: () => 'ai', content: 'Hi there', lc_kwargs: {} },
-      ];
-      
-      const mockAgent = {
-        getState: jest.fn().mockResolvedValue({
-          values: { messages: mockMessages },
-        }),
-      };
-      
-      (getAgentClient as jest.Mock).mockResolvedValue(mockAgent);
-
-      const result = await agent_messages({
-        tenant_id: mockTenantId,
-        assistant_id: mockAssistantId,
-        thread_id: mockThreadId,
-      });
-
-      expect(getAgentClient).toHaveBeenCalledWith(mockTenantId, mockAssistantId);
-      expect(mockAgent.getState).toHaveBeenCalledWith({
-        configurable: { thread_id: mockThreadId, subgraphs: false },
-      });
-      
-      expect(result).toHaveLength(2);
-      expect(result[0].role).toBe('human');
-      expect(result[1].role).toBe('ai');
-    });
-  });
-
-  describe('draw_graph', () => {
-    it('should properly await getAgentClient before calling getGraphAsync', async () => {
-      const mockGraph = {
-        drawMermaid: jest.fn().mockResolvedValue('graph TD; A-->B;'),
-      };
-      
-      const mockAgent = {
-        getGraphAsync: jest.fn().mockResolvedValue(mockGraph),
-      };
-      
-      (getAgentClient as jest.Mock).mockResolvedValue(mockAgent);
-
-      const result = await draw_graph(mockAssistantId, mockTenantId);
-
-      expect(getAgentClient).toHaveBeenCalledWith(mockTenantId, mockAssistantId);
-      expect(mockAgent.getGraphAsync).toHaveBeenCalled();
-      expect(mockGraph.drawMermaid).toHaveBeenCalled();
-      expect(result).toBe('graph TD; A-->B;');
-    });
-
-    it('should throw error when agent is not found', async () => {
-      (getAgentClient as jest.Mock).mockResolvedValue(null);
-
-      await expect(draw_graph(mockAssistantId, mockTenantId))
-        .rejects.toThrow(`Agent ${mockAssistantId} not found`);
-    });
-  });
-
-  describe('agent_stream', () => {
-    it('should properly await getAgentClient before calling stream', async () => {
-      const mockStream = async function* () {
-        yield ['updates', { messages: [{ toDict: () => ({ type: 'ai', content: 'Hello' }) }] }];
-      }();
-      
-      const mockAgent = {
-        stream: jest.fn().mockResolvedValue(mockStream),
-      };
-      
-      (getAgentClient as jest.Mock).mockResolvedValue(mockAgent);
-      (agentLatticeManager.getAgentLatticeWithTenant as jest.Mock).mockReturnValue({
-        config: { runConfig: {} },
-      });
-
-      const result = await agent_stream({
-        tenant_id: mockTenantId,
-        assistant_id: mockAssistantId,
-        thread_id: mockThreadId,
-        input: { message: 'Test' },
-      });
-
-      expect(getAgentClient).toHaveBeenCalledWith(mockTenantId, mockAssistantId);
-      expect(mockAgent.stream).toHaveBeenCalled();
-      expect(typeof result[Symbol.asyncIterator]).toBe('function');
-    });
-
-    it('should throw error when agent is not found', async () => {
-      (getAgentClient as jest.Mock).mockResolvedValue(null);
-      (agentLatticeManager.getAgentLatticeWithTenant as jest.Mock).mockReturnValue({
-        config: { runConfig: {} },
-      });
-
-      // The function will throw but also try to call chunkBuffer.abortThread
-      // which might fail in test environment, so we just check it throws
-      await expect(agent_stream({
-        tenant_id: mockTenantId,
-        assistant_id: mockAssistantId,
-        thread_id: mockThreadId,
-        input: { message: 'Test' },
-      })).rejects.toThrow();
-    });
-  });
-
-  describe('agent_invoke', () => {
-    it('should use agentLattice.client directly without calling getAgentClient', async () => {
-      const mockResult = {
-        messages: [
-          { toDict: () => ({ type: 'ai', content: 'Response' }) },
-        ],
-      };
-      
-      const mockAgent = {
-        invoke: jest.fn().mockResolvedValue(mockResult),
-      };
-      
-      (agentLatticeManager.getAgentLatticeWithTenant as jest.Mock).mockReturnValue({
-        client: mockAgent,
-        config: { runConfig: {} },
-      });
-
-      const result = await agent_invoke({
-        tenant_id: mockTenantId,
-        assistant_id: mockAssistantId,
-        thread_id: mockThreadId,
-        input: { message: 'Hello' },
-      });
-
-      // agent_invoke uses agentLattice.client directly, not getAgentClient
-      expect(getAgentClient).not.toHaveBeenCalled();
-      expect(mockAgent.invoke).toHaveBeenCalled();
-      expect(result.messages).toHaveLength(1);
-    });
-
-    it('should throw error when agent lattice client is not found', async () => {
-      (agentLatticeManager.getAgentLatticeWithTenant as jest.Mock).mockReturnValue({
-        client: null,
-        config: { runConfig: {} },
-      });
-
-      await expect(agent_invoke({
-        tenant_id: mockTenantId,
-        assistant_id: mockAssistantId,
-        thread_id: mockThreadId,
-        input: { message: 'Hello' },
-      })).rejects.toThrow(`Agent ${mockAssistantId} not found`);
-    });
-  });
-});