fix(listeners:moderationAgent): Update request parameters and response handling

2 files changed, 38 insertions, 2 deletions

diff --git a/packages/gateway/src/listeners/moderationAgent/constants.ts b/packages/gateway/src/listeners/moderationAgent/constants.ts
index 40cd82d..75109bd 100644
--- a/packages/gateway/src/listeners/moderationAgent/constants.ts
+++ b/packages/gateway/src/listeners/moderationAgent/constants.ts
@@ -7,7 +7,7 @@ export const EXCLUDED_CATEGORIES = [
 export const MODERATION_LOG_CHANNEL_ID = "1406422619934167106";
 export const MIN_MESSAGE_LENGTH = 15;
 export const MAX_SYMBOL_DENSITY = 0.6;
-export const MAX_COMPLETION_TOKENS = 2000;
+export const MAX_COMPLETION_TOKENS = 4000;
 export const MESSAGE_HISTORY_SIZE = 3;
 export const MODEL = "mistralai/mistral-nemo";
 export const SAFE_WORDS = new Set([
diff --git a/packages/gateway/src/listeners/moderationAgent/utilities.ts b/packages/gateway/src/listeners/moderationAgent/utilities.ts
index 14bf563..397afca 100644
--- a/packages/gateway/src/listeners/moderationAgent/utilities.ts
+++ b/packages/gateway/src/listeners/moderationAgent/utilities.ts
@@ -1,5 +1,5 @@
 import { Message, TextChannel, ThreadChannel } from "discord.js";
-import { MESSAGE_HISTORY_SIZE, MODEL, SERVER_RULES } from "./constants";
+import { MESSAGE_HISTORY_SIZE, MAX_COMPLETION_TOKENS, MODEL, SERVER_RULES } from "./constants";
 
 export const fetchMessageContext = async (
   channel: TextChannel | ThreadChannel,
@@ -307,6 +307,7 @@ Remember: Only enforce the exact rules provided. Do not make assumptions or inte
               content: fullContext,
             },
           ],
+          max_tokens: MAX_COMPLETION_TOKENS,
         }),
       },
     );
@@ -358,6 +359,23 @@ Remember: Only enforce the exact rules provided. Do not make assumptions or inte
         if (jsonMatch) jsonContent = jsonMatch[0];
       }
 
+      if (!jsonContent.endsWith("}")) {
+        const openBraces = (jsonContent.match(/\{/g) || []).length;
+        const closeBraces = (jsonContent.match(/\}/g) || []).length;
+        
+        if (openBraces > closeBraces) {
+          let truncatedJson = jsonContent;
+          
+          if (truncatedJson.match(/"[^"]*$/))
+            truncatedJson = truncatedJson.replace(/"[^"]*$/, '""');
+          
+          const missingBraces = openBraces - closeBraces;
+
+          truncatedJson += '}'.repeat(missingBraces);
+          jsonContent = truncatedJson;
+        }
+      }
+
       jsonContent = jsonContent
         .replace(/,\s*}/g, "}")
         .replace(/,\s*]/g, "]")
@@ -401,6 +419,23 @@ Remember: Only enforce the exact rules provided. Do not make assumptions or inte
 
         if (jsonMatch) fallbackContent = jsonMatch[0];
 
+        if (!fallbackContent.endsWith("}")) {
+          const openBraces = (fallbackContent.match(/\{/g) || []).length;
+          const closeBraces = (fallbackContent.match(/\}/g) || []).length;
+          
+          if (openBraces > closeBraces) {
+            let truncatedJson = fallbackContent;
+            
+            if (truncatedJson.match(/"[^"]*$/))
+              truncatedJson = truncatedJson.replace(/"[^"]*$/, '""');
+            
+            const missingBraces = openBraces - closeBraces;
+
+            truncatedJson += '}'.repeat(missingBraces);
+            fallbackContent = truncatedJson;
+          }
+        }
+
         fallbackContent = fallbackContent
           .replace(/,\s*}/g, "}")
           .replace(/,\s*]/g, "]")
@@ -421,6 +456,7 @@ Remember: Only enforce the exact rules provided. Do not make assumptions or inte
         }
       } catch (fallbackError) {
         console.error("Fallback parsing also failed:", fallbackError);
+
         return null;
       }
     }