fix: prevent verbose json output in router

agiresearch · Jan 7, 2025 · df56cf1 · df56cf1
1 parent 5e7c2b7
commit df56cf1
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 14 deletions.
diff --git a/aios/llm_core/adapter.py b/aios/llm_core/adapter.py
@@ -254,14 +254,17 @@ def address_syscall(
         model = self.strategy()
 
         if isinstance(model, (str, HfLocalBackend, VLLMLocalBackend, OllamaBackend)):
-            res = model(
-                messages=messages,
-                temperature=temperature,
-            ) if not isinstance(model, str) else str(completion(
-                model=model,
-                messages=messages,
-                temperature=temperature,
-            ))
+            if not isinstance(model, str):
+                res = model(
+                    messages=messages,
+                    temperature=temperature,
+                )
+            else:
+                res = completion(
+                    model=model,
+                    messages=messages,
+                    temperature=temperature,
+                ).choices[0].message.content
         else:
             raise RuntimeError(f"Unsupported model type: {type(model)}")
 

diff --git a/aios/llm_core/local.py b/aios/llm_core/local.py
@@ -35,13 +35,13 @@ def __init__(self, model_name, device="auto", max_gpu_memory=None, hostname=None
         self.tokenizer.chat_template = "{% for message in messages %}{% if message['role'] == 'user' %}{{ ' ' }}{% endif %}{{ message['content'] }}{% if not loop.last %}{{ ' ' }}{% endif %}{% endfor %}{{ eos_token }}"
 
     def inference_online(self, messages, temperature, stream=False):
-        return str(completion(
+        return completion(
             model="huggingface/" + self.model_name,
             messages=messages,
             temperature=temperature,
             api_base=self.hostname,
-        ))
-    
+        ).choices[0].message.content
+
     def __call__(
         self,
         messages,
@@ -50,7 +50,7 @@ def __call__(
     ):
         if self.hostname is not None:
             return self.inference_online(messages, temperature, stream=stream)
-        
+
         if stream:
             raise NotImplemented
 
@@ -105,12 +105,12 @@ def __init__(self, model_name, device="auto", max_gpu_memory=None, hostname=None
             print("Error loading vllm model:", err)
 
     def inference_online(self, messages, temperatures, stream=False):
-        return str(completion(
+        return completion(
             model="hosted_vllm/" + self.model_name,
             messages=messages,
             temperature=temperature,
             api_base=self.hostname,
-        ))
+        ).choices[0].message.content
 
     def __call__(
         self,