Momento/memento-note/lib/ai/providers/ollama.ts

import { AIProvider, TagSuggestion, TitleSuggestion } from '../types';

export class OllamaProvider implements AIProvider {
  private baseUrl: string;
  private modelName: string;
  private embeddingModelName: string;

  constructor(baseUrl: string, modelName: string = 'llama3', embeddingModelName?: string) {
    if (!baseUrl) {
      throw new Error('baseUrl is required for OllamaProvider')
    }
    // Ensure baseUrl ends with /api for Ollama API
    this.baseUrl = baseUrl.endsWith('/api') ? baseUrl : `${baseUrl}/api`;
    this.modelName = modelName;
    this.embeddingModelName = embeddingModelName || modelName;
  }

  async generateTags(content: string): Promise<TagSuggestion[]> {
    try {
      const response = await fetch(`${this.baseUrl}/generate`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          model: this.modelName,
          prompt: `Analyse la note suivante et extrais les concepts clés sous forme de tags courts (1-3 mots max).
        
        Règles:
        - Pas de mots de liaison (le, la, pour, et...).
        - Garde les expressions composées ensemble (ex: "semaine prochaine", "New York").
        - Normalise en minuscules sauf noms propres.
        - Maximum 5 tags.

        Réponds UNIQUEMENT sous forme de liste JSON d'objets : [{"tag": "string", "confidence": number}].
        
        Contenu de la note: "${content}"`,
          stream: false,
        }),
      });

      if (!response.ok) throw new Error(`Ollama error: ${response.statusText}`);

      const data = await response.json();
      const text = data.response;

      const jsonMatch = text.match(/\[\s*\{[\s\S]*\}\s*\]/);
      if (jsonMatch) {
        return JSON.parse(jsonMatch[0]);
      }

      // Support pour le format { "tags": [...] }
      const objectMatch = text.match(/\{\s*"tags"\s*:\s*(\[[\s\S]*\])\s*\}/);
      if (objectMatch && objectMatch[1]) {
        return JSON.parse(objectMatch[1]);
      }
      
      return [];
    } catch (e) {
      console.error('Erreur API directe Ollama:', e);
      return [];
    }
  }

  async getEmbeddings(text: string): Promise<number[]> {
    try {
      const response = await fetch(`${this.baseUrl}/embeddings`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          model: this.embeddingModelName,
          prompt: text,
        }),
      });

      if (!response.ok) throw new Error(`Ollama error: ${response.statusText}`);

      const data = await response.json();
      return data.embedding;
    } catch (e) {
      console.error('Erreur embeddings directs Ollama:', e);
      return [];
    }
  }

  async generateTitles(prompt: string): Promise<TitleSuggestion[]> {
    try {
      const response = await fetch(`${this.baseUrl}/generate`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          model: this.modelName,
          prompt: `${prompt}

Réponds UNIQUEMENT sous forme de tableau JSON : [{"title": "string", "confidence": number}]`,
          stream: false,
        }),
      });

      if (!response.ok) throw new Error(`Ollama error: ${response.statusText}`);

      const data = await response.json();
      const text = data.response;

      // Extraire le JSON de la réponse
      const jsonMatch = text.match(/\[\s*\{[\s\S]*\}\s*\]/);
      if (jsonMatch) {
        return JSON.parse(jsonMatch[0]);
      }

      return [];
    } catch (e) {
      console.error('Erreur génération titres Ollama:', e);
      return [];
    }
  }

  async generateText(prompt: string): Promise<string> {
    try {
      const response = await fetch(`${this.baseUrl}/generate`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          model: this.modelName,
          prompt: prompt,
          stream: false,
        }),
      });

      if (!response.ok) throw new Error(`Ollama error: ${response.statusText}`);

      const data = await response.json();
      return data.response.trim();
    } catch (e) {
      console.error('Erreur génération texte Ollama:', e);
      throw e;
    }
  }
}