Add tunnel support ,and implment TelegramTunnel. Now we can use telegram bot connect to installed agent/workflow.

This commit is contained in:
Liu Zhicong
2023-09-15 17:35:12 -07:00
parent c4d97942cb
commit ef9ab83e63
8 changed files with 234 additions and 228 deletions
+63
View File
@@ -0,0 +1,63 @@
from abc import ABC, abstractmethod
import logging
from typing import Coroutine
from .agent_message import AgentMsg
from .bus import AIBus
logger = logging.getLogger(__name__)
class AgentTunnel(ABC):
_all_loader = {}
_all_tunnels = {}
@classmethod
def register_loader(cls,tunnel_type:str,loader:Coroutine) -> None:
cls._all_loader[tunnel_type] = loader
@classmethod
async def load_all_tunnels_from_config(cls,config:dict) -> None:
for tunnel_config in config:
loader = cls._all_loader.get(tunnel_config["type"])
if loader is not None:
tunnel = await loader(tunnel_config)
if tunnel is not None:
cls._all_tunnels[tunnel.tunnel_id] = tunnel
tunnel.connect_to(AIBus.get_default_bus(),tunnel.target_id)
await tunnel.start()
else:
logger.error(f"load tunnel {tunnel_config['tunnel_id']} failed")
else:
logger.error(f"load tunnel {tunnel_config['type']} failed,loader not found")
def __init__(self) -> None:
super().__init__()
self.tunnel_id = None
self.target_id = None
self.target_type = None
self.ai_bus = None
self.is_connected = False
def connect_to(self, ai_bus:AIBus,target_id: str) -> None:
"""
Connect to the agent with the given id
"""
if self.is_connected:
logger.warning(f"tunnel {self.tunnel_id} is already connected to {self.target_id}")
return
self.target_id = target_id
self.target_type = "agent"
self.ai_bus = ai_bus
self.is_connected = True
@abstractmethod
async def start(self) -> bool:
pass
@abstractmethod
async def close(self) -> None:
pass
@abstractmethod
async def _process_message(self, msg: AgentMsg) -> None:
pass