v3-api/src/ai-v2/ai-v2.client.ts

import { HttpStatus, Injectable } from '@nestjs/common';
import axios, { AxiosError, AxiosRequestConfig, AxiosResponse } from 'axios';
import * as FormData from 'form-data';
import { Readable } from 'node:stream';
import { AiV2Exception } from './ai-v2.exception';
import { iterateSseEvents, readStreamToString } from './sse';

export type AiV2RunResult = {
  run_id: string;
  status: 'answered' | 'escalation_offered';
  message: string;
  suggested_questions?: string[] | null;
  escalation: { summary?: string; handoff_context?: Record<string, unknown> } | null;
};

@Injectable()
export class AiV2Client {
  private baseUrl(): string {
    const url = process.env.AI_V2_BASE_URL?.trim();
    if (!url) {
      throw new AiV2Exception(
        HttpStatus.SERVICE_UNAVAILABLE,
        'ai_v2_base_url_missing',
      );
    }
    return url.replace(/\/$/, '');
  }

  private apiKey(): string {
    const key = process.env.AI_V2_API_KEY?.trim();
    if (!key) {
      throw new AiV2Exception(
        HttpStatus.SERVICE_UNAVAILABLE,
        'ai_v2_api_key_missing',
      );
    }
    return key;
  }

  timeoutMs(): number {
    const raw = process.env.AI_V2_TIMEOUT || process.env.AI_SERVICE_TIMEOUT;
    const parsed = parseInt(String(raw || '75000'), 10);
    return Number.isFinite(parsed) && parsed > 0 ? parsed : 75_000;
  }

  isStreamToFrontendEnabled(): boolean {
    const raw = process.env.AI_V2_STREAM?.trim().toLowerCase();
    return raw === '1' || raw === 'true' || raw === 'yes' || raw === 'on';
  }

  private headers(extra: Record<string, string> = {}): Record<string, string> {
    return {
      Authorization: `Bearer ${this.apiKey()}`,
      accept: 'application/json',
      ...extra,
    };
  }

  private qs(params?: object): Record<string, unknown> {
    const out: Record<string, unknown> = {};
    if (!params) return out;
    for (const [key, value] of Object.entries(params)) {
      if (value === undefined || value === null || value === '') continue;
      out[key] = value;
    }
    return out;
  }

  private throwFromAxios(err: unknown): never {
    if (err instanceof AiV2Exception) throw err;
    const axiosErr = err as AxiosError;
    if (axiosErr?.isAxiosError) {
      const status =
        axiosErr.response?.status || HttpStatus.BAD_GATEWAY;
      const data = axiosErr.response?.data ?? axiosErr.message;
      const message =
        (data as any)?.error?.message ||
        (data as any)?.message ||
        'ai_v2_request_failed';
      throw new AiV2Exception(status, message, data);
    }
    throw new AiV2Exception(
      HttpStatus.BAD_GATEWAY,
      (err as Error)?.message || 'ai_v2_request_failed',
    );
  }

  private async request<T>(config: AxiosRequestConfig): Promise<T> {
    try {
      const response = await axios.request<T>({
        timeout: this.timeoutMs(),
        maxBodyLength: Infinity,
        ...config,
        baseURL: this.baseUrl(),
        headers: this.headers(config.headers as Record<string, string>),
      });
      return response.data;
    } catch (err) {
      this.throwFromAxios(err);
    }
  }

  listDomains(includeDisabled?: boolean) {
    return this.request({
      method: 'GET',
      url: '/api/v1/domains',
      params: this.qs({ include_disabled: includeDisabled }),
    });
  }

  createDomain(body: object) {
    return this.request({
      method: 'POST',
      url: '/api/v1/domains',
      data: body,
      headers: { 'Content-Type': 'application/json' },
    });
  }

  updateDomain(domain: string, body: object) {
    return this.request({
      method: 'PATCH',
      url: `/api/v1/domains/${encodeURIComponent(domain)}`,
      data: body,
      headers: { 'Content-Type': 'application/json' },
    });
  }

  disableDomain(domain: string) {
    return this.request({
      method: 'DELETE',
      url: `/api/v1/domains/${encodeURIComponent(domain)}`,
    });
  }

  enableDomain(domain: string) {
    return this.request({
      method: 'POST',
      url: `/api/v1/domains/${encodeURIComponent(domain)}/enable`,
    });
  }

  listDomainPoints(domain: string, query: object) {
    return this.request({
      method: 'GET',
      url: `/api/v1/domains/${encodeURIComponent(domain)}/points`,
      params: this.qs(query),
    });
  }

  async uploadFile(file: Express.Multer.File, domain: string) {
    const form = new FormData();
    form.append('file', file.buffer, {
      filename: file.originalname,
      contentType: file.mimetype || 'application/octet-stream',
    });
    form.append('domain', domain);
    return this.request({
      method: 'POST',
      url: '/api/v1/files',
      data: form,
      headers: form.getHeaders() as Record<string, string>,
      timeout: Math.max(this.timeoutMs(), 120_000),
    });
  }

  listFiles(query: object) {
    return this.request({
      method: 'GET',
      url: '/api/v1/files',
      params: this.qs(query),
    });
  }

  getFile(fileId: string) {
    return this.request({
      method: 'GET',
      url: `/api/v1/files/${encodeURIComponent(fileId)}`,
    });
  }

  deleteFile(fileId: string) {
    return this.request({
      method: 'DELETE',
      url: `/api/v1/files/${encodeURIComponent(fileId)}`,
    });
  }

  listFilePoints(fileId: string, query: object) {
    return this.request({
      method: 'GET',
      url: `/api/v1/files/${encodeURIComponent(fileId)}/points`,
      params: this.qs(query),
    });
  }

  countPoints(query: object) {
    return this.request({
      method: 'GET',
      url: '/api/v1/points/count',
      params: this.qs(query),
    });
  }

  searchPoints(query: object) {
    return this.request({
      method: 'GET',
      url: '/api/v1/points/search',
      params: this.qs(query),
    });
  }

  listPoints(query: object) {
    return this.request({
      method: 'GET',
      url: '/api/v1/points',
      params: this.qs(query),
    });
  }

  createPoint(body: object) {
    return this.request({
      method: 'POST',
      url: '/api/v1/points',
      data: body,
      headers: { 'Content-Type': 'application/json' },
    });
  }

  getPoint(pointId: string, withVectors?: boolean) {
    return this.request({
      method: 'GET',
      url: `/api/v1/points/${encodeURIComponent(pointId)}`,
      params: this.qs({ with_vectors: withVectors }),
    });
  }

  deletePoint(pointId: string) {
    return this.request({
      method: 'DELETE',
      url: `/api/v1/points/${encodeURIComponent(pointId)}`,
    });
  }

  replacePoint(pointId: string, body: object) {
    return this.request({
      method: 'PUT',
      url: `/api/v1/points/${encodeURIComponent(pointId)}`,
      data: body,
      headers: { 'Content-Type': 'application/json' },
    });
  }

  patchPointPayload(pointId: string, body: object) {
    return this.request({
      method: 'PATCH',
      url: `/api/v1/points/${encodeURIComponent(pointId)}/payload`,
      data: body,
      headers: { 'Content-Type': 'application/json' },
    });
  }

  queryRetrieval(body: object) {
    return this.request({
      method: 'POST',
      url: '/api/v1/retrieval/query',
      data: body,
      headers: { 'Content-Type': 'application/json' },
    });
  }

  createFeedback(
    threadId: string,
    runId: string,
    body: object,
  ) {
    return this.request({
      method: 'POST',
      url: `/api/v1/threads/${encodeURIComponent(threadId)}/runs/${encodeURIComponent(runId)}/feedback`,
      data: body,
      headers: { 'Content-Type': 'application/json' },
    });
  }

  private async openRunStream(
    threadId: string,
    body: { message: string; user_id: string },
  ): Promise<AxiosResponse<Readable>> {
    try {
      const response = await axios.request<Readable>({
        method: 'POST',
        baseURL: this.baseUrl(),
        url: `/api/v1/threads/${encodeURIComponent(threadId)}/runs`,
        data: body,
        headers: this.headers({
          'Content-Type': 'application/json',
          Accept: 'text/event-stream',
        }),
        responseType: 'stream',
        timeout: this.timeoutMs(),
        validateStatus: () => true,
        maxBodyLength: Infinity,
      });
      return response;
    } catch (err) {
      this.throwFromAxios(err);
    }
  }

  private async rejectStreamResponse(response: AxiosResponse<Readable>): Promise<never> {
    const raw = await readStreamToString(response.data);
    let data: unknown = raw;
    try {
      data = raw ? JSON.parse(raw) : raw;
    } catch {
      data = raw;
    }
    const message =
      (data as any)?.error?.message ||
      (data as any)?.message ||
      'ai_v2_run_failed';
    throw new AiV2Exception(response.status, message, data);
  }

  async consumeRun(
    threadId: string,
    body: { message: string; user_id: string },
    handlers?: {
      onToken?: (text: string) => void;
      onClientClose?: (abort: () => void) => void;
    },
  ): Promise<AiV2RunResult> {
    const response = await this.openRunStream(threadId, body);
    if (response.status >= 400) {
      await this.rejectStreamResponse(response);
    }

    const upstream = response.data;
    const abort = () => {
      if (typeof (upstream as any).destroy === 'function') {
        (upstream as any).destroy();
      }
    };
    handlers?.onClientClose?.(abort);

    let result: AiV2RunResult | null = null;
    for await (const event of iterateSseEvents(upstream)) {
      if (event.event === 'token') {
        const text = (event.data as any)?.text;
        if (typeof text === 'string' && text.length) {
          handlers?.onToken?.(text);
        }
        continue;
      }
      if (event.event === 'error') {
        const message =
          (event.data as any)?.message || 'ai_v2_run_failed';
        throw new AiV2Exception(HttpStatus.BAD_GATEWAY, message, event.data);
      }
      if (event.event === 'result' && event.data && typeof event.data === 'object') {
        result = event.data as AiV2RunResult;
      }
    }

    if (!result) {
      throw new AiV2Exception(
        HttpStatus.BAD_GATEWAY,
        'ai_v2_run_missing_result',
      );
    }
    return result;
  }

  async collectRun(
    threadId: string,
    body: { message: string; user_id: string },
  ): Promise<AiV2RunResult> {
    return this.consumeRun(threadId, body);
  }

  async pipeRun(
    threadId: string,
    body: { message: string; user_id: string },
    dest: NodeJS.WritableStream,
    onClientClose?: (abort: () => void) => void,
  ): Promise<void> {
    const response = await this.openRunStream(threadId, body);
    if (response.status >= 400) {
      await this.rejectStreamResponse(response);
    }

    const upstream = response.data;
    const abort = () => {
      if (typeof (upstream as any).destroy === 'function') {
        (upstream as any).destroy();
      }
    };
    onClientClose?.(abort);

    await new Promise<void>((resolve, reject) => {
      upstream.on('error', reject);
      dest.on('error', reject);
      dest.on('close', abort);
      upstream.on('end', () => resolve());
      upstream.pipe(dest, { end: true });
    });
  }
}