From 81a28178e3f4878311998529bdafbd4452ac0373 Mon Sep 17 00:00:00 2001
From: largomst <dingrui.zhen@gmail.com>
Date: Tue, 3 Dec 2024 17:34:03 +0800
Subject: [PATCH] feat: Provides support for OpenAI compatible APIs that
 support json_object

---
 package.json                        |  7 +++--
 src/plugins/openai/openai-rename.ts | 44 +++++++++++++----------------
 2 files changed, 25 insertions(+), 26 deletions(-)

diff --git a/package.json b/package.json
index 6f1d92b9..653c7c2f 100644
--- a/package.json
+++ b/package.json
@@ -48,18 +48,21 @@
   "license": "MIT",
   "dependencies": {
     "@babel/core": "^7.25.2",
+    "@babel/traverse": "^7.25.9",
     "@babel/types": "^7.25.2",
     "@google/generative-ai": "^0.20.0",
+    "@instructor-ai/instructor": "^1.5.0",
     "@types/babel__core": "^7.20.5",
     "babel-plugin-transform-beautifier": "^0.1.0",
     "commander": "^12.1.0",
     "dotenv": "^16.4.5",
     "ipull": "^3.9.0",
     "node-llama-cpp": "^3.0.0-beta.40",
-    "openai": "^4.55.1",
+    "openai": "^4.74.0",
     "tsx": "^4.16.2",
     "typescript": "^5.5.4",
-    "webcrack": "^2.13.0"
+    "webcrack": "^2.13.0",
+    "zod": "^3.23.8"
   },
   "devDependencies": {
     "@eslint/js": "^9.8.0",
diff --git a/src/plugins/openai/openai-rename.ts b/src/plugins/openai/openai-rename.ts
index f88fede2..d2036ea9 100644
--- a/src/plugins/openai/openai-rename.ts
+++ b/src/plugins/openai/openai-rename.ts
@@ -2,6 +2,8 @@ import OpenAI from "openai";
 import { visitAllIdentifiers } from "../local-llm-rename/visit-all-identifiers.js";
 import { showPercentage } from "../../progress.js";
 import { verbose } from "../../verbose.js";
+import Instructor from "@instructor-ai/instructor";
+import { z } from 'zod';
 
 export function openaiRename({
   apiKey,
@@ -14,8 +16,11 @@ export function openaiRename({
   model: string;
   contextWindowSize: number;
 }) {
-  const client = new OpenAI({ apiKey, baseURL });
-
+  const oai = new OpenAI({ apiKey, baseURL });
+  const instructor = Instructor({
+    client: oai,
+    mode: "JSON",
+  })
   return async (code: string): Promise<string> => {
     return await visitAllIdentifiers(
       code,
@@ -23,14 +28,14 @@ export function openaiRename({
         verbose.log(`Renaming ${name}`);
         verbose.log("Context: ", surroundingCode);
 
-        const response = await client.chat.completions.create(
+        const result = await instructor.chat.completions.create(
           toRenamePrompt(name, surroundingCode, model)
         );
-        const result = response.choices[0].message?.content;
+
         if (!result) {
-          throw new Error("Failed to rename", { cause: response });
+          throw new Error("Failed to rename", { cause: result });
         }
-        const renamed = JSON.parse(result).newName;
+        const renamed = result.newName;
 
         verbose.log(`Renamed to ${renamed}`);
 
@@ -46,7 +51,12 @@ function toRenamePrompt(
   name: string,
   surroundingCode: string,
   model: string
-): OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming {
+) {
+  const schema = z.object({
+    newName: z.string({
+      description: `The new name for the variable/function called \`${name}\``
+    })
+  })
   return {
     model,
     messages: [
@@ -59,23 +69,9 @@ function toRenamePrompt(
         content: surroundingCode
       }
     ],
-    response_format: {
-      type: "json_schema",
-      json_schema: {
-        strict: true,
-        name: "rename",
-        schema: {
-          type: "object",
-          properties: {
-            newName: {
-              type: "string",
-              description: `The new name for the variable/function called \`${name}\``
-            }
-          },
-          required: ["newName"],
-          additionalProperties: false
-        }
-      }
+    response_model: {
+      name: "Result",
+      schema: schema
     }
   };
 }