diff --git a/ai_agent.py b/ai_agent.py index a16b59e..ab254d7 100644 --- a/ai_agent.py +++ b/ai_agent.py @@ -21,6 +21,10 @@ PRIVATE_CHAT_SYSTEM_PROMPT = """ """ PRIVATE_CHAT_MAX_MESSAGES = 40 +OPENROUTER_HEADERS = { + 'HTTP-Referer': 'https://ultracoder.org', + 'X-Title': 'TG/VK Chat Bot' +} @dataclass() class Message: @@ -30,13 +34,14 @@ class Message: class AiAgent: - def __init__(self, api_token: str, model: str, db: BasicDatabase): + def __init__(self, api_token: str, model: str, db: BasicDatabase, platform: str): retry_config = RetryConfig(strategy="backoff", backoff=BackoffStrategy( initial_interval=2000, max_interval=8000, exponent=2, max_elapsed_time=14000), retry_connection_errors=True) self.db = db self.model = model + self.platform = platform self.client = OpenRouter(api_key=api_token, retry_config=retry_config) async def get_group_chat_reply(self, bot_id: int, chat_id: int, @@ -55,7 +60,9 @@ class AiAgent: response = await self.client.chat.send_async( model=self.model, messages=context, - max_tokens=500 + max_tokens=500, + user=f'{self.platform}_{bot_id}_{chat_id}', + http_headers=OPENROUTER_HEADERS ) # Extract AI response @@ -85,7 +92,9 @@ class AiAgent: response = await self.client.chat.send_async( model=self.model, messages=context, - max_tokens=500 + max_tokens=500, + user=f'{self.platform}_{bot_id}_{chat_id}', + http_headers=OPENROUTER_HEADERS ) # Extract AI response @@ -133,6 +142,6 @@ class AiAgent: agent: AiAgent -def create_ai_agent(api_token: str, model: str, db: BasicDatabase): +def create_ai_agent(api_token: str, model: str, db: BasicDatabase, platform: str): global agent - agent = AiAgent(api_token, model, db) + agent = AiAgent(api_token, model, db, platform) diff --git a/tg/__main__.py b/tg/__main__.py index 6edd133..782dd2e 100644 --- a/tg/__main__.py +++ b/tg/__main__.py @@ -24,9 +24,7 @@ async def main() -> None: database.create_database(config['db_connection_string']) - create_ai_agent(config['openrouter_token'], - config['openrouter_model'], - database.DB) + create_ai_agent(config['openrouter_token'], config['openrouter_model'], database.DB, 'tg') bots: list[Bot] = [] for item in database.DB.get_bots(): diff --git a/vk/__main__.py b/vk/__main__.py index fb31f06..74c58b4 100644 --- a/vk/__main__.py +++ b/vk/__main__.py @@ -24,9 +24,7 @@ if __name__ == '__main__': database.create_database(config['db_connection_string']) - create_ai_agent(config['openrouter_token'], - config['openrouter_model'], - database.DB) + create_ai_agent(config['openrouter_token'], config['openrouter_model'], database.DB, 'vk') bot = Bot(labeler=handlers.labeler)