From 640c861ad77e14e1745d2c95b8d77aa43b74c69f Mon Sep 17 00:00:00 2001 From: nullstreak <139914347+nullstreak@users.noreply.github.com> Date: Sun, 21 Jan 2024 05:07:34 +0100 Subject: DeepInfra: Fix token duplication --- g4f/Provider/DeepInfra.py | 25 ++++++++++++++++--------- 1 file 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}") -- cgit v1.2.3