src/llm/llama-cpp-llm.js

/**
 * LlamaCppLLM - node-llama-cpp wrapper as a Runnable
 *
 * @module llm/llama-cpp-llm
 */

import { Runnable } from '../core/runnable.js';
import { AIMessage, HumanMessage } from '../core/message.js';
import { getLlama, LlamaChatSession } from 'node-llama-cpp';

/**
 * LlamaCppLLM - A Runnable wrapper for node-llama-cpp
 *
 * Wraps your LLM calls from agent fundamentals into a reusable,
 * composable Runnable component.
 *
 * Key benefits over raw node-llama-cpp:
 * - Composable with other Runnables via .pipe()
 * - Supports batch processing multiple inputs
 * - Built-in streaming support
 * - Consistent interface across all LLMs
 * - Easy to swap with other LLM providers
 */
export class LlamaCppLLM extends Runnable {
  /**
   * Create a new LlamaCppLLM instance
   *
   * @param {Object} options - Configuration options
   * @param {string} options.modelPath - Path to your GGUF model file (REQUIRED)
   * @param {number} [options.temperature=0.7] - Sampling temperature (0-1)
   *   - Lower (0.1): More focused, deterministic
   *   - Higher (0.9): More creative, random
   * @param {number} [options.topP=0.9] - Nucleus sampling threshold
   * @param {number} [options.topK=40] - Top-K sampling parameter
   * @param {number} [options.maxTokens=2048] - Maximum tokens to generate
   * @param {number} [options.repeatPenalty=1.1] - Penalty for repeating tokens
   * @param {number} [options.contextSize=4096] - Context window size
   * @param {number} [options.batchSize=512] - Batch processing size
   * @param {boolean} [options.verbose=false] - Enable debug logging
   * @param {string[]} [options.stopStrings] - Strings that stop generation
   * @param {Object} [options.chatWrapper] - Custom chat wrapper instance (e.g., QwenChatWrapper)
   *   - If not provided, the library will automatically select the best wrapper for your model
   *
   * @example Basic Setup
   * ```javascript
   * const llm = new LlamaCppLLM({
   *   modelPath: './models/Meta-Llama-3.1-8B-Instruct-Q5_K_S.gguf',
   *   temperature: 0.7
   * });
   * ```
   *
   * @example With Qwen Chat Wrapper (Discourage Thoughts)
   * ```javascript
   * import { QwenChatWrapper } from 'node-llama-cpp';
   *
   * const llm = new LlamaCppLLM({
   *   modelPath: './models/Qwen3-1.7B-Q6_K.gguf',
   *   temperature: 0.7,
   *   chatWrapper: new QwenChatWrapper({
   *     thoughts: 'discourage'
   *   })
   * });
   * ```
   *
   * @example Different Configurations for Different Tasks
   * ```javascript
   * // Creative writing (higher temperature)
   * const creative = new LlamaCppLLM({
   *   modelPath: './model.gguf',
   *   temperature: 0.9,
   *   maxTokens: 1000
   * });
   *
   * // Factual responses (lower temperature)
   * const factual = new LlamaCppLLM({
   *   modelPath: './model.gguf',
   *   temperature: 0.1,
   *   maxTokens: 500
   * });
   * ```
   */
  constructor(options = {}) {
    super();

    // Validate required options
    this.modelPath = options.modelPath;
    if (!this.modelPath) {
      throw new Error(
          'modelPath is required. Example: new LlamaCppLLM({ modelPath: "./model.gguf" })'
      );
    }

    // Generation parameters
    // These control how the LLM generates text - same as in your fundamentals!
    this.temperature = options.temperature ?? 0.7;
    this.topP = options.topP ?? 0.9;
    this.topK = options.topK ?? 40;
    this.maxTokens = options.maxTokens ?? 2048;
    this.repeatPenalty = options.repeatPenalty ?? 1.1;

    // Context configuration
    this.contextSize = options.contextSize ?? 4096;
    this.batchSize = options.batchSize ?? 512;

    // Behavior
    this.verbose = options.verbose ?? false;

    // Chat wrapper configuration
    // If not provided, LlamaChatSession will auto-select the best wrapper
    this.chatWrapper = options.chatWrapper ?? 'auto';

    // Stop strings - when the model sees these, it stops generating
    // Default includes common chat separators
    this.stopStrings = options.stopStrings ?? [
      'Human:',
      'User:',
      '\n\nHuman:',
      '\n\nUser:'
    ];

    // Internal state (lazy initialized)
    this._llama = null;
    this._model = null;
    this._context = null;
    this._chatSession = null;
    this._initialized = false;
  }

  /**
   * Initialize model (lazy loading)
   *
   * This loads the model only when first needed, not at construction.
   * This pattern is useful because model loading is slow - we only
   * want to do it once and only when we actually need it.
   *
   * @private
   * @throws {Error} If model loading fails
   */
  async _initialize() {
    // Skip if already initialized
    if (this._initialized) return;

    if (this.verbose) {
      console.log(`Loading model: ${this.modelPath}`);
    }

    try {
      // Step 1: Get the llama instance
      this._llama = await getLlama();

      // Step 2: Load the model file
      this._model = await this._llama.loadModel({
        modelPath: this.modelPath
      });

      // Step 3: Create a context for generation
      this._context = await this._model.createContext({
        contextSize: this.contextSize,
        batchSize: this.batchSize
      });

      // Step 4: Create a chat session
      // This manages conversation history for us
      const contextSequence = this._context.getSequence();
      const sessionConfig = {
        contextSequence
      };

      // Add chatWrapper if specified (otherwise LlamaChatSession uses "auto")
      if (this.chatWrapper !== 'auto') {
        sessionConfig.chatWrapper = this.chatWrapper;
      }

      this._chatSession = new LlamaChatSession(sessionConfig);

      this._initialized = true;

      if (this.verbose) {
        console.log('✓ Model loaded successfully');
        if (this.chatWrapper !== 'auto') {
          console.log(`✓ Using custom chat wrapper: ${this.chatWrapper.constructor.name}`);
        } else {
          console.log('✓ Using auto-detected chat wrapper');
        }
      }
    } catch (error) {
      throw new Error(
          `Failed to initialize model at ${this.modelPath}: ${error.message}`
      );
    }
  }

  /**
   * Convert our Message objects to node-llama-cpp chat history format
   *
   * This bridges between our standardized Message types and what
   * node-llama-cpp expects. Think of it as a translator.
   *
   * @private
   * @param {Array<Message>} messages - Array of Message objects
   * @returns {Array<Object>} Chat history in llama.cpp format
   *
   * @example
   * ```javascript
   * // Input: Our messages
   * [
   *   new SystemMessage("You are helpful"),
   *   new HumanMessage("Hi"),
   *   new AIMessage("Hello!")
   * ]
   *
   * // Output: llama.cpp format
   * [
   *   { type: 'system', text: 'You are helpful' },
   *   { type: 'user', text: 'Hi' },
   *   { type: 'model', response: 'Hello!' }
   * ]
   * ```
   */
  _messagesToChatHistory(messages) {
    return messages.map(msg => {
      // System messages: instructions for the AI
      if (msg._type === 'system') {
        return { type: 'system', text: msg.content };
      }
      // Human messages: user input
      else if (msg._type === 'human') {
        return { type: 'user', text: msg.content };
      }
      // AI messages: previous AI responses
      else if (msg._type === 'ai') {
        return { type: 'model', response: msg.content };
      }
      // Tool messages: results from tool execution
      else if (msg._type === 'tool') {
        // Convert tool results to system messages
        return { type: 'system', text: `Tool Result: ${msg.content}` };
      }

      // Fallback: treat unknown types as user messages
      return { type: 'user', text: msg.content };
    });
  }

  /**
   * Clean up model response
   *
   * Sometimes models include extra prefixes or suffixes.
   * This cleans them up for a better user experience.
   *
   * @private
   * @param {string} response - Raw model response
   * @returns {string} Cleaned response
   *
   * @example
   * ```javascript
   * // Before: "Assistant: The answer is 42\n\nHuman: "
   * // After:  "The answer is 42"
   * ```
   */
  _cleanResponse(response) {
    let cleaned = response.trim();

    // Remove "Assistant:" or "AI:" prefixes
    cleaned = cleaned.replace(/^(Assistant|AI):\s*/i, '');

    // Remove any conversation continuations
    cleaned = cleaned.replace(/\n\n(Human|User):.*$/s, '');

    return cleaned.trim();
  }

  /**
   * Main generation method - this is where your LLM calls happen!
   *
   * This is the same as calling `llm.chat(messages)` in your fundamentals,
   * but wrapped to work with the Runnable interface.
   *
   * @async
   * @param {string|Array<Message>} input - User input or message array
   * @param {Object} [config={}] - Runtime configuration
   * @param {number} [config.temperature] - Override temperature for this call
   * @param {number} [config.maxTokens] - Override max tokens for this call
   * @param {boolean} [config.clearHistory=false] - Clear chat history before this call
   * @returns {Promise<AIMessage>} Generated response as AIMessage
   *
   * @example String Input (Simplest)
   * ```javascript
   * const response = await llm.invoke("What is AI?");
   * console.log(response.content); // "AI is..."
   * ```
   *
   * @example Message Array Input (Full Control)
   * ```javascript
   * const messages = [
   *   new SystemMessage("You are a helpful assistant"),
   *   new HumanMessage("What is AI?")
   * ];
   * const response = await llm.invoke(messages);
   * ```
   *
   * @example Runtime Configuration
   * ```javascript
   * // Override temperature for this specific call
   * const response = await llm.invoke(
   *   "Write a creative story",
   *   { temperature: 0.9, maxTokens: 500 }
   * );
   * ```
   *
   * @example Clear History Before Call
   * ```javascript
   * // Ensure fresh context with no prior conversation
   * const response = await llm.invoke(
   *   "What is AI?",
   *   { clearHistory: true }
   * );
   * ```
   *
   * @example In a Pipeline (Composition)
   * ```javascript
   * const pipeline = promptFormatter
   *   .pipe(llm)
   *   .pipe(outputParser);
   *
   * const result = await pipeline.invoke("user input");
   * ```
   */
  async _call(input, config = {}) {
    // Ensure model is loaded (only happens once)
    await this._initialize();

    // Clear history if requested (important for batch processing)
    if (config.clearHistory) {
      this._chatSession.setChatHistory([]);
    }

    // Handle different input types
    let messages;
    if (typeof input === 'string') {
      messages = [new HumanMessage(input)];
    } else if (Array.isArray(input)) {
      messages = input;
    } else {
      throw new Error(
          'Input must be a string or array of messages. ' +
          'Example: "Hello" or [new HumanMessage("Hello")]'
      );
    }

    // Extract system message if present
    const systemMessages = messages.filter(msg => msg._type === 'system');
    const systemPrompt = systemMessages.length > 0
        ? systemMessages[0].content
        : '';

    // Convert our Message objects to llama.cpp format
    const chatHistory = this._messagesToChatHistory(messages);
    this._chatSession.setChatHistory(chatHistory);

    // ALWAYS set system prompt (either new value or empty string to clear)
    this._chatSession.systemPrompt = systemPrompt;

    try {
      // Build prompt options
      const promptOptions = {
        temperature: config.temperature ?? this.temperature,
        topP: config.topP ?? this.topP,
        topK: config.topK ?? this.topK,
        maxTokens: config.maxTokens ?? this.maxTokens,
        repeatPenalty: config.repeatPenalty ?? this.repeatPenalty,
        customStopTriggers: config.stopStrings ?? this.stopStrings
      };

      // Add random seed if temperature > 0 and no seed specified
      // This ensures randomness works properly
      if (promptOptions.temperature > 0 && config.seed === undefined) {
        promptOptions.seed = Math.floor(Math.random() * 1000000);
      } else if (config.seed !== undefined) {
        promptOptions.seed = config.seed;
      }

      // Generate response using prompt (simpler than promptWithMeta for non-streaming)
      const response = await this._chatSession.prompt('', promptOptions);

      // Return as AIMessage for consistency
      return new AIMessage(response);
    } catch (error) {
      throw new Error(`Generation failed: ${error.message}`);
    }
  }

  /**
   * Batch processing with history isolation
   *
   * Processes multiple inputs sequentially, ensuring each gets a clean chat history.
   * Note: Local models process requests sequentially, so there's no performance
   * benefit compared to calling invoke() multiple times.
   *
   * @async
   * @param {Array<string|Array<Message>>} inputs - Array of inputs to process
   * @param {Object} [config={}] - Runtime configuration
   * @returns {Promise<Array<AIMessage>>} Array of generated responses
   *
   * @example
   * ```javascript
   * const questions = ["What is AI?", "What is ML?", "What is DL?"];
   * const answers = await llm.batch(questions);
   * ```
   */
  async batch(inputs, config = {}) {
    const results = [];
    for (const input of inputs) {
      // Clear history before each batch item to prevent contamination
      const result = await this._call(input, { ...config, clearHistory: true });
      results.push(result);
    }
    return results;
  }

  /**
   * Streaming generation - show results as they're generated!
   *
   * This is the same as _call() but yields chunks as they arrive,
   * like the typing effect you see in ChatGPT.
   *
   * @async
   * @generator
   * @param {string|Array<Message>} input - User input or message array
   * @param {Object} [config={}] - Runtime configuration
   * @yields {AIMessage} Chunks of generated text
   *
   * @example Basic Streaming
   * ```javascript
   * console.log("Response: ");
   * for await (const chunk of llm.stream("Tell me a story")) {
   *   process.stdout.write(chunk.content); // Print without newline
   * }
   * console.log("\nDone!");
   * ```
   *
   * @example Streaming in a Pipeline
   * ```javascript
   * const pipeline = promptFormatter
   *   .pipe(llm)
   *   .pipe(parser);
   *
   * // Only the last step (parser) gets streamed chunks
   * for await (const chunk of pipeline.stream(input)) {
   *   console.log(chunk);
   * }
   * ```
   *
   * @example Building a Chat UI
   * ```javascript
   * async function streamToUI(input) {
   *   let fullResponse = '';
   *
   *   for await (const chunk of llm.stream(input)) {
   *     fullResponse += chunk.content;
   *     updateUI(fullResponse); // Update your UI in real-time
   *   }
   * }
   * ```
   */
  async* stream(input, config = {}) {
    await this._initialize();

    // Clear history if requested
    if (config.clearHistory) {
      this._chatSession.setChatHistory([]);
    }

    // Handle different input types (same as _call)
    let messages;
    if (typeof input === 'string') {
      messages = [new HumanMessage(input)];
    } else if (Array.isArray(input)) {
      messages = input;
    } else {
      throw new Error(
          'Input must be a string or array of messages for streaming'
      );
    }

    // Extract system message if present
    const systemMessages = messages.filter(msg => msg._type === 'system');
    const systemPrompt = systemMessages.length > 0
        ? systemMessages[0].content
        : '';

    // Set up chat history
    const chatHistory = this._messagesToChatHistory(messages);
    this._chatSession.setChatHistory(chatHistory);

    // ALWAYS set system prompt (either new value or empty string to clear)
    this._chatSession.systemPrompt = systemPrompt;

    try {
      // Build prompt options
      const promptOptions = {
        temperature: config.temperature ?? this.temperature,
        topP: config.topP ?? this.topP,
        topK: config.topK ?? this.topK,
        maxTokens: config.maxTokens ?? this.maxTokens,
        repeatPenalty: config.repeatPenalty ?? this.repeatPenalty,
        customStopTriggers: config.stopStrings ?? this.stopStrings
      };

      // Add random seed if temperature > 0 and no seed specified
      if (promptOptions.temperature > 0 && config.seed === undefined) {
        promptOptions.seed = Math.floor(Math.random() * 1000000);
      } else if (config.seed !== undefined) {
        promptOptions.seed = config.seed;
      }

      // Use onTextChunk callback to stream chunks as they arrive
      const self = this;
      promptOptions.onTextChunk = (chunk) => {
        // This callback is synchronous, so we can't yield directly
        // We'll collect chunks and yield them after
        self._currentStreamChunks = self._currentStreamChunks || [];
        self._currentStreamChunks.push(chunk);
      };

      // Initialize chunk collection
      this._currentStreamChunks = [];

      // Start generation (this will call onTextChunk as it generates)
      const responsePromise = this._chatSession.prompt('', promptOptions);

      // Yield chunks as they become available
      let lastYieldedIndex = 0;

      // Poll for new chunks
      while (true) {
        // Yield any new chunks
        while (lastYieldedIndex < this._currentStreamChunks.length) {
          yield new AIMessage(this._currentStreamChunks[lastYieldedIndex], {
            additionalKwargs: { chunk: true }
          });
          lastYieldedIndex++;
        }

        // Check if generation is complete
        const isDone = await Promise.race([
          responsePromise.then(() => true),
          new Promise(resolve => setTimeout(() => resolve(false), 10))
        ]);

        if (isDone) {
          // Yield any remaining chunks
          while (lastYieldedIndex < this._currentStreamChunks.length) {
            yield new AIMessage(this._currentStreamChunks[lastYieldedIndex], {
              additionalKwargs: { chunk: true }
            });
            lastYieldedIndex++;
          }
          break;
        }
      }

      // Wait for the full response
      await responsePromise;

      // Clean up
      delete this._currentStreamChunks;

    } catch (error) {
      throw new Error(`Streaming failed: ${error.message}`);
    }
  }

  /**
   * Cleanup resources
   *
   * LLMs hold resources in memory. Call this when you're done
   * to free them up properly.
   *
   * @async
   * @returns {Promise<void>}
   *
   * @example
   * ```javascript
   * const llm = new LlamaCppLLM({ modelPath: './model.gguf' });
   *
   * try {
   *   const response = await llm.invoke("Hello");
   *   console.log(response.content);
   * } finally {
   *   await llm.dispose(); // Always clean up!
   * }
   * ```
   *
   * @example With Multiple Uses
   * ```javascript
   * const llm = new LlamaCppLLM({ modelPath: './model.gguf' });
   *
   * // Use it many times
   * await llm.invoke("Question 1");
   * await llm.invoke("Question 2");
   * await llm.batch(["Q3", "Q4", "Q5"]);
   *
   * // Clean up when completely done
   * await llm.dispose();
   * ```
   */
  async dispose() {
    if (this._context) {
      await this._context.dispose();
      this._context = null;
    }
    if (this._model) {
      await this._model.dispose();
      this._model = null;
    }
    this._chatSession = null;
    this._initialized = false;

    if (this.verbose) {
      console.log('✓ Model resources disposed');
    }
  }

  /**
   * String representation for debugging
   *
   * @returns {string} Human-readable representation
   *
   * @example
   * ```javascript
   * const llm = new LlamaCppLLM({ modelPath: './llama-2-7b.gguf' });
   * console.log(llm.toString());
   * // "LlamaCppLLM(model=./llama-2-7b.gguf)"
   *
   * // Useful in pipelines
   * const pipeline = formatter.pipe(llm).pipe(parser);
   * console.log(pipeline.toString());
   * // "PromptFormatter() | LlamaCppLLM(model=./llama-2-7b.gguf) | OutputParser()"
   * ```
   */
  toString() {
    return `LlamaCppLLM(model=${this.modelPath})`;
  }
}

export default LlamaCppLLM;