fix bug

piEsposito · Jul 7, 2024 · d9fa08c · d9fa08c
1 parent 130f0dd
commit d9fa08c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 20 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tiny-ai-client"
-version = "0.0.10"
+version = "0.0.11"
 description = "Tiny AI client for LLMs. As simple as it gets."
 authors = ["piEsposito <[email protected]>"]
 license = "Apache 2.0"

diff --git a/tiny_ai_client/models.py b/tiny_ai_client/models.py
@@ -85,25 +85,19 @@ def stream(
         if isinstance(images, PIL_Image.Image):
             images = [images]
         self.chat.append(Message(text=message, role="user", images=images))
+        text = ""
         for chunk in self.client_wrapper.stream(
             max_new_tokens=max_new_tokens,
             temperature=temperature,
             chat=self.chat,
             timeout=timeout,
         ):
+            text += chunk
             yield chunk
         # After streaming, update the chat history
         self.chat.append(
             Message(
-                text="".join(
-                    chunk
-                    for chunk in self.client_wrapper.stream(
-                        max_new_tokens=max_new_tokens,
-                        temperature=temperature,
-                        chat=self.chat,
-                        timeout=timeout,
-                    )
-                ),
+                text=text,
                 role="assistant",
             )
         )
@@ -196,23 +190,16 @@ async def astream(
         if isinstance(images, PIL_Image.Image):
             images = [images]
         self.chat.append(Message(text=message, role="user", images=images))
+        text = ""
         async for chunk in self.client_wrapper.astream(
             max_new_tokens=max_new_tokens,
             temperature=temperature,
             chat=self.chat,
             timeout=timeout,
         ):
+            text += chunk
             yield chunk
-        # After streaming, update the chat history
-        full_response = ""
-        async for chunk in self.client_wrapper.astream(
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            chat=self.chat,
-            timeout=timeout,
-        ):
-            full_response += chunk
-        self.chat.append(Message(text=full_response, role="assistant"))
+        self.chat.append(Message(text=text, role="assistant"))
 
 
 class ToolCall(BaseModel):