Newer
Older
TelosDB / src / llama-client.js
@楽曲作りまくりおじさん 楽曲作りまくりおじさん 8 days ago 5 KB refactor: modernize source code with improved structure and logging
import { CONFIG } from "./config.js";
import { Logger } from "./logger.js";

/**
 * llama.cpp API の基本 URL を取得
 * @returns {string} ベース URL
 */
function getBaseUrl() {
  return CONFIG.llama.baseUrl;
}

/**
 * 埋め込みモデル名を取得
 * @returns {string} モデル名
 * @throws {Error} モデルが設定されていない場合
 */
function getEmbeddingModel() {
  const model = CONFIG.llama.embeddingModel;
  if (!model) {
    throw new Error("LLAMA_CPP_EMBEDDING_MODEL environment variable is not set");
  }
  return model;
}

/**
 * 補完モデル名を取得
 * @returns {string} モデル名
 * @throws {Error} モデルが設定されていない場合
 */
function getCompletionModel() {
  const model = CONFIG.llama.completionModel;
  if (!model) {
    throw new Error("LLAMA_CPP_MODEL environment variable is not set");
  }
  return model;
}

/**
 * テキストから埋め込みベクトルを生成
 * @param {string} text - 埋め込み対象のテキスト
 * @param {number} timeout - リクエストタイムアウト(ミリ秒)
 * @returns {Promise<number[]>} 埋め込みベクトル
 * @throws {Error} API呼び出しまたはレスポンス解析に失敗した場合
 */
export async function llamaEmbedding(text, timeout = 30000) {
  if (!text || typeof text !== "string") {
    throw new Error("Text parameter must be a non-empty string");
  }

  const baseUrl = getBaseUrl();
  const model = getEmbeddingModel();

  try {
    Logger.debug("Calling llama.cpp embeddings API", {
      baseUrl,
      model,
      textLength: text.length,
    });

    const controller = new AbortController();
    const timeoutId = setTimeout(() => controller.abort(), timeout);

    const res = await fetch(`${baseUrl}/embeddings`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        model,
        input: text,
      }),
      signal: controller.signal,
    });

    clearTimeout(timeoutId);

    if (!res.ok) {
      const detail = await res.text();
      Logger.error("llama.cpp embeddings API error", {
        status: res.status,
        detail,
      });
      throw new Error(`llama.cpp embeddings error: ${res.status} ${detail}`);
    }

    const data = await res.json();

    // レスポンス形式に応じて埋め込みを抽出
    const embedding = data?.data?.[0]?.embedding ?? data?.embedding;

    if (!Array.isArray(embedding) || embedding.length === 0) {
      Logger.error("Invalid embedding response", { data });
      throw new Error("llama.cpp embeddings response is missing or invalid embedding array");
    }

    Logger.debug("Embedding generated successfully", {
      dimension: embedding.length,
    });

    return embedding;
  } catch (error) {
    if (error.name === "AbortError") {
      Logger.error("Embeddings request timeout", { timeout });
      throw new Error(`Embeddings request timed out after ${timeout}ms`);
    }
    throw error;
  }
}

/**
 * llama.cpp でテキスト補完を実行
 * @param {string} prompt - プロンプトテキスト
 * @param {Object} options - オプション設定
 * @param {number} options.n_predict - 生成するトークン数(デフォルト: 128)
 * @param {number} options.temperature - 温度パラメータ(デフォルト: 0.2)
 * @param {number} options.timeout - リクエストタイムアウト(ミリ秒)
 * @returns {Promise<string>} 生成されたテキスト
 * @throws {Error} API呼び出しまたはレスポンス解析に失敗した場合
 */
export async function llamaCompletion(prompt, options = {}) {
  if (!prompt || typeof prompt !== "string") {
    throw new Error("Prompt parameter must be a non-empty string");
  }

  const baseUrl = getBaseUrl();
  const model = getCompletionModel();
  const timeout = options.timeout ?? 60000;

  try {
    Logger.debug("Calling llama.cpp completion API", {
      baseUrl,
      model,
      promptLength: prompt.length,
      n_predict: options.n_predict ?? 128,
      temperature: options.temperature ?? 0.2,
    });

    const controller = new AbortController();
    const timeoutId = setTimeout(() => controller.abort(), timeout);

    const res = await fetch(`${baseUrl}/completion`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        model,
        prompt,
        n_predict: options.n_predict ?? 128,
        temperature: options.temperature ?? 0.2,
        stream: false,
      }),
      signal: controller.signal,
    });

    clearTimeout(timeoutId);

    if (!res.ok) {
      const detail = await res.text();
      Logger.error("llama.cpp completion API error", {
        status: res.status,
        detail,
      });
      throw new Error(`llama.cpp completion error: ${res.status} ${detail}`);
    }

    const data = await res.json();

    // レスポンス形式に応じてテキストを抽出
    const text = data?.content ?? data?.completion ?? data?.response;

    if (typeof text !== "string") {
      Logger.error("Invalid completion response", { data });
      throw new Error("llama.cpp completion response is missing or invalid text");
    }

    Logger.debug("Completion generated successfully", {
      resultLength: text.length,
    });

    return text;
  } catch (error) {
    if (error.name === "AbortError") {
      Logger.error("Completion request timeout", { timeout });
      throw new Error(`Completion request timed out after ${timeout}ms`);
    }
    throw error;
  }
}