diff options
Diffstat (limited to 'g4f/Provider/AItianhu.py')
| -rw-r--r-- | g4f/Provider/AItianhu.py | 36 |
1 files changed, 18 insertions, 18 deletions
diff --git a/g4f/Provider/AItianhu.py b/g4f/Provider/AItianhu.py index 18ab08d8..2fdcfcde 100644 --- a/g4f/Provider/AItianhu.py +++ b/g4f/Provider/AItianhu.py @@ -24,9 +24,9 @@ class AItianhu(AsyncGeneratorProvider): if not cookies: cookies = get_cookies(domain_name='www.aitianhu.com') - if not cookies: - raise RuntimeError(f"g4f.provider.{cls.__name__} requires cookies [refresh https://www.aitianhu.com on chrome]") - + if not cookies: + raise RuntimeError(f"g4f.provider.{cls.__name__} requires cookies [refresh https://www.aitianhu.com on chrome]") + data = { "prompt": format_prompt(messages), "options": {}, @@ -35,7 +35,7 @@ class AItianhu(AsyncGeneratorProvider): "top_p": 1, **kwargs } - + headers = { 'authority': 'www.aitianhu.com', 'accept': 'application/json, text/plain, */*', @@ -51,31 +51,31 @@ class AItianhu(AsyncGeneratorProvider): 'sec-fetch-site': 'same-origin', 'user-agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/118.0.0.0 Safari/537.36', } - + async with StreamSession(headers=headers, - cookies=cookies, - timeout=timeout, - proxies={"https": proxy}, - impersonate="chrome107", verify=False) as session: + cookies=cookies, + timeout=timeout, + proxies={"https": proxy}, + impersonate="chrome107", verify=False) as session: async with session.post(f"{cls.url}/api/chat-process", json=data) as response: response.raise_for_status() - + async for line in response.iter_lines(): if line == b"<script>": raise RuntimeError("Solve challenge and pass cookies") - + if b"platform's risk control" in line: raise RuntimeError("Platform's Risk Control") - + line = json.loads(line) - - if "detail" in line: - content = line["detail"]["choices"][0]["delta"].get("content") - if content: - yield content - else: + + if "detail" not in line: raise RuntimeError(f"Response: {line}") + if content := line["detail"]["choices"][0]["delta"].get( + "content" + ): + yield content @classmethod |