619 lines
15 KiB
TypeScript
619 lines
15 KiB
TypeScript
import type { ChatModel, ModelProvider } from "./modelRegistry.js";
|
|
import { apiKeyForProvider } from "./modelRegistry.js";
|
|
import { normalizeUsage, type Usage } from "./tokens.js";
|
|
|
|
export type ChatRequest = {
|
|
model: ChatModel;
|
|
provider: ModelProvider;
|
|
systemPrompt?: string;
|
|
prompt: string;
|
|
temperature: number;
|
|
maxTokens: number;
|
|
};
|
|
|
|
export type ChatResponse = {
|
|
output: string;
|
|
reasoning?: string;
|
|
usage: Usage;
|
|
rawModel: string;
|
|
};
|
|
|
|
export type StreamDelta = {
|
|
channel: "reasoning" | "output";
|
|
text: string;
|
|
};
|
|
|
|
export type ChatStreamRequest = ChatRequest & {
|
|
onDelta: (delta: StreamDelta) => void;
|
|
signal?: AbortSignal;
|
|
};
|
|
|
|
type OpenAIResponse = {
|
|
choices?: Array<{ message?: { content?: string }; text?: string }>;
|
|
usage?: {
|
|
prompt_tokens?: number;
|
|
completion_tokens?: number;
|
|
total_tokens?: number;
|
|
input_tokens?: number;
|
|
output_tokens?: number;
|
|
};
|
|
};
|
|
|
|
type OpenAIStreamDelta = {
|
|
content?: string;
|
|
reasoning?: string;
|
|
reasoning_content?: string;
|
|
thinking?: string;
|
|
reasoning_details?: Array<{ text?: string; content?: string }>;
|
|
};
|
|
|
|
type OpenAIStreamChunk = {
|
|
choices?: Array<{
|
|
delta?: OpenAIStreamDelta;
|
|
}>;
|
|
usage?: OpenAIResponse["usage"];
|
|
};
|
|
|
|
type AnthropicResponse = {
|
|
content?: Array<{ type?: string; text?: string }>;
|
|
usage?: {
|
|
input_tokens?: number;
|
|
output_tokens?: number;
|
|
};
|
|
};
|
|
|
|
type AnthropicStreamEvent = {
|
|
type?: string;
|
|
content_block?: {
|
|
type?: string;
|
|
text?: string;
|
|
thinking?: string;
|
|
};
|
|
delta?: {
|
|
type?: string;
|
|
text?: string;
|
|
thinking?: string;
|
|
};
|
|
message?: {
|
|
usage?: {
|
|
input_tokens?: number;
|
|
output_tokens?: number;
|
|
};
|
|
};
|
|
usage?: {
|
|
input_tokens?: number;
|
|
output_tokens?: number;
|
|
};
|
|
};
|
|
|
|
type GeminiResponse = {
|
|
candidates?: Array<{
|
|
content?: {
|
|
parts?: Array<{ text?: string }>;
|
|
};
|
|
}>;
|
|
usageMetadata?: {
|
|
promptTokenCount?: number;
|
|
candidatesTokenCount?: number;
|
|
totalTokenCount?: number;
|
|
};
|
|
};
|
|
|
|
type GeminiStreamChunk = GeminiResponse;
|
|
|
|
const readErrorBody = async (response: Response) => {
|
|
const text = await response.text();
|
|
if (!text) {
|
|
return response.statusText;
|
|
}
|
|
|
|
try {
|
|
const json = JSON.parse(text) as { error?: { message?: string }; message?: string };
|
|
return json.error?.message ?? json.message ?? text;
|
|
} catch {
|
|
return text;
|
|
}
|
|
};
|
|
|
|
const postJson = async <T>(
|
|
url: string,
|
|
headers: Record<string, string>,
|
|
body: unknown
|
|
): Promise<T> => {
|
|
const response = await fetch(url, {
|
|
method: "POST",
|
|
headers: {
|
|
"content-type": "application/json",
|
|
...headers
|
|
},
|
|
body: JSON.stringify(body)
|
|
});
|
|
|
|
if (!response.ok) {
|
|
throw new Error(await readErrorBody(response));
|
|
}
|
|
|
|
return (await response.json()) as T;
|
|
};
|
|
|
|
const postJsonStream = async (
|
|
url: string,
|
|
headers: Record<string, string>,
|
|
body: unknown,
|
|
signal?: AbortSignal
|
|
) => {
|
|
const response = await fetch(url, {
|
|
method: "POST",
|
|
headers: {
|
|
"content-type": "application/json",
|
|
...headers
|
|
},
|
|
body: JSON.stringify(body),
|
|
signal
|
|
});
|
|
|
|
if (!response.ok) {
|
|
throw new Error(await readErrorBody(response));
|
|
}
|
|
|
|
if (!response.body) {
|
|
throw new Error("Provider returned an empty stream");
|
|
}
|
|
|
|
return response;
|
|
};
|
|
|
|
const readSse = async (
|
|
response: Response,
|
|
onEvent: (event: { event?: string; data: string }) => void
|
|
) => {
|
|
const reader = response.body?.getReader();
|
|
if (!reader) {
|
|
throw new Error("Provider stream is not readable");
|
|
}
|
|
|
|
const decoder = new TextDecoder();
|
|
let buffer = "";
|
|
let eventName = "";
|
|
let dataLines: string[] = [];
|
|
|
|
const flush = () => {
|
|
if (!dataLines.length) {
|
|
eventName = "";
|
|
return;
|
|
}
|
|
|
|
onEvent({
|
|
event: eventName || undefined,
|
|
data: dataLines.join("\n")
|
|
});
|
|
eventName = "";
|
|
dataLines = [];
|
|
};
|
|
|
|
const handleLine = (rawLine: string) => {
|
|
const line = rawLine.endsWith("\r") ? rawLine.slice(0, -1) : rawLine;
|
|
if (!line) {
|
|
flush();
|
|
return;
|
|
}
|
|
|
|
if (line.startsWith("event:")) {
|
|
eventName = line.slice(6).trim();
|
|
return;
|
|
}
|
|
|
|
if (line.startsWith("data:")) {
|
|
dataLines.push(line.slice(5).trimStart());
|
|
return;
|
|
}
|
|
|
|
if (line.trim().startsWith("{")) {
|
|
onEvent({ data: line.trim() });
|
|
}
|
|
};
|
|
|
|
while (true) {
|
|
const { value, done } = await reader.read();
|
|
if (done) {
|
|
break;
|
|
}
|
|
|
|
buffer += decoder.decode(value, { stream: true });
|
|
const lines = buffer.split("\n");
|
|
buffer = lines.pop() ?? "";
|
|
for (const line of lines) {
|
|
handleLine(line);
|
|
}
|
|
}
|
|
|
|
buffer += decoder.decode();
|
|
if (buffer) {
|
|
handleLine(buffer);
|
|
}
|
|
flush();
|
|
};
|
|
|
|
const endpoint = (provider: ModelProvider, path: string) => {
|
|
const baseUrl = provider.baseUrl?.replace(/\/$/, "");
|
|
if (!baseUrl) {
|
|
throw new Error(`${provider.name} base URL is not configured`);
|
|
}
|
|
|
|
return `${baseUrl}${path}`;
|
|
};
|
|
|
|
const composePrompt = (systemPrompt: string | undefined, prompt: string) =>
|
|
[systemPrompt, prompt].filter(Boolean).join("\n\n");
|
|
|
|
const openAICompatibleMessages = (systemPrompt: string | undefined, prompt: string) => [
|
|
...(systemPrompt ? [{ role: "system", content: systemPrompt }] : []),
|
|
{ role: "user", content: prompt }
|
|
];
|
|
|
|
export const callOpenAICompatible = async ({
|
|
provider,
|
|
model,
|
|
systemPrompt,
|
|
prompt,
|
|
temperature,
|
|
maxTokens
|
|
}: ChatRequest): Promise<ChatResponse> => {
|
|
const apiKey = apiKeyForProvider(provider);
|
|
if (!apiKey) {
|
|
throw new Error(`${provider.name} API key is not configured (${provider.apiKeyEnv})`);
|
|
}
|
|
|
|
const json = await postJson<OpenAIResponse>(
|
|
endpoint(provider, "/chat/completions"),
|
|
{ authorization: `Bearer ${apiKey}` },
|
|
{
|
|
model: model.model,
|
|
messages: openAICompatibleMessages(systemPrompt, prompt),
|
|
temperature,
|
|
max_tokens: maxTokens
|
|
}
|
|
);
|
|
|
|
const output = json.choices?.[0]?.message?.content ?? json.choices?.[0]?.text ?? "";
|
|
const usage = normalizeUsage(composePrompt(systemPrompt, prompt), output, {
|
|
inputTokens: json.usage?.prompt_tokens ?? json.usage?.input_tokens,
|
|
outputTokens: json.usage?.completion_tokens ?? json.usage?.output_tokens,
|
|
totalTokens: json.usage?.total_tokens,
|
|
estimated: !json.usage
|
|
});
|
|
|
|
return { output, usage, rawModel: model.model };
|
|
};
|
|
|
|
const reasoningTextFromDelta = (delta?: OpenAIStreamDelta) => {
|
|
const details = delta?.reasoning_details
|
|
?.map((detail) => detail.text ?? detail.content ?? "")
|
|
.join("");
|
|
return delta?.reasoning_content ?? delta?.reasoning ?? delta?.thinking ?? details ?? "";
|
|
};
|
|
|
|
export const streamOpenAICompatible = async ({
|
|
provider,
|
|
model,
|
|
systemPrompt,
|
|
prompt,
|
|
temperature,
|
|
maxTokens,
|
|
onDelta,
|
|
signal
|
|
}: ChatStreamRequest): Promise<ChatResponse> => {
|
|
const apiKey = apiKeyForProvider(provider);
|
|
if (!apiKey) {
|
|
throw new Error(`${provider.name} API key is not configured (${provider.apiKeyEnv})`);
|
|
}
|
|
|
|
let output = "";
|
|
let reasoning = "";
|
|
let usage: OpenAIResponse["usage"] | undefined;
|
|
|
|
const response = await postJsonStream(
|
|
endpoint(provider, "/chat/completions"),
|
|
{ authorization: `Bearer ${apiKey}` },
|
|
{
|
|
model: model.model,
|
|
messages: openAICompatibleMessages(systemPrompt, prompt),
|
|
temperature,
|
|
max_tokens: maxTokens,
|
|
stream: true,
|
|
stream_options: { include_usage: true }
|
|
},
|
|
signal
|
|
);
|
|
|
|
await readSse(response, ({ data }) => {
|
|
if (data === "[DONE]") {
|
|
return;
|
|
}
|
|
|
|
const chunk = JSON.parse(data) as OpenAIStreamChunk;
|
|
usage = chunk.usage ?? usage;
|
|
const delta = chunk.choices?.[0]?.delta;
|
|
const reasoningDelta = reasoningTextFromDelta(delta);
|
|
const outputDelta = delta?.content ?? "";
|
|
|
|
if (reasoningDelta) {
|
|
reasoning += reasoningDelta;
|
|
onDelta({ channel: "reasoning", text: reasoningDelta });
|
|
}
|
|
|
|
if (outputDelta) {
|
|
output += outputDelta;
|
|
onDelta({ channel: "output", text: outputDelta });
|
|
}
|
|
});
|
|
|
|
return {
|
|
output,
|
|
reasoning,
|
|
usage: normalizeUsage(composePrompt(systemPrompt, prompt), output, {
|
|
inputTokens: usage?.prompt_tokens ?? usage?.input_tokens,
|
|
outputTokens: usage?.completion_tokens ?? usage?.output_tokens,
|
|
totalTokens: usage?.total_tokens,
|
|
estimated: !usage
|
|
}),
|
|
rawModel: model.model
|
|
};
|
|
};
|
|
|
|
export const callAnthropic = async ({
|
|
provider,
|
|
model,
|
|
systemPrompt,
|
|
prompt,
|
|
temperature,
|
|
maxTokens
|
|
}: ChatRequest): Promise<ChatResponse> => {
|
|
const apiKey = apiKeyForProvider(provider);
|
|
if (!apiKey) {
|
|
throw new Error(`${provider.name} API key is not configured (${provider.apiKeyEnv})`);
|
|
}
|
|
|
|
const json = await postJson<AnthropicResponse>(
|
|
endpoint(provider, "/messages"),
|
|
{
|
|
"x-api-key": apiKey,
|
|
"anthropic-version": "2023-06-01"
|
|
},
|
|
{
|
|
model: model.model,
|
|
system: systemPrompt || undefined,
|
|
messages: [{ role: "user", content: prompt }],
|
|
temperature,
|
|
max_tokens: maxTokens
|
|
}
|
|
);
|
|
|
|
const output = json.content?.map((part) => part.text ?? "").join("\n").trim() ?? "";
|
|
const usage = normalizeUsage(composePrompt(systemPrompt, prompt), output, {
|
|
inputTokens: json.usage?.input_tokens,
|
|
outputTokens: json.usage?.output_tokens,
|
|
estimated: !json.usage
|
|
});
|
|
|
|
return { output, usage, rawModel: model.model };
|
|
};
|
|
|
|
export const streamAnthropic = async ({
|
|
provider,
|
|
model,
|
|
systemPrompt,
|
|
prompt,
|
|
temperature,
|
|
maxTokens,
|
|
onDelta,
|
|
signal
|
|
}: ChatStreamRequest): Promise<ChatResponse> => {
|
|
const apiKey = apiKeyForProvider(provider);
|
|
if (!apiKey) {
|
|
throw new Error(`${provider.name} API key is not configured (${provider.apiKeyEnv})`);
|
|
}
|
|
|
|
let output = "";
|
|
let reasoning = "";
|
|
let inputTokens: number | undefined;
|
|
let outputTokens: number | undefined;
|
|
|
|
const response = await postJsonStream(
|
|
endpoint(provider, "/messages"),
|
|
{
|
|
"x-api-key": apiKey,
|
|
"anthropic-version": "2023-06-01"
|
|
},
|
|
{
|
|
model: model.model,
|
|
system: systemPrompt || undefined,
|
|
messages: [{ role: "user", content: prompt }],
|
|
temperature,
|
|
max_tokens: maxTokens,
|
|
stream: true
|
|
},
|
|
signal
|
|
);
|
|
|
|
await readSse(response, ({ data }) => {
|
|
const event = JSON.parse(data) as AnthropicStreamEvent;
|
|
inputTokens = event.message?.usage?.input_tokens ?? event.usage?.input_tokens ?? inputTokens;
|
|
outputTokens = event.message?.usage?.output_tokens ?? event.usage?.output_tokens ?? outputTokens;
|
|
|
|
const startText =
|
|
event.type === "content_block_start" && event.content_block?.type === "text"
|
|
? event.content_block.text ?? ""
|
|
: "";
|
|
const startThinking =
|
|
event.type === "content_block_start" && event.content_block?.type === "thinking"
|
|
? event.content_block.thinking ?? ""
|
|
: "";
|
|
const outputDelta = event.delta?.type === "text_delta" ? event.delta.text ?? "" : "";
|
|
const reasoningDelta =
|
|
event.delta?.type === "thinking_delta" ? event.delta.thinking ?? "" : "";
|
|
|
|
if (startThinking || reasoningDelta) {
|
|
const text = startThinking || reasoningDelta;
|
|
reasoning += text;
|
|
onDelta({ channel: "reasoning", text });
|
|
}
|
|
|
|
if (startText || outputDelta) {
|
|
const text = startText || outputDelta;
|
|
output += text;
|
|
onDelta({ channel: "output", text });
|
|
}
|
|
});
|
|
|
|
return {
|
|
output,
|
|
reasoning,
|
|
usage: normalizeUsage(composePrompt(systemPrompt, prompt), output, {
|
|
inputTokens,
|
|
outputTokens,
|
|
estimated: inputTokens === undefined && outputTokens === undefined
|
|
}),
|
|
rawModel: model.model
|
|
};
|
|
};
|
|
|
|
export const callGemini = async ({
|
|
provider,
|
|
model,
|
|
systemPrompt,
|
|
prompt,
|
|
temperature,
|
|
maxTokens
|
|
}: ChatRequest): Promise<ChatResponse> => {
|
|
const apiKey = apiKeyForProvider(provider);
|
|
if (!apiKey) {
|
|
throw new Error(`${provider.name} API key is not configured (${provider.apiKeyEnv})`);
|
|
}
|
|
|
|
const baseUrl = provider.baseUrl?.replace(/\/$/, "");
|
|
const url = `${baseUrl}/models/${encodeURIComponent(model.model)}:generateContent?key=${encodeURIComponent(apiKey)}`;
|
|
const json = await postJson<GeminiResponse>(
|
|
url,
|
|
{},
|
|
{
|
|
systemInstruction: systemPrompt ? { parts: [{ text: systemPrompt }] } : undefined,
|
|
contents: [{ role: "user", parts: [{ text: prompt }] }],
|
|
generationConfig: {
|
|
temperature,
|
|
maxOutputTokens: maxTokens
|
|
}
|
|
}
|
|
);
|
|
|
|
const output =
|
|
json.candidates?.[0]?.content?.parts?.map((part) => part.text ?? "").join("\n").trim() ?? "";
|
|
const usage = normalizeUsage(composePrompt(systemPrompt, prompt), output, {
|
|
inputTokens: json.usageMetadata?.promptTokenCount,
|
|
outputTokens: json.usageMetadata?.candidatesTokenCount,
|
|
totalTokens: json.usageMetadata?.totalTokenCount,
|
|
estimated: !json.usageMetadata
|
|
});
|
|
|
|
return { output, usage, rawModel: model.model };
|
|
};
|
|
|
|
export const streamGemini = async ({
|
|
provider,
|
|
model,
|
|
systemPrompt,
|
|
prompt,
|
|
temperature,
|
|
maxTokens,
|
|
onDelta,
|
|
signal
|
|
}: ChatStreamRequest): Promise<ChatResponse> => {
|
|
const apiKey = apiKeyForProvider(provider);
|
|
if (!apiKey) {
|
|
throw new Error(`${provider.name} API key is not configured (${provider.apiKeyEnv})`);
|
|
}
|
|
|
|
const baseUrl = provider.baseUrl?.replace(/\/$/, "");
|
|
const url = `${baseUrl}/models/${encodeURIComponent(model.model)}:streamGenerateContent?alt=sse&key=${encodeURIComponent(apiKey)}`;
|
|
let output = "";
|
|
let reasoning = "";
|
|
let usage: GeminiResponse["usageMetadata"] | undefined;
|
|
|
|
const response = await postJsonStream(
|
|
url,
|
|
{},
|
|
{
|
|
systemInstruction: systemPrompt ? { parts: [{ text: systemPrompt }] } : undefined,
|
|
contents: [{ role: "user", parts: [{ text: prompt }] }],
|
|
generationConfig: {
|
|
temperature,
|
|
maxOutputTokens: maxTokens
|
|
}
|
|
},
|
|
signal
|
|
);
|
|
|
|
await readSse(response, ({ data }) => {
|
|
const chunk = JSON.parse(data) as GeminiStreamChunk;
|
|
usage = chunk.usageMetadata ?? usage;
|
|
const parts = chunk.candidates?.[0]?.content?.parts ?? [];
|
|
|
|
for (const part of parts) {
|
|
const text = part.text ?? "";
|
|
if (!text) {
|
|
continue;
|
|
}
|
|
|
|
if ("thought" in part && part.thought === true) {
|
|
reasoning += text;
|
|
onDelta({ channel: "reasoning", text });
|
|
} else {
|
|
output += text;
|
|
onDelta({ channel: "output", text });
|
|
}
|
|
}
|
|
});
|
|
|
|
return {
|
|
output,
|
|
reasoning,
|
|
usage: normalizeUsage(composePrompt(systemPrompt, prompt), output, {
|
|
inputTokens: usage?.promptTokenCount,
|
|
outputTokens: usage?.candidatesTokenCount,
|
|
totalTokens: usage?.totalTokenCount,
|
|
estimated: !usage
|
|
}),
|
|
rawModel: model.model
|
|
};
|
|
};
|
|
|
|
export const callProvider = async (request: ChatRequest): Promise<ChatResponse> => {
|
|
switch (request.provider.kind) {
|
|
case "openai-compatible":
|
|
return callOpenAICompatible(request);
|
|
case "anthropic":
|
|
return callAnthropic(request);
|
|
case "gemini":
|
|
return callGemini(request);
|
|
default: {
|
|
const provider: never = request.provider.kind;
|
|
throw new Error(`Unsupported provider kind: ${provider}`);
|
|
}
|
|
}
|
|
};
|
|
|
|
export const streamProvider = async (request: ChatStreamRequest): Promise<ChatResponse> => {
|
|
switch (request.provider.kind) {
|
|
case "openai-compatible":
|
|
return streamOpenAICompatible(request);
|
|
case "anthropic":
|
|
return streamAnthropic(request);
|
|
case "gemini":
|
|
return streamGemini(request);
|
|
default: {
|
|
const provider: never = request.provider.kind;
|
|
throw new Error(`Unsupported provider kind: ${provider}`);
|
|
}
|
|
}
|
|
};
|