Update completions.py

ktransformers/server/api/ollama/completions.py

@@ -85,7 +85,11 @@ async def generate(request: Request, input: OllamaGenerateCompletionRequest):
         return check_link_response(request, inner())
     else:
         complete_response = ""
-        async for token in interface.inference(input.prompt, id):
+        async for res in interface.inference(input.prompt, id):
+            if isinstance(res, RawUsage):
+                raw_usage = res
+            else: 
+                token, finish_reason = res
                 complete_response += token
         response = OllamaGenerationResponse(
             model=config.model_name,
@@ -187,7 +191,11 @@ async def chat(request: Request, input: OllamaChatCompletionRequest):
         complete_response = ""
         eval_count = 0 
 
-        async for token in interface.inference(prompt, id):
+        async for res in interface.inference(prompt, id):
+            if isinstance(res, RawUsage):
+                raw_usage = res
+            else: 
+                token, finish_reason = res
                 complete_response += token
             eval_count += 1