去掉sdk的引入，客户端也能直连

app/api/bedrock.ts

@@ -1,131 +1,186 @@
-import { getServerSideConfig } from "../config/server";
-import { prettyObject } from "../utils/format";
 import { NextRequest, NextResponse } from "next/server";
-import { decrypt } from "../utils/encryption";
-import {
-  BedrockRuntimeClient,
-  ConverseStreamCommand,
-  ConverseStreamCommandInput,
-  Message,
-  ContentBlock,
-  ConverseStreamOutput,
-} from "@aws-sdk/client-bedrock-runtime";
+import { sign, decrypt } from "../utils/aws";
 
-const ALLOWED_PATH = new Set(["converse"]);
+const ALLOWED_PATH = new Set(["chat", "models"]);
 
-// AWS Credential Validation Function
-function validateAwsCredentials(
-  region: string,
-  accessKeyId: string,
-  secretAccessKey: string,
-): boolean {
-  const regionRegex = /^[a-z]{2}-[a-z]+-\d+$/;
-  const accessKeyRegex = /^(AKIA|A3T|ASIA)[A-Z0-9]{16}$/;
-
-  return (
-    regionRegex.test(region) &&
-    accessKeyRegex.test(accessKeyId) &&
-    secretAccessKey.length === 40
-  );
+function parseEventData(chunk: Uint8Array): any {
+  const decoder = new TextDecoder();
+  const text = decoder.decode(chunk);
+  try {
+    return JSON.parse(text);
+  } catch (e) {
+    try {
+      const base64Match = text.match(/:"([A-Za-z0-9+/=]+)"/);
+      if (base64Match) {
+        const decoded = Buffer.from(base64Match[1], "base64").toString("utf-8");
+        return JSON.parse(decoded);
+      }
+      const eventMatch = text.match(/:event-type[^\{]+({.*})/);
+      if (eventMatch) {
+        return JSON.parse(eventMatch[1]);
+      }
+    } catch (innerError) {}
+  }
+  return null;
 }
 
-export interface ConverseRequest {
-  modelId: string;
-  messages: {
-    role: "user" | "assistant" | "system";
-    content: string | any[];
-  }[];
-  inferenceConfig?: {
-    maxTokens?: number;
-    temperature?: number;
-    topP?: number;
-    stopSequences?: string[];
-  };
-  tools?: {
-    name: string;
-    description?: string;
-    input_schema: any;
-  }[];
-  stream?: boolean;
+async function* transformBedrockStream(stream: ReadableStream) {
+  const reader = stream.getReader();
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+
+      const parsed = parseEventData(value);
+      if (parsed) {
+        if (parsed.type === "content_block_delta") {
+          if (parsed.delta?.type === "text_delta") {
+            yield `data: ${JSON.stringify({
+              delta: { text: parsed.delta.text },
+            })}\n\n`;
+          } else if (parsed.delta?.type === "input_json_delta") {
+            yield `data: ${JSON.stringify(parsed)}\n\n`;
+          }
+        } else if (
+          parsed.type === "message_delta" &&
+          parsed.delta?.stop_reason
+        ) {
+          yield `data: ${JSON.stringify({
+            delta: { stop_reason: parsed.delta.stop_reason },
+          })}\n\n`;
+        } else if (
+          parsed.type === "content_block_start" &&
+          parsed.content_block?.type === "tool_use"
+        ) {
+          yield `data: ${JSON.stringify(parsed)}\n\n`;
+        } else if (parsed.type === "content_block_stop") {
+          yield `data: ${JSON.stringify(parsed)}\n\n`;
+        }
+      }
+    }
+  } finally {
+    reader.releaseLock();
+  }
 }
 
-function supportsToolUse(modelId: string): boolean {
-  return modelId.toLowerCase().includes("claude-3");
+function validateRequest(body: any, modelId: string): void {
+  if (!modelId) throw new Error("Model ID is required");
+
+  if (modelId.startsWith("anthropic.claude")) {
+    if (
+      !body.anthropic_version ||
+      body.anthropic_version !== "bedrock-2023-05-31"
+    ) {
+      throw new Error("anthropic_version must be 'bedrock-2023-05-31'");
+    }
+    if (typeof body.max_tokens !== "number" || body.max_tokens < 0) {
+      throw new Error("max_tokens must be a positive number");
+    }
+    if (modelId.startsWith("anthropic.claude-3")) {
+      if (!Array.isArray(body.messages))
+        throw new Error("messages array is required for Claude 3");
+    } else if (typeof body.prompt !== "string") {
+      throw new Error("prompt is required for Claude 2 and earlier");
+    }
+  } else if (modelId.startsWith("meta.llama")) {
+    if (!body.prompt) throw new Error("Llama requires a prompt");
+  } else if (modelId.startsWith("mistral.mistral")) {
+    if (!Array.isArray(body.messages))
+      throw new Error("Mistral requires a messages array");
+  } else if (modelId.startsWith("amazon.titan")) {
+    if (!body.inputText) throw new Error("Titan requires inputText");
+  }
 }
 
-function formatRequestBody(
-  request: ConverseRequest,
-): ConverseStreamCommandInput {
-  const messages: Message[] = request.messages.map((msg) => ({
-    role: msg.role === "system" ? "user" : msg.role,
-    content: Array.isArray(msg.content)
-      ? msg.content.map((item) => {
-          if (item.type === "tool_use") {
-            return {
-              toolUse: {
-                toolUseId: item.id,
-                name: item.name,
-                input: item.input || "{}",
-              },
-            } as ContentBlock;
-          }
-          if (item.type === "tool_result") {
-            return {
-              toolResult: {
-                toolUseId: item.tool_use_id,
-                content: [{ text: item.content || ";" }],
-                status: "success",
-              },
-            } as ContentBlock;
-          }
-          if (item.type === "text") {
-            return { text: item.text || ";" } as ContentBlock;
-          }
-          if (item.type === "image") {
-            return {
-              image: {
-                format: item.source.media_type.split("/")[1] as
-                  | "png"
-                  | "jpeg"
-                  | "gif"
-                  | "webp",
-                source: {
-                  bytes: Uint8Array.from(
-                    Buffer.from(item.source.data, "base64"),
-                  ),
-                },
-              },
-            } as ContentBlock;
-          }
-          return { text: ";" } as ContentBlock;
-        })
-      : [{ text: msg.content || ";" } as ContentBlock],
-  }));
+async function requestBedrock(req: NextRequest) {
+  const controller = new AbortController();
+  const awsRegion = req.headers.get("X-Region") ?? "";
+  const awsAccessKey = req.headers.get("X-Access-Key") ?? "";
+  const awsSecretKey = req.headers.get("X-Secret-Key") ?? "";
+  const awsSessionToken = req.headers.get("X-Session-Token");
+  const modelId = req.headers.get("X-Model-Id") ?? "";
 
-  const input: ConverseStreamCommandInput = {
-    modelId: request.modelId,
-    messages,
-    ...(request.inferenceConfig && {
-      inferenceConfig: request.inferenceConfig,
-    }),
-  };
-
-  if (request.tools?.length && supportsToolUse(request.modelId)) {
-    input.toolConfig = {
-      tools: request.tools.map((tool) => ({
-        toolSpec: {
-          name: tool.name,
-          description: tool.description,
-          inputSchema: {
-            json: tool.input_schema,
-          },
-        },
-      })),
-      toolChoice: { auto: {} },
-    };
+  if (!awsRegion || !awsAccessKey || !awsSecretKey || !modelId) {
+    throw new Error("Missing required AWS credentials or model ID");
   }
 
-  return input;
+  const decryptedAccessKey = decrypt(awsAccessKey);
+  const decryptedSecretKey = decrypt(awsSecretKey);
+  const decryptedSessionToken = awsSessionToken
+    ? decrypt(awsSessionToken)
+    : undefined;
+
+  if (!decryptedAccessKey || !decryptedSecretKey) {
+    throw new Error("Failed to decrypt AWS credentials");
+  }
+
+  const endpoint = `https://bedrock-runtime.${awsRegion}.amazonaws.com/model/${modelId}/invoke-with-response-stream`;
+  const timeoutId = setTimeout(() => controller.abort(), 10 * 60 * 1000);
+
+  try {
+    const bodyText = await req.clone().text();
+    const bodyJson = JSON.parse(bodyText);
+    validateRequest(bodyJson, modelId);
+    const canonicalBody = JSON.stringify(bodyJson);
+
+    const headers = await sign({
+      method: "POST",
+      url: endpoint,
+      region: awsRegion,
+      accessKeyId: decryptedAccessKey,
+      secretAccessKey: decryptedSecretKey,
+      sessionToken: decryptedSessionToken,
+      body: canonicalBody,
+      service: "bedrock",
+    });
+
+    const res = await fetch(endpoint, {
+      method: "POST",
+      headers,
+      body: canonicalBody,
+      redirect: "manual",
+      // @ts-ignore
+      duplex: "half",
+      signal: controller.signal,
+    });
+
+    if (!res.ok) {
+      const error = await res.text();
+      try {
+        const errorJson = JSON.parse(error);
+        throw new Error(errorJson.message || error);
+      } catch {
+        throw new Error(error);
+      }
+    }
+
+    const transformedStream = transformBedrockStream(res.body!);
+    const stream = new ReadableStream({
+      async start(controller) {
+        try {
+          for await (const chunk of transformedStream) {
+            controller.enqueue(new TextEncoder().encode(chunk));
+          }
+          controller.close();
+        } catch (err) {
+          controller.error(err);
+        }
+      },
+    });
+
+    return new Response(stream, {
+      headers: {
+        "Content-Type": "text/event-stream",
+        "Cache-Control": "no-cache",
+        Connection: "keep-alive",
+        "X-Accel-Buffering": "no",
+      },
+    });
+  } catch (e) {
+    throw e;
+  } finally {
+    clearTimeout(timeoutId);
+  }
 }
 
 export async function handle(
@@ -139,166 +194,16 @@ export async function handle(
   const subpath = params.path.join("/");
   if (!ALLOWED_PATH.has(subpath)) {
     return NextResponse.json(
-      { error: true, msg: "Path not allowed: " + subpath },
+      { error: true, msg: "you are not allowed to request " + subpath },
       { status: 403 },
     );
   }
 
-  const serverConfig = getServerSideConfig();
-  let region = serverConfig.awsRegion;
-  let accessKeyId = serverConfig.awsAccessKey;
-  let secretAccessKey = serverConfig.awsSecretKey;
-  let sessionToken = undefined;
-
-  // Attempt to get credentials from headers if not in server config
-  if (!region || !accessKeyId || !secretAccessKey) {
-    region = decrypt(req.headers.get("X-Region") ?? "");
-    accessKeyId = decrypt(req.headers.get("X-Access-Key") ?? "");
-    secretAccessKey = decrypt(req.headers.get("X-Secret-Key") ?? "");
-    sessionToken = req.headers.get("X-Session-Token")
-      ? decrypt(req.headers.get("X-Session-Token") ?? "")
-      : undefined;
-  }
-
-  // Validate AWS credentials
-  if (!validateAwsCredentials(region, accessKeyId, secretAccessKey)) {
-    return NextResponse.json(
-      {
-        error: true,
-        msg: "Invalid AWS credentials. Please check your region, access key, and secret key.",
-      },
-      { status: 401 },
-    );
-  }
-
   try {
-    const client = new BedrockRuntimeClient({
-      region,
-      credentials: {
-        accessKeyId,
-        secretAccessKey,
-        sessionToken,
-      },
-    });
-
-    const body = (await req.json()) as ConverseRequest;
-    const command = new ConverseStreamCommand(formatRequestBody(body));
-    const response = await client.send(command);
-
-    if (!response.stream) {
-      throw new Error("No stream in response");
-    }
-
-    // If stream is false, accumulate the response and return as JSON
-    if (body.stream === false) {
-      let fullResponse = {
-        content: "",
-      };
-
-      const responseStream =
-        response.stream as AsyncIterable<ConverseStreamOutput>;
-      for await (const event of responseStream) {
-        if (
-          "contentBlockDelta" in event &&
-          event.contentBlockDelta?.delta &&
-          "text" in event.contentBlockDelta.delta &&
-          event.contentBlockDelta.delta.text
-        ) {
-          fullResponse.content += event.contentBlockDelta.delta.text;
-        }
-      }
-
-      return NextResponse.json(fullResponse);
-    }
-
-    // Otherwise, return streaming response
-    const stream = new ReadableStream({
-      async start(controller) {
-        try {
-          const responseStream =
-            response.stream as AsyncIterable<ConverseStreamOutput>;
-          for await (const event of responseStream) {
-            if (
-              "contentBlockStart" in event &&
-              event.contentBlockStart?.start?.toolUse &&
-              event.contentBlockStart.contentBlockIndex !== undefined
-            ) {
-              controller.enqueue(
-                `data: ${JSON.stringify({
-                  type: "content_block",
-                  content_block: {
-                    type: "tool_use",
-                    id: event.contentBlockStart.start.toolUse.toolUseId,
-                    name: event.contentBlockStart.start.toolUse.name,
-                  },
-                  index: event.contentBlockStart.contentBlockIndex,
-                })}\n\n`,
-              );
-            } else if (
-              "contentBlockDelta" in event &&
-              event.contentBlockDelta?.delta &&
-              event.contentBlockDelta.contentBlockIndex !== undefined
-            ) {
-              const delta = event.contentBlockDelta.delta;
-
-              if ("text" in delta && delta.text) {
-                controller.enqueue(
-                  `data: ${JSON.stringify({
-                    type: "content_block_delta",
-                    delta: {
-                      type: "text_delta",
-                      text: delta.text,
-                    },
-                    index: event.contentBlockDelta.contentBlockIndex,
-                  })}\n\n`,
-                );
-              } else if ("toolUse" in delta && delta.toolUse?.input) {
-                controller.enqueue(
-                  `data: ${JSON.stringify({
-                    type: "content_block_delta",
-                    delta: {
-                      type: "input_json_delta",
-                      partial_json: delta.toolUse.input,
-                    },
-                    index: event.contentBlockDelta.contentBlockIndex,
-                  })}\n\n`,
-                );
-              }
-            } else if (
-              "contentBlockStop" in event &&
-              event.contentBlockStop?.contentBlockIndex !== undefined
-            ) {
-              controller.enqueue(
-                `data: ${JSON.stringify({
-                  type: "content_block_stop",
-                  index: event.contentBlockStop.contentBlockIndex,
-                })}\n\n`,
-              );
-            }
-          }
-          controller.close();
-        } catch (error) {
-          console.error("[Bedrock] Stream error:", error);
-          controller.error(error);
-        }
-      },
-    });
-
-    return new Response(stream, {
-      headers: {
-        "Content-Type": "text/event-stream",
-        "Cache-Control": "no-cache",
-        Connection: "keep-alive",
-      },
-    });
+    return await requestBedrock(req);
   } catch (e) {
-    console.error("[Bedrock] Error:", e);
     return NextResponse.json(
-      {
-        error: true,
-        message: e instanceof Error ? e.message : "Unknown error",
-        details: prettyObject(e),
-      },
+      { error: true, msg: e instanceof Error ? e.message : "Unknown error" },
       { status: 500 },
     );
   }