From 01f05285c7f0b09fd6bd1c4087a410539a712cac Mon Sep 17 00:00:00 2001
From: tsukasa <puotosssa@live.com>
Date: Wed, 27 Sep 2023 17:24:00 +0800
Subject: [PATCH] fix queue compute node

---
 src/aios_kernel/compute_kernel.py        |  2 +-
 src/aios_kernel/local_st_compute_node.py |  4 +--
 src/aios_kernel/open_ai_node.py          |  2 +-
 src/aios_kernel/queue_compute_node.py    | 33 ++++++++++++++----------
 src/knowledge/store.py                   |  2 +-
 5 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/src/aios_kernel/compute_kernel.py b/src/aios_kernel/compute_kernel.py
index bd3eb07..42e81fa 100644
--- a/src/aios_kernel/compute_kernel.py
+++ b/src/aios_kernel/compute_kernel.py
@@ -148,7 +148,7 @@ class ComputeKernel:
         task_result = await self._send_task(task_req)
 
         if task_req.state == ComputeTaskState.DONE:
-            return task_result.result
+            return task_result.result_str
 
         return "error!"
 
diff --git a/src/aios_kernel/local_st_compute_node.py b/src/aios_kernel/local_st_compute_node.py
index 2f36936..d4cd17e 100644
--- a/src/aios_kernel/local_st_compute_node.py
+++ b/src/aios_kernel/local_st_compute_node.py
@@ -36,7 +36,7 @@ class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode):
         except Exception as err:
             logger.error(f"load model {self.model} failed: {err}")
             return False
-
+        self.start()
         return True
 
     async def execute_task(
@@ -126,7 +126,7 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
         except Exception as err:
             logger.error(f"load model {self.model} failed: {err}")
             return False
-
+        self.start()
         return True
 
     def _load_image(self, source: Union[ObjectID, bytes]) -> Optional[Image]:
diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py
index e16d76d..6716c2b 100644
--- a/src/aios_kernel/open_ai_node.py
+++ b/src/aios_kernel/open_ai_node.py
@@ -98,7 +98,7 @@ class OpenAI_ComputeNode(ComputeNode):
                 task.state = ComputeTaskState.DONE
                 result.result_code = ComputeTaskResultCode.OK
                 result.worker_id = self.node_id
-                result.result = resp["data"][0]["embedding"]
+                result.result_str = resp["data"][0]["embedding"]
 
                 return result
             case ComputeTaskType.LLM_COMPLETION:
diff --git a/src/aios_kernel/queue_compute_node.py b/src/aios_kernel/queue_compute_node.py
index 1475672..7d88778 100644
--- a/src/aios_kernel/queue_compute_node.py
+++ b/src/aios_kernel/queue_compute_node.py
@@ -4,7 +4,7 @@ from asyncio import Queue
 import logging
 from abc import abstractmethod
 
-from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
+from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskResultCode, ComputeTaskState, ComputeTaskType
 from .compute_node import ComputeNode
 
 logger = logging.getLogger(__name__)
@@ -13,6 +13,7 @@ class Queue_ComputeNode(ComputeNode):
     def __init__(self):
         super().__init__()
         self.task_queue = Queue()
+        self.is_start = False
 
     @abstractmethod
     async def execute_task(self, task: ComputeTask) -> {
@@ -39,28 +40,32 @@ class Queue_ComputeNode(ComputeNode):
         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"]
+        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)
 
         return result
 
     async def start(self):
+        if self.is_start is True:
+            return
+        self.is_start = True
+
         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
+                logger.info(f"openai_node get task: {task.display()}")
+                self._run_task(task)
 
         asyncio.create_task(_run_task_loop())
 
diff --git a/src/knowledge/store.py b/src/knowledge/store.py
index a344719..d8b4bbb 100644
--- a/src/knowledge/store.py
+++ b/src/knowledge/store.py
@@ -62,5 +62,5 @@ class KnowledgeStore:
     
     def get_vector_store(self, model_name: str) -> VectorBase:
         if model_name not in self.vector_store:
-            self.vector_store[model_name] = ChromaVectorStore(model_name)
+            self.vector_store[model_name] = ChromaVectorStore(self.root, model_name)
         return self.vector_store[model_name]