ci/src/backends/ollama-cloud.ts

import { OllamaBaseBackend, OllamaMessage, OllamaChatResponse } from "./ollama-base.js";
import { OllamaCloudConfig, emptyBackendResult } from "./types.js";
import { ToolRegistry } from "./tool-registry.js";

const MAX_RETRIES = 3;
const BASE_BACKOFF_MS = 1000;

export class OllamaCloudBackend extends OllamaBaseBackend {
  readonly name = "ollama-cloud";

  private cloudConfig: OllamaCloudConfig;
  private apiKey: string | null;

  constructor(config?: OllamaCloudConfig) {
    super(config);
    this.cloudConfig = config || {
      base_url: "",
      api_key_env: "OLLAMA_CLOUD_API_KEY",
      model_profile: "quality",
      timeout_ms: 60000,
    };
    this.apiKey = this.resolveApiKey();
  }

  async isAvailable(): Promise<boolean> {
    if (!this.cloudConfig.base_url) return false;
    if (!this.apiKey) return false;

    try {
      const response = await fetch(`${this.cloudConfig.base_url}/v1/models`, {
        headers: this.getAuthHeaders(),
        signal: AbortSignal.timeout(10000),
      });
      return response.ok;
    } catch {
      return false;
    }
  }

  protected resolveModel(): string {
    if (this.cloudConfig.model) return this.cloudConfig.model;
    return "llama3.1:70b";
  }

  protected async callModel(
    messages: OllamaMessage[],
    model: string,
    toolRegistry: ToolRegistry
  ): Promise<OllamaChatResponse> {
    if (!this.apiKey) {
      throw new Error(`API key not found. Set ${this.cloudConfig.api_key_env} environment variable.`);
    }

    const url = `${this.cloudConfig.base_url}/v1/chat/completions`;

    const body: Record<string, unknown> = {
      model,
      messages: messages.map((m) => {
        const msg: Record<string, unknown> = { role: m.role, content: m.content };
        if (m.name) msg.name = m.name;
        if (m.tool_calls) msg.tool_calls = m.tool_calls;
        return msg;
      }),
      tools: toolRegistry.getOpenAIToolSchema(),
      stream: false,
    };

    return this.callWithRetry(url, body);
  }

  private async callWithRetry(
    url: string,
    body: Record<string, unknown>,
    attempt: number = 0
  ): Promise<OllamaChatResponse> {
    const timeout = this.cloudConfig.timeout_ms || 60000;

    try {
      const response = await fetch(url, {
        method: "POST",
        headers: {
          ...this.getAuthHeaders(),
          "Content-Type": "application/json",
        },
        body: JSON.stringify(body),
        signal: AbortSignal.timeout(timeout),
      });

      if (response.status === 429 && attempt < MAX_RETRIES) {
        const retryAfter = response.headers.get("Retry-After");
        const delay = retryAfter
          ? parseInt(retryAfter) * 1000
          : BASE_BACKOFF_MS * Math.pow(2, attempt);

        await this.sleep(delay);
        return this.callWithRetry(url, body, attempt + 1);
      }

      if (response.status === 401 || response.status === 403) {
        throw new Error(`Authentication failed. Check ${this.cloudConfig.api_key_env} environment variable.`);
      }

      if (response.status === 402) {
        throw new Error("Quota exceeded. Check your Ollama Cloud billing status.");
      }

      if (!response.ok) {
        const errorText = await response.text().catch(() => "unknown error");
        throw new Error(`Ollama Cloud API error (${response.status}): ${errorText}`);
      }

      return (await response.json()) as OllamaChatResponse;
    } catch (err) {
      if (err instanceof TypeError && err.message.includes("fetch")) {
        if (attempt < MAX_RETRIES) {
          await this.sleep(BASE_BACKOFF_MS * Math.pow(2, attempt));
          return this.callWithRetry(url, body, attempt + 1);
        }
      }
      throw err;
    }
  }

  private getAuthHeaders(): Record<string, string> {
    const headers: Record<string, string> = {};
    if (this.apiKey) {
      headers["Authorization"] = `Bearer ${this.apiKey}`;
    }
    return headers;
  }

  private resolveApiKey(): string | null {
    return process.env[this.cloudConfig.api_key_env] || null;
  }

  private sleep(ms: number): Promise<void> {
    return new Promise((resolve) => setTimeout(resolve, ms));
  }
}