"""Container-friendly multi-platform AI search script. Directly uses Playwright to start Chromium in headless mode, navigate to AI chat platforms that support anonymous access, submit a query, and collect the response. Designed for Hugging Face Docker Spaces. Supported platforms (no login required): doubao, qwen, kimi, minimaxi """ import argparse import json import sys import time from pathlib import Path from playwright.sync_api import sync_playwright SITES = { "doubao": { "url": "https://www.doubao.com/chat/", "input_selectors": ['textarea', 'div[contenteditable="true"]'], "response_container": '[class*="message"]', "submit_via": "enter", }, "qwen": { "url": "https://chat.qwen.ai/", "input_selectors": ['textarea', 'div[contenteditable="true"]'], "response_container": '[class*="message"]', "submit_via": "enter", }, "kimi": { "url": "https://www.kimi.com/", "input_selectors": ['textarea', 'div[contenteditable="true"]'], "response_container": '[class*="message"]', "submit_via": "enter", }, "minimaxi": { "url": "https://agent.minimaxi.com/", "input_selectors": ['textarea', 'div[contenteditable="true"]'], "response_container": '[class*="message"]', "submit_via": "enter", }, } def _find_input(page, selectors: list[str]): """Find the first visible input/textarea element.""" for sel in selectors: try: els = page.locator(sel).all() for el in els: try: if el.is_visible(timeout=2000): return el except Exception: continue except Exception: continue return None def _dismiss_overlays(page): """Try to close popups/modals/cookie banners.""" dismiss_texts = ["知道了", "我知道了", "关闭", "OK", "确定", "Got it", "Accept", "开始体验", "同意", "允许", "稍后再说", "以后再说", "跳过"] for text in dismiss_texts: try: btn = page.locator(f'button:has-text("{text}")').first if btn.is_visible(timeout=500): btn.click() time.sleep(0.5) except Exception: continue def _extract_response_text(page, site: str) -> str: """Extract the latest AI response from the page.""" # Strategy 1: Get all text blocks that look like responses try: # Wait a moment for content to render time.sleep(2) body_text = page.inner_text("body") return body_text except Exception: return "" def _is_response_complete(page, prev_length: int, stable_count: int) -> tuple[bool, int, int]: """Check if the response has stopped growing.""" try: body_text = page.inner_text("body") current_length = len(body_text) if current_length > prev_length + 20: return False, current_length, 0 return stable_count >= 4, current_length, stable_count + 1 except Exception: return stable_count >= 4, prev_length, stable_count + 1 def run_search(site: str, query: str, output_path: str, timeout: int = 120) -> dict: """Run a search on a single AI platform using headless Chromium.""" config = SITES.get(site) if not config: return {"ok": False, "error": f"Unknown site: {site}. Available: {list(SITES.keys())}"} with sync_playwright() as p: browser = p.chromium.launch( headless=True, executable_path="/usr/bin/chromium", args=[ "--no-sandbox", "--disable-gpu", "--disable-dev-shm-usage", "--disable-setuid-sandbox", "--no-first-run", "--no-zygote", "--single-process", "--disable-extensions", "--disable-software-rasterizer", "--disable-background-networking", ], ) context = browser.new_context( viewport={"width": 1440, "height": 960}, user_agent=( "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " "AppleWebKit/537.36 (KHTML, like Gecko) " "Chrome/125.0.0.0 Safari/537.36" ), locale="zh-CN", ) page = context.new_page() try: # Navigate print(f"[{site}] Navigating to {config['url']} ...", flush=True) page.goto(config["url"], wait_until="domcontentloaded", timeout=60000) time.sleep(3) # Dismiss overlays _dismiss_overlays(page) # Find input input_el = _find_input(page, config["input_selectors"]) if not input_el: debug_path = str(Path(output_path).parent / f"{site}_debug.png") page.screenshot(path=debug_path) return { "ok": False, "error": f"Could not find input on {site}", "debug": debug_path, "title": page.title(), "url": page.url, } # Record page text length before submitting try: pre_text = page.inner_text("body") pre_length = len(pre_text) except Exception: pre_length = 0 # Type and submit print(f"[{site}] Submitting query ...", flush=True) input_el.click() time.sleep(0.5) input_el.fill(query) time.sleep(1) input_el.press("Enter") print(f"[{site}] Query sent, waiting for response ...", flush=True) # Wait for response to complete time.sleep(5) # Initial wait prev_length = pre_length stable_count = 0 for _ in range(timeout // 3): done, prev_length, stable_count = _is_response_complete(page, prev_length, stable_count) if done: break time.sleep(3) # Extract response raw_text = _extract_response_text(page, site) # Try to extract just the new content (after the query) answer = raw_text query_pos = raw_text.rfind(query) if query_pos >= 0: after_query = raw_text[query_pos + len(query):].strip() if len(after_query) > 20: answer = after_query if answer.strip(): Path(output_path).write_text(answer.strip(), encoding="utf-8") print(f"[{site}] Got response ({len(answer)} chars)", flush=True) return {"ok": True, "answer": answer[:500], "length": len(answer)} else: debug_path = str(Path(output_path).parent / f"{site}_debug.png") page.screenshot(path=debug_path) return {"ok": False, "error": "No response captured", "debug": debug_path} except Exception as e: debug_path = str(Path(output_path).parent / f"{site}_debug.png") try: page.screenshot(path=debug_path) except Exception: pass return {"ok": False, "error": f"{type(e).__name__}: {e}"} finally: browser.close() def main(): parser = argparse.ArgumentParser(description="Container-friendly AI search (no login required)") parser.add_argument("--site", required=True, choices=sorted(SITES.keys())) parser.add_argument("--query", required=True) parser.add_argument("--output", required=True) parser.add_argument("--timeout", type=int, default=120) parser.add_argument("--test", action="store_true", help="Run browser test only") args = parser.parse_args() if args.test: # Diagnostic mode - just test if Playwright + Chromium work try: with sync_playwright() as p: browser = p.chromium.launch( headless=True, executable_path="/usr/bin/chromium", args=["--no-sandbox", "--disable-gpu", "--disable-dev-shm-usage"], ) page = browser.new_page() page.goto("https://example.com", timeout=15000) title = page.title() browser.close() print(json.dumps({"ok": True, "title": title})) return 0 except Exception as e: print(json.dumps({"ok": False, "error": f"{type(e).__name__}: {e}"})) return 1 result = run_search(args.site, args.query, args.output, args.timeout) print(json.dumps(result, ensure_ascii=False)) return 0 if result.get("ok") else 1 if __name__ == "__main__": raise SystemExit(main())