gpt4free/g4f/Provider/needs_auth/Grok.py

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)