Add maxlength command to set max response tokens

assistant/assistant.py

@@ -48,6 +48,7 @@ class Assistant(commands.Cog):
             temperature=0.7,
             system_prompt="You are a helpful assistant.",
             question_mode=False,
+            max_tokens=None,
         )
         self.config.register_global(api_key=None)
         self.history: dict[int, list[dict]] = {}
@@ -128,6 +129,7 @@ class Assistant(commands.Cog):
         model = await self.config.guild(guild).model()
         temperature = await self.config.guild(guild).temperature()
         system_prompt = await self.config.guild(guild).system_prompt()
+        max_tokens = await self.config.guild(guild).max_tokens()
 
         key = (guild.id, channel_id)
         if key not in self.history:
@@ -143,6 +145,9 @@ class Assistant(commands.Cog):
             "temperature": temperature,
         }
 
+        if max_tokens:
+            payload["max_tokens"] = max_tokens
+
         try:
             async with self.session.post(
                 f"{self.BASE_URL}/chat/completions",
@@ -309,6 +314,18 @@ class Assistant(commands.Cog):
         await self.config.guild(ctx.guild).temperature.set(temperature)
         await ctx.send(f"Temperature set to **{temperature}**.")
 
+    @assistant.command(name="maxlength")
+    async def assistant_maxlength(self, ctx: commands.Context, tokens: int = None):
+        """Set max response length in tokens. Omit to clear (unlimited)."""
+        if tokens is None:
+            await self.config.guild(ctx.guild).max_tokens.set(None)
+            await ctx.send("Max response length cleared (unlimited).")
+        elif tokens < 1:
+            await ctx.send("Must be a positive number.")
+        else:
+            await self.config.guild(ctx.guild).max_tokens.set(tokens)
+            await ctx.send(f"Max response length set to **{tokens}** tokens.")
+
     @assistant.command(name="system")
     async def assistant_system(self, ctx: commands.Context, *, prompt: str):
         """Set the system prompt."""
@@ -340,6 +357,7 @@ class Assistant(commands.Cog):
         temp = await cfg.temperature()
         sys_prompt = await cfg.system_prompt()
         q_mode = await cfg.question_mode()
+        max_tokens = await cfg.max_tokens()
         has_key = await self.config.api_key() is not None
 
         channel_str = f"<#{chat_channel}>" if chat_channel else "None"
@@ -351,6 +369,7 @@ class Assistant(commands.Cog):
         embed.add_field(name="Model", value=model, inline=True)
         embed.add_field(name="Temperature", value=str(temp), inline=True)
         embed.add_field(name="Question Mode", value=str(q_mode), inline=True)
+        embed.add_field(name="Max Tokens", value=str(max_tokens) if max_tokens else "Unlimited", inline=True)
         embed.add_field(name="System Prompt", value=sys_prompt[:200], inline=False)
 
         await ctx.send(embed=embed)