learning_ai_common_plat/services/mcp-server/src/lib/extraction-client.ts

import { config } from './config.js';

export interface ExtractionItem {
  extraction_class: string;
  extraction_text: string;
  attributes?: Record<string, string>;
}

export interface ExtractionResult {
  extractions: ExtractionItem[];
  metadata: {
    modelId: string;
    taskId?: string;
    durationMs?: number;
    cached?: boolean;
  };
}

export async function extractionRun(
  params: {
    text: string;
    taskId?: string;
    modelId?: string;
  },
  opts: { requestId?: string }
): Promise<ExtractionResult> {
  const url = `${config.EXTRACTION_SERVICE_URL}/api/extract`;
  const headers: Record<string, string> = {
    'Content-Type': 'application/json',
    ...(opts.requestId ? { 'x-request-id': opts.requestId } : {}),
  };
  const res = await fetch(url, {
    method: 'POST',
    headers,
    body: JSON.stringify(params),
    signal: AbortSignal.timeout(30_000),
  });
  if (!res.ok) {
    const body = await res.text().catch(() => '');
    throw new Error(`extraction-service POST /api/extract → ${res.status}: ${body}`);
  }
  return res.json() as Promise<ExtractionResult>;
}

export async function extractionModels(opts: { requestId?: string }): Promise<unknown> {
  const url = `${config.EXTRACTION_SERVICE_URL}/api/extract/models`;
  const headers: Record<string, string> = {
    ...(opts.requestId ? { 'x-request-id': opts.requestId } : {}),
  };
  const res = await fetch(url, { headers, signal: AbortSignal.timeout(5_000) });
  if (!res.ok) throw new Error(`extraction-service GET /api/extract/models → ${res.status}`);
  return res.json();
}

export async function extractionCacheStats(opts: { requestId?: string }): Promise<unknown> {
  const url = `${config.EXTRACTION_SERVICE_URL}/api/extract/cache-stats`;
  const headers: Record<string, string> = {
    ...(opts.requestId ? { 'x-request-id': opts.requestId } : {}),
  };
  const res = await fetch(url, { headers, signal: AbortSignal.timeout(5_000) });
  if (!res.ok) throw new Error(`extraction-service GET /api/extract/cache-stats → ${res.status}`);
  return res.json();
}

export async function extractionSidecarHealth(opts: { requestId?: string }): Promise<unknown> {
  const url = `${config.EXTRACTION_SERVICE_URL}/api/extract/sidecar-health`;
  const headers: Record<string, string> = {
    ...(opts.requestId ? { 'x-request-id': opts.requestId } : {}),
  };
  const res = await fetch(url, { headers, signal: AbortSignal.timeout(10_000) });
  return res.json();
}

export async function extractionExtractBatch(
  params: {
    inputs: Array<{
      text: string;
      taskId?: string;
      taskPrompt?: string;
    }>;
    examples?: Array<{
      text: string;
      extractions: ExtractionItem[];
    }>;
    modelId?: string;
  },
  opts: { requestId?: string }
): Promise<unknown> {
  const url = `${config.EXTRACTION_SERVICE_URL}/api/extract/batch`;
  const headers: Record<string, string> = {
    'Content-Type': 'application/json',
    ...(opts.requestId ? { 'x-request-id': opts.requestId } : {}),
  };
  const res = await fetch(url, {
    method: 'POST',
    headers,
    body: JSON.stringify(params),
    signal: AbortSignal.timeout(30_000),
  });
  if (!res.ok) {
    const body = await res.text().catch(() => '');
    throw new Error(`extraction-service POST /api/extract/batch → ${res.status}: ${body}`);
  }
  return res.json();
}

export async function extractionSubmitJob(
  params: {
    inputs: Array<{
      text: string;
      taskId?: string;
      taskPrompt?: string;
    }>;
    examples?: Array<{
      text: string;
      extractions: ExtractionItem[];
    }>;
    modelId?: string;
    productId?: string;
    webhookUrl?: string;
    webhookSecret?: string;
    webhookRetryAttempts?: number;
  },
  opts: { requestId?: string }
): Promise<unknown> {
  const url = `${config.EXTRACTION_SERVICE_URL}/api/extract/jobs`;
  const headers: Record<string, string> = {
    'Content-Type': 'application/json',
    ...(opts.requestId ? { 'x-request-id': opts.requestId } : {}),
  };
  const res = await fetch(url, {
    method: 'POST',
    headers,
    body: JSON.stringify(params),
    signal: AbortSignal.timeout(30_000),
  });
  if (!res.ok) {
    const body = await res.text().catch(() => '');
    throw new Error(`extraction-service POST /api/extract/jobs → ${res.status}: ${body}`);
  }
  return res.json();
}

export async function extractionGetJob(
  jobId: string,
  opts: { requestId?: string }
): Promise<unknown> {
  const url = `${config.EXTRACTION_SERVICE_URL}/api/extract/jobs/${jobId}`;
  const headers: Record<string, string> = {
    ...(opts.requestId ? { 'x-request-id': opts.requestId } : {}),
  };
  const res = await fetch(url, { headers, signal: AbortSignal.timeout(10_000) });
  if (!res.ok) {
    const body = await res.text().catch(() => '');
    throw new Error(`extraction-service GET /api/extract/jobs/${jobId} → ${res.status}: ${body}`);
  }
  return res.json();
}

export async function extractionListJobs(opts: { requestId?: string }): Promise<unknown> {
  const url = `${config.EXTRACTION_SERVICE_URL}/api/extract/jobs`;
  const headers: Record<string, string> = {
    ...(opts.requestId ? { 'x-request-id': opts.requestId } : {}),
  };
  const res = await fetch(url, { headers, signal: AbortSignal.timeout(10_000) });
  if (!res.ok) {
    const body = await res.text().catch(() => '');
    throw new Error(`extraction-service GET /api/extract/jobs → ${res.status}: ${body}`);
  }
  return res.json();
}

export async function extractionGetProductRateLimitStatus(
  productId?: string,
  opts?: { requestId?: string }
): Promise<unknown> {
  const url = productId
    ? `${config.EXTRACTION_SERVICE_URL}/api/extract/rate-limits/product?productId=${encodeURIComponent(productId)}`
    : `${config.EXTRACTION_SERVICE_URL}/api/extract/rate-limits/product`;
  const headers: Record<string, string> = {
    ...(opts?.requestId ? { 'x-request-id': opts.requestId } : {}),
  };
  const res = await fetch(url, { headers, signal: AbortSignal.timeout(10_000) });
  if (!res.ok) {
    const body = await res.text().catch(() => '');
    throw new Error(
      `extraction-service GET /api/extract/rate-limits/product → ${res.status}: ${body}`
    );
  }
  return res.json();
}

export async function extractionResetProductRateLimit(
  productId: string,
  opts?: { requestId?: string }
): Promise<unknown> {
  const url = `${config.EXTRACTION_SERVICE_URL}/api/extract/rate-limits/product/reset`;
  const headers: Record<string, string> = {
    'Content-Type': 'application/json',
    ...(opts?.requestId ? { 'x-request-id': opts.requestId } : {}),
  };
  const res = await fetch(url, {
    method: 'POST',
    headers,
    body: JSON.stringify({ productId }),
    signal: AbortSignal.timeout(10_000),
  });
  if (!res.ok) {
    const body = await res.text().catch(() => '');
    throw new Error(
      `extraction-service POST /api/extract/rate-limits/product/reset → ${res.status}: ${body}`
    );
  }
  return res.json();
}

export async function extractionSidecarMonitoringState(opts: {
  requestId?: string;
}): Promise<unknown> {
  const url = `${config.EXTRACTION_SERVICE_URL}/api/extract/monitoring/sidecar`;
  const headers: Record<string, string> = {
    ...(opts.requestId ? { 'x-request-id': opts.requestId } : {}),
  };
  const res = await fetch(url, { headers, signal: AbortSignal.timeout(10_000) });
  if (!res.ok) {
    const body = await res.text().catch(() => '');
    throw new Error(
      `extraction-service GET /api/extract/monitoring/sidecar → ${res.status}: ${body}`
    );
  }
  return res.json();
}