stack/apps/gateway/src/agent/adapters/openai.adapter.ts

import { Logger } from '@nestjs/common';
import OpenAI from 'openai';
import type { ModelRegistry } from '@mariozechner/pi-coding-agent';
import type {
  CompletionEvent,
  CompletionParams,
  IProviderAdapter,
  ModelInfo,
  ProviderHealth,
} from '@mosaic/types';

/**
 * OpenAI provider adapter.
 *
 * Registers OpenAI models (including Codex gpt-5.4) with the Pi ModelRegistry.
 * Configuration is driven by environment variables:
 *   OPENAI_API_KEY — OpenAI API key (required; adapter skips registration when absent)
 */
export class OpenAIAdapter implements IProviderAdapter {
  readonly name = 'openai';

  private readonly logger = new Logger(OpenAIAdapter.name);
  private registeredModels: ModelInfo[] = [];
  private client: OpenAI | null = null;

  /** Model ID used for Codex gpt-5.4 in the Pi registry. */
  static readonly CODEX_MODEL_ID = 'codex-gpt-5-4';

  constructor(private readonly registry: ModelRegistry) {}

  async register(): Promise<void> {
    const apiKey = process.env['OPENAI_API_KEY'];
    if (!apiKey) {
      this.logger.debug('Skipping OpenAI provider registration: OPENAI_API_KEY not set');
      return;
    }

    this.client = new OpenAI({ apiKey });

    const codexModel = {
      id: OpenAIAdapter.CODEX_MODEL_ID,
      name: 'Codex gpt-5.4',
      /** OpenAI-compatible completions API */
      api: 'openai-completions' as never,
      reasoning: false,
      input: ['text', 'image'] as ('text' | 'image')[],
      cost: { input: 0.003, output: 0.012, cacheRead: 0.0015, cacheWrite: 0 },
      contextWindow: 128_000,
      maxTokens: 16_384,
    };

    this.registry.registerProvider('openai', {
      apiKey,
      baseUrl: 'https://api.openai.com/v1',
      models: [codexModel],
    });

    this.registeredModels = [
      {
        id: OpenAIAdapter.CODEX_MODEL_ID,
        provider: 'openai',
        name: 'Codex gpt-5.4',
        reasoning: false,
        contextWindow: 128_000,
        maxTokens: 16_384,
        inputTypes: ['text', 'image'] as ('text' | 'image')[],
        cost: { input: 0.003, output: 0.012, cacheRead: 0.0015, cacheWrite: 0 },
      },
    ];

    this.logger.log(`OpenAI provider registered with model: ${OpenAIAdapter.CODEX_MODEL_ID}`);
  }

  listModels(): ModelInfo[] {
    return this.registeredModels;
  }

  async healthCheck(): Promise<ProviderHealth> {
    const apiKey = process.env['OPENAI_API_KEY'];
    if (!apiKey) {
      return {
        status: 'down',
        lastChecked: new Date().toISOString(),
        error: 'OPENAI_API_KEY not configured',
      };
    }

    const start = Date.now();
    try {
      // Lightweight call — list models to verify key validity
      const res = await fetch('https://api.openai.com/v1/models', {
        method: 'GET',
        headers: {
          Authorization: `Bearer ${apiKey}`,
          'Content-Type': 'application/json',
        },
        signal: AbortSignal.timeout(5000),
      });
      const latencyMs = Date.now() - start;

      if (!res.ok) {
        return {
          status: 'degraded',
          latencyMs,
          lastChecked: new Date().toISOString(),
          error: `HTTP ${res.status}`,
        };
      }

      return { status: 'healthy', latencyMs, lastChecked: new Date().toISOString() };
    } catch (err) {
      const latencyMs = Date.now() - start;
      const error = err instanceof Error ? err.message : String(err);
      return { status: 'down', latencyMs, lastChecked: new Date().toISOString(), error };
    }
  }

  /**
   * Stream a completion from OpenAI using the chat completions API.
   *
   * Maps OpenAI streaming chunks to the Mosaic CompletionEvent format.
   */
  async *createCompletion(params: CompletionParams): AsyncIterable<CompletionEvent> {
    if (!this.client) {
      throw new Error(
        'OpenAIAdapter: client not initialized. ' +
          'Ensure OPENAI_API_KEY is set and register() was called.',
      );
    }

    const stream = await this.client.chat.completions.create({
      model: params.model,
      messages: params.messages.map((m) => ({
        role: m.role,
        content: m.content,
      })),
      ...(params.temperature !== undefined && { temperature: params.temperature }),
      ...(params.maxTokens !== undefined && { max_tokens: params.maxTokens }),
      ...(params.tools &&
        params.tools.length > 0 && {
          tools: params.tools.map((t) => ({
            type: 'function' as const,
            function: {
              name: t.name,
              description: t.description,
              parameters: t.parameters,
            },
          })),
        }),
      stream: true,
      stream_options: { include_usage: true },
    });

    let inputTokens = 0;
    let outputTokens = 0;

    for await (const chunk of stream) {
      const choice = chunk.choices[0];

      // Accumulate usage when present (final chunk with stream_options.include_usage)
      if (chunk.usage) {
        inputTokens = chunk.usage.prompt_tokens;
        outputTokens = chunk.usage.completion_tokens;
      }

      if (!choice) continue;

      const delta = choice.delta;

      // Text content delta
      if (delta.content) {
        yield { type: 'text_delta', content: delta.content };
      }

      // Tool call delta — emit when arguments are complete
      if (delta.tool_calls) {
        for (const toolCallDelta of delta.tool_calls) {
          if (toolCallDelta.function?.name && toolCallDelta.function.arguments !== undefined) {
            yield {
              type: 'tool_call',
              name: toolCallDelta.function.name,
              arguments: toolCallDelta.function.arguments,
            };
          }
        }
      }

      // Stream finished
      if (choice.finish_reason === 'stop' || choice.finish_reason === 'tool_calls') {
        yield {
          type: 'done',
          usage: { inputTokens, outputTokens },
        };
        return;
      }
    }

    // Fallback done event when stream ends without explicit finish_reason
    yield { type: 'done', usage: { inputTokens, outputTokens } };
  }
}