Add embeding compute task
This commit is contained in:
@@ -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"/>
|
||||
|
||||
@@ -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}"
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user