Refactor the code directory structure to better suit the current complexity
This commit is contained in:
Liu Zhicong
@@ -0,0 +1,149 @@
|
||||
from typing import Coroutine,Dict,Any
|
||||
import asyncio
|
||||
from asyncio import Queue
|
||||
import logging
|
||||
|
||||
from ..proto.agent_msg import *
|
||||
from ..agent.agent_base import *
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class AIBusHandler:
|
||||
def __init__(self,handler:Coroutine,owner_bus,enable_defualt_proc=True) -> None:
|
||||
self.handler = handler
|
||||
self.working_task = None
|
||||
self.results = {} # recv resps
|
||||
self.queue:Queue = Queue()
|
||||
self.enable_defualt_proc = enable_defualt_proc
|
||||
self.owner_bus = owner_bus
|
||||
|
||||
async def handle_message(self,msg:AgentMsg) -> Any:
|
||||
if self.handler is None:
|
||||
return None
|
||||
|
||||
resp_msg = await self.handler(msg)
|
||||
if self.enable_defualt_proc:
|
||||
if resp_msg is not None:
|
||||
if resp_msg.msg_type == AgentMsgType.TYPE_GROUPMSG:
|
||||
await self.owner_bus.post_message(resp_msg,resp_msg.target)
|
||||
else:
|
||||
await self.owner_bus.post_message(resp_msg)
|
||||
|
||||
return resp_msg
|
||||
|
||||
class AIBus:
|
||||
_instance = None
|
||||
@classmethod
|
||||
def get_default_bus(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = AIBus()
|
||||
return cls._instance
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.handlers:Dict[AIBusHandler] = {}
|
||||
self.unhandle_handler:Coroutine = None
|
||||
|
||||
|
||||
async def post_message(self,msg:AgentMsg,target_id = None,use_unhandle=True) -> bool:
|
||||
if target_id is None:
|
||||
target_id =msg.target
|
||||
|
||||
target_id = target_id.split(".")[0]
|
||||
|
||||
handler = self.handlers.get(target_id)
|
||||
if handler:
|
||||
if msg.rely_msg_id is not None:
|
||||
handler.results[msg.rely_msg_id] = msg
|
||||
return None
|
||||
|
||||
handler.queue.put_nowait(msg)
|
||||
self.start_process(target_id)
|
||||
return True
|
||||
|
||||
if use_unhandle:
|
||||
if self.unhandle_handler is not None:
|
||||
if await self.unhandle_handler(self,target_id):
|
||||
return await self.post_message(msg,target_id,False)
|
||||
|
||||
logger.warn(f"post message to {msg.target} failed!,target not found")
|
||||
return False
|
||||
|
||||
async def resp_message(self,org_msg_id:str,resp:AgentMsg) -> None:
|
||||
assert resp.rely_msg_id == org_msg_id
|
||||
return await self.post_message(resp)
|
||||
|
||||
async def send_message(self,msg:AgentMsg,target_id = None, real_sender=None) -> AgentMsg:
|
||||
if real_sender is None:
|
||||
sender_id = msg.sender.split(".")[0]
|
||||
else:
|
||||
sender_id = real_sender.split(".")[0]
|
||||
|
||||
sender_handler = self.handlers.get(sender_id) # sender already register on bus
|
||||
if sender_handler is None:
|
||||
logger.warn(f"sender {sender_id} not register on AI_BUS!")
|
||||
return None
|
||||
|
||||
post_result = await self.post_message(msg,target_id)
|
||||
if post_result is False:
|
||||
return None
|
||||
|
||||
retry_times = 0
|
||||
while True:
|
||||
resp : AgentMsg = sender_handler.results.get(msg.msg_id)
|
||||
if resp is not None:
|
||||
msg.resp_msg = resp
|
||||
msg.status = AgentMsgStatus.RESPONSED
|
||||
del sender_handler.results[msg.msg_id]
|
||||
return resp
|
||||
|
||||
await asyncio.sleep(0.2)
|
||||
retry_times += 1
|
||||
if retry_times > 5*240: # default timeout is 240 sec
|
||||
msg.status = AgentMsgStatus.ERROR
|
||||
return None
|
||||
|
||||
return None
|
||||
|
||||
def register_unhandle_message_handler(self,handler:Any) -> Queue:
|
||||
self.unhandle_handler = handler
|
||||
|
||||
# means sub
|
||||
def register_message_handler(self,handler_name:str,handler:Any) -> Queue:
|
||||
handler_node = AIBusHandler(handler,self)
|
||||
if self.handlers.get(handler_name) is not None:
|
||||
logger.warn(f"handler {handler_name} already register on AI_BUS!")
|
||||
|
||||
self.handlers[handler_name] = handler_node
|
||||
return handler_node.queue
|
||||
|
||||
async def process_queue(self, handler:AIBusHandler):
|
||||
while True:
|
||||
# Wait for a message
|
||||
message = await handler.queue.get()
|
||||
|
||||
try:
|
||||
# Try to handle the message
|
||||
await handler.handle_message(message)
|
||||
except Exception as e:
|
||||
# If an error occurs, put the message back into the queue
|
||||
logger.error(f"handle message {message.msg_id} failed! {e}")
|
||||
logger.exception(e)
|
||||
raise e
|
||||
#self.queues[name].put_nowait(message)
|
||||
|
||||
return
|
||||
|
||||
def start_process(self,target_name):
|
||||
handler = self.handlers.get(target_name)
|
||||
if handler is None:
|
||||
logger.error(f"handler {target_name} not found!")
|
||||
return
|
||||
|
||||
if handler.handler is None:
|
||||
return
|
||||
|
||||
if handler.working_task is not None:
|
||||
logger.warn(f"handler {target_name} is already working!")
|
||||
return
|
||||
|
||||
handler.working_task = asyncio.create_task(self.process_queue(handler))
|
||||
@@ -0,0 +1,258 @@
|
||||
from abc import ABC, abstractmethod
|
||||
import random
|
||||
from typing import Optional
|
||||
import logging
|
||||
import asyncio
|
||||
import tiktoken
|
||||
from asyncio import Queue
|
||||
|
||||
from ..proto.compute_task import *
|
||||
from ..knowledge import ObjectID
|
||||
from ..agent.agent_base import AgentPrompt
|
||||
from .compute_node import ComputeNode
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# How to dispatch different computing tasks (some tasks may contain a large amount of state for correct execution)
|
||||
# to suitable computing nodes, achieving a balance of speed, cost, and power consumption,
|
||||
# is the CORE GOAL of the entire computing task schedule system (aios_kernel).
|
||||
|
||||
class ComputeKernel:
|
||||
_instance = None
|
||||
@classmethod
|
||||
def get_instance(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = ComputeKernel()
|
||||
return cls._instance
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.is_start = False
|
||||
self.task_queue = Queue()
|
||||
self.is_start = False
|
||||
self.compute_nodes = {}
|
||||
|
||||
def run(self, task: ComputeTask) -> None:
|
||||
# check there is compute node can support this task
|
||||
if self.is_task_support(task) is False:
|
||||
logger.error(
|
||||
f"task {task.display()} is not support by any compute node")
|
||||
return
|
||||
# add task to working_queue
|
||||
self.task_queue.put_nowait(task)
|
||||
|
||||
async def start(self):
|
||||
if self.is_start is True:
|
||||
logger.warn("compute_kernel is already start")
|
||||
return
|
||||
|
||||
self.is_start = True
|
||||
|
||||
async def _run_task_loop():
|
||||
while True:
|
||||
task = await self.task_queue.get()
|
||||
logger.info(f"compute_kernel get task: {task.display()}")
|
||||
c_node: ComputeNode = self._schedule(task)
|
||||
if c_node:
|
||||
await c_node.push_task(task)
|
||||
|
||||
logger.warn("compute_kernel is stoped!")
|
||||
|
||||
asyncio.create_task(_run_task_loop())
|
||||
|
||||
def _schedule(self, task) -> ComputeNode:
|
||||
# find all the node which supports this task
|
||||
support_nodes = []
|
||||
total_weights = 0
|
||||
|
||||
for node in self.compute_nodes.values():
|
||||
if node.is_support(task) is True:
|
||||
support_nodes.append({
|
||||
"pos": total_weights,
|
||||
"node": node
|
||||
})
|
||||
total_weights += node.weight()
|
||||
|
||||
if len(support_nodes) < 1:
|
||||
logger.warning(f"task {task.display()} is not support by any compute node")
|
||||
return None
|
||||
|
||||
# hit a random node with weight
|
||||
hit_pos = random.randint(0, total_weights - 1)
|
||||
for i in range(min(len(support_nodes) - 1, hit_pos), -1, -1):
|
||||
if support_nodes[i]["pos"] <= hit_pos:
|
||||
return support_nodes[i]["node"]
|
||||
|
||||
logger.warning(
|
||||
f"task {task.display()} is not support by any compute node")
|
||||
return None
|
||||
|
||||
def add_compute_node(self, node: ComputeNode):
|
||||
if self.compute_nodes.get(node.node_id) is not None:
|
||||
logger.warn(
|
||||
f"compute_node {node.display()} already in compute_kernel")
|
||||
return
|
||||
self.compute_nodes[node.node_id] = node
|
||||
logger.info(f"add compute_node {node.display()} to compute_kernel")
|
||||
|
||||
def disable_compute_node(self, node_id: str):
|
||||
node = self.compute_nodes.get(node_id)
|
||||
if node is None:
|
||||
logger.warn(f"compute_node {node_id} not in compute_kernel")
|
||||
return
|
||||
node.enable = False
|
||||
|
||||
def is_task_support(self, task: ComputeTask) -> bool:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def llm_num_tokens_from_text(text:str,model:str) -> int:
|
||||
try:
|
||||
encoding = tiktoken.encoding_for_model(model)
|
||||
except KeyError:
|
||||
logger.debug("Warning: model not found. Using cl100k_base encoding.")
|
||||
encoding = tiktoken.get_encoding("cl100k_base")
|
||||
|
||||
token_count = len(encoding.encode(text))
|
||||
return token_count
|
||||
|
||||
|
||||
# friendly interface for use:
|
||||
def llm_completion(self, prompt: AgentPrompt, resp_mode:str="text",mode_name: Optional[str] = None, max_token: int = 0,inner_functions = None):
|
||||
# craete a llm_work_task ,push on queue's end
|
||||
# then task_schedule would run this task.(might schedule some work_task to another host)
|
||||
task_req = ComputeTask()
|
||||
task_req.set_llm_params(prompt,resp_mode,mode_name, max_token,inner_functions)
|
||||
self.run(task_req)
|
||||
return task_req
|
||||
|
||||
async def _wait_task(self,task_req:ComputeTask, timeout=60)->ComputeTaskResult:
|
||||
async def check_timer():
|
||||
check_times = 0
|
||||
while True:
|
||||
if task_req.state == ComputeTaskState.DONE:
|
||||
break
|
||||
|
||||
if task_req.state == ComputeTaskState.ERROR:
|
||||
break
|
||||
|
||||
if timeout is not None and check_times >= timeout*2:
|
||||
task_req.state = ComputeTaskState.ERROR
|
||||
break
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
check_times += 1
|
||||
|
||||
await asyncio.create_task(check_timer())
|
||||
if task_req.result:
|
||||
return task_req.result
|
||||
else:
|
||||
time_out_result = ComputeTaskResult()
|
||||
time_out_result.result_code = ComputeTaskResultCode.TIMEOUT
|
||||
time_out_result.set_from_task(task_req)
|
||||
task_req.result = time_out_result
|
||||
return time_out_result
|
||||
|
||||
|
||||
async def do_llm_completion(self, prompt: AgentPrompt,resp_mode:str="text", mode_name: Optional[str]=None, max_token:int=0, inner_functions=None, timeout=60) -> str:
|
||||
task_req = self.llm_completion(prompt, resp_mode,mode_name, max_token,inner_functions)
|
||||
return await self._wait_task(task_req, timeout)
|
||||
|
||||
|
||||
def text_embedding(self,input:str,model_name:Optional[str] = None):
|
||||
task_req = ComputeTask()
|
||||
task_req.set_text_embedding_params(input,model_name)
|
||||
self.run(task_req)
|
||||
return task_req
|
||||
|
||||
async def do_text_embedding(self,input:str,model_name:Optional[str] = None) -> [float]:
|
||||
task_req = self.text_embedding(input,model_name)
|
||||
task_result = await self._wait_task(task_req)
|
||||
|
||||
if task_req.state == ComputeTaskState.DONE:
|
||||
return task_result.result.get("content")
|
||||
else:
|
||||
logging.warning(f"do_text_embedding error: {task_req.error_str},input: {input}")
|
||||
return None
|
||||
|
||||
def image_embedding(self,input:ObjectID,model_name:Optional[str] = None):
|
||||
task_req = ComputeTask()
|
||||
task_req.set_image_embedding_params(input,model_name)
|
||||
self.run(task_req)
|
||||
return task_req
|
||||
|
||||
async def do_image_embedding(self,input:ObjectID,model_name:Optional[str] = None) -> [float]:
|
||||
task_req = self.image_embedding(input,model_name)
|
||||
task_result = await self._wait_task(task_req)
|
||||
|
||||
if task_req.state == ComputeTaskState.DONE:
|
||||
return task_result.result.get("content")
|
||||
|
||||
return None
|
||||
|
||||
async def do_text_to_speech(self,
|
||||
input:str,
|
||||
language_code:Optional[str] = None,
|
||||
gender: Optional[str] = None,
|
||||
age: Optional[str] = None,
|
||||
voice_name: Optional[str] = None,
|
||||
tone: Optional[str] = None,
|
||||
model_name: Optional[str] = None):
|
||||
task_req = ComputeTask()
|
||||
task_req.params["text"] = input
|
||||
task_req.params["language_code"] = language_code
|
||||
task_req.params["gender"] = gender
|
||||
task_req.params["age"] = age
|
||||
task_req.params["voice_name"] = voice_name
|
||||
task_req.params["tone"] = tone
|
||||
task_req.params["model_name"] = model_name
|
||||
task_req.task_type = ComputeTaskType.TEXT_2_VOICE
|
||||
self.run(task_req)
|
||||
|
||||
task_result = await self._wait_task(task_req)
|
||||
|
||||
if task_req.state == ComputeTaskState.DONE:
|
||||
return task_result.result
|
||||
|
||||
async def do_speech_to_text(self,
|
||||
audio: str,
|
||||
model: str,
|
||||
prompt: Optional[str],
|
||||
response_format: Optional[str]):
|
||||
task_req = ComputeTask()
|
||||
task_req.params["file"] = audio
|
||||
task_req.params["model_name"] = model
|
||||
task_req.params["prompt"] = prompt
|
||||
task_req.params["response_format"] = response_format
|
||||
task_req.task_type = ComputeTaskType.VOICE_2_TEXT
|
||||
|
||||
self.run(task_req)
|
||||
|
||||
task_result = await self._wait_task(task_req)
|
||||
|
||||
if task_req.state == ComputeTaskState.DONE:
|
||||
return task_result
|
||||
|
||||
def text_2_image(self, prompt:str, model_name:Optional[str] = None, negative_prompt = None):
|
||||
task = ComputeTask()
|
||||
task.set_text_2_image_params(prompt,model_name, negative_prompt)
|
||||
self.run(task)
|
||||
return task
|
||||
|
||||
async def do_text_2_image(self, prompt:str, model_name:Optional[str] = None, negative_prompt = None) -> ComputeTaskResult:
|
||||
task = self.text_2_image(prompt,model_name, negative_prompt)
|
||||
task = await self._wait_task(task)
|
||||
|
||||
return task.result
|
||||
# if task_req.state == ComputeTaskState.DONE:
|
||||
# return None, task_result
|
||||
|
||||
def image_2_text(self, image_path: str, prompt:str, model_name:Optional[str] = None, negative_prompt = None):
|
||||
task = ComputeTask()
|
||||
task.set_image_2_text_params(image_path,prompt,model_name, negative_prompt)
|
||||
self.run(task)
|
||||
return task
|
||||
async def do_image_2_text(self, image_path: str, prompt:str, model_name:Optional[str] = None, negative_prompt = None) -> ComputeTaskResult:
|
||||
task = self.image_2_text(image_path,prompt, model_name, negative_prompt)
|
||||
task = await self._wait_task(task)
|
||||
return task.result
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from ..proto.compute_task import ComputeTask, ComputeTaskType
|
||||
|
||||
class ComputeNode(ABC):
|
||||
def __init__(self) -> None:
|
||||
self.node_id = "default"
|
||||
self.enable = True
|
||||
|
||||
@abstractmethod
|
||||
async def push_task(self, task: ComputeTask, proiority: int = 0):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def remove_task(self, task_id: str):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_task_state(self, task_id: str):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def display(self) -> str:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_capacity(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def is_support(self, task: ComputeTask) -> bool:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def is_local(self) -> bool:
|
||||
pass
|
||||
|
||||
# the hit weight when select this node in schedule
|
||||
def weight(self) -> int:
|
||||
return 1
|
||||
|
||||
def is_trusted(self) -> bool:
|
||||
return True
|
||||
|
||||
def get_fee_type(self) -> str:
|
||||
return "free"
|
||||
|
||||
class LocalComputeNode(ComputeNode):
|
||||
def display(self) -> str:
|
||||
return super().display()
|
||||
|
||||
def is_local(self) -> bool:
|
||||
return True
|
||||
@@ -0,0 +1,85 @@
|
||||
from typing import List
|
||||
import logging
|
||||
from datetime import datetime
|
||||
|
||||
from ..proto.agent_msg import AgentMsg
|
||||
from .tunnel import AgentTunnel
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
class Contact:
|
||||
def __init__(self, name, phone=None, email=None, telegram=None,added_by=None, tags=[], notes=""):
|
||||
self.name = name
|
||||
self.phone = phone
|
||||
self.email = email
|
||||
self.telegram = telegram
|
||||
self.added_by = added_by
|
||||
self.tags = tags
|
||||
self.notes = notes
|
||||
self.is_family_member = False
|
||||
self.active_tunnels = {}
|
||||
|
||||
def to_dict(self):
|
||||
return {
|
||||
"name": self.name,
|
||||
"phone": self.phone,
|
||||
"email": self.email,
|
||||
"telegram" : self.telegram,
|
||||
|
||||
"added_by": self.added_by,
|
||||
"tags": self.tags,
|
||||
"notes": self.notes,
|
||||
"now" : datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
||||
}
|
||||
|
||||
async def _process_msg(self,msg:AgentMsg):
|
||||
tunnel : AgentTunnel = self.get_active_tunnel(msg.sender)
|
||||
if tunnel is not None:
|
||||
await tunnel.post_message(msg)
|
||||
return None
|
||||
else:
|
||||
tunnel = await self.create_default_tunnel(msg.sender)
|
||||
if tunnel is not None:
|
||||
self.active_tunnels[msg.sender] = tunnel
|
||||
await tunnel.post_message(msg)
|
||||
return None
|
||||
|
||||
|
||||
logger.warn(f"contact {self.name} cann't get tunnel,post message failed!")
|
||||
|
||||
def get_active_tunnel(self,agent_id) -> AgentTunnel:
|
||||
tunnel = self.active_tunnels.get(agent_id)
|
||||
return tunnel
|
||||
|
||||
def set_active_tunnel(self,agent_id,tunnel:AgentTunnel):
|
||||
self.active_tunnels[agent_id] = tunnel
|
||||
|
||||
async def create_default_tunnel(self,agent_id:str) -> AgentTunnel:
|
||||
from .email_tunnel import EmailTunnel
|
||||
|
||||
result_tunnels = AgentTunnel.get_tunnel_by_agentid(agent_id)
|
||||
for tunnel in result_tunnels:
|
||||
if isinstance(tunnel,EmailTunnel):
|
||||
return tunnel
|
||||
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data):
|
||||
return Contact(data.get("name"), data.get("phone"), data.get("email"), data.get("telegram"),data.get("added_by"), data.get("tags"), data.get("notes"))
|
||||
|
||||
class FamilyMember(Contact):
|
||||
def __init__(self, name, relationship,phone=None, email=None,telegram=None):
|
||||
super().__init__(name, phone, email, telegram)
|
||||
self.name = name
|
||||
self.relationship = relationship
|
||||
self.is_family_member = True
|
||||
|
||||
def to_dict(self):
|
||||
result = super().to_dict()
|
||||
result["relationship"] = self.relationship
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data):
|
||||
return FamilyMember(data.get("name"),data.get("relationship"),data.get("phone"), data.get("email"),data.get("telegram"))
|
||||
@@ -0,0 +1,129 @@
|
||||
from typing import List
|
||||
import toml
|
||||
import time
|
||||
import logging
|
||||
|
||||
from datetime import datetime
|
||||
from ..proto.agent_msg import AgentMsg
|
||||
from .tunnel import AgentTunnel
|
||||
from .contact import Contact,FamilyMember
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ContactManager:
|
||||
_instance = None
|
||||
@classmethod
|
||||
def get_instance(cls,filename=None) -> "ContactManager":
|
||||
if cls._instance is None:
|
||||
cls._instance = ContactManager(str(filename))
|
||||
return cls._instance
|
||||
|
||||
def __init__(self, filename="contacts.toml"):
|
||||
self.filename = filename
|
||||
self.contacts = []
|
||||
self.family_members = []
|
||||
|
||||
self.is_auto_create_contact_from_telegram = True
|
||||
|
||||
def load_data(self):
|
||||
try:
|
||||
with open(self.filename, "r") as f:
|
||||
config = toml.load(f)
|
||||
return self.load_from_config(config)
|
||||
except FileNotFoundError:
|
||||
return {}
|
||||
|
||||
def load_from_config(self,config_data:dict):
|
||||
self.contacts = [Contact.from_dict(item) for item in config_data.get("contacts", [])]
|
||||
self.family_members = [FamilyMember.from_dict(item) for item in config_data.get("family_members", [])]
|
||||
|
||||
def save_data(self):
|
||||
data = {
|
||||
"contacts": [contact.to_dict() for contact in self.contacts],
|
||||
"family_members": [member.to_dict() for member in self.family_members]
|
||||
}
|
||||
with open(self.filename, "w") as f:
|
||||
toml.dump(data, f)
|
||||
|
||||
def set_contact(self, name:str, new_contact:Contact):
|
||||
assert name == new_contact.name
|
||||
for i, contact in enumerate(self.contacts):
|
||||
if contact.name == name:
|
||||
self.contacts[i] = new_contact
|
||||
self.save_data()
|
||||
return True
|
||||
for i, member in enumerate(self.family_members):
|
||||
if member.name == name:
|
||||
self.family_members[i] = new_contact
|
||||
self.save_data()
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def add_contact(self, name:str, new_contact:Contact):
|
||||
assert name == new_contact.name
|
||||
self.contacts.append(new_contact)
|
||||
self.save_data()
|
||||
|
||||
def remove_contact(self, name:str):
|
||||
self.contacts = [contact for contact in self.contacts if contact.name != name]
|
||||
self.save_data()
|
||||
|
||||
def find_contact_by_name(self, name:str):
|
||||
for contact in self.contacts:
|
||||
if contact.name == name:
|
||||
return contact
|
||||
|
||||
for member in self.family_members:
|
||||
if member.name == name:
|
||||
return member
|
||||
return None
|
||||
|
||||
def find_contact_by_telegram(self, telegram:str):
|
||||
for contact in self.contacts:
|
||||
if contact.telegram == telegram:
|
||||
return contact
|
||||
for member in self.family_members:
|
||||
if member.telegram == telegram:
|
||||
return member
|
||||
return None
|
||||
|
||||
def find_contact_by_email(self, email:str):
|
||||
for contact in self.contacts:
|
||||
if contact.email == email:
|
||||
return contact
|
||||
for member in self.family_members:
|
||||
if member.email == email:
|
||||
return member
|
||||
return None
|
||||
|
||||
def find_contact_by_phone(self, phone:str):
|
||||
for contact in self.contacts:
|
||||
if contact.phone == phone:
|
||||
return contact
|
||||
for member in self.family_members:
|
||||
if member.phone == phone:
|
||||
return member
|
||||
return None
|
||||
|
||||
|
||||
def add_family_member(self, name, new_member:FamilyMember):
|
||||
assert name == new_member.name
|
||||
self.family_members.append(new_member)
|
||||
self.save_data()
|
||||
|
||||
def list_contacts(self):
|
||||
return self.contacts
|
||||
|
||||
def list_family_members(self):
|
||||
return self.family_members
|
||||
|
||||
#def register_to_ai_bus(self, ai_bus:AIBus):
|
||||
# ai_bus.register_message_handler("contact_manager", self.process_msg)
|
||||
|
||||
|
||||
#async def process_msg(self,msg:AgentMsg):
|
||||
# # forword message to contact
|
||||
# pass
|
||||
@@ -0,0 +1,63 @@
|
||||
|
||||
import asyncio
|
||||
from asyncio import Queue
|
||||
import logging
|
||||
from abc import abstractmethod
|
||||
|
||||
from aios import ComputeTask, ComputeNode,ComputeTaskResult, ComputeTaskResultCode, ComputeTaskState, ComputeTaskType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class Queue_ComputeNode(ComputeNode):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.task_queue = Queue()
|
||||
self.is_start = False
|
||||
|
||||
@abstractmethod
|
||||
async def execute_task(self, task: ComputeTask)->ComputeTaskResult:
|
||||
pass
|
||||
|
||||
async def push_task(self, task: ComputeTask, proiority: int = 0):
|
||||
logger.info(f"{self.display()} push task: {task.display()}")
|
||||
self.task_queue.put_nowait(task)
|
||||
|
||||
async def remove_task(self, task_id: str):
|
||||
pass
|
||||
|
||||
async def _run_task(self, task: ComputeTask):
|
||||
task.state = ComputeTaskState.RUNNING
|
||||
|
||||
result = ComputeTaskResult()
|
||||
result.result_code = ComputeTaskResultCode.ERROR
|
||||
result.set_from_task(task)
|
||||
result.worker_id = self.node_id
|
||||
|
||||
real_result = await self.execute_task(task)
|
||||
|
||||
if real_result:
|
||||
if real_result.result_code == ComputeTaskResultCode.OK:
|
||||
task.state = ComputeTaskState.DONE
|
||||
else:
|
||||
task.state = ComputeTaskState.ERROR
|
||||
return real_result
|
||||
else:
|
||||
task.state = ComputeTaskState.ERROR
|
||||
return result
|
||||
|
||||
def start(self):
|
||||
if self.is_start is True:
|
||||
return
|
||||
self.is_start = True
|
||||
|
||||
async def _run_task_loop():
|
||||
while True:
|
||||
task = await self.task_queue.get()
|
||||
logger.info(f"openai_node get task: {task.display()}")
|
||||
await self._run_task(task)
|
||||
|
||||
asyncio.create_task(_run_task_loop())
|
||||
|
||||
|
||||
def get_task_state(self, task_id: str):
|
||||
pass
|
||||
@@ -0,0 +1,102 @@
|
||||
from abc import ABC, abstractmethod
|
||||
import logging
|
||||
from typing import Coroutine
|
||||
|
||||
from ..proto.agent_msg 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_id,tunnel_config in config.items():
|
||||
loader = cls._all_loader.get(tunnel_config["type"])
|
||||
tid = tunnel_config.get("tunnel_id")
|
||||
if tid is not None:
|
||||
if tunnel_id != tid:
|
||||
logger.warning(f"load tunnel {tunnel_id} error,{tunnel_id} != {tid} in config!")
|
||||
continue
|
||||
else:
|
||||
tunnel_config["tunnel_id"] = tunnel_id
|
||||
|
||||
if loader is not None:
|
||||
tunnel = await loader(tunnel_config)
|
||||
if tunnel is not None:
|
||||
cls._all_tunnels[tunnel_id] = tunnel
|
||||
tunnel.connect_to(AIBus.get_default_bus(),tunnel.target_id)
|
||||
await tunnel.start()
|
||||
else:
|
||||
logger.error(f"load tunnel {tunnel_id} failed")
|
||||
else:
|
||||
logger.error(f"load tunnel {tunnel_id} failed,loader not found")
|
||||
|
||||
@classmethod
|
||||
async def load_tunnel_from_config(cls,tunnel_config:dict):
|
||||
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()
|
||||
return True
|
||||
else:
|
||||
logger.error(f"load tunnel {tunnel_config['tunnel_id']} failed")
|
||||
else:
|
||||
logger.error(f"load tunnel {tunnel_config['type']} failed,loader not found")
|
||||
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
async def get_tunnel_by_agentid(cls,agent_id:str):
|
||||
result = []
|
||||
for tunnel in cls._all_tunnels.values():
|
||||
if tunnel.target_id == agent_id:
|
||||
result.append(tunnel)
|
||||
return result
|
||||
|
||||
|
||||
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
|
||||
def post_message(self, msg: AgentMsg) -> None:
|
||||
pass
|
||||
|
||||
|
||||
@abstractmethod
|
||||
async def start(self) -> bool:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def close(self) -> None:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def _process_message(self, msg: AgentMsg) -> None:
|
||||
pass
|
||||
Reference in New Issue
Block a user