From febe28eecc3dc9e4e7c339956dfb87597d47ef85 Mon Sep 17 00:00:00 2001
From: tsukasa <puotosssa@live.com>
Date: Thu, 31 Aug 2023 15:45:02 +0800
Subject: [PATCH] Add embeding compute task

---
 doc/mvp/workflow.drawio         |  2 +-
 src/aios_kernel/compute_task.py |  8 +++++
 src/aios_kernel/open_ai_node.py | 54 +++++++++++++++++++--------------
 3 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/doc/mvp/workflow.drawio b/doc/mvp/workflow.drawio
index c6659c1..e011651 100644
--- a/doc/mvp/workflow.drawio
+++ b/doc/mvp/workflow.drawio
@@ -123,7 +123,7 @@
         </mxGraphModel>
     </diagram>
     <diagram id="kWxfmPxtNxOAf0a73TCG" name="Page-2">
-        <mxGraphModel dx="500" dy="864" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="850" pageHeight="1100" math="0" shadow="0">
+        <mxGraphModel dx="625" dy="1080" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="850" pageHeight="1100" math="0" shadow="0">
             <root>
                 <mxCell id="0"/>
                 <mxCell id="1" parent="0"/>
diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py
index f2253e0..ea1d7dc 100644
--- a/src/aios_kernel/compute_task.py
+++ b/src/aios_kernel/compute_task.py
@@ -38,6 +38,14 @@ class ComputeTask:
             self.params["model_name"] = "gpt-4-0613"
         self.params["max_token_size"] = max_token_size
 
+    def set_embeding_params(self, model_name, input, callchain_id = None):
+        self.task_type = "embeding"
+        self.create_time = time.time()
+        self.task_id = uuid.uuid4().hex
+        self.callchain_id = callchain_id
+        self.params["model_name"] = model_name
+        self.params["input"] = input
+
     def display(self) -> str:
         return f"ComputeTask: {self.task_id} {self.task_type} {self.state}"
 
diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py
index 7f5c8e6..be0d184 100644
--- a/src/aios_kernel/open_ai_node.py
+++ b/src/aios_kernel/open_ai_node.py
@@ -47,31 +47,39 @@ class OpenAI_ComputeNode(ComputeNode):
     
     def _run_task(self,task:ComputeTask):
         task.state = ComputeTaskState.RUNNING
-        mode_name = task.params["model_name"]
-        # max_token_size = task.params["max_token_size"]
-        prompts = task.params["prompts"]
+        # switch tsak type
+        if task.task_type == "llm_completion":
+            mode_name = task.params["model_name"]
+            # max_token_size = task.params["max_token_size"]
+            prompts = task.params["prompts"]
 
-        logger.info(f"call openai {mode_name} prompts: {prompts}")
-        resp = openai.ChatCompletion.create(model=mode_name,
-                                            messages=prompts,
-                                            max_tokens=4000,
-                                            temperature=1.2)
-        logger.info(f"openai response: {resp}")
-        
-        status_code = resp["choices"][0]["finish_reason"]
-        if status_code != "stop":
-            task.state = ComputeTaskState.ERROR
-            task.error_str =f"The status code was {status_code}."
-            return None
-        
-        result = ComputeTaskResult()    
-        result.set_from_task(task)
-        result.worker_id = self.node_id
-        result.result_str = resp["choices"][0]["message"]["content"]
-        result.result = resp["choices"][0]["message"]
-        
-        return result
+	        mode_name = task.params["model_name"]
+	        # max_token_size = task.params["max_token_size"]
+	        prompts = task.params["prompts"]
 
+	        logger.info(f"call openai {mode_name} prompts: {prompts}")
+	        resp = openai.ChatCompletion.create(model=mode_name,
+	                                            messages=prompts,
+	                                            max_tokens=4000,
+	                                            temperature=1.2)
+	        logger.info(f"openai response: {resp}")
+	        
+	        status_code = resp["choices"][0]["finish_reason"]
+	        if status_code != "stop":
+	            task.state = ComputeTaskState.ERROR
+	            task.error_str =f"The status code was {status_code}."
+	            return None
+	        
+	        result = ComputeTaskResult()    
+	        result.set_from_task(task)
+	        result.worker_id = self.node_id
+	        result.result_str = resp["choices"][0]["message"]["content"]
+	        result.result = resp["choices"][0]["message"]
+	        
+	        return result
+ 		if task.task_type == "embeding":
+            pass
+            
     def start(self):
         async def _run_task_loop():
             while True: