smartai/ts/provider.mistral.ts

import * as plugins from './plugins.js';
import { MultiModalModel } from './abstract.classes.multimodal.js';
import type {
  ChatOptions,
  ChatResponse,
  ChatMessage,
  ResearchOptions,
  ResearchResponse,
  ImageGenerateOptions,
  ImageEditOptions,
  ImageResponse
} from './abstract.classes.multimodal.js';

export interface IMistralProviderOptions {
  mistralToken: string;
  chatModel?: string;       // default: 'mistral-large-latest'
  ocrModel?: string;        // default: 'mistral-ocr-latest'
  tableFormat?: 'markdown' | 'html';
}

export class MistralProvider extends MultiModalModel {
  private options: IMistralProviderOptions;
  public mistralClient: plugins.mistralai.Mistral;

  constructor(optionsArg: IMistralProviderOptions) {
    super();
    this.options = optionsArg;
  }

  async start() {
    await super.start();
    this.mistralClient = new plugins.mistralai.Mistral({
      apiKey: this.options.mistralToken,
    });
  }

  async stop() {
    await super.stop();
  }

  /**
   * Synchronous chat interaction using Mistral's chat API
   */
  public async chat(optionsArg: ChatOptions): Promise<ChatResponse> {
    // Convert message history to Mistral format
    const messages: Array<{
      role: 'system' | 'user' | 'assistant';
      content: string;
    }> = [];

    // Add system message first
    if (optionsArg.systemMessage) {
      messages.push({
        role: 'system',
        content: optionsArg.systemMessage
      });
    }

    // Add message history
    for (const msg of optionsArg.messageHistory) {
      messages.push({
        role: msg.role === 'system' ? 'system' : msg.role === 'assistant' ? 'assistant' : 'user',
        content: msg.content
      });
    }

    // Add current user message
    messages.push({
      role: 'user',
      content: optionsArg.userMessage
    });

    const result = await this.mistralClient.chat.complete({
      model: this.options.chatModel || 'mistral-large-latest',
      messages: messages,
    });

    // Extract content from response
    const choice = result.choices?.[0];
    let content = '';

    if (choice?.message?.content) {
      if (typeof choice.message.content === 'string') {
        content = choice.message.content;
      } else if (Array.isArray(choice.message.content)) {
        // Handle array of content chunks
        content = choice.message.content
          .map((chunk: any) => {
            if (typeof chunk === 'string') return chunk;
            if (chunk && typeof chunk === 'object' && 'text' in chunk) return chunk.text;
            return '';
          })
          .join('');
      }
    }

    return {
      role: 'assistant',
      message: content,
    };
  }

  /**
   * Streaming chat using Mistral's streaming API
   */
  public async chatStream(input: ReadableStream<Uint8Array>): Promise<ReadableStream<string>> {
    const decoder = new TextDecoder();
    let buffer = '';
    const mistralClient = this.mistralClient;
    const chatModel = this.options.chatModel || 'mistral-large-latest';

    const transform = new TransformStream<Uint8Array, string>({
      async transform(chunk, controller) {
        buffer += decoder.decode(chunk, { stream: true });

        // Try to parse complete JSON messages from the buffer
        while (true) {
          const newlineIndex = buffer.indexOf('\n');
          if (newlineIndex === -1) break;

          const line = buffer.slice(0, newlineIndex);
          buffer = buffer.slice(newlineIndex + 1);

          if (line.trim()) {
            try {
              const message = JSON.parse(line);

              // Build messages array
              const messages: Array<{
                role: 'system' | 'user' | 'assistant';
                content: string;
              }> = [];

              if (message.systemMessage) {
                messages.push({
                  role: 'system',
                  content: message.systemMessage
                });
              }

              messages.push({
                role: message.role === 'assistant' ? 'assistant' : 'user',
                content: message.content
              });

              // Use Mistral streaming
              const stream = await mistralClient.chat.stream({
                model: chatModel,
                messages: messages,
              });

              // Process streaming events
              for await (const event of stream) {
                const delta = event.data?.choices?.[0]?.delta;
                if (delta?.content) {
                  if (typeof delta.content === 'string') {
                    controller.enqueue(delta.content);
                  } else if (Array.isArray(delta.content)) {
                    for (const chunk of delta.content) {
                      if (typeof chunk === 'string') {
                        controller.enqueue(chunk);
                      } else if (chunk && typeof chunk === 'object' && 'text' in chunk) {
                        controller.enqueue((chunk as any).text);
                      }
                    }
                  }
                }
              }
            } catch (e) {
              console.error('Failed to parse message:', e);
            }
          }
        }
      },

      flush(controller) {
        if (buffer.trim()) {
          try {
            const message = JSON.parse(buffer);
            controller.enqueue(message.content || '');
          } catch (e) {
            console.error('Failed to parse remaining buffer:', e);
          }
        }
      }
    });

    return input.pipeThrough(transform);
  }

  /**
   * Audio generation is not supported by Mistral
   */
  public async audio(optionsArg: { message: string }): Promise<NodeJS.ReadableStream> {
    throw new Error('Audio generation is not supported by Mistral. Please use ElevenLabs or OpenAI provider for audio generation.');
  }

  /**
   * Vision using Mistral's OCR API for image analysis
   */
  public async vision(optionsArg: { image: Buffer; prompt: string }): Promise<string> {
    const base64Image = optionsArg.image.toString('base64');

    // Detect image type from buffer header
    let mimeType = 'image/jpeg';
    if (optionsArg.image[0] === 0x89 && optionsArg.image[1] === 0x50) {
      mimeType = 'image/png';
    } else if (optionsArg.image[0] === 0x47 && optionsArg.image[1] === 0x49) {
      mimeType = 'image/gif';
    } else if (optionsArg.image[0] === 0x52 && optionsArg.image[1] === 0x49) {
      mimeType = 'image/webp';
    }

    // Use OCR API with image data URL
    const ocrResult = await this.mistralClient.ocr.process({
      model: this.options.ocrModel || 'mistral-ocr-latest',
      document: {
        imageUrl: `data:${mimeType};base64,${base64Image}`,
        type: 'image_url',
      },
    });

    // Combine markdown from all pages
    const extractedText = ocrResult.pages.map(page => page.markdown).join('\n\n');

    // If a prompt is provided, use chat to analyze the extracted text
    if (optionsArg.prompt && optionsArg.prompt.trim()) {
      const chatResponse = await this.chat({
        systemMessage: 'You are an assistant analyzing image content. The following is text extracted from an image using OCR.',
        userMessage: `${optionsArg.prompt}\n\nExtracted content:\n${extractedText}`,
        messageHistory: [],
      });
      return chatResponse.message;
    }

    return extractedText;
  }

  /**
   * Document processing using Mistral's OCR API
   * PDFs are uploaded via Files API first, then processed with OCR
   */
  public async document(optionsArg: {
    systemMessage: string;
    userMessage: string;
    pdfDocuments: Uint8Array[];
    messageHistory: ChatMessage[];
  }): Promise<{ message: any }> {
    const extractedTexts: string[] = [];
    const uploadedFileIds: string[] = [];

    try {
      // Process each PDF document using Mistral OCR
      for (let i = 0; i < optionsArg.pdfDocuments.length; i++) {
        const pdfDocument = optionsArg.pdfDocuments[i];

        // Upload the PDF to Mistral's Files API first
        const uploadResult = await this.mistralClient.files.upload({
          file: {
            fileName: `document_${i + 1}.pdf`,
            content: pdfDocument,
          },
          purpose: 'ocr',
        });

        uploadedFileIds.push(uploadResult.id);

        // Now use OCR with the uploaded file
        const ocrResult = await this.mistralClient.ocr.process({
          model: this.options.ocrModel || 'mistral-ocr-latest',
          document: {
            type: 'file',
            fileId: uploadResult.id,
          },
          tableFormat: this.options.tableFormat || 'markdown',
        });

        // Combine all page markdown with page separators
        const pageTexts = ocrResult.pages.map((page, index) => {
          let pageContent = `--- Page ${index + 1} ---\n${page.markdown}`;

          // Include tables if present
          if (page.tables && page.tables.length > 0) {
            pageContent += '\n\n**Tables:**\n' + page.tables.map((t: any) => t.markdown || t.html || '').join('\n');
          }

          // Include header/footer if present
          if (page.header) {
            pageContent = `Header: ${page.header}\n${pageContent}`;
          }
          if (page.footer) {
            pageContent += `\nFooter: ${page.footer}`;
          }

          return pageContent;
        }).join('\n\n');

        extractedTexts.push(pageTexts);
      }

      // Combine all document texts
      const allDocumentText = extractedTexts.length === 1
        ? extractedTexts[0]
        : extractedTexts.map((text, i) => `=== Document ${i + 1} ===\n${text}`).join('\n\n');

      // Use chat API to process the extracted text with the user's query
      const chatResponse = await this.chat({
        systemMessage: optionsArg.systemMessage || 'You are a helpful assistant analyzing document content.',
        userMessage: `${optionsArg.userMessage}\n\n---\nDocument Content:\n${allDocumentText}`,
        messageHistory: optionsArg.messageHistory,
      });

      return {
        message: {
          role: 'assistant',
          content: chatResponse.message
        }
      };
    } finally {
      // Clean up uploaded files
      for (const fileId of uploadedFileIds) {
        try {
          await this.mistralClient.files.delete({ fileId });
        } catch (cleanupError) {
          // Ignore cleanup errors - files may have already been auto-deleted
          console.warn(`Failed to delete temporary file ${fileId}:`, cleanupError);
        }
      }
    }
  }

  /**
   * Research is not natively supported by Mistral
   */
  public async research(optionsArg: ResearchOptions): Promise<ResearchResponse> {
    throw new Error('Research/web search is not supported by Mistral. Please use Perplexity or Anthropic provider for research capabilities.');
  }

  /**
   * Image generation is not supported by Mistral
   */
  public async imageGenerate(optionsArg: ImageGenerateOptions): Promise<ImageResponse> {
    throw new Error('Image generation is not supported by Mistral. Please use OpenAI provider for image generation.');
  }

  /**
   * Image editing is not supported by Mistral
   */
  public async imageEdit(optionsArg: ImageEditOptions): Promise<ImageResponse> {
    throw new Error('Image editing is not supported by Mistral. Please use OpenAI provider for image editing.');
  }
}