polish

Author: VictorTaelin

Ask.mjs

@@ -0,0 +1,159 @@
+import fs from 'fs/promises';
+import os from 'os';
+import path from 'path';
+import { OpenAI } from "openai";
+import { Anthropic } from '@anthropic-ai/sdk';
+import { Groq } from "groq-sdk";
+import { GoogleGenerativeAI } from "@google/generative-ai";
+import { encode } from "gpt-tokenizer/esm/model/davinci-codex"; // tokenizer
+
+// Map of model shortcodes to full model names
+export const MODELS = {
+  g: 'gpt-4o', 
+  G: 'gpt-4-32k-0314',
+  h: 'claude-3-haiku-20240307',
+  s: 'claude-3-sonnet-20240229',
+  o: 'claude-3-opus-20240229',
+  l: 'llama3-8b-8192',
+  L: 'llama3-70b-8192',
+  i: 'gemini-1.5-flash-latest',
+  I: 'gemini-1.5-pro-latest'
+};
+
+// Factory function to create a stateful OpenAI chat
+export function openAIChat(clientClass) {
+  const messages = [];
+
+  async function ask(userMessage, { system, model, temperature = 0.0, max_tokens = 4096, stream = true }) {
+    model = MODELS[model] || model;
+    const client = new clientClass({ apiKey: await getToken(clientClass.name.toLowerCase()) });
+
+    if (messages.length === 0) {
+      messages.push({ role: "system", content: system });
+    }
+
+    messages.push({ role: "user", content: userMessage });
+
+    const params = { messages, model, temperature, max_tokens, stream };
+
+    let result = "";
+    const response = await client.chat.completions.create(params);
+
+    for await (const chunk of response) {
+      const text = chunk.choices[0]?.delta?.content || "";
+      process.stdout.write(text);
+      result += text;
+    }
+
+    messages.push({ role: 'assistant', content: result });
+
+    return result;
+  }
+
+  return ask;
+}
+
+// Factory function to create a stateful Anthropic chat
+export function anthropicChat(clientClass) {
+  const messages = [];
+
+  async function ask(userMessage, { system, model, temperature = 0.0, max_tokens = 4096, stream = true }) {
+    model = MODELS[model] || model;
+    const client = new clientClass({ apiKey: await getToken(clientClass.name.toLowerCase()) });
+
+    messages.push({ role: "user", content: userMessage });
+
+    const params = { system, model, temperature, max_tokens, stream };
+
+    let result = "";
+    const response = client.messages
+      .stream({ ...params, messages })
+      .on('text', (text) => {
+        process.stdout.write(text);
+        result += text;  
+      });
+    await response.finalMessage();
+
+    messages.push({ role: 'assistant', content: result });
+
+    return result;
+  }
+
+  return ask;
+}
+
+export function geminiChat(clientClass) {
+  const messages = [];
+
+  async function ask(userMessage, { system, model, temperature = 0.0, max_tokens = 4096, stream = true }) {
+    model = MODELS[model] || model;
+    const client = new clientClass(await getToken(clientClass.name.toLowerCase()));
+
+    const generationConfig = {
+      maxOutputTokens: max_tokens,
+      temperature,
+    };
+
+    const chat = client.getGenerativeModel({ model, systemInstruction: system, generationConfig })
+      .startChat({ history: messages });
+
+    messages.push({ role: "user", parts: [{ text: userMessage }] });
+
+    let result = "";
+    if (stream) {
+      const response = await chat.sendMessageStream(userMessage);
+      for await (const chunk of response.stream) {
+        const text = chunk.text();
+        process.stdout.write(text);
+        result += text;
+      }
+    } else {
+      const response = await chat.sendMessage(userMessage);
+      result = (await response.response).text();
+    }
+
+    messages.push({ role: 'model', parts: [{ text: result }] });
+
+    return result;
+  }
+
+  return ask;
+}
+
+// Generic asker function that dispatches to the correct asker based on the model name
+export function chat(model) {
+  model = MODELS[model] || model;
+  if (model.startsWith('gpt')) {
+    return openAIChat(OpenAI);
+  } else if (model.startsWith('claude')) {
+    return anthropicChat(Anthropic);
+  } else if (model.startsWith('llama')) {
+    return openAIChat(Groq);
+  } else if (model.startsWith('gemini')) {
+    return geminiChat(GoogleGenerativeAI);
+  } else {
+    throw new Error(`Unsupported model: ${model}`);
+  }
+}
+
+// Utility function to read the API token for a given vendor
+async function getToken(vendor) {
+  const tokenPath = path.join(os.homedir(), '.config', `${vendor}.token`);
+  try {
+    return (await fs.readFile(tokenPath, 'utf8')).trim();
+  } catch (err) {
+    console.error(`Error reading ${vendor}.token file:`, err.message);
+    process.exit(1);
+  }
+}
+
+export function tokenCount(inputText) {
+  // Encode the input string into tokens
+  const tokens = encode(inputText);
+
+  // Get the number of tokens 
+  const numberOfTokens = tokens.length;
+
+  // Return the number of tokens
+  return numberOfTokens;
+}

Chat.mjs

@@ -1,9 +1,10 @@
+#!/usr/bin/env node
+
 import process from "process";
 import fs from 'fs/promises';
-import { asker, MODELS } from './Ask.mjs';
+import { chat, MODELS } from './Chat.mjs';
 
-//const MODEL = "claude-3-opus-20240229";
-const MODEL = "g";
+const MODEL = process.argv[2] || "g";
 
 const SYSTEM = `
 You're a game emulator. You can emulate ANY game, but text-based. Your goal is
@@ -226,15 +227,15 @@ If the player provides feedback after a '#', use it to improve the experience.
   console.log(ASCII_ART);
 
   console.log("");
-  console.log(`\x1b[32mUsing \x1b[1m${MODEL}\x1b[0m`);
+  console.log(`\x1b[32mUsing \x1b[1m${MODELS[MODEL]||MODEL}\x1b[0m`);
   console.log("");
 
   process.stdout.write("Game: ");
   const game = (await new Promise(resolve => process.stdin.once('data', data => resolve(data.toString())))).trim();
 
   console.log(`Emulating ${game}...\n\n`);
 
-  const ask = asker();
+  const ask = chat(MODEL);
   let messages = [
     {role: "user", content: `# GAME: ${game}`},
   ];

aiemu.mjs

@@ -3,7 +3,7 @@
 import readline from 'readline';
 import { exec } from 'child_process';
 import { promisify } from 'util';
-import { asker, MODELS } from './Ask.mjs';
+import { chat, MODELS } from './Chat.mjs';
 
 const execAsync = promisify(exec);
 
@@ -43,7 +43,7 @@ const rl = readline.createInterface({
 });
 
 // Create a stateful asker
-const ask = asker();
+const ask = chat(MODEL);
 
 // Utility function to prompt the user for input
 async function prompt(query) {

chatsh.mjs

@@ -1,12 +1,10 @@
 #!/usr/bin/env node
-import { asker, MODELS, token_count } from './Ask.mjs';
+import { chat, MODELS, tokenCount } from './Chat.mjs';
 import process from "process";
 import fs from 'fs/promises';
 import os from 'os';
 import path from 'path';
 
-const ask = asker();
-
 const system = `
 You are a HOLE FILLER. You are provided with a file containing holes, formatted
 as '{{HOLE_NAME}}'. Your TASK is to complete with a string to replace this hole
@@ -98,9 +96,10 @@ function hypothenuse(a, b) {
 - Answer ONLY with the <COMPLETION/> block. Do NOT include anything outside it.
 `;
 
-var file = process.argv[2];
-var mini = process.argv[3];
+var file  = process.argv[2];
+var mini  = process.argv[3];
 var model = process.argv[4] || "g";
+var ask   = chat(model);
 
 if (!file) {
   console.log("Usage: holefill <file> [<shortened_file>] [<model_name>]");
@@ -130,7 +129,7 @@ while ((match = regex.exec(mini_code)) !== null) {
 
 await fs.writeFile(mini, mini_code, 'utf-8');
 
-var tokens = token_count(mini_code);
+var tokens = tokenCount(mini_code);
 var holes = mini_code.match(/{{\w+}}/g) || [];
 
 if (holes.length === 0 && mini_code.indexOf("??") !== -1 && (mini_code.match(/\?\?/g) || []).length == 1) {