File size: 12,308 Bytes
a4b70d9 |
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 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 |
from __future__ import annotations
import os
import json
import time
import asyncio
import uuid
from typing import Dict, Any, AsyncIterator
try:
import nodriver
except ImportError:
pass
from ...typing import Messages, AsyncResult
from ...providers.response import JsonConversation, Reasoning, ImagePreview, ImageResponse, TitleGeneration, AuthResult, RequestLogin
from ...requests import StreamSession, get_nodriver, DEFAULT_HEADERS, merge_cookies
from ...requests.raise_for_status import raise_for_status
from ...errors import MissingAuthError
from ..base_provider import AsyncAuthedProvider, ProviderModelMixin
from ..helper import format_prompt, get_last_user_message
class Conversation(JsonConversation):
def __init__(self,
conversation_id: str
) -> None:
self.conversation_id = conversation_id
class Grok(AsyncAuthedProvider, ProviderModelMixin):
label = "Grok AI"
url = "https://grok.com"
cookie_domain = ".grok.com"
assets_url = "https://assets.grok.com"
conversation_url = "https://grok.com/rest/app-chat/conversations"
needs_auth = True
working = True
# Updated to Grok 4 as default
default_model = "grok-4"
# Updated model list with latest Grok 4 and 3 models
models = [
# Grok 4 models
"grok-4",
"grok-4-heavy",
"grok-4-reasoning",
# Grok 3 models
"grok-3",
"grok-3-reasoning",
"grok-3-mini",
"grok-3-mini-reasoning",
# Legacy Grok 2 (still supported)
"grok-2",
"grok-2-image",
# Latest aliases
"grok-latest",
]
model_aliases = {
# Grok 3 aliases
"grok-3-thinking": "grok-3-reasoning",
"grok-3-r1": "grok-3-reasoning",
"grok-3-mini-thinking": "grok-3-mini-reasoning",
# Latest alias
"grok": "grok-latest",
}
@classmethod
async def on_auth_async(cls, proxy: str = None, **kwargs) -> AsyncIterator:
auth_result = AuthResult(headers=DEFAULT_HEADERS, impersonate="chrome")
auth_result.headers["referer"] = cls.url + "/"
browser, stop_browser = await get_nodriver(proxy=proxy)
yield RequestLogin(cls.__name__, os.environ.get("G4F_LOGIN_URL") or "")
try:
page = await browser.get(cls.url)
has_headers = False
def on_request(event: nodriver.cdp.network.RequestWillBeSent, page=None):
nonlocal has_headers
if event.request.url.startswith(cls.conversation_url + "/new"):
for key, value in event.request.headers.items():
auth_result.headers[key.lower()] = value
has_headers = True
await page.send(nodriver.cdp.network.enable())
page.add_handler(nodriver.cdp.network.RequestWillBeSent, on_request)
await page.reload()
auth_result.headers["user-agent"] = await page.evaluate("window.navigator.userAgent", return_by_value=True)
while True:
if has_headers:
break
textarea = await page.select("textarea", 180)
await textarea.send_keys("Hello")
await asyncio.sleep(1)
button = await page.select("button[type='submit']")
if button:
await button.click()
await asyncio.sleep(1)
auth_result.cookies = {}
for c in await page.send(nodriver.cdp.network.get_cookies([cls.url])):
auth_result.cookies[c.name] = c.value
await page.close()
finally:
stop_browser()
yield auth_result
@classmethod
async def _prepare_payload(cls, model: str, message: str) -> Dict[str, Any]:
# Map model names to API model names
api_model = "grok-latest"
if model in ["grok-4", "grok-4-heavy", "grok-4-reasoning"]:
api_model = model
elif model == "grok-3":
api_model = "grok-3"
elif model in ["grok-3-mini", "grok-3-mini-reasoning"]:
api_model = "grok-3-mini"
elif model == "grok-2":
api_model = "grok-2"
# Check if it's a reasoning model
is_reasoning = model.endswith("-reasoning") or model.endswith("-thinking") or model.endswith("-r1")
# Enable Big Brain mode for heavy models
enable_big_brain = "heavy" in model or "big-brain" in model
# Enable DeepSearch for Grok 3+ models
enable_deep_search = not model.startswith("grok-2")
return {
"temporary": True,
"modelName": api_model,
"message": message,
"fileAttachments": [],
"imageAttachments": [],
"disableSearch": False,
"enableImageGeneration": model == "grok-2-image" or model == "grok-4",
"returnImageBytes": False,
"returnRawGrokInXaiRequest": False,
"enableImageStreaming": True,
"imageGenerationCount": 2,
"forceConcise": False,
"toolOverrides": {},
"enableSideBySide": True,
"isPreset": False,
"sendFinalMetadata": True,
"customInstructions": "",
"deepsearchPreset": "enabled" if enable_deep_search else "",
"isReasoning": is_reasoning,
"enableBigBrain": enable_big_brain,
"enableLiveSearch": False, # Real-time search for Grok 4
"contextWindow": 256000 if model.startswith("grok-4") else 131072, # 256k for Grok 4, 128k for others
}
@classmethod
async def create_authed(
cls,
model: str,
messages: Messages,
auth_result: AuthResult,
conversation: Conversation = None,
**kwargs
) -> AsyncResult:
conversation_id = None if conversation is None else conversation.conversation_id
prompt = format_prompt(messages) if conversation_id is None else get_last_user_message(messages)
async with StreamSession(
**auth_result.get_dict()
) as session:
payload = await cls._prepare_payload(model, prompt)
# Add voice mode support flag (for future use)
if kwargs.get("enable_voice", False):
payload["enableVoiceMode"] = True
if conversation_id is None:
url = f"{cls.conversation_url}/new"
else:
url = f"{cls.conversation_url}/{conversation_id}/responses"
async with session.post(url, json=payload, headers={"x-xai-request-id": str(uuid.uuid4())}) as response:
if response.status == 403:
raise MissingAuthError("Invalid secrets")
auth_result.cookies = merge_cookies(auth_result.cookies, response)
await raise_for_status(response)
thinking_duration = None
deep_search_active = False
async for line in response.iter_lines():
if line:
try:
json_data = json.loads(line)
result = json_data.get("result", {})
if conversation_id is None:
conversation_id = result.get("conversation", {}).get("conversationId")
response_data = result.get("response", {})
# Handle DeepSearch status
deep_search = response_data.get("deepSearchStatus")
if deep_search:
if not deep_search_active:
deep_search_active = True
yield Reasoning(status="🔍 Deep searching...")
if deep_search.get("completed"):
deep_search_active = False
yield Reasoning(status="Deep search completed")
# Handle image generation (Aurora for Grok 3+)
image = response_data.get("streamingImageGenerationResponse", None)
if image is not None:
image_url = image.get("imageUrl")
if image_url:
yield ImagePreview(
f'{cls.assets_url}/{image_url}',
"",
{"cookies": auth_result.cookies, "headers": auth_result.headers}
)
# Handle text tokens
token = response_data.get("token", result.get("token"))
is_thinking = response_data.get("isThinking", result.get("isThinking"))
if token:
if is_thinking:
if thinking_duration is None:
thinking_duration = time.time()
# Different status for different models
if "grok-4" in model:
status = "🧠 Grok 4 is processing..."
elif "big-brain" in payload and payload["enableBigBrain"]:
status = "🧠 Big Brain mode active..."
else:
status = "🤔 Is thinking..."
yield Reasoning(status=status)
yield Reasoning(token)
else:
if thinking_duration is not None:
thinking_duration = time.time() - thinking_duration
status = f"Thought for {thinking_duration:.2f}s" if thinking_duration > 1 else ""
thinking_duration = None
yield Reasoning(status=status)
yield token
# Handle generated images
generated_images = response_data.get("modelResponse", {}).get("generatedImageUrls", None)
if generated_images:
yield ImageResponse(
[f'{cls.assets_url}/{image}' for image in generated_images],
"",
{"cookies": auth_result.cookies, "headers": auth_result.headers}
)
# Handle title generation
title = result.get("title", {}).get("newTitle", "")
if title:
yield TitleGeneration(title)
# Handle tool usage information (Grok 4)
tool_usage = response_data.get("toolUsage")
if tool_usage:
tools_used = tool_usage.get("toolsUsed", [])
if tools_used:
yield Reasoning(status=f"Used tools: {', '.join(tools_used)}")
except json.JSONDecodeError:
continue
# Return conversation ID for continuation
if conversation_id is not None and kwargs.get("return_conversation", False):
yield Conversation(conversation_id)
|