Support custom agent

src/aios_kernel/agent.py

@@ -12,7 +12,8 @@ import datetime
 import copy
 import sys
 
-from .agent_base import AgentMsg, AgentMsgStatus, AgentMsgType,FunctionItem,LLMResult,AgentPrompt,AgentReport,AgentTodo,AgentTodoResult,AgentWorkLog
+from .agent_base import AgentMsg, AgentMsgStatus, AgentMsgType, FunctionItem, LLMResult, AgentPrompt, AgentReport, \
+    AgentTodo, AgentTodoResult, AgentWorkLog, BaseAIAgent
 from .chatsession import AIChatSession
 from .compute_task import ComputeTaskResult,ComputeTaskResultCode
 from .ai_function import AIFunction
@@ -116,7 +117,7 @@ class AIAgentTemplete:
         return True
 
 
-class AIAgent:
+class AIAgent(BaseAIAgent):
     def __init__(self) -> None:
         self.role_prompt:AgentPrompt = None
         self.agent_prompt:AgentPrompt = None

src/aios_kernel/agent_base.py

@@ -1,4 +1,6 @@
+import abc
 import copy
+from abc import abstractmethod
 from datetime import datetime, timedelta
 import logging
 from enum import Enum
@@ -8,6 +10,7 @@ import re
 import shlex
 import json
 from typing import List
+
 from .ai_function import FunctionItem
 from .compute_task import ComputeTaskResult
 
@@ -548,6 +551,44 @@ class AgentWorkLog:
     def __init__(self) -> None:
         pass
 
-class BaseAIAgent:
-    def __init__(self) -> None:
+
+class BaseAIAgent(abc.ABC):
+    @abstractmethod
+    def get_id(self) -> str:
         pass
+
+    @abstractmethod
+    def get_llm_model_name(self) -> str:
+        pass
+
+    @abstractmethod
+    def get_max_token_size(self) -> int:
+        pass
+
+    @abstractmethod
+    def get_llm_learn_token_limit(self) -> int:
+        pass
+
+    @abstractmethod
+    async def _process_msg(self,msg:AgentMsg,workspace = None) -> AgentMsg:
+        pass
+
+
+class CustomAIAgent(BaseAIAgent):
+    def __init__(self, agent_id: str, llm_model_name: str, max_token_size: int, llm_learn_token_limit: int) -> None:
+        self.agent_id = agent_id
+        self.llm_model_name = llm_model_name
+        self.max_token_size = max_token_size
+        self.llm_learn_token_limit = llm_learn_token_limit
+
+    def get_id(self) -> str:
+        return self.agent_id
+
+    def get_llm_model_name(self) -> str:
+        return self.llm_model_name
+
+    def get_max_token_size(self) -> int:
+        return self.max_token_size
+
+    def get_llm_learn_token_limit(self) -> int:
+        return self.llm_learn_token_limit

src/component/agent_manager/agent_manager.py

@@ -1,11 +1,13 @@
-
+import importlib
 import logging
 import toml
 import os
+import sys
 import runpy
 from typing import Any, Callable, Dict, List, Optional, Union
 
 from aios_kernel import AIAgent,AIAgentTemplete,AIStorage,Environment
+from aios_kernel.agent_base import BaseAIAgent
 from package_manager import PackageEnv,PackageEnvManager,PackageMediaInfo,PackageInstallTask
 
 logger = logging.getLogger(__name__)
@@ -28,7 +30,7 @@ class AgentManager:
         self.agent_templete_env : PackageEnv = None
         self.agent_env : PackageEnv = None
         self.db_path : str = None
-        self.loaded_agent_instance : Dict[str,AIAgent] = None
+        self.loaded_agent_instance : Dict[str,BaseAIAgent] = None
 
     async def initial(self) -> None:
         system_app_dir = AIStorage.get_instance().get_system_app_dir()
@@ -93,7 +95,7 @@ class AgentManager:
     async def _load_templete_from_media(self,templete_media:PackageMediaInfo) -> AIAgentTemplete:
         pass
 
-    async def _load_agent_from_media(self,agent_media:PackageMediaInfo) -> AIAgent:
+    async def _load_agent_from_media(self,agent_media:PackageMediaInfo) -> BaseAIAgent:
         reader = self.agent_env._create_media_loader(agent_media)
         if reader is None:
             logger.error(f"create media loader for {agent_media} failed!")
@@ -125,8 +127,19 @@ class AgentManager:
                 return None
             return result_agent
         except Exception as e:
-            logger.error(f"read agent.toml cfg from {agent_media} failed! unexpected error occurred: {str(e)}")
-            return None
+            custom_agent = os.path.join(agent_media.full_path,"agent.py")
+            if not os.path.exists(custom_agent):
+                logger.error(f"read agent.toml cfg from {agent_media} failed! unexpected error occurred: {str(e)}")
+                return None
+
+            agent_name = os.path.split(agent_media.full_path)[1]
+            spec = importlib.util.spec_from_file_location(agent_name, custom_agent)
+            the_api = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(the_api)
+            if not hasattr(the_api,"Agent"):
+                logger.error(f"read agent.toml cfg from {agent_media} failed! unexpected error occurred: {str(e)}")
+                return None
+            return the_api.Agent()