backend/src/services/ollamaClient.js

import { config } from "../config/index.js";
import { createTtlCache } from "./simpleCache.js";


function timeoutSignal(externalSignal, timeoutMs) {
  const controller = new AbortController();
  const timer = setTimeout(() => controller.abort(new Error(`ollama_timeout:${timeoutMs}ms`)), timeoutMs);
  const onExternalAbort = () => controller.abort(externalSignal.reason);
  if (externalSignal) {
    if (externalSignal.aborted) controller.abort(externalSignal.reason);
    else externalSignal.addEventListener("abort", onExternalAbort, { once: true });
  }
  return {
    signal: controller.signal,
    cleanup() {
      clearTimeout(timer);
      externalSignal?.removeEventListener("abort", onExternalAbort);
    }
  };
}

async function ollamaFetch(path, body, { signal } = {}) {
  const { signal: combinedSignal, cleanup } = timeoutSignal(signal, config.ollama.timeoutMs);
  let res;
  try {
    res = await fetch(`${config.ollama.url}${path}`, {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify(body),
      signal: combinedSignal
    });
  } finally {
    cleanup();
  }

  if (!res.ok) {
    const text = await res.text().catch(() => "");
    let msg = `ollama_error:${res.status}:${text || res.statusText}`;
    if (
      res.status === 404 &&
      typeof body?.model === "string" &&
      typeof text === "string" &&
      text.toLowerCase().includes("model") &&
      text.toLowerCase().includes("not found")
    ) {
      let installed = [];
      try {
        const tagsRes = await fetch(`${config.ollama.url}/api/tags`);
        const tags = await tagsRes.json().catch(() => ({}));
        installed = Array.isArray(tags?.models) ? tags.models.map((m) => m?.name).filter(Boolean) : [];
      } catch {}
      const installedList = installed.length ? installed.slice(0, 10).join(",") : "none";
      msg = `ollama_model_not_found:${body.model}:installed=${installedList}:hint=ollama pull ${body.model}`;
    }
    const err = new Error(msg);
    err.statusCode = 502;
    throw err;
  }

  return res.json();
}

function embedPrefixFor(role) {
  if (role === "query") return config.ollama.embedQueryPrefix;
  if (role === "passage") return config.ollama.embedPassagePrefix;
  return "";
}


const EMBED_CACHE_TTL_MS = 15 * 60 * 1000;
const embedCache = createTtlCache(EMBED_CACHE_TTL_MS);

async function embedSingle(text, { role, model } = {}) {
  const effectiveModel = model ?? config.ollama.embeddingsModel;
  const cacheKey = `${effectiveModel}:${role ?? ""}:${text}`;
  const cached = embedCache.get(cacheKey);
  if (cached) return cached;

  const prefix = embedPrefixFor(role);
  const data = await ollamaFetch("/api/embeddings", {
    model: effectiveModel,
    prompt: `${prefix}${text}`
  });
  embedCache.set(cacheKey, data.embedding);
  return data.embedding;
}

export async function embedTexts(texts, { role, model } = {}) {
  const BATCH = 10;
  const results = [];
  for (let i = 0; i < texts.length; i += BATCH) {
    const batch = texts.slice(i, i + BATCH);
    const embeddings = await Promise.all(batch.map((text) => embedSingle(text, { role, model })));
    results.push(...embeddings);
  }
  return results;
}

export async function chatCompletion({ messages, options, format, model, signal }) {
  const data = await ollamaFetch(
    "/api/chat",
    {
      model: model || config.ollama.chatModel,
      messages,
      stream: false,
      ...(format && { format }),
      ...(options && Object.keys(options).length > 0 && { options })
    },
    { signal }
  );

  return {
    content: data?.message?.content ?? ""
  };
}

export async function chatCompletionStream({ messages, onChunk, signal, options }) {
  const idleTimeoutMs = config.ollama.streamIdleTimeoutMs;
  const idleController = new AbortController();
  let idleTimer = setTimeout(
    () => idleController.abort(new Error(`ollama_stream_idle_timeout:${idleTimeoutMs}ms`)),
    idleTimeoutMs
  );
  const resetIdleTimer = () => {
    clearTimeout(idleTimer);
    idleTimer = setTimeout(
      () => idleController.abort(new Error(`ollama_stream_idle_timeout:${idleTimeoutMs}ms`)),
      idleTimeoutMs
    );
  };
  const onExternalAbort = () => idleController.abort(signal.reason);
  if (signal) {
    if (signal.aborted) idleController.abort(signal.reason);
    else signal.addEventListener("abort", onExternalAbort, { once: true });
  }

  try {
    const res = await fetch(`${config.ollama.url}/api/chat`, {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({
        model: config.ollama.chatModel,
        messages,
        stream: true,
        ...(options && Object.keys(options).length > 0 && { options })
      }),
      signal: idleController.signal
    });

    if (!res.ok) {
      const text = await res.text().catch(() => "");
      const err = new Error(`ollama_error:${res.status}:${text || res.statusText}`);
      err.statusCode = 502;
      throw err;
    }

    const reader = res.body.getReader();
    const decoder = new TextDecoder();
    let fullContent = "";
    let buffer = "";

    const processLine = (line) => {
      if (!line.trim()) return;
      try {
        const parsed = JSON.parse(line);
        const delta = parsed?.message?.content ?? "";
        if (delta) {
          fullContent += delta;
          onChunk(delta);
        }
      } catch {}
    };

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      resetIdleTimer();

      buffer += decoder.decode(value, { stream: true });
      const lines = buffer.split("\n");
      buffer = lines.pop() ?? "";
      for (const line of lines) processLine(line);
    }

    buffer += decoder.decode();
    processLine(buffer);

    return { content: fullContent };
  } finally {
    clearTimeout(idleTimer);
    signal?.removeEventListener("abort", onExternalAbort);
  }
}

export async function visionExtractFromImage({ imageBase64, prompt }) {
  const data = await ollamaFetch("/api/chat", {
    model: config.ollama.visionModel,
    messages: [
      {
        role: "user",
        content: prompt,
        images: [imageBase64]
      }
    ],
    stream: false
  });

  return {
    content: data?.message?.content ?? ""
  };
}