From 29835d951c657c348c8f13bde987c336006d3c7e Mon Sep 17 00:00:00 2001
From: kqlio67 <kqlio67@users.noreply.github.com>
Date: Sat, 19 Oct 2024 22:33:06 +0300
Subject: fix(g4f/Provider/HuggingChat.py): handle JSON decode errors and
 response status

---
 g4f/Provider/HuggingChat.py | 36 +++++++++++++++++++++---------------
 1 file changed, 21 insertions(+), 15 deletions(-)

(limited to 'g4f/Provider')

diff --git a/g4f/Provider/HuggingChat.py b/g4f/Provider/HuggingChat.py
index 45f3a0d2..7ebbf570 100644
--- a/g4f/Provider/HuggingChat.py
+++ b/g4f/Provider/HuggingChat.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-import json, requests, re
+import json
+import requests
 
 from curl_cffi import requests as cf_reqs
 from ..typing import CreateResult, Messages
@@ -73,17 +74,18 @@ class HuggingChat(AbstractProvider, ProviderModelMixin):
                 'user-agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36',
             }
 
-            print(model)
             json_data = {
                 'model': model,
             }
 
             response = session.post('https://huggingface.co/chat/conversation', json=json_data)
-            conversationId = response.json()['conversationId']
+            if response.status_code != 200:
+                raise RuntimeError(f"Request failed with status code: {response.status_code}, response: {response.text}")
 
-            response = session.get(f'https://huggingface.co/chat/conversation/{conversationId}/__data.json?x-sveltekit-invalidated=11',)
+            conversationId = response.json().get('conversationId')
+            response = session.get(f'https://huggingface.co/chat/conversation/{conversationId}/__data.json?x-sveltekit-invalidated=11')
 
-            data: list = (response.json())["nodes"][1]["data"]
+            data: list = response.json()["nodes"][1]["data"]
             keys: list[int] = data[data[0]["messages"]]
             message_keys: dict = data[keys[0]]
             messageId: str = data[message_keys["id"]]
@@ -124,22 +126,26 @@ class HuggingChat(AbstractProvider, ProviderModelMixin):
                 files=files,
             )
 
-            first_token = True
+            full_response = ""
             for line in response.iter_lines():
-                line = json.loads(line)
+                if not line:
+                    continue
+                try:
+                    line = json.loads(line)
+                except json.JSONDecodeError as e:
+                    print(f"Failed to decode JSON: {line}, error: {e}")
+                    continue
                 
                 if "type" not in line:
                     raise RuntimeError(f"Response: {line}")
                 
                 elif line["type"] == "stream":
-                    token = line["token"]
-                    if first_token:
-                        token = token.lstrip().replace('\u0000', '')
-                        first_token = False
-                    else:
-                        token = token.replace('\u0000', '')
-
-                    yield token
+                    token = line["token"].replace('\u0000', '')
+                    full_response += token
                 
                 elif line["type"] == "finalAnswer":
                     break
+            
+            full_response = full_response.replace('<|im_end|', '').replace('\u0000', '').strip()
+
+            yield full_response
-- 
cgit v1.2.3