Add openai text embedding compute task

src/aios_kernel/compute_kernel.py

@@ -118,4 +118,34 @@ class ComputeKernel:
             return task_req.result.result_str
               
         return "error!"
+    
+    def text_embedding(self,input:str,model_name:Optional[str] = None):
+        task_req = ComputeTask()
+        task_req.set_text_embeding_params(input,model_name)
+        self.run(task_req)
+        return task_req
+    
+    async def do_text_embedding(self,input:str,model_name:Optional[str] = None) -> [float]
+        task_req = self.text_embedding(input,model_name)
+        async def check_timer():
+            check_times = 0
+            while True:
+                if task_req.state == ComputeTaskState.DONE:
+                    break
+
+                if task_req.state == ComputeTaskState.ERROR:
+                    break
+
+                if check_times >=  20:
+                    task_req.state = ComputeTaskState.ERROR
+                    break
+
+                await asyncio.sleep(0.5)
+                check_times += 1
+            
+        await asyncio.create_task(check_timer())
+        if task_req.state == ComputeTaskState.DONE:
+            return task_req.result.result
+              
+        return "error!"
     

src/aios_kernel/compute_node.py

@@ -28,7 +28,7 @@ class ComputeNode(ABC):
         pass
 
     @abstractmethod
-    def is_support(self,task_type:str) -> bool:
+    def is_support(self, task: ComputeTask) -> bool:
         pass
 
     @abstractmethod

src/aios_kernel/compute_task.py

@@ -38,12 +38,15 @@ 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"
+    def set_text_embeding_params(self, input, model_name=None, callchain_id = None):
+        self.task_type = "text_embedding"
         self.create_time = time.time()
         self.task_id = uuid.uuid4().hex
         self.callchain_id = callchain_id
-        self.params["model_name"] = model_name
+        if model_name is not None:
+            self.params["model_name"] = model_name
+        else:
+            self.params["model_name"] = "text-embedding-ada-002"
         self.params["input"] = input
 
     def display(self) -> str:

src/aios_kernel/knowledge/knowledge_base.py

@@ -7,4 +7,6 @@ class KnowledgeBase:
         pass
 
     async def query(self, prompt: AgentPrompt) -> AgentPrompt:
-        pass
+        pass
+
+    

src/aios_kernel/open_ai_node.py

@@ -53,32 +53,45 @@ class OpenAI_ComputeNode(ComputeNode):
             # max_token_size = task.params["max_token_size"]
             prompts = task.params["prompts"]
 
-	        mode_name = task.params["model_name"]
+            mode_name = task.params["model_name"]
 	        # max_token_size = task.params["max_token_size"]
-	        prompts = task.params["prompts"]
+            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
+            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 == "text_embedding":
+            model_name = task.params["model_name"]
+            input = task.params["input"]
+            logger.info(f"call openai {model_name} input: {input}")
+
+            resp = openai.Embeding.create(model=model_name,
+                                            input=input)
+            logger.info(f"openai response: {resp}")
+
+            result = ComputeTaskResult()    
+            result.set_from_task(task)
+            result.worker_id = self.node_id
+            result.result = resp["data"][0]["embedding"]
+
+            return result
             
     def start(self):
         async def _run_task_loop():
@@ -104,8 +117,13 @@ class OpenAI_ComputeNode(ComputeNode):
         pass
 
 
-    def is_support(self,task_type:str) -> bool:
-        return True
+    def is_support(self, task: ComputeTask) -> bool:
+        if task.task_type == "llm_completion":
+            return True
+        if task.task_type == "text_embedding":
+            if task.params["model_name"] == "text-embedding-ada-002":
+                return True
+        return False
 
 
     def is_local(self) -> bool: