slipstream/discordbot
slipstream/discordbot
1
1
Fork 0
Code Issues Pull Requests 3 Actions Packages Projects Releases Wiki Activity

Applying previous commit.

Browse Source
This commit is contained in:
Slipstream 2025-06-05 06:10:59 +00:00
parent 4f946ca4d0
commit 3dd6a02dfb
Signed by: slipstream
GPG Key ID: 13E498CE010AC6FD
1 changed files with 86 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

127
cogs/neru_teto_cog.py
View File

@ -4,6 +4,7 @@ from discord import app_commands
import re
import base64
import io
from typing import Optional
def strip_think_blocks(text):
# Removes all <think>...</think> blocks, including multiline
@ -18,27 +19,60 @@ _teto_conversations = {}
import os
import aiohttp
from google import genai
from google.genai import types
from google.api_core import exceptions as google_exceptions
from gurt.config import PROJECT_ID, LOCATION
STANDARD_SAFETY_SETTINGS = [
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold="BLOCK_NONE"),
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold="BLOCK_NONE"),
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, threshold="BLOCK_NONE"),
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold="BLOCK_NONE"),
]
def _get_response_text(response: Optional[types.GenerateContentResponse]) -> Optional[str]:
"""Extract text from a Vertex AI response if available."""
if not response:
return None
if hasattr(response, "text") and response.text:
return response.text
if not response.candidates:
return None
try:
candidate = response.candidates[0]
if not getattr(candidate, "content", None) or not candidate.content.parts:
return None
for part in candidate.content.parts:
if hasattr(part, "text") and isinstance(part.text, str) and part.text.strip():
return part.text
return None
except (AttributeError, IndexError, TypeError):
return None
class DmbotTetoCog(commands.Cog):
def __init__(self, bot: commands.Bot):
self.bot = bot
self._api_endpoint = "https://openrouter.ai/api/v1/chat/completions" # Default endpoint
self._ai_model = "google/gemini-2.5-flash-preview" # Default model
try:
if PROJECT_ID and LOCATION:
self.genai_client = genai.Client(
vertexai=True,
project=PROJECT_ID,
location=LOCATION,
)
else:
self.genai_client = None
except Exception:
self.genai_client = None
self._ai_model = "gemini-2.5-flash-preview-05-20" # Default model used by TetoCog
async def _teto_reply_ai_with_messages(self, messages, system_mode="reply"):
"""
Use OpenRouter AI to generate a Kasane Teto-style response.
system_mode: "reply" for replying as Kasane Teto.
"""
api_key = os.getenv("AI_API_KEY")
if not api_key:
raise RuntimeError("AI_API_KEY environment variable not set.")
url = self._api_endpoint
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
"HTTP-Referer": "https://github.com/Slipstreamm/discordbot/blob/master/cogs/teto_cog.py",
"X-Title": "Teto AI"
}
"""Use Vertex AI to generate a Kasane Teto-style response."""
if not self.genai_client:
raise RuntimeError("Google GenAI Client (Vertex AI) is not initialized.")
system_prompt = (
"For context, you speak with users via being a Discord bot in a public server. \n"
"You are Kasane Teto, a cheeky, energetic, and often mischievous UTAU character with iconic red drill-like twintails. \n"
@ -53,30 +87,42 @@ class DmbotTetoCog(commands.Cog):
"When expressing emotions, never use emojis and instead use text-based emoticons like :) and ;) as well as Japanese-style emoticons, for example >~<, ^~^, >///<, UwU, o.O, and OwO over emojis. \n"
"Reply to the user in a short, conversational manner, staying in character."
)
payload = {
"model": self._ai_model,
"messages": [{"role": "system", "content": system_prompt}] + messages,
"max_tokens": 2000
}
async with aiohttp.ClientSession() as session:
async with session.post(url, headers=headers, json=payload) as resp:
if resp.status != 200:
text = await resp.text()
raise RuntimeError(f"OpenRouter API returned error status {resp.status}: {text[:500]}")
contents = [types.Content(role="system", parts=[types.Part(text=system_prompt)])]
for msg in messages:
role = "user" if msg.get("role") == "user" else "model"
contents.append(types.Content(role=role, parts=[types.Part(text=msg.get("content", ""))]))
if resp.content_type == "application/json":
data = await resp.json()
if "choices" not in data or not data["choices"]:
raise RuntimeError(f"OpenRouter API returned unexpected response format: {data}")
return data["choices"][0]["message"]["content"]
else:
text = await resp.text()
raise RuntimeError(f"OpenRouter API returned non-JSON response (status {resp.status}): {text[:500]}")
generation_config = types.GenerateContentConfig(
temperature=1.0,
max_output_tokens=2000,
safety_settings=STANDARD_SAFETY_SETTINGS,
)
try:
response = await self.genai_client.aio.models.generate_content(
model=f"publishers/google/models/{self._ai_model}",
contents=contents,
config=generation_config,
)
except google_exceptions.GoogleAPICallError as e:
raise RuntimeError(f"Vertex AI API call failed: {e}")
ai_reply = _get_response_text(response)
if ai_reply:
return ai_reply
raise RuntimeError("Vertex AI returned no text response.")
async def _teto_reply_ai(self, text: str) -> str:
"""Replies to the text as Kasane Teto using AI via OpenRouter."""
"""Replies to the text as Kasane Teto using AI via Vertex AI."""
return await self._teto_reply_ai_with_messages([{"role": "user", "content": text}])
async def _send_followup_in_chunks(self, interaction: discord.Interaction, text: str, *, ephemeral: bool = True) -> None:
"""Send a potentially long message in chunks using followup messages."""
chunk_size = 1900
chunks = [text[i : i + chunk_size] for i in range(0, len(text), chunk_size)] or [""]
for chunk in chunks:
await interaction.followup.send(chunk, ephemeral=ephemeral)
teto = app_commands.Group(name="teto", description="Commands related to Kasane Teto.")
model = app_commands.Group(parent=teto, name="model", description="Commands related to Teto's AI model.")
endpoint = app_commands.Group(parent=teto, name="endpoint", description="Commands related to Teto's API endpoint.")
@ -121,13 +167,13 @@ class DmbotTetoCog(commands.Cog):
try:
ai_reply = await self._teto_reply_ai_with_messages(messages=convo)
ai_reply = strip_think_blocks(ai_reply)
await interaction.followup.send(ai_reply)
await self._send_followup_in_chunks(interaction, ai_reply, ephemeral=True)
convo.append({"role": "assistant", "content": ai_reply})
_teto_conversations[convo_key] = convo[-30:] # Keep last 30 messages
_teto_conversations[convo_key] = convo[-30:] # Keep last 30 messages
except Exception as e:
await interaction.followup.send(f"Teto AI reply failed: {e} desu~")
await interaction.followup.send(f"Teto AI reply failed: {e} desu~", ephemeral=True)
else:
await interaction.followup.send("Please provide a message to chat with Teto desu~")
await interaction.followup.send("Please provide a message to chat with Teto desu~", ephemeral=True)
# Context menu command must be defined at module level
@ -153,8 +199,7 @@ async def teto_context_menu_ai_reply(interaction: discord.Interaction, message:
return
ai_reply = await cog._teto_reply_ai_with_messages(messages=convo)
ai_reply = strip_think_blocks(ai_reply)
await message.reply(ai_reply)
await interaction.followup.send("Teto AI replied desu~", ephemeral=True)
await cog._send_followup_in_chunks(interaction, ai_reply, ephemeral=True)
convo.append({"role": "assistant", "content": ai_reply})
_teto_conversations[convo_key] = convo[-10:]
except Exception as e: