File size: 2,616 Bytes
5a31018 d4016bc 5a31018 d4016bc 5a31018 d4016bc 5a31018 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import { z } from "zod";
import { ANTHROPIC_API_KEY } from "$env/static/private";
import type { Endpoint } from "../endpoints";
import type { TextGenerationStreamOutput } from "@huggingface/inference";
export const endpointAnthropicParametersSchema = z.object({
weight: z.number().int().positive().default(1),
model: z.any(),
type: z.literal("anthropic"),
baseURL: z.string().url().default("https://api.anthropic.com"),
apiKey: z.string().default(ANTHROPIC_API_KEY ?? "sk-"),
defaultHeaders: z.record(z.string()).optional(),
defaultQuery: z.record(z.string()).optional(),
});
export async function endpointAnthropic(
input: z.input<typeof endpointAnthropicParametersSchema>
): Promise<Endpoint> {
const { baseURL, apiKey, model, defaultHeaders, defaultQuery } =
endpointAnthropicParametersSchema.parse(input);
let Anthropic;
try {
Anthropic = (await import("@anthropic-ai/sdk")).default;
} catch (e) {
throw new Error("Failed to import @anthropic-ai/sdk", { cause: e });
}
const anthropic = new Anthropic({
apiKey,
baseURL,
defaultHeaders,
defaultQuery,
});
return async ({ messages, preprompt, generateSettings }) => {
let system = preprompt;
if (messages?.[0]?.from === "system") {
system = messages[0].content;
}
const messagesFormatted = messages
.filter((message) => message.from !== "system")
.map((message) => ({
role: message.from,
content: message.content,
})) as unknown as {
role: "user" | "assistant";
content: string;
}[];
let tokenId = 0;
const parameters = { ...model.parameters, ...generateSettings };
return (async function* () {
const stream = anthropic.messages.stream({
model: model.id ?? model.name,
messages: messagesFormatted,
max_tokens: parameters?.max_new_tokens,
temperature: parameters?.temperature,
top_p: parameters?.top_p,
top_k: parameters?.top_k,
stop_sequences: parameters?.stop,
system,
});
while (true) {
const result = await Promise.race([stream.emitted("text"), stream.emitted("end")]);
// Stream end
if (result === undefined) {
yield {
token: {
id: tokenId++,
text: "",
logprob: 0,
special: true,
},
generated_text: await stream.finalText(),
details: null,
} satisfies TextGenerationStreamOutput;
return;
}
// Text delta
yield {
token: {
id: tokenId++,
text: result as unknown as string,
special: false,
logprob: 0,
},
generated_text: null,
details: null,
} satisfies TextGenerationStreamOutput;
}
})();
};
}
|