tsgames/src/games/storywriter/utils/llm.ts

import { formatError } from '@common/errors';
import SSE from '@common/sse';

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;
        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 ToolStringParameter {
        type: 'string';
        enum?: string[];
        description?: string;
    }

    export interface ToolNumberParameter {
        type: 'number' | 'integer';
        enum?: number[];
        description?: string;
    }

    export interface ToolBooleanParameter {
        type: 'boolean';
        enum?: boolean[];
        description?: string;
    }

    export interface ToolArrayParameter {
        type: 'array';
        description?: string;
        items: ToolParameter;
    }

    export interface ToolObjectParameter {
        type: 'object';
        description?: string;
        properties: Record<string, ToolParameter>;
        required?: string[];
    }

    export type ToolParameter = ToolStringParameter | ToolNumberParameter | ToolBooleanParameter | ToolArrayParameter | ToolObjectParameter;

    export interface Tool {
        type: 'function';
        function: {
            name: string;
            description?: string;
            parameters: {
                type: 'object';
                properties: Record<string, ToolParameter>;
                required?: string[];
            };
        };
    }

    export interface ChatCompletionRequest {
        model: string;
        messages: ChatMessage[];
        tools?: Tool[];
        temperature?: number;
        max_tokens?: number;
        stop?: string | string[];
        banned_tokens?: string[];
        top_p?: number;
        frequency_penalty?: number;
        presence_penalty?: number;
    }

    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; 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 ModelInfo {
        id: string;
        object: 'model';
        created: number;
        owned_by: string;
        support_tools: boolean;
        max_context?: number;
        max_length?: number;
    }

    export interface ModelsResponse {
        object: 'list';
        data: ModelInfo[];
    }

    export interface CountTokensRequest {
        model: string;
        input: string | ChatMessage[];
    }

    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) {
            throw new Error(`HTTP error! status: ${response.status}`);
        }
        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 getModels(connection: Connection): Promise<ModelsResponse> {
        return request<ModelsResponse>(connection, '/v1/models');
    }

    export async function countTokens(connection: Connection, body: CountTokensRequest): Promise<CountTokensResponse> {
        return request<CountTokensResponse>(connection, '/v1/responses/input_tokens', 'POST', body);
    }

    export async function* generateStream(connection: Connection, config: ChatCompletionRequest): AsyncGenerator<ChatCompletionChunk> {
        yield* streamRequest<ChatCompletionChunk>(connection, '/v1/chat/completions', 'POST', {
            ...config,
            stream: true,
        });
    }

    export async function generate(connection: Connection, config: ChatCompletionRequest): Promise<ChatCompletionResponse> {
        return request<ChatCompletionResponse>(connection, '/v1/chat/completions', 'POST', config);
    }
}

export default LLM;