persona-community-1/packages/ai-client/src/textgen.ts

/**
 * Text generation client for calling textgen API endpoints.
 */

import { TokenUsage, AIResponseMeta, createAIError, ErrorCodes } from './types';

/** Role for chat messages. */
export type MessageRole = 'system' | 'user' | 'assistant';

/** A single message in a conversation. */
export interface TextGenMessage {
  role: MessageRole;
  content: string;
}

/** Request for text generation. */
export interface TextGenRequest {
  messages: TextGenMessage[];
  model?: string;
  maxTokens?: number;
  temperature?: number;
  stream?: boolean;
}

/** Response from text generation. */
export interface TextGenResponse extends AIResponseMeta {
  text: string;
  usage?: TokenUsage;
  finishReason?: 'stop' | 'length' | 'content_filter';
}

/** Chunk received during streaming. */
export interface TextGenChunk {
  text: string;
  done: boolean;
}

/** Options for API calls. */
export interface TextGenOptions {
  timeout?: number;
  headers?: Record<string, string>;
}

/**
 * Stream text generation with Server-Sent Events.
 *
 * @param endpoint - The textgen API endpoint URL
 * @param request - The generation request
 * @param onChunk - Callback for each text chunk received
 * @param onDone - Callback when generation completes
 * @param onError - Callback for errors
 * @param options - Additional options
 * @returns Abort function to cancel the stream
 *
 * @example
 * ```ts
 * const abort = streamText(
 *   '/api/textgen',
 *   { messages: [{ role: 'user', content: 'Hello!' }] },
 *   (chunk) => console.log(chunk),
 *   (response) => console.log('Done:', response),
 *   (error) => console.error(error)
 * );
 *
 * // Cancel if needed
 * abort();
 * ```
 */
export function streamText(
  endpoint: string,
  request: TextGenRequest,
  onChunk: (chunk: string) => void,
  onDone: (response: TextGenResponse) => void,
  onError: (error: Error) => void,
  options: TextGenOptions = {}
): () => void {
  const controller = new AbortController();
  const startTime = Date.now();

  const streamRequest = { ...request, stream: true };

  fetch(endpoint, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      ...options.headers,
    },
    body: JSON.stringify(streamRequest),
    signal: controller.signal,
  })
    .then(async (response) => {
      if (!response.ok) {
        const errorData = await response.json().catch(() => ({}));
        throw createAIError(
          errorData.message || `HTTP ${response.status}`,
          response.status === 429 ? ErrorCodes.RATE_LIMITED : ErrorCodes.PROVIDER_ERROR,
          undefined,
          response.status === 429 || response.status >= 500
        );
      }

      const reader = response.body?.getReader();
      if (!reader) {
        throw createAIError('No response body', ErrorCodes.NETWORK_ERROR);
      }

      const decoder = new TextDecoder();
      let fullText = '';
      let provider = '';
      let usage: TokenUsage | undefined;

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        const chunk = decoder.decode(value, { stream: true });
        const lines = chunk.split('\n');

        for (const line of lines) {
          if (line.startsWith('data: ')) {
            const data = line.slice(6);
            if (data === '[DONE]') continue;

            try {
              const parsed = JSON.parse(data);
              if (parsed.text) {
                fullText += parsed.text;
                onChunk(parsed.text);
              }
              if (parsed.provider) provider = parsed.provider;
              if (parsed.usage) usage = parsed.usage;
            } catch {
              // Skip invalid JSON lines
            }
          }
        }
      }

      onDone({
        text: fullText,
        provider,
        usage,
        latencyMs: Date.now() - startTime,
      });
    })
    .catch((error) => {
      if (error.name === 'AbortError') return;
      onError(error);
    });

  return () => controller.abort();
}

/**
 * Generate text (non-streaming).
 *
 * @param endpoint - The textgen API endpoint URL
 * @param request - The generation request
 * @param options - Additional options
 * @returns Promise resolving to the generation response
 *
 * @example
 * ```ts
 * const response = await generateText('/api/textgen', {
 *   messages: [
 *     { role: 'system', content: 'You are a helpful assistant.' },
 *     { role: 'user', content: 'What is 2 + 2?' }
 *   ],
 *   temperature: 0.7,
 * });
 * console.log(response.text);
 * ```
 */
export async function generateText(
  endpoint: string,
  request: TextGenRequest,
  options: TextGenOptions = {}
): Promise<TextGenResponse> {
  const startTime = Date.now();
  const nonStreamRequest = { ...request, stream: false };

  const controller = new AbortController();
  const timeoutId = options.timeout
    ? setTimeout(() => controller.abort(), options.timeout)
    : undefined;

  try {
    const response = await fetch(endpoint, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        ...options.headers,
      },
      body: JSON.stringify(nonStreamRequest),
      signal: controller.signal,
    });

    if (!response.ok) {
      const errorData = await response.json().catch(() => ({}));
      throw createAIError(
        errorData.message || `HTTP ${response.status}`,
        response.status === 429 ? ErrorCodes.RATE_LIMITED : ErrorCodes.PROVIDER_ERROR,
        undefined,
        response.status === 429 || response.status >= 500
      );
    }

    const data = await response.json();

    return {
      text: data.text || '',
      provider: data.provider || '',
      model: data.model,
      usage: data.usage,
      finishReason: data.finish_reason,
      latencyMs: Date.now() - startTime,
    };
  } catch (error) {
    if (error instanceof Error && error.name === 'AbortError') {
      throw createAIError('Request timed out', ErrorCodes.TIMEOUT, undefined, true);
    }
    throw error;
  } finally {
    if (timeoutId) clearTimeout(timeoutId);
  }
}