add testing suite

2026-02-21 22:19:21 -05:00
parent 07128c3529
commit 83eee20b2c
25 changed files with 5317 additions and 1628 deletions
--- a/src/ollamaClient.js
+++ b/src/ollamaClient.js
@@ -0,0 +1,127 @@
+import { cleanText } from "./textUtils.js";
+
+const OLLAMA_API_URL = process.env.OLLAMA_API_URL;
+export let OLLAMA_MODEL = process.env.OLLAMA_MODEL || "gemma3:latest";
+
+export async function initializeModel() {
+  if (process.env.OLLAMA_MODEL) return;
+  try {
+    const apiUrl = process.env.OLLAMA_API_URL;
+    const isOpenWebUI = apiUrl?.includes("/api/chat/completions");
+    const baseUrl = apiUrl?.replace(/\/api\/.*$/, "");
+    const url = isOpenWebUI ? `${baseUrl}/api/v1/models` : `${baseUrl}/api/tags`;
+    const headers = isOpenWebUI && process.env.OLLAMA_API_KEY
+      ? { "Authorization": `Bearer ${process.env.OLLAMA_API_KEY}` }
+      : {};
+    const res = await fetch(url, { headers });
+    if (res.ok) {
+      const data = await res.json();
+      const model = isOpenWebUI 
+        ? data.data?.[0]?.id || data.data?.[0]?.name
+        : data.models?.[0]?.name;
+      if (model) {
+        OLLAMA_MODEL = model;
+        console.log(`Using default model: ${model}`);
+      }
+    }
+  } catch {
+    console.warn(`Could not fetch default model, using: ${OLLAMA_MODEL}`);
+  }
+}
+
+export { cleanText };
+
+export function inferApiType(url) {
+  if (!url) return "ollama-generate";
+  if (url.includes("/api/chat/completions")) return "open-webui";
+  if (url.includes("/api/chat")) return "ollama-chat";
+  return "ollama-generate";
+}
+
+async function sleep(ms) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+async function callOllamaBase(prompt, model, retries, stepName, apiType) {
+  const isUsingOpenWebUI = apiType === "open-webui";
+  const isUsingOllamaChat = apiType === "ollama-chat";
+
+  for (let attempt = 1; attempt <= retries; attempt++) {
+    try {
+      const promptCharCount = prompt.length;
+      const promptWordCount = prompt.split(/\s+/).length;
+
+      console.log(
+        `\n[${stepName}] Sending prompt (attempt ${attempt}/${retries})`,
+      );
+      console.log(
+        `Prompt: ${promptCharCount} chars, ~${promptWordCount} words`,
+      );
+
+      const headers = { "Content-Type": "application/json" };
+      if (isUsingOpenWebUI && process.env.OLLAMA_API_KEY) {
+        headers["Authorization"] = `Bearer ${process.env.OLLAMA_API_KEY}`;
+      }
+
+      const body = isUsingOpenWebUI || isUsingOllamaChat
+        ? { model, messages: [{ role: "user", content: prompt }] }
+        : { model, prompt, stream: false };
+
+      const response = await fetch(OLLAMA_API_URL, {
+        method: "POST",
+        headers,
+        body: JSON.stringify(body),
+      });
+
+      if (!response.ok) {
+        const errorData = await response.text().catch(() => null);
+        const errorDetails = errorData ? `: ${errorData}` : "";
+        throw new Error(
+          `Ollama request failed: ${response.status} ${response.statusText}${errorDetails}`,
+        );
+      }
+
+      const data = await response.json();
+      const rawText = isUsingOpenWebUI
+        ? data.choices?.[0]?.message?.content
+        : isUsingOllamaChat
+          ? data.message?.content
+          : data.response;
+
+      if (!rawText) throw new Error("No response from Ollama");
+
+      const cleaned = cleanText(rawText);
+      console.log(
+        `[${stepName}] Received: ${rawText.length} chars, ~${rawText.split(/\s+/).length} words`,
+      );
+
+      return cleaned;
+    } catch (err) {
+      console.warn(`[${stepName}] Attempt ${attempt} failed: ${err.message}`);
+      if (attempt === retries) throw err;
+      const delay = 1000 * Math.pow(2, attempt) + Math.random() * 500;
+      console.log(`Retrying in ${Math.round(delay / 1000)}s...`);
+      await sleep(delay);
+    }
+  }
+}
+
+export async function callOllama(
+  prompt,
+  model = OLLAMA_MODEL,
+  retries = 5,
+  stepName = "unknown",
+) {
+  const apiType = inferApiType(OLLAMA_API_URL);
+  return callOllamaBase(prompt, model, retries, stepName, apiType);
+}
+
+export async function callOllamaExplicit(
+  prompt,
+  model = OLLAMA_MODEL,
+  retries = 5,
+  stepName = "unknown",
+  apiType = "ollama-generate",
+) {
+  return callOllamaBase(prompt, model, retries, stepName, apiType);
+}