diff --git a/src/bot/modules/solaris/client.py b/src/bot/modules/solaris/client.py index 385dffc..95ed450 100644 --- a/src/bot/modules/solaris/client.py +++ b/src/bot/modules/solaris/client.py @@ -1,28 +1,22 @@ import json -from google import genai -from .structures import InputMessage, OutputMessage 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 + model="gemini-2.5-flash", config=MAIN_CONTENT_CONFIG ) - #self.tts + # self.tts + 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.fromdict(msg) - for msg in json.loads(resp.text) - ] - return output_messages \ No newline at end of file + 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.fromdict(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 a66836f..4a9e8ab 100644 --- a/src/bot/modules/solaris/content_configs.py +++ b/src/bot/modules/solaris/content_configs.py @@ -1,19 +1,18 @@ from google.genai import types -from .structures import OutputMessage +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 надо? мне просто каж что судя по всему вот эта нам + response_schema=list[ + OutputMessage + ], # ты уверен что там json надо? мне просто каж что судя по всему вот эта нам safety_settings=[ - types.SafetySetting( - category=category, - threshold=types.HarmBlockThreshold.OFF - ) + types.SafetySetting(category=category, threshold=types.HarmBlockThreshold.OFF) for category in types.HarmCategory - ] + ], ) @@ -28,15 +27,12 @@ TTS_CONTENT_CONFIG = types.GenerateContentConfig( speech_config=types.SpeechConfig( voice_config=types.VoiceConfig( prebuilt_voice_config=types.PrebuiltVoiceConfig( - voice_name='Kore', + voice_name="Kore", ) ) ), safety_settings=[ - types.SafetySetting( - category=category, - threshold=types.HarmBlockThreshold.OFF - ) + types.SafetySetting(category=category, threshold=types.HarmBlockThreshold.OFF) for category in types.HarmCategory - ] + ], ) diff --git a/src/bot/modules/solaris/structures/__init__.py b/src/bot/modules/solaris/structures/__init__.py index 807c037..57660ce 100644 --- a/src/bot/modules/solaris/structures/__init__.py +++ b/src/bot/modules/solaris/structures/__init__.py @@ -1,2 +1,2 @@ from .input_message import InputMessage -from .output_message import OutputMessage \ No newline at end of file +from .output_message import OutputMessage diff --git a/src/bot/modules/solaris/structures/input_message.py b/src/bot/modules/solaris/structures/input_message.py index 7e35d48..913c83e 100644 --- a/src/bot/modules/solaris/structures/input_message.py +++ b/src/bot/modules/solaris/structures/input_message.py @@ -1,10 +1,14 @@ -from pydantic import BaseModel from typing import Optional +from pydantic import BaseModel + + class InputMessage(BaseModel): time: str message_id: int text: str user_id: int - username: Optional[str] # хуйня я не помню зачем жт надо, наверное first_name важнее будет - reply_to: Optional[int] \ No newline at end of file + username: Optional[ + str + ] # хуйня я не помню зачем жт надо, наверное first_name важнее будет + reply_to: Optional[int] diff --git a/src/bot/modules/solaris/structures/output_message.py b/src/bot/modules/solaris/structures/output_message.py index 2563fbe..85a705c 100644 --- a/src/bot/modules/solaris/structures/output_message.py +++ b/src/bot/modules/solaris/structures/output_message.py @@ -1,16 +1,19 @@ -from pydantic import BaseModel from typing import Optional +from pydantic import BaseModel + + class OutputMessage(BaseModel): - priority: int # удалишь там есличе + priority: int # удалишь там есличе text: Optional[str] reaction: Optional[str] reply_to: Optional[int] + @classmethod def fromdict(cls, data: dict): return cls( - priority = data.get('priority'), - text = data.get('text'), - reaction = data.get('reaction'), - reply_to = data.get('reply_to') - ) \ No newline at end of file + priority=data.get("priority"), + text=data.get("text"), + reaction=data.get("reaction"), + reply_to=data.get("reply_to"), + )