Merge pull request #1 from crestalnetwork/feat/telegram-bot-mgmg-server

feat: telegram bot manager server

Author: hyacinthus

app/config/config.py

@@ -66,6 +66,12 @@ def __init__(self):
         self.cdp_api_key_name = self.load("CDP_API_KEY_NAME")
         self.cdp_api_key_private_key = self.load("CDP_API_KEY_PRIVATE_KEY")
         self.openai_api_key = self.load("OPENAI_API_KEY")
+        self.slack_token = self.load("SLACK_TOKEN")
+        self.slack_channel = self.load("SLACK_CHANNEL")
+        self.tg_base_url = self.load("TG_BASE_URL")
+        self.tg_server_host = self.load("TG_SERVER_HOST", "127.0.0.1")
+        self.tg_server_port = self.load("TG_SERVER_PORT", "8081")
+        self.tg_new_agent_poll_interval = self.load("TG_NEW_AGENT_POLL_INTERVAL", "60")
         self.twitter_endpoint_interval = int(
             self.load("TWITTER_ENDPOINT_INTERVAL", "15")
         )  # in minutes

app/entrypoints/tg.py

@@ -0,0 +1,80 @@
+import asyncio
+import logging
+import signal
+import sys
+
+from sqlmodel import Session, select
+
+from app.config.config import config
+from app.models.db import get_engine, init_db
+from app.models.agent import Agent
+from tg.bot import pool
+from tg.bot.pool import BotPool
+
+logger = logging.getLogger(__name__)
+
+
+class AgentScheduler:
+    def __init__(self, bot_pool):
+        self.bot_pool = bot_pool
+
+    def check_new_bots(self):
+        with Session(get_engine()) as db:
+            # Get all telegram agents
+            agents = db.exec(
+                select(Agent).where(
+                    Agent.telegram_enabled,
+                )
+            ).all()
+
+            new_bots = []
+            for agent in agents:
+                if agent.telegram_config["token"] not in pool._bots:
+                    agent.telegram_config["agent_id"] = agent.id
+                    new_bots.append(agent.telegram_config)
+                    logger.info("New agent with id {id} found...".format(id=agent.id))
+
+            return new_bots
+
+    async def start(self, interval):
+        logger.info("New agent addition tracking started...")
+        while True:
+            logger.info("check for new bots...")
+            await asyncio.sleep(interval)
+            if self.check_new_bots() is not None:
+                for new_bot in self.check_new_bots():
+                    await self.bot_pool.init_new_bot(
+                        new_bot["agent_id"], new_bot["kind"], new_bot["token"]
+                    )
+
+
+def run_telegram_server() -> None:
+    # Initialize database connection
+    init_db(**config.db)
+
+    # Signal handler for graceful shutdown
+    def signal_handler(signum, frame):
+        logger.info("Received termination signal. Shutting down gracefully...")
+        scheduler.shutdown()
+        sys.exit(0)
+
+    # Register signal handlers
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
+
+    logger.info("Initialize bot pool...")
+    bot_pool = BotPool(config.tg_base_url)
+
+    bot_pool.init_god_bot()
+    bot_pool.init_all_dispatchers()
+
+    scheduler = AgentScheduler(bot_pool)
+
+    loop = asyncio.new_event_loop()
+    loop.create_task(scheduler.start(int(config.tg_new_agent_poll_interval)))
+
+    bot_pool.start(loop, config.tg_server_host, int(config.tg_server_port))
+
+
+if __name__ == "__main__":
+    run_telegram_server()

app/models/agent.py

@@ -34,6 +34,11 @@ class Agent(SQLModel, table=True):
     # twitter skills require config, but not require twitter_enabled flag.
     # As long as twitter_skills is not empty, the corresponding skills will be loaded.
     twitter_skills: Optional[List[str]] = Field(sa_column=Column(ARRAY(String)))
+    telegram_enabled: bool = Field(default=False)
+    telegram_config: Optional[dict] = Field(sa_column=Column(JSONB, nullable=True))
+    # twitter skills require config, but not require twitter_enabled flag.
+    # As long as twitter_skills is not empty, the corresponding skills will be loaded.
+    telegram_skills: Optional[List[str]] = Field(sa_column=Column(ARRAY(String)))
     # crestal skills
     crestal_skills: Optional[List[str]] = Field(sa_column=Column(ARRAY(String)))
     # skills not require config

debug/create_agent.py

@@ -26,6 +26,9 @@
     twitter_enabled=False,
     twitter_config={},  # Ensure this dict structure aligns with expected config format
     twitter_skills=[],  # Confirm if no specific Twitter skills are to be enabled
+    telegram_enabled=False,
+    telegram_config={},  # Ensure this dict structure aligns with expected config format
+    telegram_skills=[],  # Confirm if no specific Telegram skills are to be enabled
     common_skills=[],  # Confirm if no common skills are to be added initially
 )
 

example.env

@@ -11,5 +11,9 @@ DB_PASSWORD=
 DB_NAME=
 DB_AUTO_MIGRATE=true
 
+TG_TOKEN_GOD_BOT=
+TG_BASE_URL=
+TG_NEW_AGENT_POLL_INTERVAL=
+
 CDP_API_KEY_NAME=
 CDP_API_KEY_PRIVATE_KEY=

poetry.lock

@@ -39,7 +39,8 @@ anyio = "^4.7.0"
 slack-sdk = "^3.34.0"
 requests = "^2.32.3"
 aws-secretsmanager-caching = "^1.1.3"
-botocore = "^1.35.81"
+botocore = "^1.35.90"
+aiogram = "^3.16.0"
 
 [tool.poetry.group.dev] 
 optional = true

pyproject.toml

@@ -0,0 +1,18 @@
+from aiogram.filters import BaseFilter
+from aiogram.types import Message
+
+
+class ChatTypeFilter(BaseFilter):
+    def __init__(self, chat_type: str | list):
+        self.chat_type = chat_type
+
+    async def __call__(self, message: Message) -> bool:
+        if isinstance(self.chat_type, str):
+            return message.chat.type == self.chat_type
+        else:
+            return message.chat.type in self.chat_type
+
+
+class GroupOnlyFilter(ChatTypeFilter):
+    def __init__(self):
+        super().__init__(["group", "supergroup"])