From b7cd2de0f3dd8f8e08f779aefe59e4bb1b66ea05 Mon Sep 17 00:00:00 2001 From: shinrei Date: Wed, 2 Jul 2025 21:46:19 +0000 Subject: [PATCH] feat(solaris): changed agents logic; made content_configs dynamic and maybe smth else.. --- src/bot/modules/solaris/agents/build.py | 0 src/bot/modules/solaris/agents/respond.py | 20 ++++++ src/bot/modules/solaris/agents/review.py | 23 +++++++ src/bot/modules/solaris/client.py | 12 ---- src/bot/modules/solaris/content_configs.py | 71 +++++++++++++--------- 5 files changed, 85 insertions(+), 41 deletions(-) create mode 100644 src/bot/modules/solaris/agents/build.py create mode 100644 src/bot/modules/solaris/agents/respond.py create mode 100644 src/bot/modules/solaris/agents/review.py diff --git a/src/bot/modules/solaris/agents/build.py b/src/bot/modules/solaris/agents/build.py new file mode 100644 index 0000000..e69de29 diff --git a/src/bot/modules/solaris/agents/respond.py b/src/bot/modules/solaris/agents/respond.py new file mode 100644 index 0000000..8d00a1c --- /dev/null +++ b/src/bot/modules/solaris/agents/respond.py @@ -0,0 +1,20 @@ +import json +from dataclasses import asdict + +from google import genai + +from ..content_configs import generate_respond_config +from ..structures import InputMessage, OutputMessage + + +class RespondAgent: + def __init__(self, api_key: str) -> None: + client = genai.Client(api_key=api_key).aio + self.chat = client.chats.create( + model="gemini-2.5-flash", config=generate_respond_config(prompt="im gay") + ) + + async def send_messages(self, messages: list[InputMessage]) -> list[OutputMessage]: + data = json.dumps([msg.model_dump() for msg in messages], ensure_ascii=True) + response = await self.chat.send_message(data) + return response.parsed diff --git a/src/bot/modules/solaris/agents/review.py b/src/bot/modules/solaris/agents/review.py new file mode 100644 index 0000000..49f124b --- /dev/null +++ b/src/bot/modules/solaris/agents/review.py @@ -0,0 +1,23 @@ +import io +import json + +from google import genai +from google.genai import types +from pydub import AudioSegment + +from ..content_configs import generate_review_config +from ..structures import InputMessage + + +class ReviewAgent: + def __init__(self, api_key: str) -> None: + self.client = genai.Client(api_key=api_key).aio + + async def review(self, messages: list[InputMessage]) -> list[InputMessage]: + data = json.dumps([msg.model_dump() for msg in messages], ensure_ascii=True) + response = await self.client.models.generate_content( + model="gemini-2.5-flash-lite-preview-06-17", # надо будет гемму + contents=data, + config=generate_review_config(prompt="fortnite balls"), + ) + return [msg for msg in messages if msg.message_id in response.parsed] diff --git a/src/bot/modules/solaris/client.py b/src/bot/modules/solaris/client.py index 41a62b2..5faae08 100644 --- a/src/bot/modules/solaris/client.py +++ b/src/bot/modules/solaris/client.py @@ -3,21 +3,9 @@ from dataclasses import asdict from google import genai -from .content_configs import MAIN_CONTENT_CONFIG, MESSAGE_CONTENT_CONFIG from .structures import InputMessage, OutputMessage class SolarisClient: def __init__(self, api_key: str) -> None: client = genai.Client(api_key=api_key).aio - self.chat = client.chats.create( - model="gemini-2.5-flash", config=MAIN_CONTENT_CONFIG - ) - - async def send_messages(self, messages: list[InputMessage]) -> list[OutputMessage]: - data = json.dumps([msg.model_dump() for msg in messages], ensure_ascii=True) - resp = await self.chat.send_message(data, config=MESSAGE_CONTENT_CONFIG) - output_messages = [ - OutputMessage.model_validate(msg) for msg in json.loads(resp.text) - ] - return output_messages diff --git a/src/bot/modules/solaris/content_configs.py b/src/bot/modules/solaris/content_configs.py index 4a9e8ab..e687d33 100644 --- a/src/bot/modules/solaris/content_configs.py +++ b/src/bot/modules/solaris/content_configs.py @@ -2,37 +2,50 @@ from google.genai import types from .structures import OutputMessage -MAIN_CONTENT_CONFIG = types.GenerateContentConfig( - system_instruction="meow meow meow", # надо где-то промпт хранить, в бд наверное хезе - thinking_config=types.ThinkingConfig(thinking_budget=0), - response_mime_type="application/json", - response_schema=list[ - OutputMessage - ], # ты уверен что там json надо? мне просто каж что судя по всему вот эта нам - safety_settings=[ - types.SafetySetting(category=category, threshold=types.HarmBlockThreshold.OFF) - for category in types.HarmCategory - ], -) +safety_settings = [ + types.SafetySetting(category=category, threshold=types.HarmBlockThreshold.OFF) + for category in types.HarmCategory +] -MESSAGE_CONTENT_CONFIG = types.GenerateContentConfig( - response_mime_type="application/json", # возможно можно скипнуть это если мы используем response_schema, надо проверить - response_schema=list[OutputMessage], -) +def generate_respond_config(prompt: str) -> types.GenerateContentConfig: + return types.GenerateContentConfig( + system_instruction=prompt, + thinking_config=types.ThinkingConfig(thinking_budget=0), + response_mime_type="application/json", + response_schema=list[ + OutputMessage + ], # ты уверен что там json надо? мне просто каж что судя по всему вот эта нам + safety_settings=safety_settings, + ) -TTS_CONTENT_CONFIG = types.GenerateContentConfig( - response_modalities=[types.Modality.AUDIO], - speech_config=types.SpeechConfig( - voice_config=types.VoiceConfig( - prebuilt_voice_config=types.PrebuiltVoiceConfig( - voice_name="Kore", +def generate_review_config(prompt: str) -> types.GenerateContentConfig: + return types.GenerateContentConfig( + system_instruction=prompt, + thinking_config=types.ThinkingConfig(thinking_budget=0), + response_mime_type="application/json", + response_schema=list[int], + safety_settings=safety_settings, + ) + + +# MESSAGE_CONTENT_CONFIG = types.GenerateContentConfig( +# response_mime_type="application/json", # возможно можно скипнуть это если мы используем response_schema, надо проверить +# response_schema=list[OutputMessage], +# ) + + +# можно было и константу оставить но хезе некрасиво чет +def generate_tts_config() -> types.GenerateContentConfig: + return types.GenerateContentConfig( + response_modalities=[types.Modality.AUDIO], + speech_config=types.SpeechConfig( + voice_config=types.VoiceConfig( + prebuilt_voice_config=types.PrebuiltVoiceConfig( + voice_name="Kore", + ) ) - ) - ), - safety_settings=[ - types.SafetySetting(category=category, threshold=types.HarmBlockThreshold.OFF) - for category in types.HarmCategory - ], -) + ), + safety_settings=safety_settings, + )