From 41927a1e0794f5a14e4c7d2758eed2de5ed2fae1 Mon Sep 17 00:00:00 2001 From: h Date: Sat, 5 Jul 2025 01:13:12 +0300 Subject: [PATCH] feat(solaris): moving to service system to support multi-chat --- src/bot/__init__.py | 1 + src/bot/handlers/start/start.py | 7 ++- src/bot/modules/solaris/agents/respond.py | 27 ++++++++--- src/bot/modules/solaris/agents/review.py | 4 +- src/bot/modules/solaris/agents/tts.py | 20 +++++++- src/bot/modules/solaris/client.py | 5 +- src/bot/modules/solaris/constants.py | 28 +++++++++++ src/bot/modules/solaris/content_configs.py | 51 -------------------- src/bot/modules/solaris/services/__init__.py | 0 src/bot/modules/solaris/services/respond.py | 22 +++++++++ src/bot/modules/solaris/tools/__init__.py | 0 src/dependencies/__init__.py | 7 ++- src/dependencies/providers/__init__.py | 3 +- src/dependencies/providers/gemini.py | 13 +++++ src/dependencies/providers/solaris.py | 22 +++++++-- src/utils/db/models/config.py | 7 +++ 16 files changed, 143 insertions(+), 74 deletions(-) create mode 100644 src/bot/modules/solaris/constants.py delete mode 100644 src/bot/modules/solaris/content_configs.py create mode 100644 src/bot/modules/solaris/services/__init__.py create mode 100644 src/bot/modules/solaris/services/respond.py create mode 100644 src/bot/modules/solaris/tools/__init__.py create mode 100644 src/dependencies/providers/gemini.py diff --git a/src/bot/__init__.py b/src/bot/__init__.py index 43ccace..c98b9f0 100644 --- a/src/bot/__init__.py +++ b/src/bot/__init__.py @@ -23,6 +23,7 @@ async def runner(): await bot.delete_webhook(drop_pending_updates=True) await dp.start_polling(bot) + def plugins(): from rich import traceback diff --git a/src/bot/handlers/start/start.py b/src/bot/handlers/start/start.py index 0605649..a8a7679 100644 --- a/src/bot/handlers/start/start.py +++ b/src/bot/handlers/start/start.py @@ -1,9 +1,12 @@ from aiogram import Router, types from aiogram.filters import CommandStart +from dishka.integrations.aiogram import FromDishka + +from bot.modules.solaris.services.respond import RespondService router = Router() @router.message(CommandStart()) -async def on_start(message: types.Message): - await message.reply("hewo everynyan") +async def on_start(message: types.Message, respond_service: FromDishka[RespondService]): + await message.reply(str(respond_service.chat_id)) diff --git a/src/bot/modules/solaris/agents/respond.py b/src/bot/modules/solaris/agents/respond.py index 5235f33..29c5e5c 100644 --- a/src/bot/modules/solaris/agents/respond.py +++ b/src/bot/modules/solaris/agents/respond.py @@ -1,18 +1,31 @@ import json -from dataclasses import asdict from google import genai +from google.genai import chats, types -from ..content_configs import generate_respond_config +from utils.config import dconfig + +from ..constants import SAFETY_SETTINGS from ..structures import InputMessage, OutputMessage -RESPOND_MODEL = "gemini-2.5-flash" - class RespondAgent: - def __init__(self, client: genai.client.AsyncClient, prompt: str) -> None: - self.chat = client.chats.create( - model=RESPOND_MODEL, config=generate_respond_config(prompt=prompt) + chat: chats.AsyncChat + + def __init__(self, client: genai.client.AsyncClient) -> None: + self.client = client + + async def load_chat(self, history: list[types.Content], system_prompt: str): + self.chat = self.client.chats.create( + model=(await dconfig()).models.respond_model, + config=types.GenerateContentConfig( + system_instruction=system_prompt, + thinking_config=types.ThinkingConfig(thinking_budget=0), + response_mime_type="application/json", + response_schema=list[OutputMessage], + safety_settings=SAFETY_SETTINGS, + ), + history=history, ) async def send_messages(self, messages: list[InputMessage]) -> list[OutputMessage]: diff --git a/src/bot/modules/solaris/agents/review.py b/src/bot/modules/solaris/agents/review.py index 754ac9c..ef392d5 100644 --- a/src/bot/modules/solaris/agents/review.py +++ b/src/bot/modules/solaris/agents/review.py @@ -2,10 +2,10 @@ import json from google import genai -from ..content_configs import generate_review_config +from ..constants import generate_review_config from ..structures import InputMessage -REVIEW_MODEL = ("gemini-2.5-flash-lite-preview-06-17",) # надо будет гемму +REVIEW_MODEL = "gemini-2.5-flash-lite-preview-06-17" # надо будет гемму class ReviewAgent: diff --git a/src/bot/modules/solaris/agents/tts.py b/src/bot/modules/solaris/agents/tts.py index 301f4cd..1cff3c1 100644 --- a/src/bot/modules/solaris/agents/tts.py +++ b/src/bot/modules/solaris/agents/tts.py @@ -4,7 +4,7 @@ from google import genai from google.genai import types from pydub import AudioSegment -from ..content_configs import generate_tts_config +from ..constants import SAFETY_SETTINGS TTS_MODEL = "gemini-2.5-flash-preview-tts" @@ -12,18 +12,34 @@ TTS_MODEL = "gemini-2.5-flash-preview-tts" class TTSAgent: def __init__(self, client: genai.client.AsyncClient) -> None: self.client = client - self.content_config = generate_tts_config() + + self.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", + ) + ) + ), + safety_settings=SAFETY_SETTINGS, + ) async def generate(self, text: str): response = await self.client.models.generate_content( model=TTS_MODEL, contents=text, config=self.content_config ) + data = response.candidates[0].content.parts[0].inline_data.data pcm_io = io.BytesIO(data) + pcm_io.seek(0) + audio = AudioSegment( pcm_io.read(), sample_width=2, frame_rate=24000, channels=1 ) + ogg_io = io.BytesIO() audio.export(ogg_io, format="ogg", codec="libopus") ogg_bytes = ogg_io.getvalue() + return ogg_bytes diff --git a/src/bot/modules/solaris/client.py b/src/bot/modules/solaris/client.py index ff8605a..4c507d2 100644 --- a/src/bot/modules/solaris/client.py +++ b/src/bot/modules/solaris/client.py @@ -4,9 +4,8 @@ from .agents import BuildAgent class SolarisClient: - def __init__(self, api_key: str) -> None: - client = genai.Client(api_key=api_key).aio - self.builder = BuildAgent(client=client) + def __init__(self, gemini_client: genai.client.AsyncClient) -> None: + self.builder = BuildAgent(client=gemini_client) async def parse_user_data(self, some_data_idk): self.reviewer, self.responder = await self.builder.build( diff --git a/src/bot/modules/solaris/constants.py b/src/bot/modules/solaris/constants.py new file mode 100644 index 0000000..be205f6 --- /dev/null +++ b/src/bot/modules/solaris/constants.py @@ -0,0 +1,28 @@ +from google.genai import types + +from .structures import OutputMessage + +SAFETY_SETTINGS = [ + types.SafetySetting(category=category, threshold=types.HarmBlockThreshold.OFF) + for category in types.HarmCategory +] + + +def generate_respond_config(prompt: str) -> types.GenerateContentConfig: + return + + +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], +# ) diff --git a/src/bot/modules/solaris/content_configs.py b/src/bot/modules/solaris/content_configs.py deleted file mode 100644 index e687d33..0000000 --- a/src/bot/modules/solaris/content_configs.py +++ /dev/null @@ -1,51 +0,0 @@ -from google.genai import types - -from .structures import OutputMessage - -safety_settings = [ - types.SafetySetting(category=category, threshold=types.HarmBlockThreshold.OFF) - for category in types.HarmCategory -] - - -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, - ) - - -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=safety_settings, - ) diff --git a/src/bot/modules/solaris/services/__init__.py b/src/bot/modules/solaris/services/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/bot/modules/solaris/services/respond.py b/src/bot/modules/solaris/services/respond.py new file mode 100644 index 0000000..4ab66c1 --- /dev/null +++ b/src/bot/modules/solaris/services/respond.py @@ -0,0 +1,22 @@ +import json + +from google import genai +from google.genai import chats, types + +from utils.config import dconfig +from utils.logging import console + +from ..agents.respond import RespondAgent +from ..constants import SAFETY_SETTINGS +from ..structures import InputMessage, OutputMessage + + +class RespondService: + def __init__(self, client: genai.client.AsyncClient, chat_id: int) -> None: + self.agent = RespondAgent(client) + self.chat_id = chat_id + + async def spawn_agent(self): + console.print(self.chat_id) + + await self.agent.load_chat(history=[], system_prompt="nya nya") diff --git a/src/bot/modules/solaris/tools/__init__.py b/src/bot/modules/solaris/tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/dependencies/__init__.py b/src/dependencies/__init__.py index 6708082..1b6fdbf 100644 --- a/src/dependencies/__init__.py +++ b/src/dependencies/__init__.py @@ -1,7 +1,10 @@ from dishka import make_async_container +from dishka.integrations.aiogram import AiogramProvider -from .providers import SolarisClientProvider +from .providers import GeminiClientProvider, SolarisProvider container = make_async_container( - SolarisClientProvider(), + AiogramProvider(), + SolarisProvider(), + GeminiClientProvider(), ) diff --git a/src/dependencies/providers/__init__.py b/src/dependencies/providers/__init__.py index d09362c..3cb50e0 100644 --- a/src/dependencies/providers/__init__.py +++ b/src/dependencies/providers/__init__.py @@ -1 +1,2 @@ -from .solaris import SolarisClientProvider +from .gemini import GeminiClientProvider +from .solaris import SolarisProvider diff --git a/src/dependencies/providers/gemini.py b/src/dependencies/providers/gemini.py new file mode 100644 index 0000000..3477154 --- /dev/null +++ b/src/dependencies/providers/gemini.py @@ -0,0 +1,13 @@ +from typing import AsyncIterable + +from dishka import Provider, Scope, provide +from google import genai + +from utils.env import env + + +class GeminiClientProvider(Provider): + @provide(scope=Scope.APP) + async def get_client(self) -> AsyncIterable[genai.client.AsyncClient]: + client = genai.Client(api_key=env.google.api_key.get_secret_value()).aio + yield client diff --git a/src/dependencies/providers/solaris.py b/src/dependencies/providers/solaris.py index 1a8cf00..8857e94 100644 --- a/src/dependencies/providers/solaris.py +++ b/src/dependencies/providers/solaris.py @@ -1,13 +1,27 @@ from typing import AsyncIterable +import aiogram.types from dishka import Provider, Scope, provide +from dishka.integrations.aiogram import AiogramMiddlewareData +from google.genai.client import AsyncClient from bot.modules.solaris.client import SolarisClient -from utils.env import env +from bot.modules.solaris.services.respond import RespondService -class SolarisClientProvider(Provider): +class SolarisProvider(Provider): @provide(scope=Scope.APP) - async def get_client(self) -> AsyncIterable[SolarisClient]: - client = SolarisClient(env.google.api_key.get_secret_value()) + async def get_solaris_client( + self, client: AsyncClient + ) -> AsyncIterable[SolarisClient]: + client = SolarisClient(gemini_client=client) yield client + + @provide(scope=Scope.REQUEST) + async def get_respond_service( + self, client: AsyncClient, middleware_data: AiogramMiddlewareData + ) -> AsyncIterable[RespondService]: + chat: aiogram.types.Chat = middleware_data["event_chat"] + service = RespondService(client=client, chat_id=chat.id) + await service.spawn_agent() + yield service diff --git a/src/utils/db/models/config.py b/src/utils/db/models/config.py index edf5319..3ec8a38 100644 --- a/src/utils/db/models/config.py +++ b/src/utils/db/models/config.py @@ -7,8 +7,15 @@ class BotConfig(BaseModel): chats_whitelist: list[int] = [] +class GeminiModelsConfig(BaseModel): + respond_model: str = "gemini-2.5-flash" + message_review_model: str = "gemini-2.5-flash-lite-preview-06-17" + tts_model: str = "gemini-2.5-flash-preview-tts" + + class DynamicConfigBase(BaseModel): bot: BotConfig = Field(default_factory=BotConfig) + models: GeminiModelsConfig = Field(default_factory=GeminiModelsConfig) class DynamicConfig(DynamicConfigBase, Document):