diff options
Diffstat (limited to 'g4f/Provider/AItianhu.py')
| -rw-r--r-- | g4f/Provider/AItianhu.py | 17 |
1 files changed, 14 insertions, 3 deletions
diff --git a/g4f/Provider/AItianhu.py b/g4f/Provider/AItianhu.py index 1d8653b7..c6e4dbad 100644 --- a/g4f/Provider/AItianhu.py +++ b/g4f/Provider/AItianhu.py @@ -4,7 +4,7 @@ import json from ..typing import AsyncGenerator from ..requests import StreamSession -from .base_provider import AsyncGeneratorProvider, format_prompt +from .base_provider import AsyncGeneratorProvider, format_prompt, get_cookies class AItianhu(AsyncGeneratorProvider): @@ -18,8 +18,12 @@ class AItianhu(AsyncGeneratorProvider): model: str, messages: list[dict[str, str]], proxy: str = None, + cookies: dict = None, + timeout: int = 30, **kwargs ) -> AsyncGenerator: + if not cookies: + cookies = get_cookies("www.aitianhu.com") data = { "prompt": format_prompt(messages), "options": {}, @@ -34,12 +38,19 @@ class AItianhu(AsyncGeneratorProvider): "Origin": cls.url, "Referer": f"{cls.url}/" } - async with StreamSession(headers=headers, proxies={"https": proxy}, impersonate="chrome107", verify=False) as session: + async with StreamSession( + headers=headers, + 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") + 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) |