Add maxlength command to set max response tokens

This commit is contained in:
2026-08-16 23:21:10 -05:00
parent ca541c2e82
commit c0ef8d0e4e
+19
View File
@@ -48,6 +48,7 @@ class Assistant(commands.Cog):
temperature=0.7, temperature=0.7,
system_prompt="You are a helpful assistant.", system_prompt="You are a helpful assistant.",
question_mode=False, question_mode=False,
max_tokens=None,
) )
self.config.register_global(api_key=None) self.config.register_global(api_key=None)
self.history: dict[int, list[dict]] = {} self.history: dict[int, list[dict]] = {}
@@ -128,6 +129,7 @@ class Assistant(commands.Cog):
model = await self.config.guild(guild).model() model = await self.config.guild(guild).model()
temperature = await self.config.guild(guild).temperature() temperature = await self.config.guild(guild).temperature()
system_prompt = await self.config.guild(guild).system_prompt() system_prompt = await self.config.guild(guild).system_prompt()
max_tokens = await self.config.guild(guild).max_tokens()
key = (guild.id, channel_id) key = (guild.id, channel_id)
if key not in self.history: if key not in self.history:
@@ -143,6 +145,9 @@ class Assistant(commands.Cog):
"temperature": temperature, "temperature": temperature,
} }
if max_tokens:
payload["max_tokens"] = max_tokens
try: try:
async with self.session.post( async with self.session.post(
f"{self.BASE_URL}/chat/completions", f"{self.BASE_URL}/chat/completions",
@@ -309,6 +314,18 @@ class Assistant(commands.Cog):
await self.config.guild(ctx.guild).temperature.set(temperature) await self.config.guild(ctx.guild).temperature.set(temperature)
await ctx.send(f"Temperature set to **{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") @assistant.command(name="system")
async def assistant_system(self, ctx: commands.Context, *, prompt: str): async def assistant_system(self, ctx: commands.Context, *, prompt: str):
"""Set the system prompt.""" """Set the system prompt."""
@@ -340,6 +357,7 @@ class Assistant(commands.Cog):
temp = await cfg.temperature() temp = await cfg.temperature()
sys_prompt = await cfg.system_prompt() sys_prompt = await cfg.system_prompt()
q_mode = await cfg.question_mode() q_mode = await cfg.question_mode()
max_tokens = await cfg.max_tokens()
has_key = await self.config.api_key() is not None has_key = await self.config.api_key() is not None
channel_str = f"<#{chat_channel}>" if chat_channel else "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="Model", value=model, inline=True)
embed.add_field(name="Temperature", value=str(temp), 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="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) embed.add_field(name="System Prompt", value=sys_prompt[:200], inline=False)
await ctx.send(embed=embed) await ctx.send(embed=embed)