Update unit test for stream completion object

NexaAI · Sep 2, 2024 · f0c9488 · f0c9488
1 parent f95d6fb
commit f0c9488
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/nexa/gguf/llama/llama.py b/nexa/gguf/llama/llama.py
@@ -1392,6 +1392,7 @@ def logit_bias_processor(
                             "choices": [
                                 {
                                     "index": 0,
+                                    "logprobs": None,
                                     "delta": {
                                         "content": ts,
                                     },

diff --git a/nexa/gguf/llama/llama_chat_format.py b/nexa/gguf/llama/llama_chat_format.py
@@ -304,6 +304,7 @@ def _convert_text_completion_chunks_to_chat(
                     }
                 ],
             }
+
         yield {
             "id": "chat" + chunk["id"],
             "model": chunk["model"],
@@ -314,7 +315,7 @@ def _convert_text_completion_chunks_to_chat(
                     "index": 0,
                     "delta": (
                         {
-                            "content": chunk["choices"][0]["text"],
+                            "content": chunk["choices"][0]["delta"]["content"],
                         }
                         if chunk["choices"][0]["finish_reason"] is None
                         else {}

diff --git a/tests/test_text_generation.py b/tests/test_text_generation.py
@@ -32,7 +32,7 @@ def test_streaming():
     )
     for chunk in output:
         if "choices" in chunk:
-            print(chunk["choices"][0]["text"], end="", flush=True)
+            print(chunk["choices"][0]["delta"]["content"], end="", flush=True)
     # TODO: add assertions here
 
 # Test conversation mode with chat format