Compare commits

..

2 commits

4 changed files with 48 additions and 12 deletions

View file

@ -1,8 +1,8 @@
from dataclasses import dataclass from dataclasses import dataclass
from typing import List, Dict, Optional from typing import List, Dict, Optional
from openrouter import OpenRouter from openrouter import OpenRouter, RetryConfig
from openrouter.utils import BackoffStrategy
SYSTEM_PROMPT = """ SYSTEM_PROMPT = """
Ты - помощник в групповом чате. Ты - помощник в групповом чате.
@ -39,7 +39,11 @@ class AiMessage:
class AiAgent: class AiAgent:
def __init__(self, api_token: str): def __init__(self, api_token: str):
self.client = OpenRouter(api_key=api_token, timeout_ms=15000) retry_config = RetryConfig(strategy="backoff",
backoff=BackoffStrategy(
initial_interval=2000, max_interval=8000, exponent=2, max_elapsed_time=14000),
retry_connection_errors=True)
self.client = OpenRouter(api_key=api_token, retry_config=retry_config)
self.chat_contexts: Dict[int, ChatContext] = {} self.chat_contexts: Dict[int, ChatContext] = {}
async def get_reply(self, chat_id: int, chat_prompt: str, async def get_reply(self, chat_id: int, chat_prompt: str,

View file

@ -1,3 +1,4 @@
from functools import partial
from typing import Optional from typing import Optional
from aiogram import Router, F, Bot from aiogram import Router, F, Bot
@ -92,8 +93,10 @@ async def any_message_handler(message: Message):
# noinspection PyUnresolvedReferences # noinspection PyUnresolvedReferences
agent = AiAgent(message.bot.config['openrouter_token']) agent = AiAgent(message.bot.config['openrouter_token'])
await message.bot.send_chat_action(chat_id, 'typing') await message.reply(
await message.reply(await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_messages)) await utils.run_with_progress(partial(agent.get_reply, chat_id, chat_prompt, ai_message, ai_fwd_messages),
partial(message.bot.send_chat_action, chat_id, 'typing'),
interval=4))
def clear_ai_chat_context(chat_id: int): def clear_ai_chat_context(chat_id: int):
@ -114,5 +117,6 @@ async def get_ai_reply(bot: Bot, chat_id, user: User, message: str, fwd_user: Us
# noinspection PyUnresolvedReferences # noinspection PyUnresolvedReferences
agent = AiAgent(bot.config['openrouter_token']) agent = AiAgent(bot.config['openrouter_token'])
await bot.send_chat_action(chat_id, 'typing') return await utils.run_with_progress(partial(agent.get_reply, chat_id, chat_prompt, ai_message, ai_fwd_messages),
return await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_messages) partial(bot.send_chat_action(chat_id, 'typing')),
interval=4)

View file

@ -1,5 +1,6 @@
import asyncio
from calendar import timegm from calendar import timegm
from typing import Optional from typing import Awaitable, Callable, Coroutine, Optional
from pymorphy3 import MorphAnalyzer from pymorphy3 import MorphAnalyzer
from time import gmtime from time import gmtime
@ -21,3 +22,26 @@ def full_name(first_name: str, last_name: Optional[str]) -> str:
if last_name is not None: if last_name is not None:
return f"{first_name} {last_name}" return f"{first_name} {last_name}"
return first_name return first_name
async def run_with_progress(main_func: Callable[[], Coroutine], progress_func: Callable[[], Awaitable], interval: int):
completion_event = asyncio.Event()
async def progress():
while not completion_event.is_set():
await progress_func()
wait_event_task = asyncio.create_task(completion_event.wait())
wait_timer_task = asyncio.create_task(asyncio.sleep(interval))
await asyncio.wait([wait_event_task, wait_timer_task],
return_when=asyncio.FIRST_COMPLETED)
if completion_event.is_set():
wait_timer_task.cancel()
progress_task = asyncio.create_task(progress())
main_task = asyncio.create_task(main_func())
result = await main_task
completion_event.set()
await progress_task
return result

View file

@ -1,4 +1,5 @@
import re import re
from functools import partial
from typing import Optional, Tuple, List from typing import Optional, Tuple, List
from vkbottle import API from vkbottle import API
@ -87,8 +88,10 @@ async def any_message_handler(message: Message):
# noinspection PyUnresolvedReferences # noinspection PyUnresolvedReferences
agent = AiAgent(message.ctx_api.config['openrouter_token']) agent = AiAgent(message.ctx_api.config['openrouter_token'])
await message.ctx_api.messages.set_activity(peer_id=chat_id, type='typing') await message.reply(
await message.reply(await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_messages)) await utils.run_with_progress(partial(agent.get_reply, chat_id, chat_prompt, ai_message, ai_fwd_messages),
partial(message.ctx_api.messages.set_activity, peer_id=chat_id, type='typing'),
interval=4))
def clear_ai_chat_context(chat_id: int): def clear_ai_chat_context(chat_id: int):
@ -112,5 +115,6 @@ async def get_ai_reply(api: API, chat_id, message: Tuple[int, str], fwd_messages
# noinspection PyUnresolvedReferences # noinspection PyUnresolvedReferences
agent = AiAgent(api.config['openrouter_token']) agent = AiAgent(api.config['openrouter_token'])
await api.messages.set_activity(peer_id=chat_id, type='typing') return await utils.run_with_progress(partial(agent.get_reply, chat_id, chat_prompt, ai_message, ai_fwd_messages),
return await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_messages) partial(api.messages.set_activity, peer_id=chat_id, type='typing'),
interval=4)