MCP-Ollama/src/client.ts

/**
 * Ollama HTTP API client.
 *
 * Wraps the Ollama REST API (default http://127.0.0.1:11434).
 * Docs: https://github.com/ollama/ollama/blob/main/docs/api.md
 */

export interface OllamaConfig {
  host: string;
}

export interface GenerateRequest {
  model: string;
  prompt: string;
  system?: string;
  temperature?: number;
  max_tokens?: number;
  format?: "json";
}

export interface ChatMessage {
  role: "system" | "user" | "assistant";
  content: string;
}

export interface ChatRequest {
  model: string;
  messages: ChatMessage[];
  temperature?: number;
  max_tokens?: number;
  format?: "json";
}

export interface ModelInfo {
  name: string;
  model: string;
  size: number;
  details: {
    parameter_size: string;
    quantization_level: string;
    family: string;
  };
}

export interface PullProgress {
  status: string;
  digest?: string;
  total?: number;
  completed?: number;
}

export class OllamaClient {
  private host: string;

  constructor(config: OllamaConfig) {
    this.host = config.host.replace(/\/+$/, "");
  }

  private async request(path: string, options?: RequestInit): Promise<Response> {
    const url = `${this.host}${path}`;
    const res = await fetch(url, options);
    if (!res.ok) {
      const body = await res.text().catch(() => "");
      throw new Error(`Ollama API ${res.status}: ${body || res.statusText}`);
    }
    return res;
  }

  /** Generate a completion (non-streaming). */
  async generate(req: GenerateRequest): Promise<string> {
    const body: Record<string, unknown> = {
      model: req.model,
      prompt: req.prompt,
      stream: false,
    };
    if (req.system) body.system = req.system;
    if (req.temperature !== undefined) body.temperature = req.temperature;
    if (req.max_tokens !== undefined) body.options = { num_predict: req.max_tokens };
    if (req.format) body.format = req.format;

    const res = await this.request("/api/generate", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify(body),
    });
    const data = await res.json() as { response: string };
    return data.response;
  }

  /** Multi-turn chat completion (non-streaming). */
  async chat(req: ChatRequest): Promise<string> {
    const body: Record<string, unknown> = {
      model: req.model,
      messages: req.messages,
      stream: false,
    };
    if (req.temperature !== undefined) body.temperature = req.temperature;
    if (req.max_tokens !== undefined) body.options = { num_predict: req.max_tokens };
    if (req.format) body.format = req.format;

    const res = await this.request("/api/chat", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify(body),
    });
    const data = await res.json() as { message: { content: string } };
    return data.message.content;
  }

  /** List locally available models. */
  async listModels(): Promise<ModelInfo[]> {
    const res = await this.request("/api/tags");
    const data = await res.json() as { models: ModelInfo[] };
    return data.models;
  }

  /** Get detailed info about a model. */
  async showModel(name: string): Promise<Record<string, unknown>> {
    const res = await this.request("/api/show", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ name }),
    });
    return await res.json() as Record<string, unknown>;
  }

  /** Pull a model (blocking — waits for completion). */
  async pullModel(name: string): Promise<string> {
    const res = await this.request("/api/pull", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ name, stream: false }),
    });
    const data = await res.json() as { status: string };
    return data.status;
  }

  /** Delete a model. */
  async deleteModel(name: string): Promise<void> {
    await this.request("/api/delete", {
      method: "DELETE",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ name }),
    });
  }

  /** Check if Ollama is reachable. */
  async health(): Promise<boolean> {
    try {
      await this.request("/");
      return true;
    } catch {
      return false;
    }
  }

  /** List running models. */
  async listRunning(): Promise<unknown[]> {
    const res = await this.request("/api/ps");
    const data = await res.json() as { models: unknown[] };
    return data.models ?? [];
  }

  /** Generate embeddings. */
  async embed(model: string, input: string | string[]): Promise<number[][]> {
    const res = await this.request("/api/embed", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ model, input }),
    });
    const data = await res.json() as { embeddings: number[][] };
    return data.embeddings;
  }
}