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

import { createOpenAI } from '@ai-sdk/openai';
import { generateObject, generateText as aiGenerateText, stepCountIs } from 'ai';
import { z } from 'zod';
import { AIProvider, TagSuggestion, TitleSuggestion, ToolUseOptions, ToolCallResult } from '../types';

export class OpenRouterProvider implements AIProvider {
  private model: any;
  private apiKey: string;
  private baseUrl: string;
  private embeddingModelName: string;

  constructor(apiKey: string, modelName: string = 'anthropic/claude-3-haiku', embeddingModelName: string = 'openai/text-embedding-3-small') {
    this.apiKey = apiKey;
    this.baseUrl = 'https://openrouter.ai/api/v1';
    this.embeddingModelName = embeddingModelName;
    // Create OpenAI-compatible client for OpenRouter
    const openrouter = createOpenAI({
      baseURL: this.baseUrl,
      apiKey: apiKey,
    });

    this.model = openrouter.chat(modelName);
  }

  private async fetchWithTimeout(url: string, options: RequestInit, timeoutMs: number = 60_000): Promise<Response> {
    const controller = new AbortController()
    const timer = setTimeout(() => controller.abort(), timeoutMs)
    try {
      return await fetch(url, { ...options, signal: controller.signal })
    } finally {
      clearTimeout(timer)
    }
  }

  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('Short tag name in lowercase'),
            confidence: z.number().min(0).max(1).describe('Confidence level between 0 and 1')
          }))
        }),
        prompt: `Analyze the following note and suggest 1 to 5 relevant tags.
        Note content: "${content}"`,
      });

      return object.tags;
    } catch (e) {
      console.error('Error generating tags (OpenRouter):', e);
      return [];
    }
  }

  async getEmbeddings(text: string): Promise<number[]> {
    try {
      const response = await this.fetchWithTimeout(`${this.baseUrl}/embeddings`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'Authorization': `Bearer ${this.apiKey}`,
          'HTTP-Referer': 'https://localhost:3000',
          'X-Title': 'Memento AI',
        },
        body: JSON.stringify({
          model: this.embeddingModelName,
          input: text,
        }),
      });

      if (!response.ok) {
        const errText = await response.text();
        throw new Error(`OpenRouter embeddings error ${response.status}: ${errText}`);
      }

      const data = await response.json();

      // OpenRouter returns { data: [{ embedding: number[] }] }
      if (data.data && Array.isArray(data.data) && data.data[0]?.embedding) {
        return data.data[0].embedding;
      }

      // Fallback: some OpenAI-compatible providers return { embedding: number[] }
      if (data.embedding && Array.isArray(data.embedding)) {
        return data.embedding;
      }

      throw new Error(`Unexpected OpenRouter embeddings response shape: ${JSON.stringify(data)}`);
    } catch (e) {
      console.error('Error generating embeddings (OpenRouter):', e);
      throw e;
    }
  }

  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('Suggested title'),
            confidence: z.number().min(0).max(1).describe('Confidence level between 0 and 1')
          }))
        }),
        prompt: prompt,
      });

      return object.titles;
    } catch (e) {
      console.error('Error generating titles (OpenRouter):', e);
      return [];
    }
  }

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

      return text.trim();
    } catch (e) {
      console.error('Error generating text (OpenRouter):', e);
      throw e;
    }
  }

  async chat(messages: any[], systemPrompt?: string): Promise<any> {
    try {
      const { text } = await aiGenerateText({
        model: this.model,
        system: systemPrompt,
        messages: messages,
      });

      return { text: text.trim() };
    } catch (e) {
      console.error('Error in chat (OpenRouter):', e);
      throw e;
    }
  }

  async generateWithTools(options: ToolUseOptions): Promise<ToolCallResult> {
    const { tools, maxSteps = 10, systemPrompt, messages, prompt } = options
    const opts: Record<string, any> = {
      model: this.model,
      tools,
      stopWhen: stepCountIs(maxSteps),
    }
    if (systemPrompt) opts.system = systemPrompt
    if (messages) opts.messages = messages
    else if (prompt) opts.prompt = prompt

    const result = await aiGenerateText(opts as any)
    return {
      toolCalls: result.toolCalls?.map((tc: any) => ({ toolName: tc.toolName, input: tc.input })) || [],
      toolResults: result.toolResults?.map((tr: any) => ({ toolName: tr.toolName, input: tr.input, output: tr.output })) || [],
      text: result.text,
      steps: result.steps?.map((step: any) => ({
        text: step.text,
        toolCalls: step.toolCalls?.map((tc: any) => ({ toolName: tc.toolName, input: tc.input })) || [],
        toolResults: step.toolResults?.map((tr: any) => ({ toolName: tr.toolName, input: tr.input, output: tr.output })) || []
      })) || []
    }
  }

  getModel() {
    return this.model;
  }
}