diff --git a/action.yml b/action.yml
index 976aed7..c7a10f8 100644
--- a/action.yml
+++ b/action.yml
@@ -11,6 +11,10 @@ inputs:
     description: "OpenAI API model."
     required: false
     default: "gpt-4"
+  MAX_TOKEN:
+    description: "Maximum number of tokens that can be generated per analysis."
+    required: false
+    default: "700"
   exclude:
     description: "Glob patterns to exclude files from the diff analysis"
     required: false
diff --git a/src/main.ts b/src/main.ts
index 13677ae..36380c7 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -8,6 +8,18 @@ import minimatch from "minimatch";
 const GITHUB_TOKEN: string = core.getInput("GITHUB_TOKEN");
 const OPENAI_API_KEY: string = core.getInput("OPENAI_API_KEY");
 const OPENAI_API_MODEL: string = core.getInput("OPENAI_API_MODEL");
+const MAX_TOKEN: number = Number(core.getInput("MAX_TOKEN"));
+
+const SUPPORTED_JSON_FORMAT_MODELS = [
+  "gpt-4o",
+  "gpt-4-turbo-preview",
+  "gpt-4-turbo",
+  "gpt-3.5-turbo",
+  "gpt-4-0125-preview",
+  "gpt-4-1106-preview",
+  "gpt-3.5-turbo-0125",
+  "gpt-3.5-turbo-1106",
+];
 
 const octokit = new Octokit({ auth: GITHUB_TOKEN });
 
@@ -117,7 +129,7 @@ async function getAIResponse(prompt: string): Promise<Array<{
   const queryConfig = {
     model: OPENAI_API_MODEL,
     temperature: 0.2,
-    max_tokens: 700,
+    max_tokens: MAX_TOKEN,
     top_p: 1,
     frequency_penalty: 0,
     presence_penalty: 0,
@@ -127,9 +139,9 @@ async function getAIResponse(prompt: string): Promise<Array<{
     const response = await openai.chat.completions.create({
       ...queryConfig,
       // return JSON if the model supports it:
-      ...(OPENAI_API_MODEL === "gpt-4-1106-preview"
-        ? { response_format: { type: "json_object" } }
-        : {}),
+      ...(SUPPORTED_JSON_FORMAT_MODELS.includes(OPENAI_API_MODEL)
+          ? { response_format: { type: "json_object" } }
+          : {}),
       messages: [
         {
           role: "system",
@@ -139,6 +151,9 @@ async function getAIResponse(prompt: string): Promise<Array<{
     });
 
     const res = response.choices[0].message?.content?.trim() || "{}";
+
+    console.log(`Trimmed Response: ${res}`);
+
     return JSON.parse(res).reviews;
   } catch (error) {
     console.error("Error:", error);