summaryrefslogtreecommitdiffstats
path: root/g4f/Provider/AItianhuSpace.py
blob: 312cb3b3e1087f8b75d9648982b9554d3fa7263d (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
from __future__ import annotations

import time
import random

from ..typing import CreateResult, Messages
from .base_provider import BaseProvider
from .helper import WebDriver, format_prompt, get_browser, get_random_string
from .. import debug

class AItianhuSpace(BaseProvider):
    url = "https://chat3.aiyunos.top/"
    working = True
    supports_stream = True
    supports_gpt_35_turbo = True
    _domains = ["aitianhu.com", "aitianhu1.top"]

    @classmethod
    def create_completion(
        cls,
        model: str,
        messages: Messages,
        stream: bool,
        domain: str = None,
        proxy: str = None,
        timeout: int = 120,
        browser: WebDriver = None,
        headless: bool = True,
        **kwargs
    ) -> CreateResult:
        if not model:
            model = "gpt-3.5-turbo"
        if not domain:
            rand = get_random_string(6)
            domain = random.choice(cls._domains)
            domain = f"{rand}.{domain}"
        if debug.logging:
            print(f"AItianhuSpace | using domain: {domain}")
        url = f"https://{domain}"
        prompt = format_prompt(messages)
        driver = browser if browser else get_browser("", headless, proxy)

        from selenium.webdriver.common.by import By
        from selenium.webdriver.support.ui import WebDriverWait
        from selenium.webdriver.support import expected_conditions as EC

        wait = WebDriverWait(driver, timeout)

        # Bypass devtools detection
        driver.get("https://blank.page/")
        wait.until(EC.visibility_of_element_located((By.ID, "sheet")))
        driver.execute_script(f"""
document.getElementById('sheet').addEventListener('click', () => {{
    window.open('{url}', '_blank');
}});
""")
        driver.find_element(By.ID, "sheet").click()
        time.sleep(10)

        original_window = driver.current_window_handle
        for window_handle in driver.window_handles:
            if window_handle != original_window:
                driver.close()
                driver.switch_to.window(window_handle)
                break

        # Wait for page load
        wait.until(EC.visibility_of_element_located((By.CSS_SELECTOR, "textarea.n-input__textarea-el")))

        try:
            # Register hook in XMLHttpRequest
            script = """
const _http_request_open = XMLHttpRequest.prototype.open;
window._last_message = window._message = "";
window._loadend = false;
XMLHttpRequest.prototype.open = function(method, url) {
    if (url == "/api/chat-process") {
        this.addEventListener("progress", (event) => {
            const lines = this.responseText.split("\\n");
            try {
                window._message = JSON.parse(lines[lines.length-1])["text"];
            } catch(e) { }
        });
        this.addEventListener("loadend", (event) => {
            window._loadend = true;
        });
    }
    return _http_request_open.call(this, method, url);
}
"""
            driver.execute_script(script)

            # Submit prompt
            driver.find_element(By.CSS_SELECTOR, "textarea.n-input__textarea-el").send_keys(prompt)
            driver.find_element(By.CSS_SELECTOR, "button.n-button.n-button--primary-type.n-button--medium-type").click()

            # Read response
            while True:
                chunk = driver.execute_script("""
if (window._message && window._message != window._last_message) {
    try {
        return window._message.substring(window._last_message.length);
    } finally {
        window._last_message = window._message;
    }
}
if (window._loadend) {
    return null;
}
return "";
""")
                if chunk:
                    yield chunk
                elif chunk != "":
                    break
                else:
                    time.sleep(0.1)
        finally:
            if not browser:
                driver.close()
                time.sleep(0.1)
                driver.quit()