測試提交

Author: Jayden Chao

api/GrokAI.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 from dataclasses import dataclass
+from typing import Literal
 import cloudscraper
 from rich.console import Console
 import json
@@ -15,11 +16,7 @@ def print_by_char_rich(console, text, delay=0.002):
 @dataclass
 class GrokAccount:
     cookies: dict
-    message: str
     headers: dict
-    NewChat: bool = True
-    ChatID: str = None
-    modelName: str = "grok-2"
 
 
 class GrokAI:
@@ -45,31 +42,33 @@ class GrokAI:
     def __init__(self,GrokAccount: GrokAccount):
         self.GrokAccount = GrokAccount
         self.Console = Console()
-        obj = self.GrokAccount
+        #obj = self.GrokAccount
 
-        has_new_chat = obj.NewChat  # 检查 NewChat 是否为 True
-        has_chat_id = obj.ChatID is not None  # 检查 ChatID 是否非 None
+       
+        
+    def Chat(self,meassge: str ,disableSearch: bool = False, enableImageGeneration: bool = True, isReasoning: bool = False,debug: bool = False,NewChat:bool = True,ChatID:str = None,ModelName: Literal['grok-3','grok-latest'] = 'grok-latest',deepsearch: Literal['deepsearch,deepersearch',None] = None) -> None:
+        has_new_chat = NewChat  # 检查 NewChat 是否为 True
+        has_chat_id = ChatID is not None  # 检查 ChatID 是否非 None
         if has_new_chat and has_chat_id:
             raise ValueError("NewChat 和 ChatID 不能同时为 True 和非 None")
 
-    def Chat(self, disableSearch: bool = False, enableImageGeneration: bool = True, imageGenerationCount: int = 2, isReasoning: bool = False,test: bool = False):
         self.cookies = self.GrokAccount.cookies
-        self.message = self.GrokAccount.message
-        self.modelName = self.GrokAccount.modelName
-        self.NewChat = self.GrokAccount.NewChat
-        self.ChatID = self.GrokAccount.ChatID
+        self.message = meassge
+        self.modelName = ModelName
+        self.NewChat = NewChat
+        self.ChatID = ChatID
         self.headers = self.GrokAccount.headers
         self.headers['referer'] = 'https://grok.com/chat/'
         self.responseUrl: str
         global has_valid_response
         has_valid_response = False
 
-        scraper = cloudscraper.create_scraper()
+        scraper = cloudscraper.create_scraper(delay=10, browser={'browser': 'chrome', 'platform': 'windows', 'mobile': False})
 
         if self.NewChat is False:
             self.url = f"https://grok.com/rest/app-chat/conversations/{self.ChatID}/responses"
             response = scraper.get(f"https://grok.com/rest/app-chat/conversations/{self.ChatID}/response-node")
-            '''if response.status_code == 200:
+            if response.status_code == 200:
                 response_data = response.json()
                 response_nodes = response_data.get("responseNodes", [])
                 if response_nodes:
@@ -81,7 +80,7 @@ def Chat(self, disableSearch: bool = False, enableImageGeneration: bool = True,
                         responseId = None
                 else:
                     print("No responseNodes found in the JSON data.")
-                    responseId = None'''
+                    responseId = None
         else:
             self.url = "https://grok.com/rest/app-chat/conversations/new"
 
@@ -96,7 +95,7 @@ def Chat(self, disableSearch: bool = False, enableImageGeneration: bool = True,
             'returnRawGrokInXaiRequest': False,
             'fileAttachments': [],
             'enableImageStreaming': True,
-            'imageGenerationCount': imageGenerationCount,
+            'imageGenerationCount': 2,
             'forceConcise': False,
             'toolOverrides': {},
             'enableSideBySide': True,
@@ -106,18 +105,21 @@ def Chat(self, disableSearch: bool = False, enableImageGeneration: bool = True,
             'disableTextFollowUps': True
         }
 
+        if deepsearch is not None:
+            self.data["deepsearchPreset"] = deepsearch
+
         console = self.Console
         # 创建 cloudscraper 实例
 
         response = scraper.post(self.url, headers=self.headers, cookies=self.cookies, json=self.data, stream=True)
-        if test is True:
+        if debug is True:
             # 這是測試代碼!!!!
             console.print(f"Request URL: {self.url}", style="bold blue")
             console.print(f"Headers: {self.headers}", style="bold blue")
             console.print(f"Cookies: {self.cookies}", style="bold blue")
             console.print(f"Data: {self.data}", style="bold blue")
 
-
+        text = response.text
         # 檢查狀態碼
         if response.status_code == 200:
             for line in response.iter_lines():
@@ -129,12 +131,13 @@ def Chat(self, disableSearch: bool = False, enableImageGeneration: bool = True,
 
                         # 提取 token
                         token = json_data.get("result", {}).get("response", {}).get("token", "")
-                        
+                        generated_image_urls = json_data.get("result", {}).get("modelResponse", {}).get("generatedImageUrls", [])
                         # 如果 token 存在，逐字輸出
-                        if token:
+                        if token or generated_image_urls:
+                            # 如果有圖片網址，打印圖片網址
 
                             has_valid_response = True
-                            print_by_char_rich(console, token, delay=0.05)
+                            print_by_char_rich(console, token if token else generated_image_urls, delay=0.05)
 
                     except json.JSONDecodeError:
                         console.print("解析 JSON 失败", style="bold red")
@@ -146,10 +149,14 @@ def Chat(self, disableSearch: bool = False, enableImageGeneration: bool = True,
         else:
             console.print("请求失败", style="bold red")
             console.print(response.status_code, style="bold red")
-            if "Just a moment..." in response.text or "<title>Just a moment...</title>" in response.text:
+            if "Just a moment..." in text or "<title>Just a moment...</title>" in text:
                 console.print("請求被 Cloudflare 防護攔截，請更新 Cookies 或 Headers。", style="bold red")
             else:
-                console.print(response.text, style="bold red")
+                console.print(text, style="bold red")
+        
+        if debug is True:
+            # 這是測試代碼!!!!
+            console.print(text, style="bold yellow")
             '''
         for line in response.iter_lines():
             if line:

test/.gitignore

@@ -0,0 +1,2 @@
+# Automatically created by ruff.
+*.txt

test/test.py

@@ -0,0 +1,43 @@
+import sys
+import os
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from api.GrokAI import GrokAI, GrokAccount
+from rich.console import Console
+
+cookies = {
+    '_ga': 'GA1.1.1691019154.1742809195',
+    'sso': 'eyJhbGciOiJIUzI1NiJ9.eyJzZXNzaW9uX2lkIjoiNGNjOTZlOTctYzNlYS00NTUyLTlhYWYtMWI1ZjJiNmY3MDNjIn0.BnY3wcBEA70vch8OHI8a_kBzhKfene7JsEYSvPuHgU0',
+    'sso-rw': 'eyJhbGciOiJIUzI1NiJ9.eyJzZXNzaW9uX2lkIjoiNGNjOTZlOTctYzNlYS00NTUyLTlhYWYtMWI1ZjJiNmY3MDNjIn0.BnY3wcBEA70vch8OHI8a_kBzhKfene7JsEYSvPuHgU0',
+    'i18nextLng': 'zh-TW',
+    'cf_clearance': 'BVyVUo3SHFMJCGaiqOdcFp0wgBTjyegtIsMJbROWumM-1743260506-1.2.1.1-93HWk2jjry2fFZMkL88eD2rgI7IapPgbrBVF7vGUNOeSX3PwiyoFTROqLKEjv18XcTqQ4dK7kuczHuKc5mQxQeBTJh23FF93IgXeSYHyNurzfqmBSqdqXw3cbkXxlbq74AmT1i9utURMwI20p3n9_.Mr8FNUX6hGpLV__H.HHhAXO.tnuhC9e2VgEGzxSYH.wCBkO76Rd12SlFYNFgNdQZ82BWk.bG791kKBLO9O3y8tXOv0okNnPXXfoqLb7Z4LmBffzSsf.4iCzif1L4QkTSfi0HxIoBL0dm6rDmpSqNb4L14EnwIB7YqvUZRAus3Z.0C01vtxtDMgB9l2aoeCnd5ImtTELwSlTBPG4NGKC_jtWlWkmydmj3xhVH4AhW23v5TudF8gWu3kDpPv7fBpmtrcjBZlZxP564Xt1DrLC7Y',
+    '_ga_8FEWB057YH': 'GS1.1.1743232889.3.1.1743233813.0.0.0',
+    '__cf_bm': 'kJlhUhv_Q5b7j8.PEBNA2oaPn3HYleB7jBVNpcQbzRQ-1743260665-1.0.1.1-YTE7Af.sWDNUsQp9USARIPeNHhiLo5TXUNoZdzs3cdk5gXc2XkcbSzPYBfyTQ3N7sSeqCGLUd4UBPW1txz0jsvWzAZ9oymP0AIureR7fubM'
+}
+account = GrokAccount(
+    cookies=cookies,
+    headers={
+        'accept': '*/*',
+        'content-type': 'application/json',
+        'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) Chrome/134.0.0.0 Safari/537.36'
+    },
+)
+
+# Initialize the GrokAI instance
+grok_ai = GrokAI(account)
+
+# Use the Chat method to send a message
+console = Console()
+console.print("Sending message to Grok AI...", style="bold green")
+
+try:
+    grok_ai.Chat(
+        disableSearch=False,
+        enableImageGeneration=True,
+        isReasoning=False,
+        NewChat=True, 
+        meassge="please generate a picture of a cat",
+        debug=True,
+        
+    )
+except Exception as e:
+    console.print(f"An error occurred: {e}", style="bold red")