1
0
Fork 0
tsgames/src/games/storywriter/utils/llm.ts

372 lines
11 KiB
TypeScript

import SSE from '@common/sse';
import type { TObject } from '@common/typebox';
namespace LLM {
export interface Connection {
url: string;
apiKey: string;
}
export interface ToolCall {
id: string;
type: 'function';
function: {
name: string;
arguments: string | Record<string, any>; // JSON string
};
}
interface ChatMessageUser {
role: 'user';
content: string;
}
interface ChatMessageAssistant {
role: 'assistant';
content: string;
reasoning_content?: string;
tool_calls?: ToolCall[];
}
interface ChatMessageSystem {
role: 'system';
content: string;
}
interface ChatMessageTool {
role: 'tool';
content: string;
tool_call_id?: string;
}
export type ChatMessage = ChatMessageUser | ChatMessageAssistant | ChatMessageSystem | ChatMessageTool;
export interface Tool {
type: 'function';
function: {
name: string;
description?: string;
parameters: TObject;
};
}
export interface ChatCompletionRequest {
model: string;
messages: ChatMessage[];
tools?: Tool[];
tool_choice?: 'none' | 'auto' | 'required' | { type: 'function'; function: { name: string } };
parallel_tool_calls?: boolean;
temperature?: number;
max_tokens?: number;
max_completion_tokens?: number;
stop?: string | string[];
banned_tokens?: string[];
top_p?: number;
top_k?: number;
min_p?: number;
frequency_penalty?: number;
repetition_penalty?: number;
reasoning?: {
effort?: 'xhigh' | 'high' | 'medium' | 'low' | 'minimal' | 'none';
exclude?: boolean;
max_tokens?: number;
};
add_generation_prompt?: boolean;
remove_last_eos?: boolean;
}
export interface ChatCompletionChoice {
index: number;
message: ChatMessage;
finish_reason: 'stop' | 'tool_calls';
}
export interface ChatCompletionResponse {
id: string;
object: 'chat.completion';
created: number;
model: string;
choices: ChatCompletionChoice[];
usage: {
prompt_tokens: number;
completion_tokens: number;
total_tokens: number;
spent_kudos?: number;
};
}
export interface ChatCompletionChunkChoice {
index: number;
delta: {
role?: string;
content?: string;
reasoning_content?: string;
tool_calls?: ToolCall[];
};
finish_reason: 'stop' | 'tool_calls' | null;
}
export interface ChatCompletionChunk {
id: string;
object: 'chat.completion.chunk';
created: number;
model: string;
choices: ChatCompletionChunkChoice[];
}
export interface ChatCompletionError {
error: string;
}
export interface ImageGenerationSettings {
width?: number;
height?: number;
negative_prompt?: string;
}
export interface ImageGenerationRequest {
model: string;
prompt: string;
n?: number;
size?: string;
quality?: 'standard' | 'hd';
output_format?: 'jpeg' | 'webp' | 'png';
image_settings?: ImageGenerationSettings;
}
export interface ImageGenerationResponse {
created: number;
data: ({ b64_json: string })[];
}
export interface ImageGenerationError {
error: string;
}
type Modality = 'text' | 'image';
interface BaseModelInfo {
id: string;
object: 'model';
created: number;
owned_by: string;
supported_parameters: string[];
architecture?: {
input_modalities: Modality[];
output_modalities: Modality[];
};
}
export interface ModelInfoText extends BaseModelInfo {
context_length: number;
top_provider: {
context_length: number;
max_completion_tokens: number;
is_moderated: boolean;
};
}
export interface ModelInfoImage extends BaseModelInfo {
}
export type ModelInfo = ModelInfoText | ModelInfoImage;
const isTextModel = (model: ModelInfo): model is ModelInfoText => ('context_length' in model);
const isImageModel = (model: ModelInfo): model is ModelInfoImage => Boolean(
!isTextModel(model) &&
model.architecture &&
(model.architecture.output_modalities).includes('image')
);
export interface ModelsResponse<T extends ModelInfo = ModelInfo> {
object: 'list';
data: T[];
}
interface CountTokensRequestString {
model: string;
input: string;
}
interface CountTokensRequestMessages {
model: string;
input: LLM.ChatMessage[];
tools?: LLM.Tool[];
add_generation_prompt?: boolean;
reasoning?: {
effort?: string;
exclude?: boolean;
max_tokens?: number;
};
}
export type CountTokensRequest = CountTokensRequestString | CountTokensRequestMessages;
export interface CountTokensResponse {
object: 'response.input_tokens';
input_tokens: number;
}
async function request<T>(connection: Connection, path: string, method: string = 'GET', body?: unknown): Promise<T> {
const headers = {
'Content-Type': 'application/json',
'Authorization': `Bearer ${connection.apiKey}`,
};
const url = new URL(connection.url);
url.pathname = path;
const response = await fetch(url, {
method,
headers,
body: body ? JSON.stringify(body) : undefined,
});
if (!response.ok) {
let text = '';
try {
text = await response.text();
} catch { }
throw new Error(`HTTP error! status: ${response.status}, text: ${text}`);
}
return response.json();
}
async function* streamRequest<T>(connection: Connection, path: string, method: string = 'GET', body?: unknown): AsyncGenerator<T> {
const headers = {
'Content-Type': 'application/json',
'Authorization': `Bearer ${connection.apiKey}`,
};
const url = new URL(connection.url);
url.pathname = path;
using sse = new SSE(url.toString(), {
headers,
method,
payload: body ? JSON.stringify(body) : undefined,
});
const readable = new ReadableStream<string>({
start: async (controller) => {
sse.addEventListener('message', (e) => {
if (isMessageEvent(e)) {
if (e.data === '[DONE]') {
controller.close();
} else {
controller.enqueue(e.data);
}
}
});
let closed = false;
const handleEnd = (_e?: unknown) => {
if (closed) return;
closed = true;
controller.close();
};
sse.addEventListener('error', handleEnd);
sse.addEventListener('abort', handleEnd);
sse.addEventListener('readystatechange', (e) => {
if (e.readyState === SSE.CLOSED) handleEnd();
});
}
});
const reader = readable.getReader();
while (true) {
const { value, done } = await reader.read();
if (done) {
break;
}
try {
yield JSON.parse(value);
} catch {
break;
}
}
await reader.closed;
sse.close();
}
function isMessageEvent(e: unknown): e is { data: string } {
return e != null && typeof e === 'object' && 'data' in e && typeof e.data === 'string';
}
export async function getTextModels(connection: Connection): Promise<ModelsResponse<ModelInfoText>> {
const response = await request<ModelsResponse>(connection, '/v1/models');
response.data = response.data.filter(isTextModel);
return response as ModelsResponse<ModelInfoText>;
}
export async function getImageModels(connection: Connection): Promise<ModelsResponse<ModelInfoImage>> {
const response = await request<ModelsResponse>(connection, '/v1/models');
response.data = response.data.filter(isImageModel);
return response as ModelsResponse<ModelInfoImage>;
}
export async function countTokens(connection: Connection, body: CountTokensRequest) {
return request<CountTokensResponse>(connection, '/v1/responses/input_tokens', 'POST', body);
}
export async function* generateStream(connection: Connection, config: ChatCompletionRequest) {
yield* streamRequest<ChatCompletionChunk | ChatCompletionError>(connection, '/v1/chat/completions', 'POST', {
...config,
stream: true,
});
}
export async function generate(connection: Connection, config: ChatCompletionRequest) {
return request<ChatCompletionResponse | ChatCompletionError>(connection, '/v1/chat/completions', 'POST', {
...config,
stream: false,
});
}
export async function generateImage(connection: Connection, config: ImageGenerationRequest) {
return request<ImageGenerationResponse | ImageGenerationError>(connection, '/v1/images/generations', 'POST', config);
}
const SUMMARIZATION_PROMPT = `Summarize the following text concisely while preserving key information and meaning. {level}
Text:
{text}
Provide a clear and coherent summary:`;
export type SummarizationLevel = 'sentence' | 'paragraph' | 'arbitrary';
const LEVEL_INSTRUCTIONS: Record<SummarizationLevel, string> = {
sentence: 'Summarize in exactly one sentence.',
paragraph: 'Summarize in exactly one paragraph (2-4 sentences).',
arbitrary: 'Summarize in a way you think is appropriate for the text length and complexity.',
};
export async function summarize(connection: Connection, model: string, text: string, level: SummarizationLevel = 'arbitrary'): Promise<string> {
const prompt = SUMMARIZATION_PROMPT
.replace('{text}', text)
.replace('{level}', LEVEL_INSTRUCTIONS[level]);
const response = await generate(connection, {
model,
messages: [{
role: 'user',
content: prompt,
}],
temperature: 0.3,
max_tokens: 500,
});
if ('error' in response) {
throw new Error(response.error);
}
return response.choices[0]?.message.content ?? '';
}
}
export default LLM;