Refactor chatsession impl, fix bug.

src/aios_kernel/bus.py

@@ -1,5 +1,5 @@
-from typing import Any
-from .agent_message import AgentMsg,AgentMsgState
+from typing import Coroutine,Dict,Any
+from .agent_message import AgentMsg,AgentMsgStatus
 import asyncio
 from asyncio import Queue
 
@@ -8,19 +8,32 @@ import logging
 logger = logging.getLogger(__name__)
 
 class AIBusHandler:
-    def __init__(self,handler:Any) -> None:
+    def __init__(self,handler:Coroutine,owner_bus,enable_defualt_proc=True) -> None:
         self.handler = handler
         self.working_task = None
-        self.results = {}
+        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
         
-        return await self.handler(msg)
+        if self.enable_defualt_proc:
+            # do default process
+            if msg.rely_msg_id is not None:
+                self.results[msg.rely_msg_id] = msg
+                return None
+            
 
+        resp_msg = await self.handler(msg)
+        if self.enable_defualt_proc:
+            if resp_msg is not None:
+                await self.owner_bus.post_message(resp_msg,False)
 
+        return resp_msg
+            
 class AIBus:
     _instance = None
     @classmethod
@@ -30,11 +43,13 @@ class AIBus:
         return cls._instance
 
     def __init__(self) -> None:
-        self.handlers = {}
-        self.unhandle_handler = None
-        
-    async def post_message(self,target_id,msg:AgentMsg,use_unhandle=True) -> bool:
-        target_id = target_id.split(".")[0]
+        self.handlers:Dict[AIBusHandler] = {}
+        self.unhandle_handler:Coroutine = None
+
+
+    async def post_message(self,msg:AgentMsg,use_unhandle=True) -> bool:
+        target_id = msg.target.split(".")[0]
+
         handler = self.handlers.get(target_id)
         if handler:
             handler.queue.put_nowait(msg)
@@ -44,46 +59,39 @@ class AIBus:
         if use_unhandle:
             if self.unhandle_handler is not None:
                 if await self.unhandle_handler(self,msg):
-                    return await self.post_message(target_id,msg,False)
+                    return await self.post_message(msg,False)
       
         logger.warn(f"post message to {msg.target} failed!,target not found")
         return False
 
-    def resp_message(self,my_id:str,org_msg_id:str,resp:AgentMsg) -> None:
-        handler = self.handlers.get(my_id)
-        if handler is None:
-            return None
-        handler.results[org_msg_id] = resp
+    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 get_message_resp(self,name:str,msg_id:str) -> AgentMsg:
-        handler = self.handlers.get(name)
-        if handler is None:
+    async def send_message(self,msg:AgentMsg) -> AgentMsg:
+        sender_id = msg.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
         
-        return handler.results.get(msg_id)
-
-    async def send_message(self,target_id:str,msg:AgentMsg) -> AgentMsg:
-        target_id = target_id.split(".")[0]
-        post_result = await self.post_message(target_id,msg)
+        post_result = await self.post_message(msg)
         if post_result is False:
             return None
-        
-        handler = self.handlers.get(target_id)
-        if handler is None:
-            return None
-                
+
         retry_times = 0
         while True:
-            resp = handler.results.get(msg.id)
+            resp = sender_handler.results.get(msg.msg_id)
             if resp is not None:
                 msg.resp_msg = resp
-                msg.state = AgentMsgState.RESPONSED
+                msg.status = AgentMsgStatus.RESPONSED
+                del sender_handler.results[msg.msg_id]
                 return resp
             
             await asyncio.sleep(0.2)
             retry_times += 1
             if retry_times > 100:
-                msg.state = AgentMsgState.ERROR
+                msg.status = AgentMsgStatus.ERROR
                 return None
             
         return None
@@ -93,7 +101,7 @@ class AIBus:
 
     # means sub
     def register_message_handler(self,handler_name:str,handler:Any) -> Queue:        
-        handler_node =  AIBusHandler(handler) 
+        handler_node =  AIBusHandler(handler,self) 
         self.handlers[handler_name] = handler_node
         return handler_node.queue
     
@@ -102,13 +110,12 @@ class AIBus:
             # Wait for a message
             message = await handler.queue.get()
 
-            try:
+            #try:
                 # Try to handle the message
-                result = await handler.handle_message(message)
-                handler.results[message.id] = result
-            except Exception as e:
+            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.id} failed! {e}")
+            #    logger.error(f"handle message {message.msg_id} failed! {e}")
                 #self.queues[name].put_nowait(message)
         
         return
@@ -126,14 +133,4 @@ class AIBus:
             logger.warn(f"handler {target_name} is already working!")
             return
         
-        handler.working_task = asyncio.create_task(self.process_queue(handler))
-
-    
-
-#send message to target logic:
-# find target handler:
-#   process_message(msg):
-#       session = get_session(msg.sender,msg.target)
-
-
-# history: open(sender,target,topic)
+        handler.working_task = asyncio.create_task(self.process_queue(handler))