diff --git a/src/aios_kernel/compute_kernel.py b/src/aios_kernel/compute_kernel.py index 1d05afd..5829bbc 100644 --- a/src/aios_kernel/compute_kernel.py +++ b/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!" \ No newline at end of file diff --git a/src/aios_kernel/compute_node.py b/src/aios_kernel/compute_node.py index 6fe0059..ef8f933 100644 --- a/src/aios_kernel/compute_node.py +++ b/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 diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py index ea1d7dc..35a1afd 100644 --- a/src/aios_kernel/compute_task.py +++ b/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: diff --git a/src/aios_kernel/knowledge/knowledge_base.py b/src/aios_kernel/knowledge/knowledge_base.py index 59e4f25..a576280 100644 --- a/src/aios_kernel/knowledge/knowledge_base.py +++ b/src/aios_kernel/knowledge/knowledge_base.py @@ -7,4 +7,6 @@ class KnowledgeBase: pass async def query(self, prompt: AgentPrompt) -> AgentPrompt: - pass \ No newline at end of file + pass + + \ No newline at end of file diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py index be0d184..27177f0 100644 --- a/src/aios_kernel/open_ai_node.py +++ b/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: