anya/tools/ask.ts

import OpenAI from "openai";
import { saveApiUsage } from "../usage";
import axios from "axios";
import fs from "fs";
import path from "path";
import { RunnableToolFunctionWithParse } from "openai/lib/RunnableFunction.mjs";
import {
  ChatCompletion,
  ChatCompletionAssistantMessageParam,
  ChatCompletionMessageParam,
} from "openai/resources/index.mjs";
import { send_sys_log } from "../interfaces/log";
import { pathInDataDir } from "../config";

const ai_token = process.env.OPENAI_API_KEY?.trim();
const groq_token = process.env.GROQ_API_KEY?.trim();
const groq_baseurl = process.env.GROQ_BASE_URL?.trim();

// Messages saving implementation

interface MessageHistory {
  messages: ChatCompletionMessageParam[];
  timeout: NodeJS.Timer;
}

const seedMessageHistories: Map<string, MessageHistory> = new Map();
const HISTORY_TIMEOUT_MS = 10 * 60 * 1000;

/**
 * Retrieves the message history for a given seed.
 * If it doesn't exist, initializes a new history.
 * Resets the timeout each time it's accessed.
 *
 * @param seed - The seed identifier for the message history
 * @returns The message history array
 */
function getMessageHistory(seed: string): ChatCompletionMessageParam[] {
  const existingHistory = seedMessageHistories.get(seed);

  if (existingHistory) {
    // Reset the timeout
    clearTimeout(existingHistory.timeout);
    existingHistory.timeout = setTimeout(() => {
      seedMessageHistories.delete(seed);
      console.log(`Cleared message history for seed: ${seed}`);
      send_sys_log(`Cleared message history for seed: ${seed}`);
    }, HISTORY_TIMEOUT_MS);

    return existingHistory.messages;
  } else {
    // Initialize new message history
    const messages: ChatCompletionMessageParam[] = [];
    const timeout = setTimeout(() => {
      seedMessageHistories.delete(seed);
      console.log(`Cleared message history for seed: ${seed}`);
      send_sys_log(`Cleared message history for seed: ${seed}`);
    }, HISTORY_TIMEOUT_MS);

    seedMessageHistories.set(seed, { messages, timeout });
    return messages;
  }
}

/**
 * Sets the entire message history for a given seed.
 *
 * @param seed - The seed identifier for the message history
 * @param messages - The complete message history to set
 */
function setMessageHistory(
  seed: string,
  messages: ChatCompletionMessageParam[]
): void {
  const existingHistory = seedMessageHistories.get(seed);
  if (existingHistory) {
    clearTimeout(existingHistory.timeout);
    existingHistory.messages = messages;
    existingHistory.timeout = setTimeout(() => {
      seedMessageHistories.delete(seed);
      console.log(`Cleared message history for seed: ${seed}`);
      send_sys_log(`Cleared message history for seed: ${seed}`);
    }, HISTORY_TIMEOUT_MS);
  } else {
    const timeout = setTimeout(() => {
      seedMessageHistories.delete(seed);
      console.log(`Cleared message history for seed: ${seed}`);
      send_sys_log(`Cleared message history for seed: ${seed}`);
    }, HISTORY_TIMEOUT_MS);
    seedMessageHistories.set(seed, { messages, timeout });
  }
}

/**
 * Appends a message to the message history for a given seed.
 *
 * @param seed - The seed identifier for the message history
 * @param message - The message to append
 */
function appendMessage(
  seed: string,
  message: ChatCompletionMessageParam
): void {
  console.log(
    "Appending message",
    message.content,
    "tool_calls" in message && message.tool_calls
  );

  const history = seedMessageHistories.get(seed);
  if (history) {
    history.messages.push(message);
    // Reset the timeout
    clearTimeout(history.timeout);
    history.timeout = setTimeout(() => {
      seedMessageHistories.delete(seed);
      send_sys_log(`Cleared message history for seed: ${seed}`);
      console.log(`Cleared message history for seed: ${seed}`);
    }, HISTORY_TIMEOUT_MS);
  }
}

/**
 * The updated ask function with support for persistent message history via a seed.
 * Separates system prompt and user message to prevent duplication.
 *
 * @param params - The parameters for the ask function
 * @returns The response from the LLM API
 */
export async function ask({
  model = "gpt-4o-mini",
  prompt, // System prompt
  message, // User input message (optional)
  name,
  tools,
  seed,
  json,
  image_url,
}: {
  model?: string;
  prompt: string;
  message?: string;
  image_url?: string;
  name?: string;
  tools?: RunnableToolFunctionWithParse<any>[];
  seed?: string;
  json?: boolean;
}): Promise<ChatCompletion> {
  // Initialize OpenAI instances
  const openai = new OpenAI({
    apiKey: ai_token,
  });

  const groq = new OpenAI({
    apiKey: groq_token,
    baseURL: groq_baseurl,
  });

  // Initialize messages array with the system prompt
  let messages: ChatCompletionMessageParam[] = [
    {
      role: "system",
      content: prompt,
    },
  ];

  if (seed && message) {
    // Retrieve existing message history
    const history = getMessageHistory(seed);

    // Combine system prompt with message history and new user message
    messages = [
      {
        role: "system",
        content: prompt,
      },
      ...history,
      {
        role: "user",
        content: image_url
          ? [
              {
                type: "text",
                text: message,
              },
              {
                type: "image_url",
                image_url: {
                  url: image_url,
                },
              },
            ]
          : message,
        name,
      },
    ];
    console.log("got image:", image_url?.slice(0, 20));
  } else if (seed && !message) {
    // If seed is provided but no new message, just retrieve history
    const history = getMessageHistory(seed);
    messages = [
      {
        role: "system",
        content: prompt,
      },
      ...history,
    ];
  } else if (!seed && message) {
    // If no seed but message is provided, send system prompt and user message without history
    messages.push({
      role: "user",
      content: image_url
        ? [
            {
              type: "text",
              text: message,
            },
            {
              type: "image_url",
              image_url: {
                url: image_url,
              },
            },
          ]
        : message,
      name,
    });
  }

  let res: ChatCompletion;

  if (model === "groq-small") {
    res = await groq.chat.completions.create({
      model: "llama-3.1-8b-instant",
      messages,
    });

    if (res.usage) {
      saveApiUsage(
        new Date().toISOString().split("T")[0],
        model,
        res.usage.prompt_tokens,
        res.usage.completion_tokens
      );
    } else {
      console.log("No usage data");
    }

    // Handle response with seed
    if (seed && res.choices && res.choices.length > 0) {
      appendMessage(seed, res.choices[0].message);
    }

    return res;
  }

  if (tools?.length) {
    // Create a new runner with the current messages and tools
    const runner = openai.beta.chat.completions
      .runTools({
        model,
        messages,
        tools,
        response_format: json ? { type: "json_object" } : undefined,
      })
      .on("functionCall", (functionCall) => {
        send_sys_log(`ASK Function call: ${JSON.stringify(functionCall)}`);
        console.log("ASK Function call:", functionCall);
      })
      .on("message", (message) => {
        // remove empty tool_calls array
        if (
          "tool_calls" in message &&
          message.tool_calls &&
          message.tool_calls.length === 0
        ) {
          message.tool_calls = undefined;
          delete message.tool_calls;
        }
        seed && appendMessage(seed, message);
      })
      .on("totalUsage", (usage) => {
        send_sys_log(
          `ASK Total usage: ${usage.prompt_tokens} prompt tokens, ${usage.completion_tokens} completion tokens`
        );
        console.log("ASK Total usage:", usage);
        saveApiUsage(
          new Date().toISOString().split("T")[0],
          model,
          usage.prompt_tokens,
          usage.completion_tokens
        );
      });

    // Await the final chat completion
    res = await runner.finalChatCompletion();

    return res;
  }

  // Default behavior without tools
  res = await openai.chat.completions.create({
    model,
    messages,
  });

  if (res.usage) {
    saveApiUsage(
      new Date().toISOString().split("T")[0],
      model,
      res.usage.prompt_tokens,
      res.usage.completion_tokens
    );
  } else {
    console.log("No usage data");
  }

  // Handle response with seed
  if (seed && res.choices && res.choices.length > 0) {
    const assistantMessage = res.choices[0].message;
    appendMessage(seed, assistantMessage);
  }

  return res;
}

const transcriptionCacheFile = pathInDataDir("transcription_cache.json");

export async function get_transcription(
  input: string | File, // Accept either a file URL (string) or a File object
  binary?: boolean,
  key?: string
) {
  // const openai = new OpenAI({
  //   apiKey: ai_token,
  // });

  const openai = new OpenAI({
    apiKey: groq_token,
    baseURL: groq_baseurl,
  });

  // Step 1: Check if the transcription for this input (file_url or File) is already cached
  let transcriptionCache: Record<string, string> = {};

  // Try to read the cache file if it exists
  if (fs.existsSync(transcriptionCacheFile)) {
    const cacheData = fs.readFileSync(transcriptionCacheFile, "utf-8");
    transcriptionCache = JSON.parse(cacheData);
  }

  let filePath: string;
  let isAudio = false;
  let fileExtension: string;

  // Determine if the input is a File or URL and handle accordingly
  if (input instanceof File) {
    // Check the MIME type for audio validation
    if (!input.type.startsWith("audio/")) {
      throw new Error("The provided file is not an audio file.");
    }
    isAudio = true;

    // Set file extension based on the MIME type
    fileExtension = getExtensionFromMimeType(input.type) ?? "ogg";
    if (!fileExtension) {
      throw new Error(`Unsupported audio file type: ${input.type}`);
    }

    // Write the file to the filesystem temporarily with the correct extension
    filePath = `/tmp/audio${Date.now()}.${fileExtension}`;
    const buffer = await input.arrayBuffer();
    fs.writeFileSync(filePath, new Uint8Array(buffer));
  } else if (typeof input === "string") {
    if (binary) {
      // If input is binary data
      const binaryData = Buffer.from(input, "base64");
      if (key && transcriptionCache[key]) {
        console.log("Transcription found in cache:", transcriptionCache[key]);
        return transcriptionCache[key];
      }
      filePath = `/tmp/audio${Date.now()}.ogg`; // Default to .ogg for binary input
      fs.writeFileSync(filePath, new Uint8Array(binaryData));
    } else {
      // Treat input as a file URL and extract the file extension
      fileExtension = path.extname(input).slice(1).toLowerCase();
      if (!["mp3", "ogg", "wav", "m4a"].includes(fileExtension)) {
        throw new Error(
          "The provided URL does not point to a valid audio file."
        );
      }
      isAudio = true;

      // Step 2: Download the file from the URL
      const response = await axios({
        url: input,
        method: "GET",
        responseType: "stream",
      });

      filePath = `/tmp/audio${Date.now()}.${fileExtension}`;

      // Save the downloaded file locally
      const writer = fs.createWriteStream(filePath);
      response.data.pipe(writer);

      await new Promise((resolve, reject) => {
        writer.on("finish", resolve);
        writer.on("error", reject);
      });
    }
  } else {
    throw new Error(
      "Invalid input type. Must be either a file URL or a File object."
    );
  }

  try {
    // Step 3: Send the file to OpenAI's Whisper model for transcription
    const transcription = await openai.audio.transcriptions.create({
      // model: "whisper-1",
      model: "distil-whisper-large-v3-en",
      file: fs.createReadStream(filePath),
      language: "en", // Optional
      temperature: 0.0, // Optional
    });

    // Delete the temp file
    fs.unlinkSync(filePath);

    // Step 4: Save the transcription to the cache
    if (key) {
      transcriptionCache[key] = transcription.text;
    } else if (typeof input === "string") {
      transcriptionCache[input] = transcription.text;
    }
    fs.writeFileSync(
      transcriptionCacheFile,
      JSON.stringify(transcriptionCache, null, 2)
    );

    console.log("Transcription:", transcription);
    return transcription.text;
  } catch (error) {
    console.error("Error transcribing audio:", error);
    throw error;
  }
}

// Helper function to get file extension based on MIME type
function getExtensionFromMimeType(mimeType: string): string | null {
  const mimeTypesMap: Record<string, string> = {
    "audio/mpeg": "mp3",
    "audio/ogg": "ogg",
    "audio/wav": "wav",
    "audio/x-wav": "wav",
    "audio/x-m4a": "m4a",
    "audio/m4a": "m4a",
    // Add other audio types as necessary
  };
  return mimeTypesMap[mimeType] || null;
}