Merge pull request #53 from streetycat/MVP

add compute node for llama

.gitignore

@@ -2,4 +2,9 @@
 *.pyc
 rootfs/email/config.local.toml
 rootfs/data
-venv
+venv
+
+aios_shell.log
+history.txt
+math_school_env.db
+workflows.db

rootfs/agents/math_teacher/agent.toml

@@ -1,5 +1,6 @@
 instance_id = "math_teacher"
 fullname = "the one"
+llm_model_name = "gpt-4-0613"
 [[prompt]]
 role = "system"
 content = "你是精通数学的老师"

src/aios_kernel/__init__.py

@@ -9,6 +9,7 @@ from .role import AIRole,AIRoleGroup
 from .workflow import Workflow
 from .bus import AIBus
 from .workflow_env import WorkflowEnvironment,CalenderEnvironment,CalenderEvent
+from .local_llama_compute_node import LocalLlama_ComputeNode
 from .whisper_node import WhisperComputeNode
 from .google_text_to_speech_node import GoogleTextToSpeechNode
 from .tunnel import AgentTunnel
@@ -17,3 +18,4 @@ from .email_tunnel import EmailTunnel
 from .storage import ResourceLocation,AIStorage,UserConfig,UserConfigItem
 
 AIOS_Version = "0.5.1, build 2023-9-17"
+

src/aios_kernel/compute_kernel.py

@@ -60,7 +60,7 @@ class ComputeKernel:
 
     def _schedule(self, task) -> ComputeNode:
         for node in self.compute_nodes.values():
-            if node.is_support(task.task_type) is True:
+            if node.is_support(task) is True:
                 return node
         logger.warning(
             f"task {task.display()} is not support by any compute node")

src/aios_kernel/compute_node.py

@@ -28,7 +28,7 @@ class ComputeNode(ABC):
         pass
 
     @abstractmethod
-    def is_support(self, task_type: ComputeTaskType) -> bool:
+    def is_support(self, task: ComputeTask) -> bool:
         pass
 
     @abstractmethod
@@ -41,10 +41,9 @@ class ComputeNode(ABC):
     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
+        return True

src/aios_kernel/environment.py

@@ -109,7 +109,7 @@ class Environment:
     def get_value(self,key:str) -> Optional[str]:
         handler = self.get_handlers.get(key)
         if handler is not None:
-            return handler(key)
+            return handler()
     
         s = self.values.get(key)
         if isinstance(s,str):

src/aios_kernel/local_llama_compute_node.py

@@ -0,0 +1,91 @@
+
+import logging
+import requests
+from typing import Optional, List
+from pydantic import BaseModel
+
+from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType
+from .queue_compute_node import Queue_ComputeNode
+
+logger = logging.getLogger(__name__)
+
+"""
+This is a custom implementation, it should be redesigned.
+"""
+
+class LocalLlama_ComputeNode(Queue_ComputeNode):
+    async def execute_task(self, task: ComputeTask) -> {
+        "content": str,
+        "message": str,
+        "state": ComputeTaskState,
+        "error": {
+            "code": int,
+            "message": str,
+        }
+    }:
+        class GenerateResponse(BaseModel):
+            error: Optional[int]
+            msg: Optional[str]
+            results: Optional[List[str]]
+
+        try:
+            prompt_msgs = []
+            for prompt in task.params["prompts"]:
+                prompt_msgs.append(prompt["content"])
+                
+            body = {
+                "prompts": prompt_msgs
+            }
+            
+            response = requests.post("http://aigc:7880/generate", json = body, verify=False, headers={"Content-Type": "application/json"})
+            response.close()
+
+            logger.info(f"LocalLlama_ComputeNode task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
+
+            if response.status_code != 200:
+                return {
+                    "state": ComputeTaskState.ERROR,
+                    "error": {
+                        "code": response.status_code,
+                        "message": "http request failed: " + response.status_code
+                    }
+                }
+            else:
+                resp = response.json()
+                if "error" in resp:
+                    return {
+                        "state": ComputeTaskState.ERROR,
+                        "error": {
+                            "code": resp["error"],
+                            "message": "local llama failed:" + resp["msg"]
+                        }
+                    }
+                else:
+                    return {
+                        "state": ComputeTaskState.DONE,
+                        "content": str(resp["results"]),
+                        "message": str(resp["results"])
+                    }
+        except Exception as err:
+            import traceback
+            logger.error(f"{traceback.format_exc()}, error: {err}")
+            
+            return {
+                "state": ComputeTaskState.ERROR,
+                "error": {
+                    "code": -1,
+                    "message": "unknown exception: " + str(err)
+                }
+            }
+
+    def display(self) -> str:
+        return f"LocalLlama_ComputeNode: {self.node_id}"
+
+    def get_capacity(self):
+        pass
+
+    def is_support(self, task: ComputeTask) -> bool:
+        return task.task_type == ComputeTaskType.LLM_COMPLETION and (not task.params["model_name"] or task.params["model_name"] == "llama")
+
+    def is_local(self) -> bool:
+        return True

src/aios_kernel/open_ai_node.py

@@ -64,6 +64,7 @@ class OpenAI_ComputeNode(ComputeNode):
         prompts = task.params["prompts"]
 
         logger.info(f"call openai {mode_name} prompts: {prompts}")
+
         if task.params.get("inner_functions") is None:
             resp = openai.ChatCompletion.create(model=mode_name,
                                             messages=prompts,
@@ -75,6 +76,7 @@ class OpenAI_ComputeNode(ComputeNode):
                                                 functions=task.params["inner_functions"],
                                                 max_tokens=task.params["max_token_size"],
                                                 temperature=0.7) # TODO: add temperature to task params?
+
         
         logger.info(f"openai response: {resp}")
         
@@ -123,8 +125,8 @@ class OpenAI_ComputeNode(ComputeNode):
     def get_capacity(self):
         pass
 
-    def is_support(self, task_type: ComputeTaskType) -> bool:
-        return task_type == ComputeTaskType.LLM_COMPLETION
+    def is_support(self, task: ComputeTask) -> bool:
+        return task.task_type == ComputeTaskType.LLM_COMPLETION and (not task.params["model_name"] or task.params["model_name"] == "gpt-4-0613")
 
     def is_local(self) -> bool:
         return False

src/aios_kernel/queue_compute_node.py

@@ -0,0 +1,69 @@
+
+import asyncio
+from asyncio import Queue
+import logging
+from abc import abstractmethod
+
+from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
+from .compute_node import ComputeNode
+
+logger = logging.getLogger(__name__)
+
+class Queue_ComputeNode(ComputeNode):
+    def __init__(self):
+        super().__init__()
+        self.task_queue = Queue()
+
+    @abstractmethod
+    async def execute_task(self, task: ComputeTask) -> {
+        "content": str,
+        "message": str,
+        "state": ComputeTaskState,
+        "error": {
+            "code": int,
+            "message": str,
+        }
+    }:
+        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
+        
+        resp = await self.execute_task(task)
+
+        result = ComputeTaskResult()
+        result.set_from_task(task)
+
+        task.state = resp["state"]
+
+        if task.state == ComputeTaskState.ERROR:
+            task.error_str = resp["error"]["message"]
+
+
+        result.worker_id = self.node_id
+        result.result_str = resp["content"]
+        result.result_message = resp["message"]
+
+        return result
+
+    def start(self):
+        async def _run_task_loop():
+            while True:
+                task = await self.task_queue.get()
+                logger.info(f"{self.display()} get task: {task.display()}")
+                result = await self._run_task(task)
+                if result is not None:
+                    task.result = result
+
+        asyncio.create_task(_run_task_loop())
+
+
+    def get_task_state(self, task_id: str):
+        pass

src/aios_kernel/stability_node.py

@@ -134,8 +134,8 @@ class Stability_ComputeNode(ComputeNode):
     def get_capacity(self):
         pass
 
-    def is_support(self, task_type: ComputeTaskType) -> bool:
-        return task_type == ComputeTaskType.TEXT_2_IMAGE
+    def is_support(self, task: ComputeTask) -> bool:
+        return task.task_type == ComputeTaskType.TEXT_2_IMAGE
 
     def is_local(self) -> bool:
         return False

src/aios_kernel/workflow_env.py

@@ -58,12 +58,8 @@ class CalenderEnvironment(Environment):
     def stop(self):
         self.is_run = False
 
-    def get_now(self,key)->str:
-        now = datetime.now()
-        formatted_time = now.strftime('%Y-%m-%d %H:%M:%S')
-        return formatted_time    
-    
     async def _get_now(self) -> str:
+
         now = datetime.now()
         formatted_time = now.strftime('%Y-%m-%d %H:%M:%S')
         return formatted_time

src/service/aios_shell/aios_shell.py

@@ -20,7 +20,9 @@ from prompt_toolkit.styles import Style
 
 directory = os.path.dirname(__file__)
 sys.path.append(directory + '/../../')
-from aios_kernel import AIOS_Version,UserConfigItem,AIStorage,Workflow,AIAgent,AgentMsg,AgentMsgStatus,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession,AgentTunnel,TelegramTunnel,CalenderEnvironment,Environment,EmailTunnel
+
+from aios_kernel import AIOS_Version,UserConfigItem,AIStorage,Workflow,AIAgent,AgentMsg,AgentMsgStatus,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession,AgentTunnel,TelegramTunnel,CalenderEnvironment,Environment,EmailTunnel,LocalLlama_ComputeNode
+
 
 sys.path.append(directory + '/../../component/')
 from agent_manager import AgentManager
@@ -87,6 +89,11 @@ class AIOS_Shell:
             return False
         
         ComputeKernel.get_instance().add_compute_node(open_ai_node)
+        
+        llama_ai_node = LocalLlama_ComputeNode()
+        llama_ai_node.start()
+        ComputeKernel().add_compute_node(llama_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)