1 files changed, 16 insertions, 9 deletions

diff --git a/g4f/Provider/DeepInfra.py b/g4f/Provider/DeepInfra.py
index 96e3a680..12bd49c7 100644
--- a/g4f/Provider/DeepInfra.py
+++ b/g4f/Provider/DeepInfra.py
@@ -58,16 +58,23 @@ class DeepInfra(AsyncGeneratorProvider):
                 response.raise_for_status()
                 first = True
                 async for line in response.iter_lines():
+                    if not line.startswith(b"data: "):
+                        continue
+
                     try:
-                        if line.startswith(b"data: [DONE]"):
+                        decoded_line = line.decode().lstrip("data: ")
+                        json_line = json.loads(decoded_line)
+
+                        choices = json_line.get("choices", [{}])
+                        finish_reason = choices[0].get("finish_reason", "")
+                        if finish_reason:
                             break
-                        elif line.startswith(b"data: "):
-                            chunk = json.loads(line[6:])["choices"][0]["delta"].get("content")
-                        if chunk:
+                        token = choices[0].get("delta", {}).get("content", "")
+
+                        if token:
                             if first:
-                                chunk = chunk.lstrip()
-                                if chunk:
-                                    first = False
-                            yield chunk
+                                token = token.lstrip()
+                                first = False
+                            yield token
                     except Exception:
-                        raise RuntimeError(f"Response: {line}")
\ No newline at end of file
+                        raise RuntimeError(f"Response: {line}")