add compute node for llama

src/aios_kernel/local_llama_compute_node.py

@@ -1,7 +1,7 @@
 
 import logging
 import requests
-from typing import Optional
+from typing import Optional, List
 from pydantic import BaseModel
 
 from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType
@@ -26,14 +26,18 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
         class GenerateResponse(BaseModel):
             error: Optional[int]
             msg: Optional[str]
-            results: Optional[str]
+            results: Optional[List[str]]
 
-        try:    
+        try:
+            prompt_msgs = []
+            for prompt in task.params["prompts"]:
+                prompt_msgs.append(prompt["content"])
+                
             body = {
-                "prompts": task.params["prompts"]
+                "prompts": prompt_msgs
             }
             
-            response = requests.post("http://aigc:7880/generate", data = body, verify=False)
+            response = requests.post("http://aigc:7880/generate", json = body, verify=False, headers={"Content-Type": "application/json"})
             response.close()
 
             logger.info(f"LocalLlama_ComputeNode task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
@@ -47,19 +51,20 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
                     }
                 }
             else:
-                resp = GenerateResponse.parse_raw(response.content.decode("utf-8"))
-                if resp.error:
+                resp = response.json()
+                if "error" in resp:
                     return {
                         "state": ComputeTaskState.ERROR,
                         "error": {
-                            "code": resp.error,
-                            "message": "local llama failed:" + resp.msg
+                            "code": resp["error"],
+                            "message": "local llama failed:" + resp["msg"]
                         }
                     }
                 else:
                     return {
-                        "content": str(resp.results),
-                        "message": {}
+                        "state": ComputeTaskState.DONE,
+                        "content": str(resp["results"]),
+                        "message": str(resp["results"])
                     }
         except Exception as err:
             import traceback