Momento/memento-note/lib/ai/providers/custom-openai.ts

import { createOpenAI } from '@ai-sdk/openai';
import { generateObject, generateText, embed } from 'ai';
import { z } from 'zod';
import { AIProvider, TagSuggestion, TitleSuggestion } from '../types';

export class CustomOpenAIProvider implements AIProvider {
  private model: any;
  private embeddingModel: any;

  constructor(
    apiKey: string,
    baseUrl: string,
    modelName: string = 'gpt-4o-mini',
    embeddingModelName: string = 'text-embedding-3-small'
  ) {
    // Create OpenAI-compatible client with custom base URL
    const customClient = createOpenAI({
      baseURL: baseUrl,
      apiKey: apiKey,
    });

    this.model = customClient(modelName);
    this.embeddingModel = customClient.embedding(embeddingModelName);
  }

  async generateTags(content: string): Promise<TagSuggestion[]> {
    try {
      const { object } = await generateObject({
        model: this.model,
        schema: z.object({
          tags: z.array(z.object({
            tag: z.string().describe('Le nom du tag, court et en minuscules'),
            confidence: z.number().min(0).max(1).describe('Le niveau de confiance entre 0 et 1')
          }))
        }),
        prompt: `Analyse la note suivante et suggère entre 1 et 5 tags pertinents.
        Contenu de la note: "${content}"`,
      });

      return object.tags;
    } catch (e) {
      console.error('Erreur génération tags Custom OpenAI:', e);
      return [];
    }
  }

  async getEmbeddings(text: string): Promise<number[]> {
    try {
      const { embedding } = await embed({
        model: this.embeddingModel,
        value: text,
      });
      return embedding;
    } catch (e) {
      console.error('Erreur embeddings Custom OpenAI:', e);
      return [];
    }
  }

  async generateTitles(prompt: string): Promise<TitleSuggestion[]> {
    try {
      const { object } = await generateObject({
        model: this.model,
        schema: z.object({
          titles: z.array(z.object({
            title: z.string().describe('Le titre suggéré'),
            confidence: z.number().min(0).max(1).describe('Le niveau de confiance entre 0 et 1')
          }))
        }),
        prompt: prompt,
      });

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

  async generateText(prompt: string): Promise<string> {
    try {
      const { text } = await generateText({
        model: this.model,
        prompt: prompt,
      });

      return text.trim();
    } catch (e) {
      console.error('Erreur génération texte Custom OpenAI:', e);
      throw e;
    }
  }
}