From 6a00d097012b6e6890e154657f4b1ce91daff0eb Mon Sep 17 00:00:00 2001
From: streetycat <305190374@qq.com>
Date: Wed, 29 Nov 2023 07:46:07 +0000
Subject: [PATCH] Adapt to the function of llama.cpp

---
 src/aios_kernel/local_llama_compute_node.py | 25 +++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/src/aios_kernel/local_llama_compute_node.py b/src/aios_kernel/local_llama_compute_node.py
index 7f8ad0a..3b1fcb8 100644
--- a/src/aios_kernel/local_llama_compute_node.py
+++ b/src/aios_kernel/local_llama_compute_node.py
@@ -113,12 +113,27 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
         llm_inner_functions = task.params.get("inner_functions")
         if max_token_size is None:
             max_token_size = max_token_size
-            
+
         body = {
             "messages": [],
+            "functions": llm_inner_functions,
+            "tools": [],
+            "tool_choices": [],
             "max_tokens": 4000
         }
 
+        for fun in llm_inner_functions:
+            body["tools"].append({
+                "type": "function",
+                "function": fun,
+            })
+            body["tool_choices"].append({
+                "type": "function",
+                "function": {
+                    "name": fun["name"]
+                }
+            })
+
         for prompt in prompts:
             body["messages"].append({
                 "role": prompt["role"],
@@ -126,6 +141,8 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
             })
         
         try:
+            logger.info(f"will post http request to {self.url}/v1/chat/completions, body: {body}")
+
             response = requests.post(self.url + "/v1/chat/completions", json = body, verify=False, headers={"Content-Type": "application/json"})
             response.close()
 
@@ -138,8 +155,12 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
                 token_usage = resp["usage"]
 
                 match status_code:
-                    case "function_call":
+                    case "tool_calls":
                         task.state = ComputeTaskState.DONE
+                        # rebuild the function name
+                        fun_name = resp["choices"][0]["message"]["function_call"]["name"]
+                        if len(llm_inner_functions) == 1 and (fun_name is None or fun_name == ""):
+                            resp["choices"][0]["message"]["function_call"]["name"] = llm_inner_functions[0]["name"]
                     case "stop":
                         task.state = ComputeTaskState.DONE
                     case _: