OmniRoute/open-sse/executors/pollinations.ts

import { BaseExecutor } from "./base.ts";
import { PROVIDERS } from "../config/constants.ts";
import { DEFAULT_POOL_CONFIG } from "../services/sessionPool/types.ts";
import type { ExecuteInput } from "./base.ts";

/** Premium Pollinations models — upstream answers 401 UNAUTHORIZED without a key. */
const PREMIUM_MODELS = new Set([
  "claude",
  "claude-fast",
  "claude-large",
  "gemini",
  "gemini-fast",
  "midijourney",
  "midijourney-large",
]);

/** Build the actionable 401 error shown when a premium model is used without a key. */
function premiumModelRequiresKeyError(model: string): Error {
  const enhanced = new Error(
    `Pollinations model "${model}" requires an API key. ` +
      `Free keyless models: openai, openai-fast, openai-large, qwen-coder, mistral, deepseek, grok, gemini-flash-lite-3.1, perplexity-fast, perplexity-reasoning. ` +
      `Get a Pollinations API key at https://enter.pollinations.ai and add it in Settings → API Keys.`
  );
  (enhanced as any).status = 401;
  (enhanced as any).type = "authentication_error";
  return enhanced;
}

export class PollinationsExecutor extends BaseExecutor {
  constructor() {
    super("pollinations", PROVIDERS["pollinations"] || { format: "openai" });
    this.poolConfig = DEFAULT_POOL_CONFIG;
  }

  buildUrl(_model: string, _stream: boolean, urlIndex = 0, _credentials = null): string {
    const baseUrls = this.getBaseUrls();
    return baseUrls[urlIndex] || baseUrls[0] || "https://gen.pollinations.ai/v1/chat/completions";
  }

  buildHeaders(credentials: any, stream = true): Record<string, string> {
    const key = credentials?.apiKey || credentials?.accessToken;

    const headers: Record<string, string> = {
      "Content-Type": "application/json",
    };

    if (key) {
      headers.Authorization = `Bearer ${key}`;
    }

    if (stream) {
      headers["Accept"] = "text/event-stream";
    }

    return headers;
  }

  transformRequest(model: string, body: any, stream: boolean, _credentials: any): any {
    if (typeof body === "object" && body !== null) {
      body.model = model;
      body.stream = stream;
      // #3981: Pollinations treats jsonMode=true as "the model MUST return JSON"
      // and rejects (HTTP 400) any request whose messages don't mention "json".
      // Only enable it when the caller actually asked for JSON output.
      const responseFormatType = body.response_format?.type;
      if (responseFormatType === "json_object" || responseFormatType === "json_schema") {
        body.jsonMode = true;
      }
    }
    return body;
  }

  async execute(input: ExecuteInput) {
    const isAnonymous = !input.credentials?.apiKey && !input.credentials?.accessToken;

    if (!isAnonymous) {
      return super.execute(input);
    }

    // #9827 — premium models require a key upstream (verified: 401 UNAUTHORIZED).
    // Fail fast with guidance instead of dispatching an anonymous request whose
    // 401 would be recorded against the keyless connection's health and flip the
    // anonymous pool to "all accounts unavailable".
    const requestedModel = input.model || "";
    if (PREMIUM_MODELS.has(requestedModel)) {
      throw premiumModelRequiresKeyError(requestedModel);
    }

    const pool = this.getPool();

    // Use acquireBlocking for anonymous requests to wait for available session
    let session;
    try {
      session = pool ? await pool.acquireBlocking(10_000) : null;
    } catch {
      // Pool exhausted — fall through to direct request without fingerprint
      session = null;
    }

    if (session) {
      const fpHeaders = session.buildHeaders();
      input.upstreamExtraHeaders = {
        ...fpHeaders,
        ...input.upstreamExtraHeaders,
      };
    }

    try {
      const result = await super.execute(input);

      if (session && pool) {
        // execute() contracts for `Response | { response, ... }`; both arms carry the
        // status this pool bookkeeping needs.
        const status = (result instanceof Response ? result : result.response).status;
        if (status === 429) {
          pool.reportCooldown(session);
        } else if (status >= 500) {
          pool.reportDead(session);
        } else {
          pool.reportSuccess(session);
        }
      }

      return result;
    } catch (err: any) {
      if (session && pool) {
        pool.reportCooldown(session);
      }
      // Enhance 401 errors with actionable guidance
      if (err?.status === 401 || err?.statusCode === 401) {
        const model = input.model || "";
        if (PREMIUM_MODELS.has(model)) {
          throw premiumModelRequiresKeyError(model);
        }
      }
      throw err;
    } finally {
      session?.release();
    }
  }
}

export default PollinationsExecutor;