local llama

rootfs/agents/math_teacher/agent.toml

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

src/aios_kernel/local_llama_compute_node.py

@@ -1,10 +1,12 @@
 
+import json
 import logging
 import requests
 from typing import Optional, List
 from pydantic import BaseModel
+from llama_cpp import Llama
 
-from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType
+from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
 from .queue_compute_node import Queue_ComputeNode
 
 logger = logging.getLogger(__name__)
@@ -14,69 +16,64 @@ This is a custom implementation, it should be redesigned.
 """
 
 class LocalLlama_ComputeNode(Queue_ComputeNode):
-    async def execute_task(self, task: ComputeTask) -> {
+    def __init__(self, model_path: str, model_name: str):
-        "content": str,
+        super().__init__()
-        "message": str,
+        self.model_path = model_path
-        "state": ComputeTaskState,
+        self.model_name = model_name
-        "error": {
+        self.llm = Llama(model_path=model_path)
-            "code": int,
-            "message": str,
-        }
-    }:
-        class GenerateResponse(BaseModel):
-            error: Optional[int]
-            msg: Optional[str]
-            results: Optional[List[str]]
 
-        try:
+    async def execute_task(self, task: ComputeTask) -> ComputeTaskResult:
-            prompt_msgs = []
+        match task.task_type:
-            for prompt in task.params["prompts"]:
+            case ComputeTaskType.TEXT_EMBEDDING:
-                prompt_msgs.append(prompt["content"])
+                model_name = task.params["model_name"]
+                input = task.params["input"]
+                logger.info(f"call openai {model_name} input: {input}")
 
-            body = {
+                embedding = self.llm.embed(input=input)
-                "prompts": prompt_msgs
-            }
                 
-            response = requests.post("http://aigc:7880/generate", json = body, verify=False, headers={"Content-Type": "application/json"})
+                logger.info(f"local-llama({self.model_path}) response: {resp}")
-            response.close()
 
-            logger.info(f"LocalLlama_ComputeNode task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
+                result = ComputeTaskResult()    
+                result.set_from_task(task)
+                result.result = embedding
 
-            if response.status_code != 200:
+                return result
-                return {
+            case ComputeTaskType.LLM_COMPLETION:
-                    "state": ComputeTaskState.ERROR,
+                mode_name = task.params["model_name"]
-                    "error": {
+                prompts = task.params["prompts"]
-                        "code": response.status_code,
+                max_token_size = task.params.get("max_token_size")
-                        "message": "http request failed: " + str(response.status_code)
+                llm_inner_functions = task.params.get("inner_functions")
-                    }
+                if max_token_size is None:
-                }
+                    max_token_size = 4000
-            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 {
+                logger.info(f"local-llama({self.model_path}) prompts: {prompts}")
-                "state": ComputeTaskState.ERROR,
+
-                "error": {
+                resp = self.llm.create_chat_completion(model=mode_name,
-                    "code": -1,
+                                                messages=prompts,
-                    "message": "unknown exception: " + str(err)
+                                                functions=llm_inner_functions, # function has not support?
-                }
+                                                max_tokens=max_token_size,
-            }
+                                                temperature=0.7) # TODO: add temperature to task params?
+
+            
+                logger.info(f"local-llama({self.model_path}) response: {json.dumps(resp, indent=4)}")
+
+                result = ComputeTaskResult()
+                result.set_from_task(task)
+
+                status_code = resp["choices"][0]["finish_reason"]
+                match status_code:
+                    case "function_call":
+                        task.state = ComputeTaskState.DONE
+                    case "stop":
+                        task.state = ComputeTaskState.DONE
+                    case _:
+                        task.state = ComputeTaskState.ERROR
+                        task.error_str = f"The status code was {status_code}."
+                        return None
+
+                result.result_str = resp["choices"][0]["message"]["content"]
+                result.result_message = resp["choices"][0]["message"]
+                return result
 
     async def initial(self) -> bool:
         return True
@@ -88,7 +85,7 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
         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")
+        return (task.task_type == ComputeTaskType.TEXT_EMBEDDING or task.task_type == ComputeTaskType.LLM_COMPLETION) and (not task.params["model_name"] or task.params["model_name"] == self.model_name)
 
     def is_local(self) -> bool:
         return True

src/aios_kernel/queue_compute_node.py

@@ -16,15 +16,7 @@ class Queue_ComputeNode(ComputeNode):
         self.is_start = False
 
     @abstractmethod
-    async def execute_task(self, task: ComputeTask) -> {
+    async def execute_task(self, task: ComputeTask) -> ComputeTaskResult:
-        "content": str,
-        "message": str,
-        "state": ComputeTaskState,
-        "error": {
-            "code": int,
-            "message": str,
-        }
-    }:
         pass
 
     async def push_task(self, task: ComputeTask, proiority: int = 0):
@@ -37,22 +29,12 @@ class Queue_ComputeNode(ComputeNode):
     async def _run_task(self, task: ComputeTask):
         task.state = ComputeTaskState.RUNNING
         
-        resp = await self.execute_task(task)
+        result = await self.execute_task(task)
-
+        if result is not None:
-        result = ComputeTaskResult()
-
-        result.worker_id = self.node_id
-        task.state = resp["state"]
-        
-        if task.state == ComputeTaskState.ERROR:
-            result.result_code = ComputeTaskResultCode.ERROR
-            task.error_str = resp["error"]["message"]
-        else:
-            result.result_code = ComputeTaskResultCode.OK
-            result.result_str = resp["content"]
-            result.result_message = resp["message"]
-
             result.set_from_task(task)
+            result.worker_id = self.node_id
+        else:
+            task.state = ComputeTaskState.ERROR
             
         return result