js/ahmad-jit-engine.js

/**
 * Ahmad JIT Engine
 * Browser-based LLM inference using WebLLM (ES module)
 * Model: TinyLlama-1.1B-Chat-v1.0-q4f32_1-MLC
 */

class AhmadJITEngine {
    constructor() {
        this.engine = null;
        this.state = 'OFFLINE';
        this.conversationHistory = [];
        this.notebookCells = [];
        this.abortController = null;
        this.modelId = 'TinyLlama-1.1B-Chat-v1.0-q4f32_1-MLC';  // Real WebLLM model
    }

    async initialize() {
        try {
            this.state = 'CHECKING_WEBGPU';
            if (!navigator.gpu) {
                console.warn('WebGPU not available, will use CPU (slower)');
            } else {
                console.log('✓ WebGPU available for acceleration');
            }

            this.state = 'LOADING';
            console.log('Dynamically importing WebLLM ES module...');

            // Dynamic import for ES module (not a global script)
            let webllm;
            try {
                webllm = await import('https://esm.run/@mlc-ai/web-llm');
                console.log('✓ WebLLM ES module imported successfully');
            } catch (importError) {
                console.error('Failed to import WebLLM:', importError);
                throw new Error(`Failed to import WebLLM ES module: ${importError.message}`);
            }

            // Verify CreateMLCEngine is available
            if (!webllm.CreateMLCEngine) {
                throw new Error('WebLLM loaded, but CreateMLCEngine is unavailable. Check module export.');
            }

            console.log('✓ CreateMLCEngine available from WebLLM module');

            // Create engine using CreateMLCEngine (not MLCEngine constructor)
            this.state = 'DOWNLOADING_MODEL';
            console.log(`Creating MLCEngine for model: ${this.modelId}`);

            try {
                this.engine = await webllm.CreateMLCEngine(this.modelId, {
                    initProgressCallback: (progress) => {
                        console.log('Model loading progress:', progress);
                    },
                });
                console.log('✓ MLCEngine created and model initialized successfully');
            } catch (engineError) {
                console.error('Failed to create engine:', engineError);
                throw new Error(`Failed to create MLCEngine: ${engineError.message}`);
            }

            this.state = 'INDEXING';
            this.buildNotebookIndex();

            this.state = 'READY';
            console.log('✓ Ahmad Bot ready!');
            return true;
        } catch (error) {
            this.state = 'ERROR';
            console.error('Ahmad JIT Engine initialization failed:', error);
            throw error;
        }
    }

    buildNotebookIndex() {
        this.notebookCells = [];
        const cellElements = document.querySelectorAll(
            '[data-cell-id], .notebook-cell, .cell, [class*="cell"]'
        );

        let cellIndex = 0;
        cellElements.forEach((elem) => {
            if (elem.contains(document.getElementById('ahmad-jit-box'))) {
                return; // Skip Ahmad Bot's own UI
            }

            const cellId = elem.dataset.cellId || `ROWM-${String(cellIndex + 1).padStart(3, '0')}`;
            elem.dataset.cellId = cellId;

            const text = elem.textContent.substring(0, 500).trim();
            if (text.length > 0) {
                this.notebookCells.push({
                    id: cellId,
                    text: text,
                    element: elem,
                });
                cellIndex++;
            }
        });

        console.log(`Indexed ${this.notebookCells.length} notebook cells`);
    }

    findRelevantCells(question, maxCells = 3) {
        const questionTerms = question.toLowerCase().split(/\s+/);
        const scored = this.notebookCells.map((cell) => {
            const cellLower = cell.text.toLowerCase();
            const matches = questionTerms.filter((term) => cellLower.includes(term)).length;
            return { cell, score: matches };
        });

        return scored
            .filter((s) => s.score > 0)
            .sort((a, b) => b.score - a.score)
            .slice(0, maxCells)
            .map((s) => s.cell);
    }

    buildSystemPrompt(question) {
        const relevant = this.findRelevantCells(question);
        let context = 'NOTEBOOK CONTEXT:\n\n';

        relevant.forEach((cell) => {
            context += `[${cell.id}]\n${cell.text}\n\n`;
        });

        return `You are Ahmad Bot, the embedded local assistant for the ROWM Polymorphic Notebook.

Answer from the supplied notebook cells.

Cite notebook cells using their real identifiers.

Do not invent missing information.

Preserve Unicode exactly.

Do not claim a hash is a signature.

Do not claim visual animation is formal verification.

${context}

USER QUESTION:
${question}`;
    }

    async generate(userMessage, onToken) {
        if (this.state !== 'READY') {
            throw new Error(`Engine not ready: ${this.state}`);
        }

        if (!userMessage || userMessage.trim().length === 0) {
            throw new Error('User message cannot be empty');
        }

        this.state = 'GENERATING';
        this.abortController = new AbortController();
        let fullResponse = '';

        try {
            const systemPrompt = this.buildSystemPrompt(userMessage);
            const messages = [
                { role: 'system', content: systemPrompt },
                ...this.conversationHistory,
                { role: 'user', content: userMessage },
            ];

            console.log('Generating response with WebLLM...');

            // WebLLM completion() is the correct streaming method
            const response = await this.engine.completion({
                messages: messages,
                temperature: 0.2,
                top_p: 0.9,
                max_tokens: 600,
                stream: true,
            });

            // Stream tokens from response
            for await (const chunk of response) {
                if (this.abortController.signal.aborted) {
                    break;
                }

                // WebLLM chunk has format: { choices: [{ delta: { content: "token" } }] }
                const content = chunk.choices?.[0]?.delta?.content ?? '';
                if (content && typeof content === 'string') {
                    fullResponse += content;
                    if (onToken && typeof onToken === 'function') {
                        try {
                            onToken(content);
                        } catch (callbackError) {
                            console.error('onToken callback error:', callbackError);
                        }
                    }
                }
            }

            // Store in conversation history
            this.conversationHistory.push({ role: 'user', content: userMessage });
            this.conversationHistory.push({ role: 'assistant', content: fullResponse });

            this.state = 'READY';
            return fullResponse;
        } catch (error) {
            if (error.name === 'AbortError') {
                this.state = 'STOPPED';
                return fullResponse;
            }
            this.state = 'ERROR';
            console.error('Generation error:', error);
            throw error;
        }
    }

    stop() {
        if (this.abortController) {
            this.abortController.abort();
        }
    }

    resetConversation() {
        this.conversationHistory = [];
    }

    unload() {
        if (this.engine) {
            this.engine.terminate?.();
        }
        this.engine = null;
        this.state = 'OFFLINE';
    }

    getState() {
        return this.state;
    }
}

// Global instance
window.ahmadEngine = null;

// Export
if (typeof module !== 'undefined' && module.exports) {
    module.exports = AhmadJITEngine;
}