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
+6
View File
@@ -0,0 +1,6 @@
[[tunnels]]
tunnel_id = "MyRoobot"
type="TelegramTunnel"
target="agent_1"
token=""
+3 -1
View File
@@ -8,4 +8,6 @@ from .open_ai_node import OpenAI_ComputeNode
from .role import AIRole,AIRoleGroup
from .workflow import Workflow
from .bus import AIBus
from .workflow_env import WorkflowEnvironment,CalenderEnvironment,CalenderEvent
from .workflow_env import WorkflowEnvironment,CalenderEnvironment,CalenderEvent
from .tunnel import AgentTunnel
from .tg_tunnel import TelegramTunnel
+2 -2
View File
@@ -243,8 +243,8 @@ class AIAgent:
result_prompt = AgentPrompt()
for msg in reversed(messages):
if msg.target == chatsession.owner_id:
result_prompt.messages.append({"role":"user","content":f"{msg.sender}:{msg.body}"})
if msg.sender == chatsession.owner_id:
result_prompt.messages.append({"role":"user","content":msg.body})
elif msg.sender == chatsession.owner_id:
result_prompt.messages.append({"role":"assistant","content":msg.body})
return result_prompt
+34
View File
@@ -0,0 +1,34 @@
from typing import List
class Contact:
def __init__(self,name:str) -> None:
self.name = name
self.tags = []
def is_zone_owner(self,zone_id=None) -> bool:
return True
def get_tags(self)->List[str]:
return self.tags
def get_name(self)->str:
return self.name
class ContactManager:
_instance = None
@classmethod
def get_instance(cls):
if cls._instance is None:
cls._instance = ContactManager()
return cls._instance
def __init__(self) -> None:
self.contacts = {}
self.contacts["liuzhicong"] = Contact("liuzhicong")
#def get_by_addr(self,addr:str) -> Contact:
# pass
def get_by_name(self,name:str) -> Contact:
return self.contacts.get(name)
+116
View File
@@ -0,0 +1,116 @@
import logging
import threading
import asyncio
import uuid
from typing import Callable
from telegram import ForceReply, Update
from telegram.ext import Application, CommandHandler, ContextTypes, MessageHandler, filters
from .tunnel import AgentTunnel
from .contact_manager import ContactManager
from .agent_message import AgentMsg
logger = logging.getLogger(__name__)
class TelegramTunnel(AgentTunnel):
@classmethod
def register_to_loader(cls):
async def load_tg_tunnel(config:dict) -> AgentTunnel:
result_tunnel = TelegramTunnel("")
if await result_tunnel.load_from_config(config):
return result_tunnel
else:
return None
AgentTunnel.register_loader("TelegramTunnel",load_tg_tunnel)
async def load_from_config(self,config:dict)->bool:
self.tg_token = config["token"]
self.target_id = config["target"]
self.tunnel_id = config["tunnel_id"]
self.type = "TelegramTunnel"
return True
def dump_to_config(self) -> dict:
pass
def __init__(self,tg_token:str) -> None:
super().__init__()
self.is_start = False
self.tg_token = tg_token
#self.tunnel_id = "tg_tunnel#" + self.app.bot.id
async def start(self) -> bool:
if self.is_start:
logger.warning(f"tunnel {self.tunnel_id} is already started")
return False
self.is_start = True
self.app:Application = Application.builder().token(self.tg_token).build()
self.app.add_handler(MessageHandler(filters.TEXT, self.on_message))
def _run_app():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
self.app.run_polling(allowed_updates=Update.ALL_TYPES)
self.poll_thread = threading.Thread(target=_run_app)
self.poll_thread.start()
return True
async def close(self) -> None:
pass
async def _process_message(self, msg: AgentMsg) -> None:
logger.warn(f"process message {msg.msg_id} from {msg.sender} to {msg.target}")
async def conver_tg_msg_to_agent_msg(self,update:Update) -> AgentMsg:
agent_msg = AgentMsg()
agent_msg.topic = "_telegram"
agent_msg.msg_id = "tg_msg#" + str(update.message.message_id) + "#" + uuid.uuid4().hex
agent_msg.target = self.target_id
agent_msg.body = update.message.text
agent_msg.create_time = update.message.date.timestamp()
#if update.message.photo is not None:
# agent_msg.body_mime = "image"
# agent_msg.body = update.message.photo[-1].get_file().download()
return agent_msg
async def on_message(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
cm = ContactManager.get_instance()
reomte_user_name = f"{update.effective_user.id}@telegram"
#contact = cm.get_by_name(update.effective_user.username)
#if contact is not None:
# reomte_user_name = contact.get_name()
#if contact is None:
# update.message.reply_text(f"{self.target_id} process message error, unknown user!")
#if not contact.is_zone_owner():
# update.message.reply_text(f"{self.target_id} process message error, you are not my owner!")
agent_msg = await self.conver_tg_msg_to_agent_msg(update)
agent_msg.sender = reomte_user_name
self.ai_bus.register_message_handler(reomte_user_name, self._process_message)
resp_msg = await self.ai_bus.send_message(agent_msg)
if resp_msg is None:
await update.message.reply_text(f"{self.target_id} process message error")
else:
if resp_msg.body_mime is None:
await update.message.reply_text(resp_msg.body)
else:
if resp_msg.body_mime.startswith("image"):
photo_file = open(resp_msg.body,"rb")
if photo_file:
await update.message.reply_photo(resp_msg.body)
else:
await update.message.reply_text(resp_msg.body)
else:
await update.message.reply_text(resp_msg.body)
+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
+10 -2
View File
@@ -4,6 +4,7 @@ import sys
import os
import logging
import re
import toml
from typing import Any, Optional, TypeVar, Tuple, Sequence
import argparse
@@ -26,12 +27,12 @@ shell_style = Style.from_dict({
directory = os.path.dirname(__file__)
sys.path.append(directory + '/../../')
from aios_kernel import Workflow,AIAgent,AgentMsg,AgentMsgStatus,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession
from aios_kernel import Workflow,AIAgent,AgentMsg,AgentMsgStatus,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession,AgentTunnel,TelegramTunnel,CalenderEnvironment,Environment
sys.path.append(directory + '/../../component/')
from agent_manager import AgentManager
from workflow_manager import WorkflowManager
from aios_kernel import CalenderEnvironment,Environment
class AIOS_Shell:
@@ -74,6 +75,13 @@ class AIOS_Shell:
ComputeKernel().add_compute_node(open_ai_node)
AIBus().get_default_bus().register_unhandle_message_handler(self._handle_no_target_msg)
AIBus().get_default_bus().register_message_handler(self.username,self._user_process_msg)
TelegramTunnel.register_to_loader()
tunnels_config_path = os.path.abspath(directory + "/../../../rootfs/tunnels.cfg.toml")
tunnel_config = toml.load(tunnels_config_path)
await AgentTunnel.load_all_tunnels_from_config(tunnel_config["tunnels"])
return True
-223
View File
@@ -1,223 +0,0 @@
import asyncio
import sys
import os
import logging
import re
from typing import Any, Optional, TypeVar, Tuple, Sequence
import argparse
from telegram import ForceReply, Update
from telegram.ext import Application, CommandHandler, ContextTypes, MessageHandler, filters
directory = os.path.dirname(__file__)
sys.path.append(directory + '/../../')
from aios_kernel import Workflow,AIAgent,AgentMsg,AgentMsgStatus,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession
sys.path.append(directory + '/../../component/')
from agent_manager import AgentManager
from workflow_manager import WorkflowManager
from aios_kernel import CalenderEnvironment,Environment
import logging
# Enable logging
logging.basicConfig(
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
)
# set higher logging level for httpx to avoid all GET and POST requests being logged
logging.getLogger("httpx").setLevel(logging.WARNING)
logger = logging.getLogger(__name__)
class AIOS_Shell:
def __init__(self,username:str) -> None:
self.username = username
self.current_target = "_"
self.current_topic = "default"
async def _handle_no_target_msg(self,bus:AIBus,msg:AgentMsg) -> bool:
target_id = msg.target.split(".")[0]
agent : AIAgent = await AgentManager().get(target_id)
if agent is not None:
bus.register_message_handler(target_id,agent._process_msg)
return True
a_workflow = await WorkflowManager().get_workflow(target_id)
if a_workflow is not None:
bus.register_message_handler(target_id,a_workflow._process_msg)
return True
return False
async def is_agent(self,target_id:str) -> bool:
agent : AIAgent = await AgentManager().get(target_id)
if agent is not None:
return True
else:
return False
async def initial(self) -> bool:
cal_env = CalenderEnvironment("calender")
cal_env.start()
Environment.set_env_by_id("calender",cal_env)
AgentManager().initial(os.path.abspath(directory + "/../../../rootfs/"))
WorkflowManager().initial(os.path.abspath(directory + "/../../../rootfs/workflows/"))
open_ai_node = OpenAI_ComputeNode()
open_ai_node.start()
ComputeKernel().add_compute_node(open_ai_node)
AIBus().get_default_bus().register_unhandle_message_handler(self._handle_no_target_msg)
return True
def get_version(self) -> str:
return "0.0.1"
async def send_msg(self,msg:str,target_id:str,topic:str,sender:str = None) -> str:
agent_msg = AgentMsg()
agent_msg.set(sender,target_id,msg)
agent_msg.topic = topic
resp = await AIBus.get_default_bus().send_message(target_id,agent_msg)
if resp is not None:
return resp.body
else:
return "error!"
async def install_workflow(self,workflow_id:Workflow) -> None:
pass
async def call_func(self,func_name, args):
match func_name:
case 'send':
target_id = args[0]
msg_content = args[1]
topic = args[2]
resp = await self.send_msg(msg_content,target_id,topic,self.username)
return resp
case 'open':
if len(args) >= 1:
target_id = args[0]
if len(args) >= 2:
topic = args[1]
self.current_target = target_id
self.current_topic = topic
show_text = f"current session switch to {topic}@{target_id}"
return show_text
case 'history':
num = 10
offset = 0
if args is not None:
if len(args) >= 1:
num = args[0]
if len(args) >= 2:
offset = args[1]
db_path = ""
if await self.is_agent(self.current_target):
db_path = AgentManager().db_path
else:
db_path = WorkflowManager().db_file
chatsession:AIChatSession = AIChatSession.get_session(self.current_target,f"{self.username}#{self.current_topic}",db_path,False)
if chatsession is not None:
msgs = chatsession.read_history(num,offset)
result_str = ""
for msg in reversed(msgs):
result_str += f"{msg.sender} >>> {msg.body}"
result_str += f"\n-------------------\n"
return result_str
return f"chatsession not found"
case 'help':
return "help info "
shell:AIOS_Shell = None
async def _init_shell():
global shell
if shell is None:
shell = AIOS_Shell("user")
await shell.initial()
return
def parse_function_call(func_string):
match = re.search(r'\s*(\w+)\s*\(\s*(.*)\s*\)\s*', func_string)
if not match:
return None
func_name = match.group(1)
params_string = match.group(2).strip()
params = re.split(r'\s*,\s*(?=(?:[^"]*"[^"]*")*[^"]*$)', params_string)
params = [param.strip('"') for param in params]
if len(params[0]) == 0:
params = None
return func_name, params
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
"""Send a message when the command /start is issued."""
await _init_shell()
user = update.effective_user
await update.message.reply_html(
rf"Hi {user.mention_html()}!",
reply_markup=ForceReply(selective=True),
)
async def on_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
await _init_shell()
cmd_msg = update.message.text[1:]
logger.info("command:" + cmd_msg)
func_call = parse_function_call(cmd_msg)
if func_call:
shell.username = update.effective_user
show_text = await shell.call_func(func_call[0], func_call[1])
await update.message.reply_text(show_text)
else:
await update.message.reply_text("command not found!")
async def on_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
await _init_shell()
shell.username = update.effective_user
resp = await shell.send_msg(update.message.text,shell.current_target,shell.current_topic,shell.username)
await update.message.reply_text(resp)
def main() -> None:
"""Start the bot."""
# Create the Application and pass it your bot's token.
tg_robot_token = os.getenv("TG_ROBOT_TOKEN")
if tg_robot_token is None:
print("TG_ROBOT_TOKEN not found!")
tg_robot_token = input("input your telegram robot token:")
if tg_robot_token is None:
return
application = Application.builder().token(tg_robot_token).build()
# on different commands - answer in Telegram
application.add_handler(CommandHandler("start", start))
application.add_handler(CommandHandler("help", on_command))
application.add_handler(CommandHandler("send", on_command))
application.add_handler(CommandHandler("open", on_command))
application.add_handler(CommandHandler("history", on_command))
# on non command i.e message - echo the message on Telegram
application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, on_message))
# Run the bot until the user presses Ctrl-C
application.run_polling(allowed_updates=Update.ALL_TYPES)
if __name__ == "__main__":
main()