diff --git a/gemini_worker/src/gemini_sdk/constants.py b/gemini_worker/src/gemini_sdk/constants.py index 66718a7..e483a61 100644 --- a/gemini_worker/src/gemini_sdk/constants.py +++ b/gemini_worker/src/gemini_sdk/constants.py @@ -1,3 +1,3 @@ -GEMINI_TOKENS_LIMIT: int = 10_000 \ No newline at end of file +GEMINI_TOKENS_LIMIT: int = 3_000 \ No newline at end of file diff --git a/gemini_worker/src/gemini_sdk/promt.py b/gemini_worker/src/gemini_sdk/promt.py index aff718d..777b707 100644 --- a/gemini_worker/src/gemini_sdk/promt.py +++ b/gemini_worker/src/gemini_sdk/promt.py @@ -7,7 +7,7 @@ ANALYTIC_PROMT = """ Ты получаешь json с такими полями { chats: [ - "chat_id": integer + "slice_id": UUID4 "messages": [ { "user_id: integer, @@ -38,7 +38,7 @@ messages - это срез диалога в чате телеграмма. success: [ { "user_id": integer, - "chat_id": integer, + "slice_id": integer, "reason": string } ] diff --git a/gemini_worker/src/gemini_sdk/schemas.py b/gemini_worker/src/gemini_sdk/schemas.py index 78a7938..49496dc 100644 --- a/gemini_worker/src/gemini_sdk/schemas.py +++ b/gemini_worker/src/gemini_sdk/schemas.py @@ -1,10 +1,10 @@ -from pydantic import BaseModel, PositiveInt, NegativeInt +from pydantic import UUID4, PositiveInt, NegativeInt from src.schemas import MessageFromChatSchema, BaseModelWithSerializeDatetime class ChatMessageSchema(BaseModelWithSerializeDatetime): - chat_id: int + slice_id: UUID4 messages: list[MessageFromChatSchema] @@ -14,7 +14,7 @@ class FullRequestForGeminiSchema(BaseModelWithSerializeDatetime): class SuccessChatFromAiSchema(BaseModelWithSerializeDatetime): user_id: PositiveInt | None - chat_id: NegativeInt | None + slice_id: UUID4 | None reason: str | None diff --git a/gemini_worker/src/rmq/connect.py b/gemini_worker/src/rmq/connect.py index 7229ece..7bea38b 100644 --- a/gemini_worker/src/rmq/connect.py +++ b/gemini_worker/src/rmq/connect.py @@ -1,6 +1,7 @@ from faststream.rabbit import RabbitBroker, RabbitExchange, ExchangeType, RabbitQueue +from src.gemini_sdk.schemas import ResponseFromGeminiSchema url = "amqp://test:test@localhost:9072/" @@ -19,6 +20,9 @@ base_exchange = RabbitExchange( base_queue = RabbitQueue( name="base_queue", ) +success_queue = RabbitQueue( + name="success_queue", +) message_subscriber = broker.subscriber( @@ -26,7 +30,15 @@ message_subscriber = broker.subscriber( exchange=base_exchange, ) +success_publisher = broker.publisher( + queue=success_queue, + exchange=base_exchange, + schema=ResponseFromGeminiSchema, +) + + async def init_queue_and_publisher(): await broker.declare_exchange(base_exchange) await broker.declare_queue(base_queue) + await broker.declare_queue(success_queue) diff --git a/gemini_worker/src/schemas.py b/gemini_worker/src/schemas.py index 536b6ec..566b17b 100644 --- a/gemini_worker/src/schemas.py +++ b/gemini_worker/src/schemas.py @@ -32,4 +32,5 @@ class MessageFromChatSchema(BaseModelWithSerializeDatetime): class MessagesForSendToWorkersSchema(BaseModelWithSerializeDatetime): + slice_id: UUID4 messages: list[MessageFromChatSchema] \ No newline at end of file diff --git a/gemini_worker/src/service/rmq_msg_handler.py b/gemini_worker/src/service/rmq_msg_handler.py index 7aca9c5..ff71489 100644 --- a/gemini_worker/src/service/rmq_msg_handler.py +++ b/gemini_worker/src/service/rmq_msg_handler.py @@ -1,7 +1,8 @@ from time import sleep -from src.gemini_sdk.schemas import ChatMessageSchema +from src.gemini_sdk.schemas import ChatMessageSchema, ResponseFromGeminiSchema from src.gemini_sdk.service import gemini_helper +from src.rmq.connect import success_publisher from src.schemas import MessagesForSendToWorkersSchema from src.service.limter_checker import check_current_token_limit from src.service.storage import MESSAGES_STORAGE_SCHEMA @@ -9,9 +10,9 @@ from src.service.storage import MESSAGES_STORAGE_SCHEMA async def handle_message_service( message_from_rmq: MessagesForSendToWorkersSchema, -): +) -> ResponseFromGeminiSchema | None: chat_message_schema = ChatMessageSchema( - chat_id=message_from_rmq.messages[0].chat_id, + slice_id=message_from_rmq.slice_id, messages=message_from_rmq.messages ) MESSAGES_STORAGE_SCHEMA.chats.append(chat_message_schema) @@ -24,9 +25,10 @@ async def handle_message_service( result = gemini_helper.create_request_ai( messages=MESSAGES_STORAGE_SCHEMA ) - print(result) sleep(60) MESSAGES_STORAGE_SCHEMA.chats = [] + if result.success is not None: + await success_publisher.publish(result)