/**
* Model Download Manager Service
*
* Centralized service responsible for orchestrating model downloads.
* Manages download state in the CuratedModels section of FMDM.
* Prevents duplicate downloads and tracks progress globally.
*/
import { ILoggingService } from '../../di/interfaces.js';
import { IFMDMService } from './fmdm-service.js';
import { CuratedModelInfo } from '../models/fmdm.js';
import { PythonEmbeddingService } from '../../infrastructure/embeddings/python-embedding-service.js';
import { ONNXDownloader } from '../../infrastructure/embeddings/onnx/onnx-downloader.js';
import { readFileSync } from 'fs';
import { join, dirname } from 'path';
import { fileURLToPath } from 'url';
/**
* Download request tracking
*/
interface DownloadRequest {
modelId: string;
startTime: Date;
promise: Promise<void>;
}
/**
* Model Download Manager interface
*/
export interface IModelDownloadManager {
/**
* Request a model download (idempotent - won't download if already downloading)
*/
requestModelDownload(modelId: string): Promise<void>;
/**
* Check if a model is available (installed or downloading)
*/
isModelAvailable(modelId: string): Promise<boolean>;
/**
* Check if a model is currently downloading
*/
isModelDownloading(modelId: string): boolean;
/**
* Get current download progress for a model
*/
getDownloadProgress(modelId: string): number | undefined;
/**
* Set the Python embedding service (for dependency injection)
*/
setPythonEmbeddingService(service: PythonEmbeddingService): void;
/**
* Set the ONNX downloader (for dependency injection)
*/
setONNXDownloader(downloader: ONNXDownloader): void;
/**
* Set cache update callback (called when model download completes successfully)
*/
setCacheUpdateCallback(callback: (modelId: string) => Promise<void>): void;
/**
* Set download complete callback (called ONLY when a model download completes successfully)
* This callback is used to trigger folder indexing after model downloads.
* Note: The callback is NOT invoked on download failures.
*/
setDownloadCompleteCallback(callback: (modelId: string) => void): void;
}
/**
* Model Download Manager implementation
*/
export class ModelDownloadManager implements IModelDownloadManager {
private activeDownloads: Map<string, DownloadRequest> = new Map();
private pythonEmbeddingService?: PythonEmbeddingService;
private onnxDownloader?: ONNXDownloader;
private cacheUpdateCallback?: (modelId: string) => Promise<void>;
private downloadCompleteCallback?: (modelId: string) => void;
private maxProgress: Map<string, number> = new Map(); // Track max progress to prevent jumps
constructor(
private logger: ILoggingService,
private fmdmService: IFMDMService
) {}
/**
* Set the Python embedding service (injected after construction to avoid circular deps)
*/
setPythonEmbeddingService(service: PythonEmbeddingService): void {
this.pythonEmbeddingService = service;
}
/**
* Set the ONNX downloader (injected after construction to avoid circular deps)
*/
setONNXDownloader(downloader: ONNXDownloader): void {
this.onnxDownloader = downloader;
}
/**
* Set cache update callback (called when model download completes successfully)
*/
setCacheUpdateCallback(callback: (modelId: string) => Promise<void>): void {
this.cacheUpdateCallback = callback;
}
/**
* Set download complete callback (called when any model download finishes)
* This is used by the orchestrator to trigger folder indexing for all folders
* waiting for this model
*/
setDownloadCompleteCallback(callback: (modelId: string) => void): void {
this.downloadCompleteCallback = callback;
}
async requestModelDownload(modelId: string): Promise<void> {
// Check if already downloading
const existingDownload = this.activeDownloads.get(modelId);
if (existingDownload) {
this.logger.debug(`Model ${modelId} already downloading, waiting for completion`);
return existingDownload.promise;
}
// Check if already installed
const modelInfo = this.getModelInfo(modelId);
if (modelInfo?.installed) {
this.logger.debug(`Model ${modelId} already installed`);
return;
}
// Start new download
const downloadPromise = this.downloadModel(modelId);
const request: DownloadRequest = {
modelId,
startTime: new Date(),
promise: downloadPromise
};
this.activeDownloads.set(modelId, request);
try {
await downloadPromise;
} finally {
this.activeDownloads.delete(modelId);
}
}
async isModelAvailable(modelId: string): Promise<boolean> {
const modelInfo = this.getModelInfo(modelId);
return modelInfo?.installed || this.isModelDownloading(modelId) || false;
}
isModelDownloading(modelId: string): boolean {
return this.activeDownloads.has(modelId);
}
getDownloadProgress(modelId: string): number | undefined {
const modelInfo = this.getModelInfo(modelId);
return modelInfo?.downloadProgress;
}
/**
* Get model info from FMDM
*/
private getModelInfo(modelId: string): CuratedModelInfo | undefined {
const fmdm = this.fmdmService.getFMDM();
return fmdm.curatedModels.find(m => m.id === modelId);
}
/**
* Perform the actual model download
*/
private async downloadModel(modelId: string): Promise<void> {
this.logger.debug(`[MODEL-DOWNLOAD] downloadModel called for ${modelId}`);
const modelInfo = this.getModelInfo(modelId);
if (!modelInfo) {
this.logger.error(`[MODEL-DOWNLOAD] Model ${modelId} not found in curated models`);
throw new Error(`Model ${modelId} not found in curated models`);
}
this.logger.debug(`[MODEL-DOWNLOAD] Model info for ${modelId}: type=${modelInfo.type}, installed=${modelInfo.installed}`);
try {
// Update FMDM to show downloading state
this.logger.debug(`[MODEL-DOWNLOAD] Setting download state to downloading for ${modelId}`);
this.updateModelDownloadState(modelId, true, 0);
if (modelInfo.type === 'gpu') {
this.logger.debug(`[MODEL-DOWNLOAD] Attempting GPU model download for ${modelId}`);
await this.downloadGPUModel(modelId);
} else if (modelInfo.type === 'cpu') {
this.logger.debug(`[MODEL-DOWNLOAD] Attempting CPU model download for ${modelId}`);
await this.downloadCPUModel(modelId);
} else {
this.logger.error(`[MODEL-DOWNLOAD] Unknown model type for ${modelId}: ${modelInfo.type}`);
throw new Error(`Unknown model type for ${modelId}`);
}
// Update FMDM to show completion
this.logger.debug(`[MODEL-DOWNLOAD] Download succeeded, marking as installed for ${modelId}`);
this.updateModelDownloadState(modelId, false, 100, true);
this.logger.info(`[MODEL-DOWNLOAD] Successfully downloaded model ${modelId}`);
// Update cache if callback is available
if (this.cacheUpdateCallback) {
try {
await this.cacheUpdateCallback(modelId);
this.logger.debug(`[MODEL-DOWNLOAD] Cache updated for ${modelId}`);
} catch (error) {
// Cache update is non-critical
this.logger.warn(`[MODEL-DOWNLOAD] Failed to update cache for ${modelId}:`, error);
}
}
// Trigger download complete callback to notify orchestrator
// This allows folders waiting for this model to start indexing
if (this.downloadCompleteCallback) {
this.logger.debug(`[MODEL-DOWNLOAD] Triggering download complete callback for ${modelId}`);
// Wrap in Promise.resolve to catch both sync and async callback errors
Promise.resolve(this.downloadCompleteCallback(modelId))
.catch(error => {
this.logger.error(
`[MODEL-DOWNLOAD] Error in download complete callback for ${modelId}:`,
error instanceof Error ? error : new Error(String(error))
);
});
}
} catch (error) {
// Update FMDM to show error
const errorMessage = error instanceof Error ? error.message : String(error);
this.logger.error(`[MODEL-DOWNLOAD] Download failed for ${modelId}: ${errorMessage}`);
this.updateModelDownloadState(modelId, false, 0, false, errorMessage);
this.logger.error(`[MODEL-DOWNLOAD] Failed to download model ${modelId}`, error instanceof Error ? error : new Error(String(error)));
throw error;
}
}
/**
* Download a GPU model using Python embedding service
*/
private async downloadGPUModel(modelId: string): Promise<void> {
this.logger.debug(`[MODEL-DOWNLOAD] downloadGPUModel called for ${modelId}`);
if (!this.pythonEmbeddingService) {
this.logger.error(`[MODEL-DOWNLOAD] Python embedding service not available for ${modelId}`);
throw new Error('Python embedding service not available for GPU model download');
}
// Get the HuggingFace model ID from curated models
this.logger.debug(`[MODEL-DOWNLOAD] Looking up HuggingFace ID for ${modelId}`);
// Load the curated models configuration
let huggingfaceId: string | undefined;
try {
// Get the path to curated-models.json relative to this file
const __dirname = dirname(fileURLToPath(import.meta.url));
const curatedModelsPath = join(__dirname, '../../config/curated-models.json');
const curatedModelsContent = readFileSync(curatedModelsPath, 'utf-8');
const curatedModels = JSON.parse(curatedModelsContent);
// Find the model configuration in GPU models
const modelConfig = curatedModels.gpuModels?.models?.find((m: any) => m.id === modelId);
if (!modelConfig || !modelConfig.huggingfaceId) {
this.logger.error(`[MODEL-DOWNLOAD] No HuggingFace ID found for model ${modelId}`);
throw new Error(`No HuggingFace ID mapping found for model ${modelId}`);
}
huggingfaceId = modelConfig.huggingfaceId;
this.logger.info(`[MODEL-DOWNLOAD] Found HuggingFace ID for ${modelId}: ${huggingfaceId}`);
} catch (error) {
this.logger.error(`[MODEL-DOWNLOAD] Failed to load curated models config`, error instanceof Error ? error : new Error(String(error)));
throw new Error(`Failed to find HuggingFace ID for model ${modelId}`);
}
// Use the Python embedding service to download the model
this.logger.info(`[MODEL-DOWNLOAD] Starting download of ${huggingfaceId} via Python embedding service`);
try {
// Initialize max progress for this download
this.maxProgress.set(modelId, 0);
// Set up real-time progress callback with monotonic progress (never decrease)
this.pythonEmbeddingService!.setDownloadProgressCallback((progress: number) => {
const currentMax = this.maxProgress.get(modelId) || 0;
const monotonicProgress = Math.max(currentMax, progress);
if (monotonicProgress > currentMax) {
this.maxProgress.set(modelId, monotonicProgress);
this.logger.debug(`[MODEL-DOWNLOAD] Progress for ${modelId}: ${monotonicProgress}% (raw: ${progress}%)`);
this.updateModelDownloadState(modelId, true, monotonicProgress);
}
});
try {
// Call the Python service to download the model
const result = await this.pythonEmbeddingService.downloadModel(huggingfaceId!);
this.logger.debug(`[MODEL-DOWNLOAD] Download result for ${modelId}: success=${result.success}, progress=${result.progress}`);
if (!result.success) {
const errorMsg = result.error || 'Model download failed';
this.logger.error(`[MODEL-DOWNLOAD] Download failed for ${modelId}: ${errorMsg}`);
this.updateModelDownloadState(modelId, false, 0);
throw new Error(errorMsg);
}
// Update progress to 100% on success
this.logger.info(`[MODEL-DOWNLOAD] Successfully downloaded ${modelId} (${huggingfaceId})`);
this.updateModelDownloadState(modelId, false, 100);
} catch (downloadError) {
// Re-throw error - cleanup handled in finally block
throw downloadError;
} finally {
// Always clear the progress callback after download completes/fails
this.pythonEmbeddingService!.setDownloadProgressCallback(() => {});
// Clean up max progress tracking
this.maxProgress.delete(modelId);
}
} catch (error) {
const errorMsg = error instanceof Error ? error.message : String(error);
this.logger.error(`[MODEL-DOWNLOAD] Failed to download ${modelId}: ${errorMsg}`);
throw error;
}
}
/**
* Download a CPU model using ONNX downloader
*/
private async downloadCPUModel(modelId: string): Promise<void> {
if (!this.onnxDownloader) {
throw new Error('ONNX downloader not available for CPU model download');
}
// Check if model is already available
const isAvailable = await this.onnxDownloader.isModelAvailable(modelId);
if (!isAvailable) {
this.logger.info(`[MODEL-DOWNLOAD] Starting ONNX model download: ${modelId}`);
// Initialize max progress for this download
this.maxProgress.set(modelId, 0);
try {
// Download with progress callback for real-time updates with monotonic progress
await this.onnxDownloader.downloadModel(modelId, {
onProgress: (progress) => {
const currentMax = this.maxProgress.get(modelId) || 0;
const monotonicProgress = Math.max(currentMax, progress.progress);
if (monotonicProgress > currentMax) {
this.maxProgress.set(modelId, monotonicProgress);
this.logger.debug(`[MODEL-DOWNLOAD] ONNX progress for ${modelId}: ${monotonicProgress}% (raw: ${progress.progress}%)`);
this.updateModelDownloadState(modelId, progress.status === 'downloading', monotonicProgress);
}
},
verifySize: true
});
// Update to 100% on success
this.logger.info(`[MODEL-DOWNLOAD] Successfully downloaded ONNX model ${modelId}`);
this.updateModelDownloadState(modelId, false, 100);
} finally {
// Clean up max progress tracking
this.maxProgress.delete(modelId);
}
} else {
this.logger.debug(`[MODEL-DOWNLOAD] ONNX model ${modelId} already cached`);
}
}
/**
* Update model download state in FMDM
*/
private updateModelDownloadState(
modelId: string,
downloading: boolean,
progress: number,
installed?: boolean,
error?: string
): void {
const fmdm = this.fmdmService.getFMDM();
const modelIndex = fmdm.curatedModels.findIndex(m => m.id === modelId);
if (modelIndex === -1) {
this.logger.warn(`Cannot update download state for unknown model ${modelId}`);
return;
}
// Update the model info
const currentModel = fmdm.curatedModels[modelIndex];
if (!currentModel) {
this.logger.warn(`Model at index ${modelIndex} is undefined`);
return;
}
const updatedModel: CuratedModelInfo = {
id: currentModel.id,
installed: currentModel.installed,
type: currentModel.type,
downloading,
downloadProgress: progress,
lastChecked: new Date().toISOString()
};
if (installed !== undefined) {
updatedModel.installed = installed;
}
if (error) {
updatedModel.downloadError = error;
} else {
delete updatedModel.downloadError;
}
// Update FMDM
const updatedModels = [...fmdm.curatedModels];
updatedModels[modelIndex] = updatedModel;
// Use FMDMService to update and broadcast
this.fmdmService.setCuratedModelInfo(updatedModels, fmdm.modelCheckStatus!);
// Also update any folders using this model
this.fmdmService.updateModelDownloadStatus(
modelId,
downloading ? 'downloading' : (installed ? 'completed' : 'failed'),
progress
);
}
}
