diff --git a/.github/workflows/test-gemini.yml b/.github/workflows/test-gemini.yml new file mode 100644 index 0000000..f73f6f9 --- /dev/null +++ b/.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}} diff --git a/README.md b/README.md index 3a28e26..746bbda 100644 --- a/README.md +++ b/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 diff --git a/package-lock.json b/package-lock.json index 3f93170..de7cdfb 100644 --- a/package-lock.json +++ b/package-lock.json @@ -11,6 +11,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", @@ -915,6 +916,15 @@ "node": "^18.18.0 || ^20.9.0 || >=21.1.0" } }, + "node_modules/@google/generative-ai": { + "version": "0.17.1", + "resolved": "https://registry.npmjs.org/@google/generative-ai/-/generative-ai-0.17.1.tgz", + "integrity": "sha512-TgWz02c5l2XJlEDys81UVat5+Qg9xqmYah7tQt6xlsBwFvzIFPz64aZFGd1av2sxT22NsssqLATjNsusAIJICA==", + "license": "Apache-2.0", + "engines": { + "node": ">=18.0.0" + } + }, "node_modules/@huggingface/jinja": { "version": "0.2.2", "resolved": "https://registry.npmjs.org/@huggingface/jinja/-/jinja-0.2.2.tgz", diff --git a/package.json b/package.json index 8c2c33c..dbcf7f8 100644 --- a/package.json +++ b/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", diff --git a/src/commands/gemini.ts b/src/commands/gemini.ts new file mode 100644 index 0000000..f28d2a8 --- /dev/null +++ b/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 ", "The model to use", "gemini-1.5-flash") + .option("-o, --outputDir ", "The output directory", "output") + .option( + "-k, --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 + ]); + }); diff --git a/src/index.ts b/src/index.ts index 8a8647c..d14a0d5 100755 --- a/src/index.ts +++ b/src/index.ts @@ -4,6 +4,7 @@ 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") @@ -11,5 +12,6 @@ cli() .version(version) .addCommand(local) .addCommand(openai) + .addCommand(azure) .addCommand(download()) .parse(process.argv); diff --git a/src/plugins/gemini-rename.ts b/src/plugins/gemini-rename.ts new file mode 100644 index 0000000..69b2224 --- /dev/null +++ b/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 => { + 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"] + } + } + }; +} diff --git a/src/test/e2e.geminitest.ts b/src/test/e2e.geminitest.ts new file mode 100644 index 0000000..84a0d8d --- /dev/null +++ b/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`) + ); +});