KnowledgeBase/server/src/vision-pipeline/vision-pipeline-cost-aware.service.ts

import { Injectable, Logger } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import * as fs from 'fs/promises';
import * as path from 'path';
import { LibreOfficeService } from '../libreoffice/libreoffice.service';
import { Pdf2ImageService } from '../pdf2image/pdf2image.service';
import { VisionService } from '../vision/vision.service';
import { ElasticsearchService } from '../elasticsearch/elasticsearch.service';
import { ModelConfigService } from '../model-config/model-config.service';
import { PreciseModeOptions, PipelineResult, ProcessingStatus, ModeRecommendation } from './vision-pipeline.interface';
import { VisionModelConfig, VisionAnalysisResult } from '../vision/vision.interface';
import { CostControlService } from './cost-control.service';

@Injectable()
export class VisionPipelineCostAwareService {
  private readonly logger = new Logger(VisionPipelineCostAwareService.name);

  constructor(
    private libreOffice: LibreOfficeService,
    private pdf2Image: Pdf2ImageService,
    private vision: VisionService,
    private elasticsearch: ElasticsearchService,
    private modelConfigService: ModelConfigService,
    private configService: ConfigService,
    private costControl: CostControlService, 
  ) { }

  
  async processPreciseMode(
    filePath: string,
    options: PreciseModeOptions,
  ): Promise<PipelineResult> {
    const startTime = Date.now();
    const results: VisionAnalysisResult[] = [];
    let processedPages = 0;
    let failedPages = 0;
    let totalCost = 0;
    let pdfPath = filePath;
    let imagesToProcess: any[] = [];

    this.logger.log(
      `Starting precise mode processing for ${options.fileName} (user: ${options.userId})`
    );

    try {
      
      this.updateStatus('converting', 10, 'Converting document format...');
      pdfPath = await this.convertToPDF(filePath);

      
      this.updateStatus('splitting', 30, 'Converting PDF to image...');
      const conversionResult = await this.pdf2Image.convertToImages(pdfPath, {
        density: 300,
        quality: 85,
        format: 'jpeg',
      });

      if (conversionResult.images.length === 0) {
        throw new Error('PDF to image conversion failed. No image was generated');
      }

      
      imagesToProcess = options.maxPages
        ? conversionResult.images.slice(0, options.maxPages)
        : conversionResult.images;

      const pageCount = imagesToProcess.length;

      
      this.updateStatus('checking', 40, 'Checking quotas and estimating costs...');
      const costEstimate = this.costControl.estimateCost(pageCount);
      this.logger.log(
        `Estimated cost: $${costEstimate.estimatedCost.toFixed(2)}, Estimated time: ${this.costControl.formatTime(costEstimate.estimatedTime)}`
      );

      
      const quotaCheck = await this.costControl.checkQuota(
        options.userId,
        costEstimate.estimatedCost,
      );

      if (!quotaCheck.allowed) {
        throw new Error(quotaCheck.reason);
      }

      
      const warning = await this.costControl.checkWarningThreshold(options.userId);
      if (warning.shouldWarn) {
        this.logger.warn(warning.message);
      }

      
      const modelConfig = await this.getVisionModelConfig(options.userId, options.modelId, options.tenantId);

      
      this.updateStatus('analyzing', 50, 'Analyzing the page using the vision model...');
      const batchResult = await this.vision.batchAnalyze(
        imagesToProcess.map(img => img.path),
        modelConfig,
        {
          startIndex: 1,
          skipQualityCheck: options.skipQualityCheck,
        }
      );

      totalCost = batchResult.estimatedCost;
      processedPages = batchResult.successCount;
      failedPages = batchResult.failedCount;
      results.push(...batchResult.results);

      
      if (totalCost > 0) {
        await this.costControl.deductQuota(options.userId, totalCost);
        this.logger.log(`Actual deducted cost: $${totalCost.toFixed(2)}`);
      }

      
      this.updateStatus('completed', 100, 'Processing completed. Cleaning up temporary files...');
      await this.pdf2Image.cleanupImages(imagesToProcess);

      
      if (pdfPath !== filePath) {
        try {
          await fs.unlink(pdfPath);
        } catch (error) {
          this.logger.warn(`Failed to cleanup converted PDF: ${error.message}`);
        }
      }

      const duration = (Date.now() - startTime) / 1000;

      this.logger.log(
        `Precise mode completed: ${processedPages} pages processed, ` +
        `cost: $${totalCost.toFixed(2)}, duration: ${duration.toFixed(1)}s`
      );

      return {
        success: true,
        fileId: options.fileId,
        fileName: options.fileName,
        totalPages: conversionResult.totalPages,
        processedPages,
        failedPages,
        results,
        cost: totalCost,
        duration,
        mode: 'precise',
      };
    } catch (error) {
      this.logger.error(`Precise mode failed: ${error.message}`);

      
      try {
        if (pdfPath !== filePath && pdfPath !== filePath) {
          await fs.unlink(pdfPath);
        }
        if (imagesToProcess.length > 0) {
          await this.pdf2Image.cleanupImages(imagesToProcess);
        }
      } catch { }

      return {
        success: false,
        fileId: options.fileId,
        fileName: options.fileName,
        totalPages: 0,
        processedPages,
        failedPages,
        results: [],
        cost: totalCost,
        duration: (Date.now() - startTime) / 1000,
        mode: 'precise',
      };
    }
  }

  
  private async getVisionModelConfig(userId: string, modelId: string, tenantId?: string): Promise<VisionModelConfig> {
    const config = await this.modelConfigService.findOne(modelId, userId, tenantId || 'default');

    if (!config) {
      throw new Error(`Model configuration not found: ${modelId}`);
    }

    // API key is optional - allow local models

    return {
      baseUrl: config.baseUrl || '',
      apiKey: config.apiKey || '',
      modelId: config.modelId,
    };
  }

  
  private async convertToPDF(filePath: string): Promise<string> {
    const ext = path.extname(filePath).toLowerCase();

    
    if (ext === '.pdf') {
      return filePath;
    }

    
    return await this.libreOffice.convertToPDF(filePath);
  }

  
  async recommendMode(filePath: string): Promise<ModeRecommendation> {
    const ext = path.extname(filePath).toLowerCase();
    const stats = await fs.stat(filePath);
    const sizeMB = stats.size / (1024 * 1024);

    const supportedFormats = ['.pdf', '.doc', '.docx', '.ppt', '.pptx', '.xls', '.xlsx'];
    const preciseFormats = ['.pdf', '.doc', '.docx', '.ppt', '.pptx'];

    if (!supportedFormats.includes(ext)) {
      return {
        recommendedMode: 'fast',
        reason: `Unsupported file format: ${ext}`,
        warnings: ['Using Fast Mode (text extraction only)'],
      };
    }

    if (!preciseFormats.includes(ext)) {
      return {
        recommendedMode: 'fast',
        reason: `Format ${ext} does not support Precise Mode`,
        warnings: ['Using Fast Mode (text extraction only)'],
      };
    }

    
    const estimatedPages = Math.max(1, Math.ceil(sizeMB * 2));
    const costEstimate = this.costControl.estimateCost(estimatedPages);

    
    if (sizeMB > 50) {
      return {
        recommendedMode: 'precise',
        reason: 'Due to large files, Precise Mode is recommended to retain complete information',
        estimatedCost: costEstimate.estimatedCost,
        estimatedTime: costEstimate.estimatedTime,
        warnings: ['Processing time may be longer', 'API charges may apply'],
      };
    }

    
    return {
      recommendedMode: 'precise',
      reason: 'Precise Mode is available. Can hold mixed content of text and images',
      estimatedCost: costEstimate.estimatedCost,
      estimatedTime: costEstimate.estimatedTime,
      warnings: ['API charges will apply'],
    };
  }

  
  async getUserQuotaInfo(userId: string) {
    const quota = await this.costControl.getUserQuota(userId);
    const report = await this.costControl.getCostReport(userId);

    return {
      ...quota,
      report,
      warnings: await this.costControl.checkWarningThreshold(userId),
    };
  }

  
  private updateStatus(status: ProcessingStatus['status'], progress: number, message: string): void {
    this.logger.log(`[${status}] ${progress}% - ${message}`);
  }
}