diff --git a/cogs/neru_teto_cog.py b/cogs/neru_teto_cog.py index bda5269..60a2856 100644 --- a/cogs/neru_teto_cog.py +++ b/cogs/neru_teto_cog.py @@ -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 ... 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: