stack/apps/api/src/ollama/ollama.service.ts

import { Injectable, Inject, HttpException, HttpStatus } from "@nestjs/common";
import type {
  GenerateOptionsDto,
  GenerateResponseDto,
  ChatMessage,
  ChatOptionsDto,
  ChatResponseDto,
  EmbedResponseDto,
  ListModelsResponseDto,
  HealthCheckResponseDto,
} from "./dto";

/**
 * Configuration for Ollama service
 */
export interface OllamaConfig {
  mode: "local" | "remote";
  endpoint: string;
  model: string;
  timeout: number;
}

/**
 * Service for interacting with Ollama API
 * Supports both local and remote Ollama instances
 */
@Injectable()
export class OllamaService {
  constructor(
    @Inject("OLLAMA_CONFIG")
    private readonly config: OllamaConfig
  ) {}

  /**
   * Generate text from a prompt
   * @param prompt - The text prompt to generate from
   * @param options - Generation options (temperature, max_tokens, etc.)
   * @param model - Optional model override (defaults to config model)
   * @returns Generated text response
   */
  async generate(
    prompt: string,
    options?: GenerateOptionsDto,
    model?: string
  ): Promise<GenerateResponseDto> {
    const url = `${this.config.endpoint}/api/generate`;

    const requestBody = {
      model: model || this.config.model,
      prompt,
      stream: false,
      ...(options && {
        options: this.mapGenerateOptions(options),
      }),
    };

    try {
      const controller = new AbortController();
      const timeoutId = setTimeout(() => controller.abort(), this.config.timeout);

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

      clearTimeout(timeoutId);

      if (!response.ok) {
        throw new HttpException(
          `Ollama API error: ${response.statusText}`,
          response.status
        );
      }

      const data = await response.json();
      return data as GenerateResponseDto;
    } catch (error: unknown) {
      if (error instanceof HttpException) {
        throw error;
      }

      const errorMessage =
        error instanceof Error ? error.message : "Unknown error";

      throw new HttpException(
        `Failed to connect to Ollama: ${errorMessage}`,
        HttpStatus.SERVICE_UNAVAILABLE
      );
    }
  }

  /**
   * Complete a chat conversation
   * @param messages - Array of chat messages
   * @param options - Chat options (temperature, max_tokens, etc.)
   * @param model - Optional model override (defaults to config model)
   * @returns Chat completion response
   */
  async chat(
    messages: ChatMessage[],
    options?: ChatOptionsDto,
    model?: string
  ): Promise<ChatResponseDto> {
    const url = `${this.config.endpoint}/api/chat`;

    const requestBody = {
      model: model || this.config.model,
      messages,
      stream: false,
      ...(options && {
        options: this.mapChatOptions(options),
      }),
    };

    try {
      const controller = new AbortController();
      const timeoutId = setTimeout(() => controller.abort(), this.config.timeout);

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

      clearTimeout(timeoutId);

      if (!response.ok) {
        throw new HttpException(
          `Ollama API error: ${response.statusText}`,
          response.status
        );
      }

      const data = await response.json();
      return data as ChatResponseDto;
    } catch (error: unknown) {
      if (error instanceof HttpException) {
        throw error;
      }

      const errorMessage =
        error instanceof Error ? error.message : "Unknown error";

      throw new HttpException(
        `Failed to connect to Ollama: ${errorMessage}`,
        HttpStatus.SERVICE_UNAVAILABLE
      );
    }
  }

  /**
   * Generate embeddings for text
   * @param text - The text to generate embeddings for
   * @param model - Optional model override (defaults to config model)
   * @returns Embedding vector
   */
  async embed(text: string, model?: string): Promise<EmbedResponseDto> {
    const url = `${this.config.endpoint}/api/embeddings`;

    const requestBody = {
      model: model || this.config.model,
      prompt: text,
    };

    try {
      const controller = new AbortController();
      const timeoutId = setTimeout(() => controller.abort(), this.config.timeout);

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

      clearTimeout(timeoutId);

      if (!response.ok) {
        throw new HttpException(
          `Ollama API error: ${response.statusText}`,
          response.status
        );
      }

      const data = await response.json();
      return data as EmbedResponseDto;
    } catch (error: unknown) {
      if (error instanceof HttpException) {
        throw error;
      }

      const errorMessage =
        error instanceof Error ? error.message : "Unknown error";

      throw new HttpException(
        `Failed to connect to Ollama: ${errorMessage}`,
        HttpStatus.SERVICE_UNAVAILABLE
      );
    }
  }

  /**
   * List available models
   * @returns List of available Ollama models
   */
  async listModels(): Promise<ListModelsResponseDto> {
    const url = `${this.config.endpoint}/api/tags`;

    try {
      const controller = new AbortController();
      const timeoutId = setTimeout(() => controller.abort(), this.config.timeout);

      const response = await fetch(url, {
        method: "GET",
        signal: controller.signal,
      });

      clearTimeout(timeoutId);

      if (!response.ok) {
        throw new HttpException(
          `Ollama API error: ${response.statusText}`,
          response.status
        );
      }

      const data = await response.json();
      return data as ListModelsResponseDto;
    } catch (error: unknown) {
      if (error instanceof HttpException) {
        throw error;
      }

      const errorMessage =
        error instanceof Error ? error.message : "Unknown error";

      throw new HttpException(
        `Failed to connect to Ollama: ${errorMessage}`,
        HttpStatus.SERVICE_UNAVAILABLE
      );
    }
  }

  /**
   * Check health and connectivity of Ollama instance
   * @returns Health check status
   */
  async healthCheck(): Promise<HealthCheckResponseDto> {
    try {
      const controller = new AbortController();
      const timeoutId = setTimeout(() => controller.abort(), 5000); // 5s timeout for health check

      const response = await fetch(`${this.config.endpoint}/api/tags`, {
        method: "GET",
        signal: controller.signal,
      });

      clearTimeout(timeoutId);

      if (response.ok) {
        return {
          status: "healthy",
          mode: this.config.mode,
          endpoint: this.config.endpoint,
          available: true,
        };
      } else {
        return {
          status: "unhealthy",
          mode: this.config.mode,
          endpoint: this.config.endpoint,
          available: false,
          error: `HTTP ${response.status}: ${response.statusText}`,
        };
      }
    } catch (error: unknown) {
      const errorMessage =
        error instanceof Error ? error.message : "Unknown error";

      return {
        status: "unhealthy",
        mode: this.config.mode,
        endpoint: this.config.endpoint,
        available: false,
        error: errorMessage,
      };
    }
  }

  /**
   * Map GenerateOptionsDto to Ollama API options format
   */
  private mapGenerateOptions(
    options: GenerateOptionsDto
  ): Record<string, unknown> {
    const mapped: Record<string, unknown> = {};

    if (options.temperature !== undefined) {
      mapped.temperature = options.temperature;
    }
    if (options.top_p !== undefined) {
      mapped.top_p = options.top_p;
    }
    if (options.max_tokens !== undefined) {
      mapped.num_predict = options.max_tokens;
    }
    if (options.stop !== undefined) {
      mapped.stop = options.stop;
    }

    return mapped;
  }

  /**
   * Map ChatOptionsDto to Ollama API options format
   */
  private mapChatOptions(options: ChatOptionsDto): Record<string, unknown> {
    const mapped: Record<string, unknown> = {};

    if (options.temperature !== undefined) {
      mapped.temperature = options.temperature;
    }
    if (options.top_p !== undefined) {
      mapped.top_p = options.top_p;
    }
    if (options.max_tokens !== undefined) {
      mapped.num_predict = options.max_tokens;
    }
    if (options.stop !== undefined) {
      mapped.stop = options.stop;
    }

    return mapped;
  }
}