summaryrefslogtreecommitdiffstats
path: root/g4f/Provider
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--g4f/Provider/DeepInfra.py25
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}")