gyt/opendan
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #70 from streetycat/MVP

Browse Source 
Add the llama nodes at local dynamically
This commit is contained in:
Liu Zhicong
2023-09-28 09:17:39 -07:00
committed by GitHub
parent 14a76476ef 861096fb0c
commit 8ccab139aa
5 changed files with 290 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/aios_kernel/compute_node_config.py
+87
View File
@@ -0,0 +1,87 @@
"""
Configuration for nodes:
```
├── nodes
│ └── llama
| └── 0
| | └── url
| | └── model_name
| └── 1
| └── url
| └── model_name
```
"""
import logging
from typing import List
import os
import toml
from .local_llama_compute_node import LocalLlama_ComputeNode
from .storage import AIStorage
# define singleton class knowledge pipline
class ComputeNodeConfig:
_instance = None
@classmethod
def get_instance(cls):
if cls._instance is None:
cls._instance = ComputeNodeConfig()
cls._instance.__singleton_init__()
return cls._instance
def initial(self) -> List[LocalLlama_ComputeNode]:
config_path = self.__config_path()
logging.info(f"initial nodes from {config_path}")
if os.path.exists(config_path):
self.config = toml.load(self.__config_path())
if self.config is None:
return []
nodes = []
llama_nodes_cfg = self.config["llama"]
if llama_nodes_cfg is not None:
for cfg in llama_nodes_cfg:
node = LocalLlama_ComputeNode(url=cfg["url"], model_name=cfg["model_name"])
nodes.append(node)
return nodes
return []
def save(self):
with open(self.__config_path(), "w") as f:
toml.dump(self.config, f)
def add_node(self, model_type: str, url: str, model_name: str):
if model_type == "llama":
llama_nodes_cfg = self.config.get("llama") or []
for cfg in llama_nodes_cfg:
if url == cfg["url"] and model_name == cfg["model_name"]:
return
llama_nodes_cfg.append({"url": url, "model_name": model_name})
self.config["llama"] = llama_nodes_cfg
def remove_node(self, model_type: str, url: str, model_name: str):
if model_type == "llama":
llama_nodes_cfg = self.config.get("llama") or []
for i in range(0, len(llama_nodes_cfg)):
cfg = llama_nodes_cfg[i]
if url == cfg["url"] and model_name == cfg["model_name"]:
llama_nodes_cfg.pop(i)
def list(self) -> str:
return toml.dumps(self.config)
def __singleton_init__(self):
self.config = {}
@classmethod
def __config_path(cls) -> str:
user_data_dir = AIStorage.get_instance().get_myai_dir()
return os.path.abspath(f"{user_data_dir}/etc/compute_nodes.cfg.toml")
src/aios_kernel/local_llama_compute_node.py
+133 -63
View File
@@ -1,11 +1,13 @@
import json
import logging
import requests
from typing import Optional, List
from pydantic import BaseModel
from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType
from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskResultCode, ComputeTaskState, ComputeTaskType
from .queue_compute_node import Queue_ComputeNode
from .storage import AIStorage,UserConfig
logger = logging.getLogger(__name__)
@@ -14,81 +16,149 @@ 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]]
def __init__(self, url: str, model_name: str):
super().__init__()
self.url = url
self.model_name = model_name
try:
prompt_msgs = []
for prompt in task.params["prompts"]:
prompt_msgs.append(prompt["content"])
async def execute_task(self, task: ComputeTask, result: ComputeTaskResult):
match task.task_type:
case ComputeTaskType.TEXT_EMBEDDING:
model_name = task.params["model_name"]
input = task.params["input"]
logger.info(f"call local-llama ({self.url}, {self.model_name}) {model_name} input: {input}")
self.embedding(input, result)
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: " + str(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"]
}
}
if result.result_code == ComputeTaskResultCode.OK:
task.state = ComputeTaskState.DONE
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)
}
}
task.state = ComputeTaskState.ERROR
task.error_str = result.error_str
return result
case ComputeTaskType.LLM_COMPLETION:
mode_name = task.params["model_name"]
prompts = task.params["prompts"]
logger.info(f"local-llama({self.url}, {self.model_name}) prompts: {prompts}")
self.completion(task, result)
if result.result_code == ComputeTaskResultCode.OK:
task.state = ComputeTaskState.DONE
else:
task.state = ComputeTaskState.ERROR
task.error_str = result.error_str
case _:
task.state = ComputeTaskState.ERROR
result.result_code = ComputeTaskResultCode.ERROR
task.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
result.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
return None
async def initial(self) -> bool:
return True
def display(self) -> str:
return f"LocalLlama_ComputeNode: {self.node_id}"
return f"local-llama: {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")
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
def embedding(self, input: str, result: ComputeTaskResult):
body = {
"input": input
}
try:
response = requests.post(self.url + "/v1/embeddings", json = body, verify=False, headers={"Content-Type": "application/json"})
response.close()
logger.info(f"local-llama({self.url}, {self.model_name}) task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
if response.status_code == 200:
resp = response.json()
result.result = resp["data"][0]["embedding"]
elif response.status_code == 422:
resp = response.json()
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = "http request failed: " + str(resp["detail"][0]["msg"])
else:
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = "http request failed: " + str(response.status_code)
except Exception as e:
logger.error(f"call local-llama({self.url}, {self.model_name}) run TEXT_EMBEDDING task error: {e}")
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = str(e)
return result
def completion(self, task: ComputeTask, result: ComputeTaskResult):
mode_name = task.params["model_name"]
prompts = task.params["prompts"]
max_token_size = task.params.get("max_token_size")
llm_inner_functions = task.params.get("inner_functions")
if max_token_size is None:
max_token_size = max_token_size
body = {
"messages": [],
"max_tokens": 4000
}
for prompt in prompts:
body["messages"].append({
"role": prompt["role"],
"content": prompt["content"]
})
try:
response = requests.post(self.url + "/v1/chat/completions", json = body, verify=False, headers={"Content-Type": "application/json"})
response.close()
logger.info(f"local-llama({self.url}, {self.model_name}) task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
if response.status_code == 200:
resp = response.json()
status_code = resp["choices"][0]["finish_reason"]
token_usage = resp["usage"]
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}."
result.error_str = f"The status code was {status_code}."
result.result_code = ComputeTaskResultCode.ERROR
return None
result.result_code = ComputeTaskResultCode.OK
result.result_str = resp["choices"][0]["message"]["content"]
result.result_message = resp["choices"][0]["message"]
if token_usage:
result.result_refers["token_usage"] = token_usage
logger.info(f"local-llama({self.url}, {self.model_name}) success response: {result.result_str}")
elif response.status_code == 422:
resp = response.json()
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = "http request failed: " + str(resp["detail"][0]["msg"])
else:
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = "http request failed: " + str(response.status_code)
except Exception as e:
logger.error(f"call local-llama({self.url}, {self.model_name}) run LLM_COMPLETION task error: {e}")
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = str(e)
return result
src/aios_kernel/queue_compute_node.py
+6 -23
View File
@@ -16,15 +16,7 @@ class Queue_ComputeNode(ComputeNode):
self.is_start = False
@abstractmethod
async def execute_task(self, task: ComputeTask) -> {
"content": str,
"message": str,
"state": ComputeTaskState,
"error": {
"code": int,
"message": str,
}
}:
async def execute_task(self, task: ComputeTask, result: ComputeTaskResult):
pass
async def push_task(self, task: ComputeTask, proiority: int = 0):
@@ -36,24 +28,15 @@ class Queue_ComputeNode(ComputeNode):
async def _run_task(self, task: ComputeTask):
task.state = ComputeTaskState.RUNNING
resp = await self.execute_task(task)
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.result_code = ComputeTaskResultCode.ERROR
result.set_from_task(task)
result.worker_id = self.node_id
await self.execute_task(task, result)
return result
def start(self):
src/requirements.txt
+1 -1
View File
@@ -137,4 +137,4 @@ python-telegram-bot
pydub
stability_sdk
sentence-transformers==2.2.2
tiktoken
tiktoken
src/service/aios_shell/aios_shell.py
+63 -8
View File
@@ -20,14 +20,15 @@ from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.styles import Style
directory = os.path.dirname(__file__)
sys.path.append(directory + '/../../')
from aios_kernel import AIOS_Version,AgentMsgType,UserConfigItem,AIStorage,Workflow,AIAgent,AgentMsg,AgentMsgStatus,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession,AgentTunnel,TelegramTunnel,CalenderEnvironment,Environment,EmailTunnel,LocalLlama_ComputeNode,Local_Stability_ComputeNode,Stability_ComputeNode,PaintEnvironment
from aios_kernel import ContactManager,Contact
import proxy
from aios_kernel import *
from aios_kernel.compute_node_config import ComputeNodeConfig
sys.path.append(directory + '/../../component/')
from agent_manager import AgentManager
@@ -140,11 +141,15 @@ class AIOS_Shell:
return False
ComputeKernel.get_instance().add_compute_node(open_ai_node)
nodes = ComputeNodeConfig.get_instance().initial()
for node in nodes:
node.start()
ComputeKernel.get_instance().add_compute_node(node)
if await AIStorage.get_instance().is_feature_enable("llama"):
llama_ai_node = LocalLlama_ComputeNode()
if await llama_ai_node.initial() is True:
await llama_ai_node.start()
llama_ai_node.start()
ComputeKernel.get_instance().add_compute_node(llama_ai_node)
else:
logger.error("llama node initial failed!")
@@ -164,9 +169,7 @@ class AIOS_Shell:
# if await stability_api_node.initial() is not True:
# logger.error("stability api node initial failed!")
# ComputeKernel.get_instance().add_compute_node(stability_api_node)
local_st_text_compute_node = LocalSentenceTransformer_Text_ComputeNode()
if local_st_text_compute_node.initial() is not True:
logger.error("local sentence transformer text embedding node initial failed!")
@@ -179,7 +182,6 @@ class AIOS_Shell:
else:
ComputeKernel.get_instance().add_compute_node(local_st_image_compute_node)
await ComputeKernel.get_instance().start()
AIBus().get_default_bus().register_unhandle_message_handler(self._handle_no_target_msg)
@@ -358,6 +360,54 @@ class AIOS_Shell:
journals = [str(journal) for journal in KnowledgePipline.get_instance().get_latest_journals(topn)]
print_formatted_text("\r\n".join(journals))
if sub_cmd == "query":
if len(args) < 2:
return show_text
prompt = AgentPrompt()
prompt.messages.append({"role": "user", "content":" ".join(args[1:])})
result = await KnowledgeBase().query_prompt(prompt)
print_formatted_text(result.as_str())
async def handle_node_commands(self, args):
show_text = FormattedText([("class:title", "sub command not support!\n"
"/node add llama $model_name $url\n"
"/node rm llama $model_name $url\n"
"/node list\n")])
if len(args) < 1:
return show_text
sub_cmd = args[0]
if sub_cmd == "add":
if len(args) < 2:
return show_text
if args[1] == "llama":
if len(args) < 4:
return show_text
model_name = args[2]
url = args[3]
ComputeNodeConfig.get_instance().add_node("llama", url, model_name)
ComputeNodeConfig.get_instance().save()
node = LocalLlama_ComputeNode(url, model_name)
node.start()
ComputeKernel.get_instance().add_compute_node(node)
else:
return show_text
elif sub_cmd == "rm":
if len(args) < 2:
return show_text
if args[1] == "llama":
if len(args) < 4:
return show_text
model_name = args[3]
url = args[4]
ComputeNodeConfig.get_instance().remove_node("llama", url, model_name)
ComputeNodeConfig.get_instance().save()
else:
return show_text
elif sub_cmd == "list":
print_formatted_text(ComputeNodeConfig.get_instance().list())
async def call_func(self,func_name, args):
match func_name:
case 'send':
@@ -483,6 +533,8 @@ class AIOS_Shell:
format_texts.append(("",f"\n-------------------\n"))
return FormattedText(format_texts)
return FormattedText([("class:title", f"chatsession not found")])
case 'node':
return await self.handle_node_commands(args)
case 'exit':
os._exit(0)
case 'help':
@@ -671,6 +723,9 @@ async def main():
'/enable $feature',
'/disable $feature',
'/list_config',
'/node add llama $model_name $url',
'/node rm llama $model_name $url',
'/node list',
'/show',
'/exit',
'/help'], ignore_case=True)