Explorar o código

feat(core): use native transport for custom anthropic chat

Ma hai 6 meses
pai
achega
5f10f209b2
Modificáronse 2 ficheiros con 203 adicións e 1 borrados
  1. 82 0
      packages/core/src/__tests__/provider.test.ts
  2. 121 1
      packages/core/src/llm/provider.ts

+ 82 - 0
packages/core/src/__tests__/provider.test.ts

@@ -286,6 +286,88 @@ describe("chatCompletion via pi-ai", () => {
 
     vi.unstubAllGlobals();
   });
+
+  it("uses native fetch transport for custom anthropic-compatible non-stream chat", async () => {
+    const fetchMock = vi.fn().mockResolvedValue({
+      ok: true,
+      json: async () => ({
+        content: [{ type: "text", text: "你好,Anthropic!" }],
+        usage: { input_tokens: 5, output_tokens: 3 },
+      }),
+    });
+    vi.stubGlobal("fetch", fetchMock);
+
+    const client = makeClient(0.7, {
+      provider: "anthropic",
+      service: "custom",
+      stream: false,
+      _piModel: {
+        ...MOCK_PI_MODEL,
+        provider: "anthropic",
+        api: "anthropic-messages" as Api,
+        baseUrl: "https://gateway.example",
+      },
+    });
+    const result = await chatCompletion(client, "claude-sonnet-4-6", [{ role: "user", content: "nihao" }]);
+
+    expect(result.content).toBe("你好,Anthropic!");
+    expect(result.usage.promptTokens).toBe(5);
+    expect(result.usage.completionTokens).toBe(3);
+    expect(fetchMock).toHaveBeenCalledOnce();
+    expect(mockCompleteSimple).not.toHaveBeenCalled();
+    expect(mockStreamSimple).not.toHaveBeenCalled();
+
+    vi.unstubAllGlobals();
+  });
+
+  it("uses native fetch transport for custom anthropic-compatible stream chat", async () => {
+    const encoder = new TextEncoder();
+    const sse = [
+      "event: message_start\n",
+      "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":4}}}\n\n",
+      "event: content_block_delta\n",
+      "data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"你\"}}\n\n",
+      "event: content_block_delta\n",
+      "data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"好\"}}\n\n",
+      "event: message_delta\n",
+      "data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":2}}\n\n",
+      "event: message_stop\n",
+      "data: {\"type\":\"message_stop\"}\n\n",
+    ].join("");
+    const fetchMock = vi.fn().mockResolvedValue({
+      ok: true,
+      body: new ReadableStream({
+        start(controller) {
+          controller.enqueue(encoder.encode(sse));
+          controller.close();
+        },
+      }),
+    });
+    vi.stubGlobal("fetch", fetchMock);
+
+    const client = makeClient(0.7, {
+      provider: "anthropic",
+      service: "custom",
+      stream: true,
+      _piModel: {
+        ...MOCK_PI_MODEL,
+        provider: "anthropic",
+        api: "anthropic-messages" as Api,
+        baseUrl: "https://gateway.example",
+      },
+    });
+    const result = await chatCompletion(client, "claude-sonnet-4-6", [{ role: "user", content: "nihao" }]);
+
+    expect(result.content).toBe("你好");
+    expect(result.usage.promptTokens).toBe(4);
+    expect(result.usage.completionTokens).toBe(2);
+    expect(result.usage.totalTokens).toBe(6);
+    expect(fetchMock).toHaveBeenCalledOnce();
+    expect(mockCompleteSimple).not.toHaveBeenCalled();
+    expect(mockStreamSimple).not.toHaveBeenCalled();
+
+    vi.unstubAllGlobals();
+  });
 });
 
 describe("chatCompletion fixed-temperature clamp (thinking models)", () => {

+ 121 - 1
packages/core/src/llm/provider.ts

@@ -308,7 +308,7 @@ function wrapLLMError(error: unknown, context?: { readonly baseUrl?: string; rea
 }
 
 function shouldUseNativeCustomTransport(client: LLMClient): boolean {
-  return client.provider === "openai" && client.service === "custom";
+  return client.service === "custom" && (client.provider === "openai" || client.provider === "anthropic");
 }
 
 function buildCustomHeaders(client: LLMClient): Record<string, string> {
@@ -335,6 +335,15 @@ function buildChatMessages(messages: ReadonlyArray<LLMMessage>): Array<{ role: s
     }));
 }
 
+function buildAnthropicMessages(messages: ReadonlyArray<LLMMessage>): Array<{ role: "user" | "assistant"; content: string }> {
+  return messages
+    .filter((message): message is Readonly<LLMMessage> & { role: "user" | "assistant" } => message.role === "user" || message.role === "assistant")
+    .map((message) => ({
+      role: message.role,
+      content: message.content,
+    }));
+}
+
 function buildResponsesInput(messages: ReadonlyArray<LLMMessage>): Array<{ role: string; content: Array<{ type: "input_text"; text: string }> }> {
   return messages
     .filter((message) => message.role !== "system")
@@ -415,6 +424,114 @@ function extractResponsesContent(json: any): string {
     .join("");
 }
 
+function extractAnthropicContent(json: any): string {
+  const content = Array.isArray(json?.content) ? json.content : [];
+  return content
+    .map((part: any) => typeof part?.text === "string" ? part.text : "")
+    .join("");
+}
+
+async function chatCompletionViaCustomAnthropicCompatible(
+  client: LLMClient,
+  model: string,
+  messages: ReadonlyArray<LLMMessage>,
+  resolved: { readonly temperature: number; readonly maxTokens: number; readonly extra: Record<string, unknown> },
+  onStreamProgress?: OnStreamProgress,
+  onTextDelta?: (text: string) => void,
+): Promise<LLMResponse> {
+  const baseUrl = client._piModel?.baseUrl ?? "";
+  const errorCtx = { baseUrl, model };
+  const monitor = createStreamMonitor(onStreamProgress);
+  const extra = stripReservedKeys(resolved.extra);
+  const payload: Record<string, unknown> = {
+    model,
+    messages: buildAnthropicMessages(messages),
+    stream: client.stream,
+    max_tokens: resolved.maxTokens,
+    temperature: resolved.temperature,
+    ...extra,
+  };
+  const system = joinSystemPrompt(messages);
+  if (system) payload.system = system;
+
+  const response = await fetch(`${baseUrl.replace(/\/$/, "")}/messages`, {
+    method: "POST",
+    headers: {
+      "x-api-key": client._apiKey ?? "",
+      "anthropic-version": "2023-06-01",
+      "Content-Type": "application/json",
+      Authorization: `Bearer ${client._apiKey ?? ""}`,
+      ...(client._piModel?.headers ?? {}),
+    },
+    body: JSON.stringify(payload),
+  });
+
+  if (!response.ok) {
+    throw wrapLLMError(new Error(await readErrorResponse(response)), errorCtx);
+  }
+
+  if (!client.stream) {
+    const json = await response.json() as any;
+    const content = extractAnthropicContent(json);
+    if (!content) {
+      throw wrapLLMError(new Error("LLM returned empty response"), errorCtx);
+    }
+    return {
+      content,
+      usage: {
+        promptTokens: json?.usage?.input_tokens ?? 0,
+        completionTokens: json?.usage?.output_tokens ?? 0,
+        totalTokens: (json?.usage?.input_tokens ?? 0) + (json?.usage?.output_tokens ?? 0),
+      },
+    };
+  }
+
+  const reader = response.body?.getReader();
+  if (!reader) throw wrapLLMError(new Error("Streaming body unavailable"), errorCtx);
+  const decoder = new TextDecoder();
+  let buffer = "";
+  let content = "";
+  let usage = { promptTokens: 0, completionTokens: 0, totalTokens: 0 };
+
+  try {
+    while (true) {
+      const { value, done } = await reader.read();
+      if (done) break;
+      buffer += decoder.decode(value, { stream: true });
+      const parsed = parseSseEvents(buffer);
+      buffer = parsed.rest;
+      for (const event of parsed.events) {
+        if (!event.data) continue;
+        const json = JSON.parse(event.data);
+        if (json.type === "message_start" && json.message?.usage) {
+          usage.promptTokens = json.message.usage.input_tokens ?? usage.promptTokens;
+        }
+        if (json.type === "content_block_delta" && json.delta?.type === "text_delta" && typeof json.delta.text === "string") {
+          content += json.delta.text;
+          monitor.onChunk(json.delta.text);
+          onTextDelta?.(json.delta.text);
+        }
+        if (json.type === "message_delta" && json.usage) {
+          usage.completionTokens = json.usage.output_tokens ?? usage.completionTokens;
+        }
+        if (json.type === "message_stop") {
+          usage.totalTokens = usage.promptTokens + usage.completionTokens;
+        }
+      }
+    }
+  } finally {
+    monitor.stop();
+  }
+
+  if (!content) {
+    throw wrapLLMError(new Error("LLM returned empty response from stream"), errorCtx);
+  }
+  if (!usage.totalTokens) {
+    usage.totalTokens = usage.promptTokens + usage.completionTokens;
+  }
+  return { content, usage };
+}
+
 async function chatCompletionViaCustomOpenAICompatible(
   client: LLMClient,
   model: string,
@@ -423,6 +540,9 @@ async function chatCompletionViaCustomOpenAICompatible(
   onStreamProgress?: OnStreamProgress,
   onTextDelta?: (text: string) => void,
 ): Promise<LLMResponse> {
+  if (client.provider === "anthropic") {
+    return chatCompletionViaCustomAnthropicCompatible(client, model, messages, resolved, onStreamProgress, onTextDelta);
+  }
   const baseUrl = client._piModel?.baseUrl ?? "";
   const headers = buildCustomHeaders(client);
   const errorCtx = { baseUrl, model };