SwagManager MCP

// server/handlers/llm-providers.ts — Multi-provider LLM with automatic failover
// Providers: Anthropic, OpenAI, Google, AWS Bedrock, Ollama (local)

import type { SupabaseClient } from "@supabase/supabase-js";
import Anthropic from "@anthropic-ai/sdk";
import OpenAI from "openai";
import { GoogleGenAI } from "@google/genai";
import { BedrockRuntimeClient, InvokeModelCommand } from "@aws-sdk/client-bedrock-runtime";

// ============================================================================
// TYPES
// ============================================================================

type Provider = "anthropic" | "openai" | "google" | "bedrock" | "ollama";

interface ProviderResult {
  provider: Provider;
  model: string;
  text: string;
  tokens: { input: number; output: number };
}

interface ProviderCredentials {
  anthropic?: { apiKey: string };
  openai?: { apiKey: string };
  google?: { apiKey: string };
  bedrock?: { accessKeyId: string; secretAccessKey: string; region: string };
  ollama?: { available: boolean };
}

// ============================================================================
// STATIC MODEL LISTS
// ============================================================================

const ANTHROPIC_MODELS = [
  { model_id: "claude-sonnet-4-20250514", name: "Claude Sonnet 4", context_window: 200000 },
  { model_id: "claude-opus-4-20250514", name: "Claude Opus 4", context_window: 200000 },
  { model_id: "claude-haiku-4-20250514", name: "Claude Haiku 4", context_window: 200000 },
  { model_id: "claude-3-5-sonnet-20241022", name: "Claude 3.5 Sonnet", context_window: 200000 },
  { model_id: "claude-3-5-haiku-20241022", name: "Claude 3.5 Haiku", context_window: 200000 },
];

const OPENAI_MODELS = [
  { model_id: "gpt-4o", name: "GPT-4o", context_window: 128000 },
  { model_id: "gpt-4o-mini", name: "GPT-4o Mini", context_window: 128000 },
  { model_id: "gpt-4-turbo", name: "GPT-4 Turbo", context_window: 128000 },
  { model_id: "o1", name: "o1", context_window: 200000 },
  { model_id: "o1-mini", name: "o1-mini", context_window: 128000 },
  { model_id: "o3-mini", name: "o3-mini", context_window: 200000 },
];

const BEDROCK_MODELS = [
  { model_id: "us.anthropic.claude-sonnet-4-20250514-v1:0", name: "Claude Sonnet 4 (Bedrock)", context_window: 200000 },
  { model_id: "us.anthropic.claude-sonnet-4-5-20250929-v1:0", name: "Claude Sonnet 4.5 (Bedrock)", context_window: 200000 },
  { model_id: "us.anthropic.claude-haiku-4-5-20251001-v1:0", name: "Claude Haiku 4.5 (Bedrock)", context_window: 200000 },
  { model_id: "us.anthropic.claude-3-5-haiku-20241022-v1:0", name: "Claude 3.5 Haiku (Bedrock)", context_window: 200000 },
  { model_id: "us.meta.llama3-1-70b-instruct-v1:0", name: "Llama 3.1 70B (Bedrock)", context_window: 128000 },
  { model_id: "us.amazon.nova-pro-v1:0", name: "Amazon Nova Pro (Bedrock)", context_window: 300000 },
];

const GOOGLE_MODELS = [
  { model_id: "gemini-3-pro-preview", name: "Gemini 3 Pro", context_window: 1048576 },
  { model_id: "gemini-3-flash-preview", name: "Gemini 3 Flash", context_window: 1048576 },
  { model_id: "gemini-2.5-pro", name: "Gemini 2.5 Pro", context_window: 1048576 },
  { model_id: "gemini-2.5-flash", name: "Gemini 2.5 Flash", context_window: 1048576 },
  { model_id: "gemini-2.5-flash-lite", name: "Gemini 2.5 Flash Lite", context_window: 1048576 },
  { model_id: "gemini-2.0-flash", name: "Gemini 2.0 Flash", context_window: 1048576 },
  { model_id: "gemini-2.0-flash-lite", name: "Gemini 2.0 Flash Lite", context_window: 1048576 },
];

const OLLAMA_MODELS = [
  { model_id: "llama3.2", name: "Llama 3.2 (local)", context_window: 128000 },
  { model_id: "mistral", name: "Mistral (local)", context_window: 32000 },
  { model_id: "codellama", name: "Code Llama (local)", context_window: 16000 },
  { model_id: "phi3", name: "Phi-3 (local)", context_window: 128000 },
];

// ============================================================================
// CREDENTIAL RESOLUTION
// ============================================================================

async function getCredentials(sb: SupabaseClient, storeId: string): Promise<ProviderCredentials> {
  const creds: ProviderCredentials = {};

  // Anthropic — available via process.env on the server
  const anthropicKey = process.env.ANTHROPIC_API_KEY;
  if (anthropicKey) {
    creds.anthropic = { apiKey: anthropicKey };
  }

  // OpenAI — from user_tool_secrets (encrypted)
  try {
    const { data: openaiKey } = await sb.rpc("decrypt_secret", { p_name: "OPENAI_API_KEY", p_store_id: storeId });
    if (openaiKey) {
      creds.openai = { apiKey: openaiKey as string };
    }
  } catch { /* not configured */ }

  // Google AI — from user_tool_secrets (encrypted)
  try {
    const { data: googleKey } = await sb.rpc("decrypt_secret", { p_name: "GOOGLE_AI_API_KEY", p_store_id: storeId });
    if (googleKey) {
      creds.google = { apiKey: googleKey as string };
    }
  } catch { /* not configured */ }

  // AWS Bedrock — from user_tool_secrets
  try {
    const [accessKeyRes, secretKeyRes, regionRes] = await Promise.all([
      sb.rpc("decrypt_secret", { p_name: "AWS_ACCESS_KEY_ID", p_store_id: storeId }),
      sb.rpc("decrypt_secret", { p_name: "AWS_SECRET_ACCESS_KEY", p_store_id: storeId }),
      sb.rpc("decrypt_secret", { p_name: "AWS_REGION", p_store_id: storeId }),
    ]);
    if (accessKeyRes.data && secretKeyRes.data) {
      creds.bedrock = {
        accessKeyId: accessKeyRes.data as string,
        secretAccessKey: secretKeyRes.data as string,
        region: (regionRes.data as string) || "us-east-1",
      };
    }
  } catch { /* not configured */ }

  // Ollama — check if local server is available
  try {
    const controller = new AbortController();
    const timer = setTimeout(() => controller.abort(), 2000);
    const resp = await fetch("http://localhost:11434/api/version", { signal: controller.signal });
    clearTimeout(timer);
    if (resp.ok) {
      creds.ollama = { available: true };
    }
  } catch { /* not available */ }

  return creds;
}

// ============================================================================
// PROVIDER IMPLEMENTATIONS
// ============================================================================

const PROVIDER_TIMEOUT_MS = 30_000;

async function callAnthropic(
  apiKey: string,
  prompt: string,
  model: string | undefined,
  maxTokens: number,
  temperature: number,
  systemPrompt?: string,
): Promise<ProviderResult> {
  const client = new Anthropic({ apiKey });
  const selectedModel = model || "claude-sonnet-4-20250514";

  const controller = new AbortController();
  const timer = setTimeout(() => controller.abort(), PROVIDER_TIMEOUT_MS);

  try {
    const resp = await client.messages.create({
      model: selectedModel,
      max_tokens: maxTokens,
      temperature,
      ...(systemPrompt ? { system: systemPrompt } : {}),
      messages: [{ role: "user", content: prompt }],
    }, { signal: controller.signal });

    clearTimeout(timer);

    const text = resp.content
      .filter((b): b is Anthropic.TextBlock => b.type === "text")
      .map((b) => b.text)
      .join("\n");

    return {
      provider: "anthropic",
      model: selectedModel,
      text,
      tokens: {
        input: resp.usage.input_tokens,
        output: resp.usage.output_tokens,
      },
    };
  } catch (err) {
    clearTimeout(timer);
    throw err;
  }
}

async function callOpenAI(
  apiKey: string,
  prompt: string,
  model: string | undefined,
  maxTokens: number,
  temperature: number,
  systemPrompt?: string,
): Promise<ProviderResult> {
  const client = new OpenAI({ apiKey });
  const selectedModel = model || "gpt-4o";

  const controller = new AbortController();
  const timer = setTimeout(() => controller.abort(), PROVIDER_TIMEOUT_MS);

  try {
    const messages: OpenAI.ChatCompletionMessageParam[] = [];
    if (systemPrompt) {
      messages.push({ role: "system" as const, content: systemPrompt });
    }
    messages.push({ role: "user" as const, content: prompt });

    const resp = await client.chat.completions.create({
      model: selectedModel,
      max_tokens: maxTokens,
      temperature,
      messages,
    }, { signal: controller.signal });

    clearTimeout(timer);

    const text = resp.choices[0]?.message?.content || "";
    return {
      provider: "openai",
      model: selectedModel,
      text,
      tokens: {
        input: resp.usage?.prompt_tokens || 0,
        output: resp.usage?.completion_tokens || 0,
      },
    };
  } catch (err) {
    clearTimeout(timer);
    throw err;
  }
}

async function callGoogle(
  apiKey: string,
  prompt: string,
  model: string | undefined,
  maxTokens: number,
  temperature: number,
  systemPrompt?: string,
): Promise<ProviderResult> {
  const client = new GoogleGenAI({ apiKey });
  const selectedModel = model || "gemini-2.5-flash";

  // Thinking models (Pro/2.5+) use thinking tokens that count against maxOutputTokens.
  // Enforce a minimum of 2048 so thinking doesn't consume the entire budget.
  const isThinkingModel = /pro|2\.5|3-/.test(selectedModel);
  const effectiveMaxTokens = isThinkingModel ? Math.max(maxTokens, 2048) : maxTokens;

  const response = await client.models.generateContent({
    model: selectedModel,
    contents: prompt,
    config: {
      maxOutputTokens: effectiveMaxTokens,
      temperature,
      ...(systemPrompt ? { systemInstruction: systemPrompt } : {}),
    },
  });

  const text = response.text ?? "";
  const usage = response.usageMetadata;

  return {
    provider: "google",
    model: selectedModel,
    text,
    tokens: {
      input: usage?.promptTokenCount || 0,
      output: usage?.candidatesTokenCount || 0,
    },
  };
}

async function callBedrock(
  accessKeyId: string,
  secretAccessKey: string,
  region: string,
  prompt: string,
  model: string | undefined,
  maxTokens: number,
  temperature: number,
  systemPrompt?: string,
): Promise<ProviderResult> {
  const client = new BedrockRuntimeClient({
    region: region || "us-east-1",
    credentials: { accessKeyId, secretAccessKey },
  });
  const selectedModel = model || "us.anthropic.claude-3-5-haiku-20241022-v1:0";

  const body = JSON.stringify({
    anthropic_version: "bedrock-2023-05-31",
    max_tokens: maxTokens,
    temperature,
    ...(systemPrompt ? { system: systemPrompt } : {}),
    messages: [{ role: "user", content: prompt }],
  });

  const command = new InvokeModelCommand({
    modelId: selectedModel,
    contentType: "application/json",
    accept: "application/json",
    body: new TextEncoder().encode(body),
  });

  // AbortController for timeout
  const controller = new AbortController();
  const timer = setTimeout(() => controller.abort(), PROVIDER_TIMEOUT_MS);

  try {
    const resp = await client.send(command, { abortSignal: controller.signal });
    clearTimeout(timer);

    const respBody = JSON.parse(new TextDecoder().decode(resp.body));
    const text = respBody.content?.[0]?.text || "";

    return {
      provider: "bedrock",
      model: selectedModel,
      text,
      tokens: {
        input: respBody.usage?.input_tokens || 0,
        output: respBody.usage?.output_tokens || 0,
      },
    };
  } catch (err) {
    clearTimeout(timer);
    throw err;
  }
}

async function callOllama(
  prompt: string,
  model: string | undefined,
  systemPrompt?: string,
): Promise<ProviderResult> {
  const selectedModel = model || "llama3.2";

  const controller = new AbortController();
  const timer = setTimeout(() => controller.abort(), PROVIDER_TIMEOUT_MS);

  try {
    const resp = await fetch("http://localhost:11434/api/generate", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        model: selectedModel,
        prompt,
        ...(systemPrompt ? { system: systemPrompt } : {}),
        stream: false,
      }),
      signal: controller.signal,
    });

    clearTimeout(timer);

    if (!resp.ok) {
      const errText = await resp.text();
      throw new Error(`Ollama HTTP ${resp.status}: ${errText.substring(0, 200)}`);
    }

    const data = await resp.json() as {
      response: string;
      prompt_eval_count?: number;
      eval_count?: number;
    };

    return {
      provider: "ollama",
      model: selectedModel,
      text: data.response || "",
      tokens: {
        input: data.prompt_eval_count || 0,
        output: data.eval_count || 0,
      },
    };
  } catch (err) {
    clearTimeout(timer);
    throw err;
  }
}

// ============================================================================
// FAILOVER LOGIC
// ============================================================================

const FAILOVER_ORDER: Provider[] = ["anthropic", "openai", "google", "bedrock", "ollama"];

async function completeWithFailover(
  creds: ProviderCredentials,
  prompt: string,
  model: string | undefined,
  provider: Provider | undefined,
  maxTokens: number,
  temperature: number,
  systemPrompt?: string,
): Promise<{ success: boolean; data?: ProviderResult; error?: string }> {
  const providers = provider ? [provider] : FAILOVER_ORDER;
  const errors: string[] = [];

  for (const p of providers) {
    try {
      switch (p) {
        case "anthropic": {
          if (!creds.anthropic) {
            errors.push("anthropic: no API key configured");
            continue;
          }
          const result = await callAnthropic(creds.anthropic.apiKey, prompt, model, maxTokens, temperature, systemPrompt);
          return { success: true, data: result };
        }
        case "openai": {
          if (!creds.openai) {
            errors.push("openai: no API key configured");
            continue;
          }
          const result = await callOpenAI(creds.openai.apiKey, prompt, model, maxTokens, temperature, systemPrompt);
          return { success: true, data: result };
        }
        case "google": {
          if (!creds.google) {
            errors.push("google: no API key configured");
            continue;
          }
          const result = await callGoogle(creds.google.apiKey, prompt, model, maxTokens, temperature, systemPrompt);
          return { success: true, data: result };
        }
        case "bedrock": {
          if (!creds.bedrock) {
            errors.push("bedrock: no AWS credentials configured");
            continue;
          }
          const result = await callBedrock(
            creds.bedrock.accessKeyId,
            creds.bedrock.secretAccessKey,
            creds.bedrock.region,
            prompt, model, maxTokens, temperature, systemPrompt,
          );
          return { success: true, data: result };
        }
        case "ollama": {
          if (!creds.ollama?.available) {
            errors.push("ollama: local server not available");
            continue;
          }
          const result = await callOllama(prompt, model, systemPrompt);
          return { success: true, data: result };
        }
        default:
          errors.push(`${p}: unknown provider`);
      }
    } catch (err) {
      const msg = err instanceof Error ? err.message : String(err);
      errors.push(`${p}: ${msg}`);
      console.error(`[llm-failover] ${p} failed:`, msg);
    }
  }

  return {
    success: false,
    error: `All providers failed. ${errors.join(" | ")}`,
  };
}

// ============================================================================
// HANDLER
// ============================================================================

export async function handleLLM(
  sb: SupabaseClient,
  args: Record<string, unknown>,
  storeId?: string,
): Promise<{ success: boolean; data?: unknown; error?: string }> {
  const action = args.action as string;

  if (!storeId) {
    return { success: false, error: "store_id is required for LLM provider operations" };
  }

  switch (action) {
    // ================================================================
    // COMPLETE — send a completion request with failover
    // ================================================================
    case "complete": {
      const prompt = args.prompt as string;
      if (!prompt) {
        return { success: false, error: "prompt is required for complete action" };
      }

      const model = args.model as string | undefined;
      const provider = args.provider as Provider | undefined;
      const maxTokens = (args.max_tokens as number) || 1024;
      const temperature = (args.temperature as number) ?? 0.7;
      const systemPrompt = args.system as string | undefined;

      if (provider && !FAILOVER_ORDER.includes(provider)) {
        return { success: false, error: `Invalid provider: ${provider}. Must be one of: ${FAILOVER_ORDER.join(", ")}` };
      }

      const creds = await getCredentials(sb, storeId);
      return completeWithFailover(creds, prompt, model, provider, maxTokens, temperature, systemPrompt);
    }

    // ================================================================
    // LIST_MODELS — list available models across configured providers
    // ================================================================
    case "list_models": {
      const creds = await getCredentials(sb, storeId);
      const models: Array<{ provider: string; model_id: string; name: string; context_window: number }> = [];

      if (creds.anthropic) {
        models.push(...ANTHROPIC_MODELS.map((m) => ({ provider: "anthropic", ...m })));
      }
      if (creds.openai) {
        models.push(...OPENAI_MODELS.map((m) => ({ provider: "openai", ...m })));
      }
      if (creds.google) {
        models.push(...GOOGLE_MODELS.map((m) => ({ provider: "google", ...m })));
      }
      if (creds.bedrock) {
        models.push(...BEDROCK_MODELS.map((m) => ({ provider: "bedrock", ...m })));
      }
      if (creds.ollama?.available) {
        // Try to get actual installed models from Ollama
        try {
          const controller = new AbortController();
          const timer = setTimeout(() => controller.abort(), 3000);
          const resp = await fetch("http://localhost:11434/api/tags", { signal: controller.signal });
          clearTimeout(timer);
          if (resp.ok) {
            const data = await resp.json() as { models?: Array<{ name: string; size?: number }> };
            if (data.models?.length) {
              models.push(...data.models.map((m) => ({
                provider: "ollama",
                model_id: m.name,
                name: `${m.name} (local)`,
                context_window: 0, // Ollama doesn't report this via API
              })));
            } else {
              models.push(...OLLAMA_MODELS.map((m) => ({ provider: "ollama", ...m })));
            }
          }
        } catch {
          models.push(...OLLAMA_MODELS.map((m) => ({ provider: "ollama", ...m })));
        }
      }

      return { success: true, data: { models, total: models.length } };
    }

    // ================================================================
    // FAILOVER_STATUS — show provider configuration and health
    // ================================================================
    case "failover_status": {
      const creds = await getCredentials(sb, storeId);

      const providers = [
        {
          provider: "anthropic",
          configured: !!creds.anthropic,
          source: creds.anthropic ? "process.env.ANTHROPIC_API_KEY" : null,
          priority: 1,
        },
        {
          provider: "openai",
          configured: !!creds.openai,
          source: creds.openai ? "user_tool_secrets (encrypted)" : null,
          priority: 2,
        },
        {
          provider: "google",
          configured: !!creds.google,
          source: creds.google ? "user_tool_secrets (encrypted)" : null,
          priority: 3,
        },
        {
          provider: "bedrock",
          configured: !!creds.bedrock,
          source: creds.bedrock ? "user_tool_secrets (encrypted)" : null,
          region: creds.bedrock?.region || null,
          priority: 4,
        },
        {
          provider: "ollama",
          configured: !!creds.ollama?.available,
          source: creds.ollama?.available ? "localhost:11434" : null,
          priority: 5,
        },
      ];

      const configuredCount = providers.filter((p) => p.configured).length;

      return {
        success: true,
        data: {
          providers,
          failover_order: FAILOVER_ORDER,
          configured_count: configuredCount,
          total_providers: providers.length,
          note: configuredCount === 0
            ? "No LLM providers configured. Add API keys to user_tool_secrets."
            : `${configuredCount} provider(s) ready. Failover will try in order: ${FAILOVER_ORDER.filter((p) => providers.find((pr) => pr.provider === p)?.configured).join(" → ")}`,
        },
      };
    }

    default:
      return { success: false, error: `Unknown LLM action: ${action}. Valid actions: complete, list_models, failover_status` };
  }
}