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

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

export class AnthropicProvider implements AIProvider {
  private model: any;

  /**
   * @param baseURL Optional Messages API root (no trailing slash). The SDK calls `{baseURL}/messages`.
   *   MiniMax: `https://api.minimax.io/anthropic` (China: `https://api.minimaxi.com/anthropic`).
   */
  constructor(apiKey: string, modelName: string = 'claude-sonnet-4-20250514', baseURL?: string) {
    const trimmedBase = baseURL?.trim().replace(/\/+$/, '');
    const zdrHeaders = { 'Anthropic-No-Train': '1' };
    const anthropicClient = createAnthropic(
      trimmedBase
        ? { apiKey, baseURL: trimmedBase, headers: zdrHeaders }
        : { apiKey, headers: zdrHeaders }
    );
    this.model = anthropicClient.chat(modelName);
  }

  async generateTags(content: string): Promise<TagSuggestion[]> {
    try {
      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 (err) {
        console.warn('Anthropic generateObject tags failed, falling back to generateText:', err);
        const { text } = await aiGenerateText({
          model: this.model,
          prompt: `Analyze the following note and suggest 1 to 5 relevant tags.
          Note content: "${content.substring(0, 1500)}"
          Return ONLY a JSON array of tag objects, like: [{"tag":"example","confidence":0.9}]`,
        });
        const cleaned = text.replace(/<think>[\s\S]*?<\/think>/gi, '').replace(/^```json\n?/, '').replace(/\n?```$/, '').trim();
        const parsed = JSON.parse(cleaned);
        const arr = Array.isArray(parsed) ? parsed : (parsed.tags || parsed.suggestions || []);
        return arr.map((t: any) => ({
          tag: t.tag || t.label || t.name || '',
          confidence: t.confidence || t.score || 0.7,
        }));
      }
    } catch (e) {
      console.error('Error generating tags (Anthropic):', e);
      return [];
    }
  }

  async getEmbeddings(_text: string): Promise<number[]> {
    throw new Error(
      'Anthropic does not expose embedding models in Memento. Choose another provider for embeddings (e.g. Ollama or OpenAI).'
    );
  }

  async generateTitles(prompt: string): Promise<TitleSuggestion[]> {
    try {
      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,
        });

        return object.titles;
      } catch (err) {
        console.warn('Anthropic generateObject titles failed, falling back to generateText:', err);
        const { text } = await aiGenerateText({
          model: this.model,
          prompt: prompt + '\n\nRespond ONLY as a JSON array of title suggestions: [{"title": "Suggested title", "confidence": 0.9}]',
        });
        const cleaned = text.replace(/<think>[\s\S]*?<\/think>/gi, '').replace(/^```json\n?/, '').replace(/\n?```$/, '').trim();
        const parsed = JSON.parse(cleaned);
        const arr = Array.isArray(parsed) ? parsed : (parsed.titles || parsed.suggestions || []);
        return arr.map((t: any) => ({
          title: typeof t === 'string' ? t : t.title || t.name || '',
          confidence: typeof t === 'number' ? t : (t.confidence || t.score || 0.8),
        }));
      }
    } catch (e) {
      console.error('Error generating titles (Anthropic):', e);
      return [];
    }
  }

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

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

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

      return { text: text.trim() };
    } catch (e) {
      console.error('Error in chat (Anthropic):', 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;
  }
}