add agents to support route image

src/index.ts

@@ -16,7 +16,13 @@ import createWriteStream from "pino-rotating-file-stream";
 import { HOME_DIR } from "./constants";
 import { configureLogging } from "./utils/log";
 import { sessionUsageCache } from "./utils/cache";
-import Stream from "node:stream";
+import {SSEParserTransform} from "./utils/SSEParser.transform";
+import {SSESerializerTransform} from "./utils/SSESerializer.transform";
+import {rewriteStream} from "./utils/rewriteStream";
+import JSON5 from "json5";
+import { IAgent } from "./agents/type";
+import agentsManager from "./agents";
+
 
 async function initializeClaudeConfig() {
   const homeDir = homedir();
@@ -58,7 +64,7 @@ async function run(options: RunOptions = {}) {
   // Configure logging based on config
   configureLogging(config);
 
-  let HOST = config.HOST;
+  let HOST = config.HOST || "127.0.0.1";
 
   if (config.HOST && !config.APIKEY) {
     HOST = "127.0.0.1";
@@ -82,7 +88,6 @@ async function run(options: RunOptions = {}) {
     cleanupPidFile();
     process.exit(0);
   });
-  console.log(HOST);
 
   // Use port from environment variable if set (for background process)
   const servicePort = process.env.SERVICE_PORT
@@ -131,12 +136,135 @@ async function run(options: RunOptions = {}) {
   });
   server.addHook("preHandler", async (req, reply) => {
     if (req.url.startsWith("/v1/messages")) {
-      router(req, reply, config);
+      const useAgents = []
+
+      for (const agent of agentsManager.getAllAgents()) {
+        if (agent.shouldHandle(req, config)) {
+          // 设置agent标识
+          useAgents.push(agent.name)
+
+          // change request body
+          agent.reqHandler(req, config);
+
+          // append agent tools
+          if (agent.tools.size) {
+            req.body.tools.unshift(...Array.from(agent.tools.values()).map(item => {
+              return {
+                name: item.name,
+                description: item.description,
+                input_schema: item.input_schema
+              }
+            }))
+          }
+        }
+      }
+
+      if (useAgents.length) {
+        req.agents = useAgents;
+      }
+      await router(req, reply, config);
     }
   });
-  server.addHook("onSend", (req, reply, payload, done) => {
+  server.addHook("onSend", async (req, reply, payload) => {
     if (req.sessionId && req.url.startsWith("/v1/messages")) {
       if (payload instanceof ReadableStream) {
+        if (req.agents) {
+          const eventStream = payload.pipeThrough(new SSEParserTransform())
+          let currentAgent: undefined | IAgent;
+          let currentToolIndex = -1
+          let currentToolName = ''
+          let currentToolArgs = ''
+          let currentToolId = ''
+          const toolMessages: any[] = []
+          const assistantMessages: any[] = []
+          // 存储Anthropic格式的消息体，区分文本和工具类型
+          return rewriteStream(eventStream, async (data, controller) => {
+            // 检测工具调用开始
+            if (data.event === 'content_block_start' && data?.data?.content_block?.name) {
+              const agent = req.agents.find((name: string) => agentsManager.getAgent(name)?.tools.get(data.data.content_block.name))
+              if (agent) {
+                currentAgent = agentsManager.getAgent(agent)
+                currentToolIndex = data.data.index
+                currentToolName = data.data.content_block.name
+                currentToolId = data.data.content_block.id
+                return undefined;
+              }
+            }
+
+            // 收集工具参数
+            if (currentToolIndex > -1 && data.data.index === currentToolIndex && data.data?.delta?.type === 'input_json_delta') {
+              currentToolArgs += data.data?.delta?.partial_json;
+              return undefined;
+            }
+
+            // 工具调用完成，处理agent调用
+            if (currentToolIndex > -1 && data.data.index === currentToolIndex && data.data.type === 'content_block_stop') {
+              try {
+                const args = JSON5.parse(currentToolArgs);
+                assistantMessages.push({
+                  type: "tool_use",
+                  id: currentToolId,
+                  name: currentToolName,
+                  input: args
+                })
+                const toolResult = await currentAgent?.tools.get(currentToolName)?.handler(args, {
+                  req,
+                  config
+                });
+                toolMessages.push({
+                  "tool_use_id": currentToolId,
+                  "type": "tool_result",
+                  "content": toolResult
+                })
+                currentAgent = undefined
+                currentToolIndex = -1
+                currentToolName = ''
+                currentToolArgs = ''
+                currentToolId = ''
+              } catch (e) {
+                console.log(e);
+              }
+              return undefined;
+            }
+
+            if (data.event === 'message_delta' && toolMessages.length) {
+              req.body.messages.push({
+                role: 'assistant',
+                content: assistantMessages
+              })
+              req.body.messages.push({
+                role: 'user',
+                content: toolMessages
+              })
+              const response = await fetch(`http://127.0.0.1:${config.PORT}/v1/messages`, {
+                method: "POST",
+                headers: {
+                  'x-api-key': config.APIKEY,
+                  'content-type': 'application/json',
+                },
+                body: JSON.stringify(req.body),
+              })
+              if (!response.ok) {
+                return undefined;
+              }
+              const stream = response.body!.pipeThrough(new SSEParserTransform())
+              const reader = stream.getReader()
+              while (true) {
+                const {value, done} = await reader.read();
+                if (done) {
+                  break;
+                }
+                if (['message_start', 'message_stop'].includes(value.event)) {
+                  continue
+                }
+                controller.enqueue(value)
+              }
+              return undefined
+            }
+            return data
+          }).pipeThrough(new SSESerializerTransform())
+        }
+
         const [originalStream, clonedStream] = payload.tee();
         const read = async (stream: ReadableStream) => {
           const reader = stream.getReader();
@@ -156,29 +284,13 @@ async function run(options: RunOptions = {}) {
           }
         }
         read(clonedStream);
-        done(null, originalStream)
-      } else {
-        req.log.debug({payload}, 'onSend Hook')
-        sessionUsageCache.put(req.sessionId, payload.usage);
-        if (payload instanceof Buffer || payload instanceof Response) {
-          done(null, payload);
-        } else if(typeof payload === "object") {
-          done(null, JSON.stringify(payload));
-        } else {
-          done(null, payload);
-        }
-      }
-    } else {
-      if(payload instanceof Buffer || payload instanceof Response || payload === null || payload instanceof ReadableStream || payload instanceof Stream) {
-        done(null, payload);
-      } else if(typeof payload === "object") {
-        req.log.debug({payload}, 'onSend Hook')
-        done(null, JSON.stringify(payload));
-      } else {
-        done(null, payload);
+        return originalStream
       }
+      sessionUsageCache.put(req.sessionId, payload.usage);
     }
+    return payload;
   });
+
   server.start();
 }