from __future__ import annotations
import os
import json
import uuid
import time
import base64
import sys
import inspect
import secrets
import asyncio
import contextvars
import functools
import re
import threading
from dataclasses import dataclass, field
from typing import List, Optional, Any, Dict, Union, Tuple, Iterator
from pathlib import Path
from contextlib import contextmanager
from concurrent.futures import ThreadPoolExecutor
from enum import Enum
import requests
from flask import (
Flask,
request,
Response,
jsonify,
stream_with_context,
render_template,
redirect,
session,
)
from curl_cffi import requests as curl_requests
from werkzeug.middleware.proxy_fix import ProxyFix
from playwright.async_api import async_playwright, Browser, BrowserContext
class PlaywrightStatsigManager:
"""
x-statsig-id capture using Playwright (adapted from Grok3API driver.py)
This approach captures authentic x-statsig-id headers by:
1. Patching window.fetch to intercept grok.com's own API calls
2. Triggering a real request on grok.com to generate authentic headers
3. Capturing and storing the real x-statsig-id for reuse
"""
def __init__(self, proxy_url: Optional[str] = None):
self._cached_statsig_id: Optional[str] = None
self._cache_timestamp: Optional[int] = None
self._cache_duration = 300
self._context: Optional[BrowserContext] = None
self._playwright = None
self._lock = threading.Lock()
self._base_url = "https://grok.com/"
self._proxy_url = proxy_url
def _run_async(self, coro):
"""Run async function in thread-safe manner"""
try:
loop = asyncio.get_event_loop()
if loop.is_running():
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(asyncio.run, coro)
return future.result(timeout=300)
else:
return loop.run_until_complete(coro)
except RuntimeError:
return asyncio.run(coro)
async def _ensure_browser(self):
"""Ensure browser is available and ready"""
if not self._context:
self._playwright = await async_playwright().start()
context_options = {
"viewport": {"width": 1920, "height": 1080},
"user_agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36",
}
if self._proxy_url:
context_options["proxy"] = {"server": self._proxy_url}
self._context = await self._playwright.chromium.launch_persistent_context(
user_data_dir="./data/chrome",
headless=True,
no_viewport=True,
channel="chrome",
args=[
"--no-first-run",
"--force-color-profile=srgb",
"--metrics-recording-only",
"--password-store=basic",
"--no-default-browser-check",
"--no-sandbox",
"--disable-dev-shm-usage",
"--disable-gpu",
"--disable-web-security",
"--disable-features=VizDisplayCompositor",
"--user-agent=Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36",
],
**context_options,
)
async def check_real_ip(self) -> str:
"""Check the real IP address using Playwright browser"""
try:
await self._ensure_browser()
page = await self._context.new_page() # type: ignore
try:
print("Checking real IP address via ipify API")
await page.goto("https://api.ipify.org?format=json", timeout=30000)
content = await page.content()
ip_info = await page.evaluate(
"""
() => {
try {
const bodyText = document.body.textContent || document.body.innerText;
return JSON.parse(bodyText);
} catch (e) {
return null;
}
}
"""
)
if ip_info and ip_info.get("ip"):
ip_address = ip_info["ip"]
print(f"Playwright real IP address: {ip_address}")
return ip_address
else:
print("Failed to parse IP from ipify response")
return "unknown"
except Exception as e:
print(f"Error checking IP address: {e}")
return "error"
finally:
await page.close()
except Exception as e:
print(f"Failed to check real IP address: {e}")
return "failed"
async def _cleanup(self):
"""Clean up browser resources"""
if self._context:
await self._context.close()
self._context = None
if self._playwright:
await self._playwright.stop()
self._playwright = None
def cleanup(self):
"""Synchronous cleanup wrapper"""
if self._context:
self._run_async(self._cleanup())
async def _patch_fetch_for_statsig(self, page):
"""Patch window.fetch to intercept x-statsig-id headers (adapted from driver.py)"""
result = await page.evaluate(
"""
(() => {
if (window.__fetchPatched) {
return "fetch already patched";
}
window.__fetchPatched = false;
const originalFetch = window.fetch;
window.__xStatsigId = null;
window.fetch = async function(...args) {
console.log("Intercepted fetch call with args:", args);
const response = await originalFetch.apply(this, args);
try {
const req = args[0];
const opts = args[1] || {};
const url = typeof req === 'string' ? req : req.url;
const headers = opts.headers || {};
const targetUrl = "https://grok.com/rest/app-chat/conversations/new";
if (url === targetUrl) {
let id = null;
if (headers["x-statsig-id"]) {
id = headers["x-statsig-id"];
} else if (typeof opts.headers?.get === "function") {
id = opts.headers.get("x-statsig-id");
}
if (id) {
window.__xStatsigId = id;
console.log("Captured x-statsig-id:", id);
} else {
console.warn("x-statsig-id not found in headers");
}
} else {
console.log("Skipped fetch, URL doesn't match target:", url);
}
} catch (e) {
console.warn("Error capturing x-statsig-id:", e);
}
return response;
};
window.__fetchPatched = true;
return "fetch successfully patched";
})()
"""
)
print(f"Fetch patching result: {result}")
async def _initiate_answer(self, page):
"""Trigger a real request to grok.com to capture x-statsig-id"""
try:
await page.wait_for_selector("div.relative.z-10 textarea", timeout=30000)
import random
import string
random_char = random.choice(string.ascii_lowercase)
await page.fill("div.relative.z-10 textarea", random_char)
await page.press("div.relative.z-10 textarea", "Enter")
print(f"Triggered request with character: {random_char}")
except Exception as e:
print(f"Error triggering answer: {e}")
title = await page.title()
url = page.url
print(f"Page title: {title}, URL: {url}")
raise
async def _capture_statsig_id_async(
self, restart_session: bool = False
) -> Optional[str]:
"""Capture x-statsig-id from real grok.com interaction"""
try:
await self._ensure_browser()
page = await self._context.new_page() # type: ignore
try:
print("Navigating to grok.com")
await page.goto(
self._base_url, wait_until="domcontentloaded", timeout=30000
)
await self._patch_fetch_for_statsig(page)
captcha_visible = await page.evaluate(
"""
(() => {
const elements = document.querySelectorAll("p");
for (const el of elements) {
if (el.textContent.includes("Making sure you're human")) {
const style = window.getComputedStyle(el);
if (style.visibility !== 'hidden' && style.display !== 'none') {
return true;
}
}
}
return false;
})()
"""
)
if captcha_visible:
print("Captcha detected, cannot capture x-statsig-id")
return None
await self._initiate_answer(page)
try:
await page.locator("div.message-bubble p[dir='auto']").or_(
page.locator("div.w-full.max-w-\\[48rem\\]")
).or_(
page.locator("p", has_text="Making sure you're human")
).wait_for(
timeout=20000
)
except:
print("No response elements found within timeout")
error_elements = await page.query_selector_all(
"div.w-full.max-w-\\[48rem\\]"
)
if error_elements:
print("Authentication error detected")
return None
captcha_elements = await page.query_selector_all(
"p:has-text('Making sure you\\'re human')"
)
if captcha_elements:
print("Captcha appeared during request")
return None
statsig_id = await page.evaluate("window.__xStatsigId")
if statsig_id:
print(f"Successfully captured x-statsig-id: {statsig_id[:30]}...")
return statsig_id
else:
print("No x-statsig-id was captured")
return None
finally:
await page.close()
except Exception as e:
print(f"Error capturing x-statsig-id: {e}")
return None
def capture_statsig_id(self, restart_session: bool = False) -> Optional[str]:
"""Capture x-statsig-id (sync wrapper)"""
with self._lock:
return self._run_async(self._capture_statsig_id_async(restart_session))
def check_real_ip_sync(self) -> str:
"""Check real IP address (sync wrapper)"""
with self._lock:
return self._run_async(self.check_real_ip())
def generate_xai_request_id(self) -> str:
"""Generate x-xai-request-id (simple UUID)"""
return str(uuid.uuid4())
def get_dynamic_headers(
self, method: str = "POST", pathname: str = "/rest/app-chat/conversations/new"
) -> Dict[str, str]:
"""Get dynamic headers including captured x-statsig-id and x-xai-request-id"""
headers = {}
current_time = int(time.time())
with self._lock:
if (
self._cached_statsig_id
and self._cache_timestamp
and (current_time - self._cache_timestamp) < self._cache_duration
):
print("Using cached x-statsig-id")
headers["x-statsig-id"] = self._cached_statsig_id
if "x-statsig-id" not in headers:
print("Capturing fresh x-statsig-id")
statsig_id = self.capture_statsig_id()
if statsig_id:
with self._lock:
self._cached_statsig_id = statsig_id
self._cache_timestamp = current_time
headers["x-statsig-id"] = statsig_id
else:
print("Failed to capture x-statsig-id, using fallback")
headers["x-statsig-id"] = (
"ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk="
)
headers["x-xai-request-id"] = self.generate_xai_request_id()
print(f"Generated dynamic headers: {list(headers.keys())}")
return headers
_global_statsig_manager: Optional[PlaywrightStatsigManager] = None
def initialize_statsig_manager(proxy_url: Optional[str] = None) -> None:
"""Initialize the global StatsigManager instance with configuration"""
global _global_statsig_manager
if _global_statsig_manager is None:
_global_statsig_manager = PlaywrightStatsigManager(proxy_url=proxy_url)
def get_statsig_manager() -> PlaywrightStatsigManager:
"""Get or create the global StatsigManager instance"""
global _global_statsig_manager
if _global_statsig_manager is None:
_global_statsig_manager = PlaywrightStatsigManager()
return _global_statsig_manager
class ModelType(Enum):
"""Supported Grok model types with new architecture."""
GROK_3 = "grok-3"
GROK_4 = "grok-4"
GROK_AUTO = "grok-auto" # grok-3 + MODEL_MODE_AUTO
GROK_FAST = "grok-fast" # grok-3 + MODEL_MODE_FAST
GROK_EXPERT = "grok-expert" # grok-4 + MODEL_MODE_EXPERT
GROK_SEARCH = "grok-deepsearch" # grok-4 + MODEL_MODE_EXPERT + workspaceIds
GROK_IMAGE = "grok-image" # grok-4 + MODEL_MODE_EXPERT + enableImageGeneration
class TokenType(Enum):
"""Token privilege levels."""
NORMAL = "normal"
SUPER = "super"
class ResponseState(Enum):
"""Response processing states."""
IDLE = "idle"
THINKING = "thinking"
GENERATING_IMAGE = "generating_image"
COMPLETE = "complete"
MESSAGE_LENGTH_LIMIT = 40000
MAX_FILE_ATTACHMENTS = 4
DEFAULT_REQUEST_TIMEOUT = 120000
MAX_RETRY_ATTEMPTS = 3
BASE_RETRY_DELAY = 1.0
BASE_HEADERS = {
"Accept": "*/*",
"Accept-Language": "zh-CN,zh;q=0.9",
"Accept-Encoding": "gzip, deflate, br, zstd",
"Content-Type": "text/plain;charset=UTF-8",
"Connection": "keep-alive",
"Origin": "https://grok.com",
"Priority": "u=1, i",
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36",
"Sec-Ch-Ua": '"Not(A:Brand";v="99", "Google Chrome";v="133", "Chromium";v="133"',
"Sec-Ch-Ua-Mobile": "?0",
"Sec-Ch-Ua-Platform": '"macOS"',
"Sec-Fetch-Dest": "empty",
"Sec-Fetch-Mode": "cors",
"Sec-Fetch-Site": "same-origin",
"Baggage": "sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c",
}
def get_dynamic_headers(
method: str = "POST",
pathname: str = "/rest/app-chat/conversations/new",
config: Optional["ConfigurationManager"] = None,
) -> Dict[str, str]:
"""
Get headers with dynamic x-statsig-id and x-xai-request-id or fallback headers
Args:
method: HTTP method for the request
pathname: Request pathname for statsig generation
config: Configuration manager to check if dynamic headers are disabled
Returns:
Dictionary with all headers including dynamic ones or fallback
"""
try:
headers = BASE_HEADERS.copy()
if config and config.get("API.DISABLE_DYNAMIC_HEADERS", False):
print("Dynamic headers disabled, using fallback headers")
headers["x-xai-request-id"] = str(uuid.uuid4())
headers["x-statsig-id"] = (
"ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk="
)
return headers
statsig_manager = get_statsig_manager()
dynamic_headers = statsig_manager.get_dynamic_headers(method, pathname)
headers.update(dynamic_headers)
print(f"Generated dynamic headers for {method} {pathname}")
return headers
except Exception as e:
print(f"Error generating dynamic headers: {e}")
headers = BASE_HEADERS.copy()
headers["x-xai-request-id"] = str(uuid.uuid4())
headers["x-statsig-id"] = (
"ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk="
)
return headers
class GrokApiException(Exception):
"""Base exception for Grok API errors."""
def __init__(self, message: str, error_code: str = "UNKNOWN_ERROR"):
super().__init__(message)
self.error_code = error_code
class TokenException(GrokApiException):
"""Token-related exceptions."""
pass
class ValidationException(GrokApiException):
"""Input validation exceptions."""
pass
class RateLimitException(GrokApiException):
"""Rate limiting exceptions."""
pass
@dataclass
class TokenCredential:
"""Represents a token credential with validation."""
sso_token: str
token_type: TokenType
def __post_init__(self):
"""Validate token format."""
if not self.sso_token or not self.sso_token.strip():
raise ValidationException("SSO token cannot be empty")
if "sso=" not in self.sso_token:
raise ValidationException("Invalid SSO token format")
try:
parts = self.sso_token.split("sso=")
if len(parts) < 2 or not parts[1]:
raise ValidationException(
"Invalid SSO token format: missing value after 'sso='"
)
except Exception as e:
raise ValidationException(f"Invalid SSO token format: {e}")
@classmethod
def from_raw_token(
cls, raw_token: str, token_type: TokenType = TokenType.NORMAL
) -> "TokenCredential":
"""Create TokenCredential from raw SSO value."""
if not raw_token or not raw_token.strip():
raise ValidationException("Raw token cannot be empty")
sanitized_token = raw_token.strip()
if ";" in sanitized_token:
raise ValidationException("Raw token contains invalid character (';')")
formatted_token = f"sso-rw={sanitized_token};sso={sanitized_token}"
return cls(formatted_token, token_type)
def extract_sso_value(self) -> str:
"""Extract the SSO value from the token."""
try:
return self.sso_token.split("sso=")[1].split(";")[0]
except (IndexError, AttributeError) as e:
raise TokenException(f"Failed to parse SSO token: {self.sso_token}") from e
@dataclass
class GeneratedImage:
"""Represents a generated image with metadata."""
url: str
base_url: str = "https://assets.grok.com"
cookies: List[Dict[str, Any]] = field(default_factory=list)
def __post_init__(self):
"""Validate image data."""
if not self.url:
raise ValidationException("Image URL cannot be empty")
@dataclass
class ProcessingState:
"""Immutable state for response processing."""
is_thinking: bool = False
is_generating_image: bool = False
image_generation_phase: int = 0
def with_thinking(self, thinking: bool) -> "ProcessingState":
"""Return new state with updated thinking status."""
return ProcessingState(
thinking, self.is_generating_image, self.image_generation_phase
)
def with_image_generation(
self, generating: bool, phase: int = 0
) -> "ProcessingState":
"""Return new state with updated image generation status."""
return ProcessingState(self.is_thinking, generating, phase)
@dataclass
class ModelResponse:
"""Enhanced model response with proper validation and transformation."""
response_id: str
message: str
sender: str
create_time: str
parent_response_id: str
manual: bool
partial: bool
shared: bool
query: str
query_type: str
web_search_results: List[Any] = field(default_factory=list)
xpost_ids: List[Any] = field(default_factory=list)
xposts: List[Any] = field(default_factory=list)
generated_images: List[GeneratedImage] = field(default_factory=list)
image_attachments: List[Any] = field(default_factory=list)
file_attachments: List[Any] = field(default_factory=list)
card_attachments_json: List[Any] = field(default_factory=list)
file_uris: List[Any] = field(default_factory=list)
file_attachments_metadata: List[Any] = field(default_factory=list)
is_control: bool = False
steps: List[Any] = field(default_factory=list)
media_types: List[Any] = field(default_factory=list)
@classmethod
def from_api_response(
cls, data: Dict[str, Any], enable_artifact_files: bool = False
) -> "ModelResponse":
"""Create ModelResponse from API response data with validation."""
try:
response_id = str(data.get("responseId", ""))
sender = str(data.get("sender", ""))
create_time = str(data.get("createTime", ""))
parent_response_id = str(data.get("parentResponseId", ""))
query = str(data.get("query", ""))
query_type = str(data.get("queryType", ""))
manual = bool(data.get("manual", False))
partial = bool(data.get("partial", False))
shared = bool(data.get("shared", False))
is_control = bool(data.get("isControl", False))
raw_message = data.get("message", "")
processed_message = cls._transform_xai_artifacts(str(raw_message))
generated_images = []
for image_url in data.get("generatedImageUrls", []):
if image_url:
generated_images.append(GeneratedImage(url=str(image_url)))
return cls(
response_id=response_id,
message=processed_message,
sender=sender,
create_time=create_time,
parent_response_id=parent_response_id,
manual=manual,
partial=partial,
shared=shared,
query=query,
query_type=query_type,
web_search_results=data.get("webSearchResults", []),
xpost_ids=data.get("xpostIds", []),
xposts=data.get("xposts", []),
generated_images=generated_images,
image_attachments=data.get("imageAttachments", []),
file_attachments=data.get("fileAttachments", []),
card_attachments_json=data.get("cardAttachmentsJson", []),
file_uris=data.get("fileUris", []),
file_attachments_metadata=data.get("fileAttachmentsMetadata", []),
is_control=is_control,
steps=data.get("steps", []),
media_types=data.get("mediaTypes", []),
)
except Exception as e:
print(f"Failed to create ModelResponse: {e}")
return cls(
response_id="",
message="Error processing response",
sender="system",
create_time=str(int(time.time())),
parent_response_id="",
manual=False,
partial=False,
shared=False,
query="",
query_type="",
)
@staticmethod
def _transform_xai_artifacts(text: str) -> str:
"""
Transform xaiArtifact blocks to proper markdown code blocks.
Comprehensive version that handles all xaiArtifact formats including:
1. blocks → ```\ncode\n```
2. ```x-src format → ```
3. ```x- format → ```
4. Any xaiArtifact with artifact_id, title, etc.
5. Self-closing xaiArtifact tags
"""
if not text:
return text
def replace_artifact_with_content(match):
full_match = match.group(0)
content = match.group(1).strip() if match.group(1) else ""
content_type_match = re.search(r'contentType="([^"]+)"', full_match)
if content_type_match:
content_type = content_type_match.group(1).strip()
if "/" in content_type:
lang = content_type.split("/")[-1]
else:
lang = content_type
if content:
return f"```{lang}\n{content}\n```"
else:
return ""
else:
return content
text = re.sub(
r"]*?>(.*?)",
replace_artifact_with_content,
text,
flags=re.DOTALL,
)
text = re.sub(r"]*?/>", "", text)
text = re.sub(r"]*>", "", text)
text = re.sub(r"", "", text)
text = re.sub(
r"```x-([a-zA-Z0-9_+-]+)src\b", lambda m: f"```{m.group(1)}", text
)
text = re.sub(
r"```x-([a-zA-Z0-9_+-]+)\b(?![a-zA-Z0-9_-]*src)",
lambda m: f"```{m.group(1)}",
text,
)
return text
@dataclass
class GrokResponse:
"""Complete Grok API response wrapper."""
model_response: ModelResponse
is_thinking: bool = False
is_soft_stop: bool = False
response_id: str = ""
conversation_id: Optional[str] = None
title: Optional[str] = None
conversation_create_time: Optional[str] = None
conversation_modify_time: Optional[str] = None
temporary: Optional[bool] = None
error: Optional[str] = None
error_code: Optional[Union[int, str]] = None
@classmethod
def from_api_response(
cls, data: Dict[str, Any], enable_artifact_files: bool = False
) -> "GrokResponse":
"""Create GrokResponse from API response data."""
try:
error = data.get("error")
error_code = data.get("error_code")
result = data.get("result", {})
response_data = result.get("response", {})
model_response = ModelResponse.from_api_response(
response_data.get("modelResponse", {}), enable_artifact_files
)
is_thinking = bool(response_data.get("isThinking", False))
is_soft_stop = bool(response_data.get("isSoftStop", False))
response_id = str(response_data.get("responseId", ""))
conversation_id = response_data.get("conversationId")
new_title = result.get("newTitle") or result.get("title")
title = new_title if new_title else None
conversation_create_time = response_data.get("createTime")
conversation_modify_time = response_data.get("modifyTime")
temporary = response_data.get("temporary")
return cls(
model_response=model_response,
is_thinking=is_thinking,
is_soft_stop=is_soft_stop,
response_id=response_id,
conversation_id=conversation_id,
title=title,
conversation_create_time=conversation_create_time,
conversation_modify_time=conversation_modify_time,
temporary=temporary,
error=error,
error_code=error_code,
)
except Exception as e:
error_msg = str(e)
return cls(
model_response=ModelResponse.from_api_response({}),
error=error_msg,
error_code="RESPONSE_PARSING_ERROR",
)
class ConfigurationManager:
"""Centralized configuration management with validation."""
def __init__(self):
"""Initialize configuration with environment variables."""
self.data_dir = Path("/data")
self.data_dir.mkdir(parents=True, exist_ok=True)
self._config = self._load_configuration()
self._validate_configuration()
def _load_configuration(self) -> Dict[str, Any]:
"""Load configuration from environment variables."""
model_mapping = {}
for model in ModelType:
alias = model.value
if alias in ["grok-3", "grok-auto"]:
model_mapping[alias] = "grok-3"
elif alias in ["grok-fast"]:
model_mapping[alias] = "grok-3"
elif alias in ["grok-4", "grok-expert"]:
model_mapping[alias] = "grok-4"
elif alias in ["grok-search", "grok-deepsearch"]:
model_mapping[alias] = "grok-4"
elif alias in ["grok-image", "grok-draw"]:
model_mapping[alias] = "grok-4"
else:
model_mapping[alias] = alias
return {
"MODELS": model_mapping,
"API": {
"IS_TEMP_CONVERSATION": self._get_bool_env(
"IS_TEMP_CONVERSATION", True
),
"IS_CUSTOM_SSO": self._get_bool_env("IS_CUSTOM_SSO", False),
"BASE_URL": "https://grok.com",
"API_KEY": os.environ.get("API_KEY", "sk-123456"),
"RETRY_TIME": 1000,
"PROXY": os.environ.get("PROXY"),
"DISABLE_DYNAMIC_HEADERS": self._get_bool_env(
"DISABLE_DYNAMIC_HEADERS", False
),
},
"ADMIN": {
"MANAGER_SWITCH": os.environ.get("MANAGER_SWITCH"),
"PASSWORD": os.environ.get("ADMINPASSWORD"),
},
"SERVER": {
"CF_CLEARANCE": os.environ.get("CF_CLEARANCE"),
"PORT": int(os.environ.get("PORT", 5200)),
},
"RETRY": {
"RETRYSWITCH": False,
"MAX_ATTEMPTS": MAX_RETRY_ATTEMPTS,
},
"TOKEN_STATUS_FILE": str(self.data_dir / "token_status.json"),
"SHOW_THINKING": self._get_bool_env("SHOW_THINKING", False),
"SHOW_SEARCH_RESULTS": self._get_bool_env("SHOW_SEARCH_RESULTS", True),
"IS_SUPER_GROK": self._get_bool_env("IS_SUPER_GROK", False),
"FILTERED_TAGS": self._get_list_env(
"FILTERED_TAGS",
[
"xaiartifact",
"xai:tool_usage_card",
"grok:render",
"details",
"summary",
],
),
"TAG_CONFIG": self._get_tag_config(),
"CONTENT_TYPE_MAPPINGS": self._get_content_type_mappings(),
}
def _get_bool_env(self, key: str, default: bool = False) -> bool:
"""Get boolean environment variable."""
return os.environ.get(key, str(default)).lower() == "true"
def _get_list_env(self, key: str, default: List[str]) -> List[str]:
"""Get comma-separated list from environment variable."""
value = os.environ.get(key)
if not value:
return default
return [tag.strip() for tag in value.split(",") if tag.strip()]
def _get_content_type_mappings(self) -> Dict[str, Dict[str, str]]:
"""Get content type mappings from environment or defaults."""
mappings_env = os.environ.get("CONTENT_TYPE_MAPPINGS")
if mappings_env:
try:
return json.loads(mappings_env)
except json.JSONDecodeError:
print("Invalid CONTENT_TYPE_MAPPINGS JSON, using defaults")
return {
"text/plain": {"stag": "```", "etag": "```"},
"text/markdown": {"stag": "", "etag": ""},
"application/json": {"stag": "```json\n", "etag": "\n```"},
"text/javascript": {"stag": "```javascript\n", "etag": "\n```"},
"text/python": {"stag": "```python\n", "etag": "\n```"},
"text/html": {"stag": "```html\n", "etag": "\n```"},
"text/css": {"stag": "```css\n", "etag": "\n```"},
"text/xml": {"stag": "```xml\n", "etag": "\n```"},
"application/xml": {"stag": "```xml\n", "etag": "\n```"},
"text/yaml": {"stag": "```yaml\n", "etag": "\n```"},
"application/yaml": {"stag": "```yaml\n", "etag": "\n```"},
"text/x-yaml": {"stag": "```yaml\n", "etag": "\n```"},
"text/sql": {"stag": "```sql\n", "etag": "\n```"},
"application/sql": {"stag": "```sql\n", "etag": "\n```"},
"text/x-sql": {"stag": "```sql\n", "etag": "\n```"},
"text/typescript": {"stag": "```typescript\n", "etag": "\n```"},
"application/typescript": {"stag": "```typescript\n", "etag": "\n```"},
"text/x-typescript": {"stag": "```typescript\n", "etag": "\n```"},
"text/jsx": {"stag": "```jsx\n", "etag": "\n```"},
"text/x-jsx": {"stag": "```jsx\n", "etag": "\n```"},
"text/tsx": {"stag": "```tsx\n", "etag": "\n```"},
"text/x-tsx": {"stag": "```tsx\n", "etag": "\n```"},
"text/java": {"stag": "```java\n", "etag": "\n```"},
"application/java": {"stag": "```java\n", "etag": "\n```"},
"text/x-java": {"stag": "```java\n", "etag": "\n```"},
"text/csharp": {"stag": "```csharp\n", "etag": "\n```"},
"text/x-csharp": {"stag": "```csharp\n", "etag": "\n```"},
"application/x-csharp": {"stag": "```csharp\n", "etag": "\n```"},
"text/cpp": {"stag": "```cpp\n", "etag": "\n```"},
"text/x-c++": {"stag": "```cpp\n", "etag": "\n```"},
"application/x-cpp": {"stag": "```cpp\n", "etag": "\n```"},
"text/c": {"stag": "```c\n", "etag": "\n```"},
"text/x-c": {"stag": "```c\n", "etag": "\n```"},
"text/go": {"stag": "```go\n", "etag": "\n```"},
"text/x-go": {"stag": "```go\n", "etag": "\n```"},
"application/x-go": {"stag": "```go\n", "etag": "\n```"},
"text/rust": {"stag": "```rust\n", "etag": "\n```"},
"text/x-rust": {"stag": "```rust\n", "etag": "\n```"},
"application/x-rust": {"stag": "```rust\n", "etag": "\n```"},
"text/php": {"stag": "```php\n", "etag": "\n```"},
"application/x-php": {"stag": "```php\n", "etag": "\n```"},
"text/ruby": {"stag": "```ruby\n", "etag": "\n```"},
"application/x-ruby": {"stag": "```ruby\n", "etag": "\n```"},
"text/swift": {"stag": "```swift\n", "etag": "\n```"},
"text/x-swift": {"stag": "```swift\n", "etag": "\n```"},
"text/kotlin": {"stag": "```kotlin\n", "etag": "\n```"},
"text/x-kotlin": {"stag": "```kotlin\n", "etag": "\n```"},
"text/scala": {"stag": "```scala\n", "etag": "\n```"},
"text/x-scala": {"stag": "```scala\n", "etag": "\n```"},
"text/bash": {"stag": "```bash\n", "etag": "\n```"},
"text/x-bash": {"stag": "```bash\n", "etag": "\n```"},
"application/x-bash": {"stag": "```bash\n", "etag": "\n```"},
"text/shell": {"stag": "```bash\n", "etag": "\n```"},
"text/x-shell": {"stag": "```bash\n", "etag": "\n```"},
"application/x-shell": {"stag": "```bash\n", "etag": "\n```"},
"text/powershell": {"stag": "```powershell\n", "etag": "\n```"},
"text/x-powershell": {"stag": "```powershell\n", "etag": "\n```"},
"application/x-powershell": {"stag": "```powershell\n", "etag": "\n```"},
"text/dockerfile": {"stag": "```dockerfile\n", "etag": "\n```"},
"text/x-dockerfile": {"stag": "```dockerfile\n", "etag": "\n```"},
"application/x-dockerfile": {"stag": "```dockerfile\n", "etag": "\n```"},
"text/toml": {"stag": "```toml\n", "etag": "\n```"},
"application/toml": {"stag": "```toml\n", "etag": "\n```"},
"text/ini": {"stag": "```ini\n", "etag": "\n```"},
"text/x-ini": {"stag": "```ini\n", "etag": "\n```"},
"application/x-ini": {"stag": "```ini\n", "etag": "\n```"},
"text/properties": {"stag": "```properties\n", "etag": "\n```"},
"text/x-properties": {"stag": "```properties\n", "etag": "\n```"},
"text/csv": {"stag": "```csv\n", "etag": "\n```"},
"application/csv": {"stag": "```csv\n", "etag": "\n```"},
"text/x-csv": {"stag": "```csv\n", "etag": "\n```"},
"text/log": {"stag": "```\n", "etag": "\n```"},
"application/x-log": {"stag": "```\n", "etag": "\n```"},
"text/x-log": {"stag": "```\n", "etag": "\n```"},
"application/x-httpd-php": {"stag": "```php\n", "etag": "\n```"},
"text/x-python": {"stag": "```python\n", "etag": "\n```"},
"application/x-python": {"stag": "```python\n", "etag": "\n```"},
"text/x-javascript": {"stag": "```javascript\n", "etag": "\n```"},
"application/javascript": {"stag": "```javascript\n", "etag": "\n```"},
"text/ecmascript": {"stag": "```javascript\n", "etag": "\n```"},
"application/ecmascript": {"stag": "```javascript\n", "etag": "\n```"},
"text/jscript": {"stag": "```javascript\n", "etag": "\n```"},
"application/x-javascript": {"stag": "```javascript\n", "etag": "\n```"},
"text/vbscript": {"stag": "```vbscript\n", "etag": "\n```"},
"application/x-vbscript": {"stag": "```vbscript\n", "etag": "\n```"},
"text/x-markdown": {"stag": "", "etag": ""},
"application/x-markdown": {"stag": "", "etag": ""},
"text/x-web-markdown": {"stag": "", "etag": ""},
"code/python": {"stag": "```python\n", "etag": "\n```"},
"code/javascript": {"stag": "```javascript\n", "etag": "\n```"},
"code/typescript": {"stag": "```typescript\n", "etag": "\n```"},
"code/html": {"stag": "```html\n", "etag": "\n```"},
"code/css": {"stag": "```css\n", "etag": "\n```"},
"code/json": {"stag": "```json\n", "etag": "\n```"},
"code/xml": {"stag": "```xml\n", "etag": "\n```"},
"code/yaml": {"stag": "```yaml\n", "etag": "\n```"},
"code/sql": {"stag": "```sql\n", "etag": "\n```"},
"code/bash": {"stag": "```bash\n", "etag": "\n```"},
"code/shell": {"stag": "```bash\n", "etag": "\n```"},
"code/dockerfile": {"stag": "```dockerfile\n", "etag": "\n```"},
"text/code": {"stag": "```\n", "etag": "\n```"},
"application/code": {"stag": "```\n", "etag": "\n```"},
"text/source": {"stag": "```\n", "etag": "\n```"},
"application/source": {"stag": "```\n", "etag": "\n```"},
}
def _get_tag_config(self) -> Dict[str, Dict[str, Any]]:
"""Get tag configuration from environment or defaults."""
tag_config_env = os.environ.get("TAG_CONFIG")
if tag_config_env:
try:
return json.loads(tag_config_env)
except json.JSONDecodeError:
print("Invalid TAG_CONFIG JSON, using defaults")
filtered_tags = self._get_list_env(
"FILTERED_TAGS",
["xaiartifact", "xai:tool_usage_card", "grok:render", "details", "summary"],
)
default_config = {}
for tag in filtered_tags:
if tag.lower() in ["xai:tool_usage_card", "grok:render"]:
default_config[tag.lower()] = {"behavior": "remove_all"}
else:
default_config[tag.lower()] = {"behavior": "preserve_content"}
if not default_config:
default_config = {
"xaiartifact": {"behavior": "preserve_content"},
"xai:tool_usage_card": {"behavior": "remove_all"},
"grok:render": {"behavior": "remove_all"},
"details": {"behavior": "preserve_content"},
"summary": {"behavior": "preserve_content"},
}
return default_config
def _validate_configuration(self) -> None:
"""Validate configuration settings."""
issues = []
if not os.environ.get("API_KEY"):
issues.append("Missing required environment variable: API_KEY")
if not self._config["API"]["IS_CUSTOM_SSO"]:
sso_env = os.environ.get("SSO", "")
sso_super_env = os.environ.get("SSO_SUPER", "")
if not sso_env and not sso_super_env:
issues.append(
"No SSO tokens configured. Set SSO or SSO_SUPER environment variables."
)
proxy = self._config["API"]["PROXY"]
if proxy and not any(
proxy.startswith(p) for p in ["http://", "https://", "socks5://"]
):
issues.append(f"Invalid proxy format: {proxy}")
if issues:
for issue in issues:
print(f"Configuration issue: {issue}")
else:
print("Configuration validation passed")
def get(self, key_path: str, default: Any = None) -> Any:
"""Get configuration value using dot notation."""
keys = key_path.split(".")
value = self._config
for key in keys:
if isinstance(value, dict) and key in value:
value = value[key]
else:
return default
return value
def set(self, key_path: str, value: Any) -> None:
"""Set configuration value using dot notation."""
keys = key_path.split(".")
config = self._config
for key in keys[:-1]:
if key not in config:
config[key] = {}
config = config[key]
config[keys[-1]] = value
@property
def models(self) -> Dict[str, str]:
"""Get supported models mapping."""
return self._config["MODELS"]
@property
def data_directory(self) -> Path:
"""Get data directory path."""
return self.data_dir
class UtilityFunctions:
"""Collection of utility functions for common operations."""
@staticmethod
def get_proxy_configuration(proxy_url: Optional[str]) -> Dict[str, Any]:
"""Get proxy configuration for requests."""
if not proxy_url:
return {}
print(f"Using proxy: {proxy_url}")
if proxy_url.startswith("socks5://"):
proxy_config: Dict[str, Any] = {"proxy": proxy_url}
if "@" in proxy_url:
auth_part = proxy_url.split("@")[0].split("://")[1]
if ":" in auth_part:
username, password = auth_part.split(":", 1)
proxy_config["proxy_auth"] = (username, password)
return proxy_config
else:
return {"proxies": {"https": proxy_url, "http": proxy_url}}
@staticmethod
def organize_search_results(search_results: Dict[str, Any]) -> str:
"""Format search results for display."""
if not search_results or "results" not in search_results:
return ""
results = search_results["results"]
formatted_results = []
for index, result in enumerate(results):
title = result.get("title", "未知标题")
url = result.get("url", "#")
preview = result.get("preview", "无预览内容")
formatted_result = (
f"\r\n资料[{index}]: {title}
\r\n"
f"{preview}\r\n\n[Link]({url})\r\n