diff --git a/src/aios_kernel/agent.py b/src/aios_kernel/agent.py index 136ac0d..5d0566f 100644 --- a/src/aios_kernel/agent.py +++ b/src/aios_kernel/agent.py @@ -340,7 +340,9 @@ class AIAgent: org_msg.inner_call_chain.append(ineternal_call_record) if stack_limit > 0: - inner_func_call_node = task_result.result_message.get("function_call") + result_message = task_result.result.get("message") + if result_message: + inner_func_call_node = result_message.get("function_call") if inner_func_call_node: return await self._execute_func(inner_func_call_node,prompt,org_msg,stack_limit-1) @@ -395,7 +397,9 @@ class AIAgent: final_result = task_result.result_str - inner_func_call_node = task_result.result_message.get("function_call") + result_message = task_result.result.get("message") + if result_message: + inner_func_call_node = result_message.get("function_call") if inner_func_call_node: #TODO to save more token ,can i use msg_prompt? final_result,error_code = await self._execute_func(inner_func_call_node,prompt,msg) diff --git a/src/aios_kernel/compute_kernel.py b/src/aios_kernel/compute_kernel.py index 0cfbe9b..0315c21 100644 --- a/src/aios_kernel/compute_kernel.py +++ b/src/aios_kernel/compute_kernel.py @@ -107,7 +107,7 @@ class ComputeKernel: self.run(task_req) return task_req - async def _send_task(self,task_req:ComputeTask)->ComputeTaskResult: + async def _wait_task(self,task_req:ComputeTask)->ComputeTaskResult: async def check_timer(): check_times = 0 while True: @@ -135,7 +135,7 @@ class ComputeKernel: async def do_llm_completion(self, prompt: AgentPrompt, mode_name: Optional[str] = None, max_token: int = 0, inner_functions = None) -> str: task_req = self.llm_completion(prompt, mode_name, max_token,inner_functions) - return await self._send_task(task_req) + return await self._wait_task(task_req) def text_embedding(self,input:str,model_name:Optional[str] = None): @@ -146,12 +146,13 @@ class ComputeKernel: async def do_text_embedding(self,input:str,model_name:Optional[str] = None) -> [float]: task_req = self.text_embedding(input,model_name) - task_result = await self._send_task(task_req) + task_result = await self._wait_task(task_req) if task_req.state == ComputeTaskState.DONE: - return task_result.result_str - - return "error!" + return task_result.result.get("content") + else: + logging.warning(f"do_text_embedding error: {task_req.error_str},input: {input}") + return None def image_embedding(self,input:ObjectID,model_name:Optional[str] = None): task_req = ComputeTask() @@ -161,12 +162,12 @@ class ComputeKernel: async def do_image_embedding(self,input:ObjectID,model_name:Optional[str] = None) -> [float]: task_req = self.image_embedding(input,model_name) - task_result = await self._send_task(task_req) + task_result = await self._wait_task(task_req) if task_req.state == ComputeTaskState.DONE: - return task_result.result_str + return task_result.result.get("content") - return "error!" + return None async def do_text_to_speech(self, input:str, @@ -185,7 +186,7 @@ class ComputeKernel: task_req.task_type = ComputeTaskType.TEXT_2_VOICE self.run(task_req) - task_result = await self._send_task(task_req) + task_result = await self._wait_task(task_req) if task_req.state == ComputeTaskState.DONE: return task_result.result @@ -199,7 +200,7 @@ class ComputeKernel: async def do_text_2_image(self, prompt:str, model_name:Optional[str] = None, negative_prompt = None) -> ComputeTaskResult: task = self.text_2_image(prompt,model_name, negative_prompt) - task = await self._send_task(task) + task = await self._wait_task(task) return task.result # if task_req.state == ComputeTaskState.DONE: diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py index 6ab9dd2..ea7c431 100644 --- a/src/aios_kernel/compute_task.py +++ b/src/aios_kernel/compute_task.py @@ -110,7 +110,9 @@ class ComputeTaskResult: self.error_str : str = None self.result_code: int = 0 self.result_str: str = None # easy to use,can read from result - self.result_message: dict = {} + + self.result : dict = {} + self.result_refers: dict = {} self.pading_data: bytearray = None diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index 2945bb5..0863c76 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -33,7 +33,8 @@ class KnowledgeBase: text = chunk.read().decode("utf-8") vector = await self.compute_kernel.do_text_embedding(text, self._default_text_model) - await self.store.get_vector_store(self._default_text_model).insert(vector, chunk_id) + if vector: + await self.store.get_vector_store(self._default_text_model).insert(vector, chunk_id) async def __embedding_image(self, image: ImageObject): # desc = {} @@ -45,7 +46,8 @@ class KnowledgeBase: # desc["tags"] = image.get_tags() # vector = await self.compute_kernel.do_text_embedding(json.dumps(desc), self._default_text_model) vector = await self.compute_kernel.do_image_embedding(image.calculate_id(), self._default_image_model) - await self.store.get_vector_store(self._default_image_model).insert(vector, image.calculate_id()) + if vector: + await self.store.get_vector_store(self._default_image_model).insert(vector, image.calculate_id()) async def __embedding_video(self, vedio: VideoObject): desc = {} diff --git a/src/aios_kernel/local_llama_compute_node.py b/src/aios_kernel/local_llama_compute_node.py index 52e6af4..7f8ad0a 100644 --- a/src/aios_kernel/local_llama_compute_node.py +++ b/src/aios_kernel/local_llama_compute_node.py @@ -21,7 +21,11 @@ class LocalLlama_ComputeNode(Queue_ComputeNode): self.url = url self.model_name = model_name - async def execute_task(self, task: ComputeTask, result: ComputeTaskResult): + async def execute_task(self, task: ComputeTask)->ComputeTaskResult: + result = ComputeTaskResult() + result.result_code = ComputeTaskResultCode.ERROR + result.set_from_task(task) + result.worker_id = self.node_id match task.task_type: case ComputeTaskType.TEXT_EMBEDDING: model_name = task.params["model_name"] @@ -56,7 +60,9 @@ class LocalLlama_ComputeNode(Queue_ComputeNode): 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 + return result + + return result async def initial(self) -> bool: return True @@ -145,7 +151,8 @@ class LocalLlama_ComputeNode(Queue_ComputeNode): result.result_code = ComputeTaskResultCode.OK result.result_str = resp["choices"][0]["message"]["content"] - result.result_message = resp["choices"][0]["message"] + result.result["message"] = resp["choices"][0]["message"] + if token_usage: result.result_refers["token_usage"] = token_usage diff --git a/src/aios_kernel/local_st_compute_node.py b/src/aios_kernel/local_st_compute_node.py index 848b6b9..4142d5a 100644 --- a/src/aios_kernel/local_st_compute_node.py +++ b/src/aios_kernel/local_st_compute_node.py @@ -5,14 +5,12 @@ from pydantic import BaseModel from typing import Union from PIL import Image import io -from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType +from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType,ComputeTaskResult,ComputeTaskResultCode from .queue_compute_node import Queue_ComputeNode from knowledge import ObjectID logger = logging.getLogger(__name__) - - class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode): # For valid pretrained models, see https://www.sbert.net/docs/pretrained_models.html def __init__(self, model_name: str = "all-MiniLM-L6-v2"): @@ -39,18 +37,11 @@ class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode): self.start() return True - async def execute_task( - self, task: ComputeTask - ) -> { - "task_type": str, - "content": str, - "message": str, - "state": ComputeTaskState, - "error": { - "code": int, - "message": str, - }, - }: + async def execute_task(self, task: ComputeTask) : + result = ComputeTaskResult() + result.result_code = ComputeTaskResultCode.ERROR + result.set_from_task(task) + result.worker_id = self.node_id try: # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task: {task}") if task.task_type == ComputeTaskType.TEXT_EMBEDDING: @@ -60,25 +51,19 @@ class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode): ) sentence_embeddings = self.model.encode(input, show_progress_bar=False).tolist() # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}") - return { - "state": ComputeTaskState.DONE, - "content": sentence_embeddings, - "message": None, - } + result.result_code = ComputeTaskResultCode.OK + result.result["content"] = sentence_embeddings + else: - return { - "state": ComputeTaskState.ERROR, - "error": {"code": -1, "message": "unsupport embedding task type"}, - } + result.error_str = f"unsupport embedding task type: {task.task_type}" except Exception as err: import traceback logger.error(f"{traceback.format_exc()}, error: {err}") + result.error_str = f"{traceback.format_exc()}, error: {err}" + + return result - return { - "state": ComputeTaskState.ERROR, - "error": {"code": -1, "message": "unknown exception: " + str(err)}, - } def display(self) -> str: return f"LocalSentenceTransformer_Text_ComputeNode: {self.node_id}, {self.model_name}" @@ -170,16 +155,11 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode): async def execute_task( self, task: ComputeTask - ) -> { - "task_type": str, - "content": str, - "message": str, - "state": ComputeTaskState, - "error": { - "code": int, - "message": str, - }, - }: + ) -> ComputeTaskResult: + result = ComputeTaskResult() + result.result_code = ComputeTaskResultCode.ERROR + result.set_from_task(task) + result.worker_id = self.node_id try: # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task: {task}") if task.task_type == ComputeTaskType.TEXT_EMBEDDING: @@ -189,11 +169,9 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode): ) sentence_embeddings = self.multi_model.encode(input, show_progress_bar=False).tolist() # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}") - return { - "state": ComputeTaskState.DONE, - "content": sentence_embeddings, - "message": None, - } + result.result_code = ComputeTaskResultCode.OK + result.result["content"] = sentence_embeddings + elif task.task_type == ComputeTaskType.IMAGE_EMBEDDING: input = task.params["input"] logger.debug( @@ -202,33 +180,22 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode): img = self._load_image(input) if img is None: - return { - "state": ComputeTaskState.ERROR, - "error": {"code": -1, "message": "load image failed"}, - } + result.error_str = f"load image failed: {input}" + return result sentence_embeddings = self.model.encode(img, show_progress_bar=False).tolist() - - # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}") - return { - "state": ComputeTaskState.DONE, - "content": sentence_embeddings, - "message": None, - } + result.result_code = ComputeTaskResultCode.OK + result.result["content"] = sentence_embeddings else: - return { - "state": ComputeTaskState.ERROR, - "error": {"code": -1, "message": "unsupport embedding task type"}, - } + result.error_str = f"unsupport embedding task type: {task.task_type}" 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)}, - } + result.error_str = f"{traceback.format_exc()}, error: {err}" + + + return result def display(self) -> str: return f"LocalSentenceTransformer_Image_ComputeNode: {self.node_id}, {self.model_name}" diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py index 6716c2b..b9d0e44 100644 --- a/src/aios_kernel/open_ai_node.py +++ b/src/aios_kernel/open_ai_node.py @@ -150,7 +150,8 @@ class OpenAI_ComputeNode(ComputeNode): result.result_code = ComputeTaskResultCode.OK result.worker_id = self.node_id result.result_str = resp["choices"][0]["message"]["content"] - result.result_message = resp["choices"][0]["message"] + result.result["message"] = resp["choices"][0]["message"] + if token_usage: result.result_refers["token_usage"] = token_usage logger.info(f"openai success response: {result.result_str}") diff --git a/src/aios_kernel/queue_compute_node.py b/src/aios_kernel/queue_compute_node.py index 3c1bf3e..4eb22f4 100644 --- a/src/aios_kernel/queue_compute_node.py +++ b/src/aios_kernel/queue_compute_node.py @@ -16,7 +16,7 @@ class Queue_ComputeNode(ComputeNode): self.is_start = False @abstractmethod - async def execute_task(self, task: ComputeTask, result: ComputeTaskResult): + async def execute_task(self, task: ComputeTask)->ComputeTaskResult: pass async def push_task(self, task: ComputeTask, proiority: int = 0): @@ -29,15 +29,22 @@ class Queue_ComputeNode(ComputeNode): async def _run_task(self, task: ComputeTask): task.state = ComputeTaskState.RUNNING - result = ComputeTaskResult() result.result_code = ComputeTaskResultCode.ERROR result.set_from_task(task) result.worker_id = self.node_id - await self.execute_task(task, result) - - return result + real_result = await self.execute_task(task) + + if real_result: + if real_result.result_code == ComputeTaskResultCode.OK: + task.state = ComputeTaskState.DONE + else: + task.state = ComputeTaskState.ERROR + return real_result + else: + task.state = ComputeTaskState.ERROR + return result def start(self): if self.is_start is True: diff --git a/src/aios_kernel/workflow.py b/src/aios_kernel/workflow.py index 69dcbdd..4944ce7 100644 --- a/src/aios_kernel/workflow.py +++ b/src/aios_kernel/workflow.py @@ -425,7 +425,10 @@ class Workflow: ineternal_call_record.done_time = time.time() org_msg.inner_call_chain.append(ineternal_call_record) if stack_limit > 0: - inner_func_call_node = task_result.result_message.get("function_call") + result_message = task_result.result.get("message") + if result_message: + inner_func_call_node = result_message.get("function_call") + if inner_func_call_node: return await self._role_execute_func(the_role,inner_func_call_node,prompt,org_msg,stack_limit-1) else: @@ -466,7 +469,9 @@ class Workflow: result_str = task_result.result_str logger.info(f"{the_role.role_id} process {msg.sender}:{msg.body},llm str is :{result_str}") - inner_func_call_node = task_result.result_message.get("function_call") + result_message = task_result.result.get("message") + if result_message: + inner_func_call_node = result_message.get("function_call") if inner_func_call_node: #TODO to save more token ,can i use msg_prompt? diff --git a/src/knowledge/vector/chroma_store.py b/src/knowledge/vector/chroma_store.py index 260a8e8..8cd3084 100644 --- a/src/knowledge/vector/chroma_store.py +++ b/src/knowledge/vector/chroma_store.py @@ -28,7 +28,8 @@ class ChromaVectorStore(VectorBase): self.collection = collection async def insert(self, vector: [float], id: ObjectID): - logging.info(f"will insert vector: {vector} id: {str(id)}") + logging.info(f"will insert vector: {len(vector)} id: {str(id)}") + logging.debug(f"vector is {vector}") self.collection.add( embeddings=vector, ids=str(id),