plugin-ai-api 1.0.24 → 1.0.25

package/src/server/services/__tests__/file-processor.test.ts

@@ -0,0 +1,184 @@
+import { describe, expect, it, vi } from 'vitest';
+import {
+  FileProcessorService,
+  base64FileForwarder,
+  httpFileUrlFetcher,
+  fetchFileAsBase64,
+  pdfFileProcessor,
+  type PdfToImageRenderer,
+} from '../file-processor';
+
+describe('FileProcessorService', () => {
+  it('forwards file blocks with base64 data unchanged', async () => {
+    const service = new FileProcessorService();
+    service.register(base64FileForwarder);
+
+    const block = { type: 'file', file: { file_data: 'data:application/pdf;base64,JVBERi0=' } };
+    const result = await service.process(block, { ctx: {} as any });
+
+    expect(result).toEqual(block);
+  });
+
+  it('returns the block unchanged when no processor can handle it', async () => {
+    const service = new FileProcessorService();
+    const block = { type: 'text', text: 'hello' };
+    const result = await service.process(block, { ctx: {} as any });
+
+    expect(result).toEqual(block);
+  });
+
+  it('allows custom processors to override default behavior', async () => {
+    const service = new FileProcessorService();
+    service.register(base64FileForwarder);
+
+    const customProcessor = {
+      name: 'customFileProcessor',
+      canHandle: (block: { type: string }) => block.type === 'file',
+      process: vi.fn().mockResolvedValue({ type: 'file', file: { file_data: 'data:text/plain;base64,SGVsbG8=' } }),
+    };
+    service.register(customProcessor);
+
+    const block = { type: 'file', file: { file_data: 'data:application/pdf;base64,JVBERi0=' } };
+    const result = await service.process(block, { ctx: {} as any });
+
+    expect(customProcessor.process).toHaveBeenCalledWith(block, { ctx: {} as any });
+    expect(result).toEqual({ type: 'file', file: { file_data: 'data:text/plain;base64,SGVsbG8=' } });
+  });
+
+  it('unregister removes a processor by name', () => {
+    const service = new FileProcessorService();
+    service.register(base64FileForwarder);
+    expect(service.list()).toHaveLength(1);
+
+    service.unregister(base64FileForwarder.name);
+    expect(service.list()).toHaveLength(0);
+  });
+});
+
+describe('httpFileUrlFetcher', () => {
+  it('fetches a file from an http(s) URL and converts it to a file block', async () => {
+    const originalFetch = globalThis.fetch;
+    const fileBuffer = Buffer.from('hello world');
+    globalThis.fetch = vi.fn().mockResolvedValue({
+      ok: true,
+      status: 200,
+      headers: new Headers({
+        'content-type': 'text/plain',
+        'content-length': String(fileBuffer.length),
+      }),
+      arrayBuffer: () =>
+        Promise.resolve(fileBuffer.buffer.slice(fileBuffer.byteOffset, fileBuffer.byteOffset + fileBuffer.length)),
+    } as unknown as Response);
+
+    const result = await httpFileUrlFetcher.process(
+      { type: 'file_url', file_url: { url: 'https://example.com/hello.txt' } },
+      { ctx: {} as any },
+    );
+
+    expect(result).toMatchObject({
+      type: 'file',
+      file: {
+        file_data: expect.stringContaining('data:text/plain;base64,'),
+        mime_type: 'text/plain',
+        filename: 'hello.txt',
+      },
+    });
+
+    globalThis.fetch = originalFetch;
+  });
+
+  it('rejects unsupported protocols', async () => {
+    await expect(
+      httpFileUrlFetcher.process(
+        { type: 'file_url', file_url: { url: 'ftp://example.com/file.txt' } },
+        { ctx: {} as any },
+      ),
+    ).rejects.toThrow('protocol');
+  });
+});
+
+describe('fetchFileAsBase64', () => {
+  it('validates URL protocols', async () => {
+    await expect(fetchFileAsBase64('ftp://example.com/file.txt')).rejects.toThrow('protocol');
+  });
+
+  it('rejects malformed URLs', async () => {
+    await expect(fetchFileAsBase64('not a url')).rejects.toThrow('valid URL');
+  });
+});
+
+describe('pdfFileProcessor', () => {
+  const pdfBuffer = Buffer.from('%PDF-1.4\n1 0 obj\n<<\n>>\nendobj\n', 'binary');
+  const pdfDataUrl = `data:application/pdf;base64,${pdfBuffer.toString('base64')}`;
+
+  function createContext(config: { pdfRenderPagesAsImages?: boolean }, renderer?: PdfToImageRenderer | null) {
+    const service = new FileProcessorService();
+    if (renderer) {
+      service.registerPdfRenderer(renderer);
+    }
+    return {
+      ctx: {
+        db: {
+          getRepository: vi.fn((name: string) => {
+            if (name === 'aiApiConfig') {
+              return {
+                findOne: vi.fn().mockResolvedValue({
+                  get: (key: string) => (key === 'pdfRenderPagesAsImages' ? config.pdfRenderPagesAsImages : undefined),
+                  pdfRenderPagesAsImages: config.pdfRenderPagesAsImages,
+                }),
+              };
+            }
+            return { findOne: vi.fn() };
+          }),
+        },
+        app: {
+          pm: {
+            get: vi.fn().mockReturnValue({
+              fileProcessorService: service,
+            }),
+          },
+        },
+        log: { warn: vi.fn() },
+      },
+    } as any;
+  }
+
+  it('forwards a PDF file block when pdfRenderPagesAsImages is disabled', async () => {
+    const block = { type: 'file', file: { file_data: pdfDataUrl } } as any;
+    const result = await pdfFileProcessor.process(block, createContext({ pdfRenderPagesAsImages: false }));
+    expect(result).toBe(block);
+  });
+
+  it('forwards a PDF and warns when rendering is enabled but no renderer is registered', async () => {
+    const block = { type: 'file', file: { file_data: pdfDataUrl } } as any;
+    const context = createContext({ pdfRenderPagesAsImages: true });
+    const result = await pdfFileProcessor.process(block, context);
+
+    expect(result).toBe(block);
+    expect(context.ctx.log.warn).toHaveBeenCalledWith(
+      expect.stringContaining('pdfRenderPagesAsImages is enabled but no PdfToImageRenderer is registered'),
+    );
+  });
+
+  it('converts a PDF file block into image_url blocks when a renderer is registered', async () => {
+    const renderer: PdfToImageRenderer = {
+      name: 'mockRenderer',
+      render: vi.fn().mockResolvedValue([Buffer.from('page1'), Buffer.from('page2')]),
+    };
+    const block = { type: 'file', file: { file_data: pdfDataUrl } } as any;
+    const result = await pdfFileProcessor.process(block, createContext({ pdfRenderPagesAsImages: true }, renderer));
+
+    expect(Array.isArray(result)).toBe(true);
+    expect(result).toHaveLength(2);
+    expect((result as any[])[0]).toMatchObject({ type: 'image_url' });
+    expect((result as any[])[0].image_url.url).toMatch(/^data:image\/png;base64,/);
+  });
+
+  it('does not convert non-PDF file blocks', async () => {
+    const block = {
+      type: 'file',
+      file: { file_data: 'data:text/plain;base64,SGVsbG8=', filename: 'hello.txt' },
+    } as any;
+    expect(pdfFileProcessor.canHandle(block)).toBe(false);
+  });
+});

package/src/server/services/file-processor.ts

@@ -0,0 +1,323 @@
+import type { Context } from '@nocobase/actions';
+import { basename } from 'path';
+
+export interface FileContentBlock {
+  type: string;
+  [key: string]: unknown;
+}
+
+export interface FileProcessorContext {
+  ctx: Context;
+}
+
+export interface FileProcessor {
+  name: string;
+  canHandle(block: FileContentBlock): boolean;
+  process(block: FileContentBlock, context: FileProcessorContext): Promise<FileContentBlock | FileContentBlock[]>;
+}
+
+export interface PdfToImageRenderer {
+  name: string;
+  /**
+   * Render each page of a PDF into a PNG image buffer.
+   * @param buffer Raw PDF bytes.
+   * @returns Array of PNG buffers, one per page.
+   */
+  render(buffer: Buffer): Promise<Buffer[]>;
+}
+
+export type FileProcessorErrorCode =
+  | 'invalid_url'
+  | 'unsupported_protocol'
+  | 'fetch_failed'
+  | 'file_too_large'
+  | 'content_type_not_allowed'
+  | 'missing_url';
+
+export class FileProcessorError extends Error {
+  constructor(
+    readonly code: FileProcessorErrorCode,
+    message: string,
+  ) {
+    super(message);
+    this.name = 'FileProcessorError';
+  }
+}
+
+export interface FetchFileOptions {
+  maxSizeBytes?: number;
+  timeoutMs?: number;
+  allowedProtocols?: string[];
+  allowedContentTypes?: string[];
+}
+
+const DEFAULT_MAX_FILE_SIZE = 50 * 1024 * 1024; // 50 MB
+const DEFAULT_TIMEOUT_MS = 30_000;
+const ALLOWED_PROTOCOLS = new Set(['http:', 'https:']);
+
+export class FileProcessorService {
+  private processors: FileProcessor[] = [];
+  private pdfRenderer: PdfToImageRenderer | null = null;
+
+  register(processor: FileProcessor): void {
+    this.unregister(processor.name);
+    this.processors.push(processor);
+  }
+
+  unregister(name: string): void {
+    this.processors = this.processors.filter((p) => p.name !== name);
+  }
+
+  async process(
+    block: FileContentBlock,
+    context: FileProcessorContext,
+  ): Promise<FileContentBlock | FileContentBlock[]> {
+    // Search from the end so the most recently registered processor wins.
+    // This makes it easy for a custom plugin to override a default processor.
+    const processor = [...this.processors].reverse().find((p) => p.canHandle(block));
+    if (!processor) {
+      return block;
+    }
+    return processor.process(block, context);
+  }
+
+  list(): ReadonlyArray<FileProcessor> {
+    return [...this.processors];
+  }
+
+  /**
+   * Register a renderer used to convert PDF pages into PNG images when
+   * `pdfRenderPagesAsImages` is enabled in the AI API config.
+   */
+  registerPdfRenderer(renderer: PdfToImageRenderer): void {
+    this.pdfRenderer = renderer;
+  }
+
+  unregisterPdfRenderer(): void {
+    this.pdfRenderer = null;
+  }
+
+  getPdfRenderer(): PdfToImageRenderer | null {
+    return this.pdfRenderer;
+  }
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === 'object' && value !== null && !Array.isArray(value);
+}
+
+function getUrlString(value: unknown): string | undefined {
+  if (typeof value === 'string') return value;
+  if (isRecord(value) && typeof value.url === 'string') return value.url;
+  return undefined;
+}
+
+function extractFilename(url: string, contentDisposition: string | null): string | undefined {
+  if (contentDisposition) {
+    const match = contentDisposition.match(/filename="?([^"]+)"?/);
+    if (match) return match[1];
+  }
+  try {
+    const pathname = new URL(url).pathname;
+    if (pathname) return basename(pathname);
+  } catch {
+    // ignore malformed URL
+  }
+  return undefined;
+}
+
+export async function fetchFileAsBase64(
+  url: string,
+  options: FetchFileOptions = {},
+): Promise<{ fileData: string; mimeType: string | undefined; filename: string | undefined }> {
+  const maxSize = options.maxSizeBytes ?? DEFAULT_MAX_FILE_SIZE;
+  const timeoutMs = options.timeoutMs ?? DEFAULT_TIMEOUT_MS;
+  const allowedProtocols = options.allowedProtocols ? new Set(options.allowedProtocols) : ALLOWED_PROTOCOLS;
+
+  let protocol: string;
+  try {
+    protocol = new URL(url).protocol;
+  } catch {
+    throw new FileProcessorError('invalid_url', `File URL '${url}' is not a valid URL.`);
+  }
+
+  if (!allowedProtocols.has(protocol)) {
+    throw new FileProcessorError('unsupported_protocol', `File URL protocol '${protocol}' is not allowed.`);
+  }
+
+  const controller = new AbortController();
+  const timeout = setTimeout(() => controller.abort(), timeoutMs);
+
+  try {
+    const response = await fetch(url, {
+      signal: controller.signal,
+      redirect: 'follow',
+    });
+
+    if (!response.ok) {
+      throw new FileProcessorError(
+        'fetch_failed',
+        `Failed to fetch file from '${url}': ${response.status} ${response.statusText}`,
+      );
+    }
+
+    const contentLength = response.headers.get('content-length');
+    if (contentLength && Number(contentLength) > maxSize) {
+      throw new FileProcessorError('file_too_large', `File at '${url}' exceeds maximum allowed size.`);
+    }
+
+    const contentType = response.headers.get('content-type') || undefined;
+    if (
+      options.allowedContentTypes &&
+      contentType &&
+      !options.allowedContentTypes.some((type) => contentType.includes(type))
+    ) {
+      throw new FileProcessorError('content_type_not_allowed', `File content type '${contentType}' is not allowed.`);
+    }
+
+    const buffer = Buffer.from(await response.arrayBuffer());
+    if (buffer.length > maxSize) {
+      throw new FileProcessorError('file_too_large', `File at '${url}' exceeds maximum allowed size.`);
+    }
+
+    const mimeType = contentType?.split(';')[0].trim() ?? 'application/octet-stream';
+    const contentDisposition = response.headers.get('content-disposition');
+    const filename = extractFilename(url, contentDisposition) ?? 'file';
+
+    return {
+      fileData: `data:${mimeType};base64,${buffer.toString('base64')}`,
+      mimeType,
+      filename,
+    };
+  } finally {
+    clearTimeout(timeout);
+  }
+}
+
+/**
+ * Default processor: forwards base64 `file` blocks unchanged.
+ * Handles `{ type: 'file', file: { file_data: 'data:...;base64,...' } }`.
+ */
+export const base64FileForwarder: FileProcessor = {
+  name: 'base64FileForwarder',
+  canHandle(block: FileContentBlock): boolean {
+    if (block.type !== 'file') return false;
+    const file = isRecord(block.file) ? block.file : undefined;
+    if (!file) return false;
+    const fileData = String(file.file_data ?? '');
+    return fileData.startsWith('data:') && fileData.includes(';base64,');
+  },
+  async process(block: FileContentBlock): Promise<FileContentBlock> {
+    return block;
+  },
+};
+
+/**
+ * Default processor: downloads an http(s) `file_url` and converts it into a `file` block.
+ * Handles `{ type: 'file_url', file_url: { url: 'https://...' } }`.
+ */
+export const httpFileUrlFetcher: FileProcessor = {
+  name: 'httpFileUrlFetcher',
+  canHandle(block: FileContentBlock): boolean {
+    if (block.type !== 'file_url') return false;
+    const fileUrl = isRecord(block.file_url) ? block.file_url : undefined;
+    if (!fileUrl) return false;
+    const url = String(fileUrl.url ?? '');
+    return url.startsWith('http://') || url.startsWith('https://');
+  },
+  async process(block: FileContentBlock): Promise<FileContentBlock> {
+    const fileUrl = isRecord(block.file_url) ? block.file_url : undefined;
+    const url = String(fileUrl?.url ?? '');
+    if (!url) {
+      throw new FileProcessorError('missing_url', "file_url block requires a 'url' property.");
+    }
+    const { fileData, mimeType, filename } = await fetchFileAsBase64(url);
+    return {
+      type: 'file',
+      file: {
+        file_data: fileData,
+        mime_type: mimeType,
+        filename: filename || 'file',
+      },
+    };
+  },
+};
+
+function decodeBase64DataUrl(url: string): { mimeType: string; buffer: Buffer } | undefined {
+  const match = /^data:([^;]+);base64,([A-Za-z0-9+/]+=*)$/.exec(url);
+  if (!match) return undefined;
+  try {
+    const buffer = Buffer.from(match[2], 'base64');
+    return { mimeType: match[1].toLowerCase(), buffer };
+  } catch {
+    return undefined;
+  }
+}
+
+function isPdfBuffer(buffer: Buffer): boolean {
+  return buffer.length >= 4 && buffer.toString('binary', 0, 4) === '%PDF';
+}
+
+function isPdfFileBlock(block: FileContentBlock): boolean {
+  if (block.type !== 'file') return false;
+  const file = isRecord(block.file) ? block.file : undefined;
+  if (!file) return false;
+  const fileData = String(file.file_data ?? '');
+  if (!fileData.startsWith('data:')) return false;
+  const mimeType = String(file.mime_type ?? '').toLowerCase();
+  if (mimeType === 'application/pdf') return true;
+  if (fileData.startsWith('data:application/pdf')) return true;
+  const decoded = decodeBase64DataUrl(fileData);
+  if (decoded && isPdfBuffer(decoded.buffer)) return true;
+  return false;
+}
+
+function getPluginFromContext(context: FileProcessorContext): any | undefined {
+  // @ts-expect-error app may not be typed on Context in test mocks.
+  return context.ctx.app?.pm?.get?.('plugin-ai-api');
+}
+
+/**
+ * Optional processor: converts a PDF `file` block into a series of `image_url`
+ * blocks when `pdfRenderPagesAsImages` is enabled and a PdfToImageRenderer is
+ * registered. When disabled, unavailable, or the block is not a PDF, the block
+ * is forwarded unchanged.
+ */
+export const pdfFileProcessor: FileProcessor = {
+  name: 'pdfFileProcessor',
+  canHandle(block: FileContentBlock): boolean {
+    return isPdfFileBlock(block);
+  },
+  async process(
+    block: FileContentBlock,
+    context: FileProcessorContext,
+  ): Promise<FileContentBlock | FileContentBlock[]> {
+    const config = await context.ctx.db.getRepository('aiApiConfig').findOne();
+    if (!config?.pdfRenderPagesAsImages) {
+      return block;
+    }
+
+    const plugin = getPluginFromContext(context);
+    const renderer = plugin?.fileProcessorService?.getPdfRenderer?.();
+    if (!renderer) {
+      context.ctx.log?.warn?.(
+        '[pdfFileProcessor] pdfRenderPagesAsImages is enabled but no PdfToImageRenderer is registered. ' +
+          'Forwarding PDF as a file block.',
+      );
+      return block;
+    }
+
+    const file = isRecord(block.file) ? block.file : undefined;
+    const fileData = String(file?.file_data ?? '');
+    const decoded = decodeBase64DataUrl(fileData);
+    if (!decoded) {
+      return block;
+    }
+
+    const pages = await renderer.render(decoded.buffer);
+    return pages.map((buffer) => ({
+      type: 'image_url',
+      image_url: { url: `data:image/png;base64,${buffer.toString('base64')}` },
+    }));
+  },
+};

package/src/server/usage.ts

@@ -6,6 +6,7 @@ export type Usage = {
   prompt_tokens: number | null;
   completion_tokens: number | null;
   total_tokens: number | null;
+  prompt_cache_tokens?: number | null;
 };
 
 export type AiApiAuthType = 'apiKey' | 'bearer' | 'oidc' | 'unknown';
@@ -48,6 +49,33 @@ function normalizeTokenCount(value: unknown): number | null {
   return typeof value === 'number' && Number.isSafeInteger(value) && value >= 0 ? value : null;
 }
 
+/**
+ * Extract cached prompt tokens from a provider usage or response metadata object.
+ * Tolerates both LangChain-style (`input_token_details.cache_read`) and raw
+ * OpenAI-style (`prompt_tokens_details.cached_tokens`) shapes.
+ */
+function extractPromptCacheTokens(source: Record<string, unknown>): number | null {
+  const inputDetails = source.input_token_details;
+  if (inputDetails && typeof inputDetails === 'object') {
+    const cacheRead = (inputDetails as Record<string, unknown>).cache_read;
+    if (typeof cacheRead === 'number') return cacheRead;
+    const cachedTokens = (inputDetails as Record<string, unknown>).cached_tokens;
+    if (typeof cachedTokens === 'number') return cachedTokens;
+  }
+
+  const promptDetails = source.prompt_tokens_details;
+  if (promptDetails && typeof promptDetails === 'object') {
+    const cachedTokens = (promptDetails as Record<string, unknown>).cached_tokens;
+    if (typeof cachedTokens === 'number') return cachedTokens;
+  }
+
+  if (typeof source.cached_tokens === 'number') return source.cached_tokens;
+  if (typeof source.cache_read_tokens === 'number') return source.cache_read_tokens;
+  if (typeof source.cache_read === 'number') return source.cache_read;
+
+  return null;
+}
+
 export function normalizeUsage(value: unknown): Usage | undefined {
   if (!value || typeof value !== 'object') return undefined;
   const source = value as Record<string, unknown>;
@@ -65,6 +93,7 @@ export function normalizeUsage(value: unknown): Usage | undefined {
     prompt_tokens: prompt,
     completion_tokens: completion,
     total_tokens: total,
+    prompt_cache_tokens: extractPromptCacheTokens(source),
   };
 }
 
@@ -72,8 +101,24 @@ export function setAiApiUsageResult(
   ctx: Context,
   value: unknown,
   metadata: Pick<AiApiUsageResult, 'gatewayResponseId' | 'providerRequestId'> = {},
+  responseMetadata?: unknown,
 ): Usage | undefined {
-  const usage = normalizeUsage(value);
+  let usage = normalizeUsage(value);
+
+  // Some providers only surface cached token details on response_metadata; fall
+  // back there when the usage object itself does not carry them.
+  if (usage?.prompt_cache_tokens === null && responseMetadata && typeof responseMetadata === 'object') {
+    const responseMetaRecord = responseMetadata as Record<string, unknown>;
+    let cacheTokens = extractPromptCacheTokens(responseMetaRecord);
+    // Providers sometimes nest the raw provider usage under response_metadata.usage.
+    if (cacheTokens === null && responseMetaRecord.usage && typeof responseMetaRecord.usage === 'object') {
+      cacheTokens = extractPromptCacheTokens(responseMetaRecord.usage as Record<string, unknown>);
+    }
+    if (cacheTokens !== null) {
+      usage = { ...usage, prompt_cache_tokens: cacheTokens };
+    }
+  }
+
   getAiApiState(ctx).aiApiUsageResult = usage
     ? { source: 'provider', usage, ...metadata }
     : { source: 'unavailable', ...metadata };
@@ -174,6 +219,7 @@ export async function finishUsageRecord(ctx: Context, id: unknown, startedAt: nu
     inputTokens: usage?.prompt_tokens ?? null,
     outputTokens: usage?.completion_tokens ?? null,
     totalTokens: usage?.total_tokens ?? null,
+    promptCacheTokens: usage?.prompt_cache_tokens ?? null,
     resolvedService: state.aiApiLlmBilling?.resolution?.service ?? null,
     resolvedProvider: state.aiApiLlmBilling?.resolution?.provider ?? null,
     resolvedModel: state.aiApiLlmBilling?.resolution?.model ?? null,