Merge pull request #66 from jehna/add/gemini

Add Google Gemini/AIStudio plugin

.github/workflows/test-gemini.yml

@@ -0,0 +1,28 @@
+name: Gemini tests
+
+on: pull_request
+
+jobs:
+  test-gemini:
+    name: Run tests
+    runs-on: macos-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Cache test model
+        id: cache-model
+        uses: actions/cache@v4
+        with:
+          path: /Users/runner/.humanifyjs/models/
+          key: models-phi
+      - name: Checkout code
+        uses: actions/checkout@v3
+      - uses: actions/setup-node@v4
+        with:
+          node-version: 20
+          cache: "npm"
+      - run: npm install --ci
+      - run: npm run build
+      - run: npm run download-ci-model
+      - run: npm run test:gemini
+        env:
+          GEMINI_API_KEY: ${{secrets.GEMINI_API_KEY}}

README.md

@@ -92,10 +92,10 @@ expect the tool to be installed globally, but they should work by replacing
 
 ### Usage
 
-Next you'll need to decide whether to use `openai` or `local` mode. In a
+Next you'll need to decide whether to use `openai`, `gemini` or `local` mode. In a
 nutshell:
 
-* `openai` mode
+* `openai` or `gemini` mode
   * Runs on someone else's computer that's specifically optimized for this kind
     of things
   * Costs money depending on the length of your code
@@ -115,14 +115,25 @@ https://openai.com/.
 
 There are several ways to provide the API key to the tool:
 ```shell
-echo "OPENAI_API_KEY=your-token" > .env && humanify openai obfuscated-file.js
-export OPENAI_API_KEY="your-token" && humanify openai obfuscated-file.js
-OPENAI_TOKEN=your-token humanify openai obfuscated-file.js
-humanify --apiKey="your-token" obfuscated-file.js
+humanify openai --apiKey="your-token" obfuscated-file.js
 ```
 
-Use your preferred way to provide the API key. Use `humanify --help` to see
-all available options.
+Alternatively you can also use an environment variable `OPENAI_API_KEY`. Use
+`humanify --help` to see all available options.
+
+### Gemini mode
+
+You'll need a Google AI Studio key. You can get one by signing up at
+https://aistudio.google.com/.
+
+You need to provice the API key to the tool:
+
+```shell
+humanify gemini --apiKey="your-token" obfuscated-file.js
+```
+
+Alternatively you can also use an environment variable `GEMINI_API_KEY`. Use
+`humanify --help` to see all available options.
 
 ### Local mode
 

package.json

@@ -16,6 +16,7 @@
     "test:e2e": "npm run build && find src -name '*.e2etest.ts' | xargs tsx --test --test-concurrency=1",
     "test:llm": "find src -name '*.llmtest.ts' | xargs tsx --test --test-concurrency=1",
     "test:openai": "npm run build && find src -name '*.openaitest.ts' | xargs tsx --test",
+    "test:gemini": "npm run build && find src -name '*.geminitest.ts' | xargs tsx --test",
     "lint": "npm run lint:prettier && npm run lint:eslint",
     "lint:prettier": "prettier --check src/* src/**/*",
     "lint:eslint": "eslint src/* src/**/*",
@@ -48,6 +49,7 @@
   "dependencies": {
     "@babel/core": "^7.25.2",
     "@babel/types": "^7.25.2",
+    "@google/generative-ai": "^0.17.1",
     "@types/babel__core": "^7.20.5",
     "babel-plugin-transform-beautifier": "^0.1.0",
     "commander": "^12.1.0",

src/commands/gemini.ts

@@ -0,0 +1,31 @@
+import { cli } from "../cli.js";
+import prettier from "../plugins/prettier.js";
+import { unminify } from "../unminify.js";
+import babel from "../plugins/babel/babel.js";
+import { verbose } from "../verbose.js";
+import { geminiRename } from "../plugins/gemini-rename.js";
+import { env } from "../env.js";
+
+export const azure = cli()
+  .name("gemini")
+  .description("Use Google Gemini/AIStudio API to unminify code")
+  .option("-m, --model <model>", "The model to use", "gemini-1.5-flash")
+  .option("-o, --outputDir <output>", "The output directory", "output")
+  .option(
+    "-k, --apiKey <apiKey>",
+    "The Google Gemini/AIStudio API key. Alternatively use GEMINI_API_KEY environment variable"
+  )
+  .option("--verbose", "Show verbose output")
+  .argument("input", "The input minified Javascript file")
+  .action(async (filename, opts) => {
+    if (opts.verbose) {
+      verbose.enabled = true;
+    }
+
+    const apiKey = opts.apiKey ?? env("GEMINI_API_KEY");
+    await unminify(filename, opts.outputDir, [
+      babel,
+      geminiRename({ apiKey, model: opts.model }),
+      prettier
+    ]);
+  });

src/index.ts

@@ -4,12 +4,14 @@ import { download } from "./commands/download.js";
 import { local } from "./commands/local.js";
 import { openai } from "./commands/openai.js";
 import { cli } from "./cli.js";
+import { azure } from "./commands/gemini.js";
 
 cli()
   .name("humanify")
   .description("Unminify code using OpenAI's API or a local LLM")
   .version(version)
   .addCommand(local)
   .addCommand(openai)
+  .addCommand(azure)
   .addCommand(download())
   .parse(process.argv);

src/plugins/gemini-rename.ts

@@ -0,0 +1,64 @@
+import { visitAllIdentifiers } from "./local-llm-rename/visit-all-identifiers.js";
+import { verbose } from "../verbose.js";
+import { showPercentage } from "../progress.js";
+import {
+  GoogleGenerativeAI,
+  ModelParams,
+  SchemaType
+} from "@google/generative-ai";
+
+export function geminiRename({
+  apiKey,
+  model: modelName
+}: {
+  apiKey: string;
+  model: string;
+}) {
+  const client = new GoogleGenerativeAI(apiKey);
+
+  return async (code: string): Promise<string> => {
+    return await visitAllIdentifiers(
+      code,
+      async (name, surroundingCode) => {
+        verbose.log(`Renaming ${name}`);
+        verbose.log("Context: ", surroundingCode);
+
+        const model = client.getGenerativeModel(
+          toRenameParams(name, modelName)
+        );
+
+        const result = await model.generateContent(surroundingCode);
+
+        const renamed = JSON.parse(result.response.text()).newName;
+
+        verbose.log(`Renamed to ${renamed}`);
+
+        return renamed;
+      },
+      showPercentage
+    );
+  };
+}
+
+function toRenameParams(name: string, model: string): ModelParams {
+  return {
+    model,
+    systemInstruction: `Rename Javascript variables/function \`${name}\` to have descriptive name based on their usage in the code."`,
+    generationConfig: {
+      responseMimeType: "application/json",
+      responseSchema: {
+        nullable: false,
+        description: "The new name for the variable/function",
+        type: SchemaType.OBJECT,
+        properties: {
+          newName: {
+            type: SchemaType.STRING,
+            nullable: false,
+            description: `The new name for the variable/function called \`${name}\``
+          }
+        },
+        required: ["newName"]
+      }
+    }
+  };
+}

src/test/e2e.geminitest.ts

@@ -0,0 +1,47 @@
+import test from "node:test";
+import { readFile, rm } from "node:fs/promises";
+import { testPrompt } from "./test-prompt.js";
+import { gbnf } from "../plugins/local-llm-rename/gbnf.js";
+import assert from "node:assert";
+import { humanify } from "../test-utils.js";
+
+const TEST_OUTPUT_DIR = "test-output";
+
+test.afterEach(async () => {
+  await rm(TEST_OUTPUT_DIR, { recursive: true, force: true });
+});
+
+test("Unminifies an example file successfully", async () => {
+  const fileIsMinified = async (filename: string) => {
+    const prompt = await testPrompt();
+    return await prompt(
+      `Your job is to read the following code and rate it's readabillity and variable names. Answer "EXCELLENT", "GOOD" or "UNREADABLE".`,
+      await readFile(filename, "utf-8"),
+      gbnf`${/("EXCELLENT" | "GOOD" | "UNREADABLE") [^.]+/}.`
+    );
+  };
+  const expectStartsWith = (expected: string[], actual: string) => {
+    assert(
+      expected.some((e) => actual.startsWith(e)),
+      `Expected "${expected}" but got ${actual}`
+    );
+  };
+
+  await expectStartsWith(
+    ["UNREADABLE"],
+    await fileIsMinified(`fixtures/example.min.js`)
+  );
+
+  await humanify(
+    "gemini",
+    "fixtures/example.min.js",
+    "--verbose",
+    "--outputDir",
+    TEST_OUTPUT_DIR
+  );
+
+  await expectStartsWith(
+    ["EXCELLENT", "GOOD"],
+    await fileIsMinified(`${TEST_OUTPUT_DIR}/deobfuscated.js`)
+  );
+});