cscsj/geminiweb

1import asyncio2import base643import hashlib4import hmac5import importlib.metadata6import io7import json8import logging9import os10import re11import secrets12import tempfile13import time14import uuid15from contextlib import asynccontextmanager16from datetime import datetime, timezone17from pathlib import Path18from typing import Dict, List, Optional, Union19from urllib.parse import quote, urlparse20 21import httpx22import numpy as np23from fastapi import Depends, FastAPI, Header, HTTPException, Request, Response24from fastapi.middleware.cors import CORSMiddleware25from fastapi.responses import JSONResponse, StreamingResponse26from gemini_webapi import GeminiClient, set_log_level27from gemini_webapi.constants import Model28from PIL import Image29from pydantic import BaseModel30 31# Configure logging32logging.basicConfig(level=logging.INFO)33logger = logging.getLogger(__name__)34set_log_level("INFO")35 36gemini_client = None37gemini_client_lock = asyncio.Lock()38 39 40@asynccontextmanager41async def lifespan(app: FastAPI):42	"""Initialize the Gemini client during startup and close it on shutdown."""43	await get_gemini_client()44	try:45		yield46	finally:47		global gemini_client48		if gemini_client is not None:49			try:50				await gemini_client.close()51			except Exception as e:52				logger.warning(f"Failed to close Gemini client during shutdown: {e}")53			finally:54				gemini_client = None55 56 57app = FastAPI(title="Gemini API FastAPI Server", lifespan=lifespan)58 59 60def get_gemini_webapi_version() -> str:61	"""Return the installed gemini-webapi package version for runtime diagnostics."""62	try:63		return importlib.metadata.version("gemini-webapi")64	except importlib.metadata.PackageNotFoundError:65		return "unknown"66 67 68def get_cached_1psidts_path(psid: str) -> str:69	"""Return the cache path for a rotated 1PSIDTS value."""70	if not psid or not re.match("^[\\w\\-\\.]+$", psid):71		return ""72	return os.path.join(GEMINI_COOKIE_PATH, f".cached_1psidts_{psid}.txt")73 74 75def load_cached_1psidts(psid: str) -> str:76	"""Load a cached rotated 1PSIDTS value for the given 1PSID."""77	cached_file_path = get_cached_1psidts_path(psid)78	if not cached_file_path:79		return ""80 81	if os.path.exists(cached_file_path):82		try:83			content = Path(cached_file_path).read_text().strip()84			if content:85				return content86		except Exception as e:87			logger.warning(f"Error reading cache file {cached_file_path}: {e}")88 89	return ""90 91 92def get_cookie_value(cookies, name: str) -> str:93	"""Safely read a cookie value from an httpx cookie jar or mapping."""94	if not cookies:95		return ""96 97	for domain in (".google.com", ".googleusercontent.com", None):98		try:99			value = cookies.get(name, domain=domain) if domain is not None else cookies.get(name)100		except TypeError:101			value = cookies.get(name)102		except Exception:103			value = ""104 105		if value:106			return value107 108	return ""109 110 111# Add CORS middleware112app.add_middleware(113	CORSMiddleware,114	allow_origins=["*"],115	allow_credentials=True,116	allow_methods=["*"],117	allow_headers=["*"],118)119 120# Authentication credentials121SECURE_1PSID = os.environ.get("SECURE_1PSID", "")122SECURE_1PSIDTS = os.environ.get("SECURE_1PSIDTS", "")123API_KEY = os.environ.get("API_KEY", "")124ENABLE_THINKING = os.environ.get("ENABLE_THINKING", "false").lower() == "true"125TEMPORARY_CHAT = os.environ.get("TEMPORARY_CHAT", "false").lower() == "true"126AUTO_DELETE_CHAT = os.environ.get("AUTO_DELETE_CHAT", "true").lower() == "true" and not TEMPORARY_CHAT127PUBLIC_BASE_URL = os.environ.get("PUBLIC_BASE_URL", "").rstrip("/")128SECRET_FILE_PATH = os.path.join(os.path.dirname(__file__), "secrets", "proxy_secret")129GEMINI_COOKIE_PATH = os.path.join(os.path.dirname(__file__), "secrets")130SESSION_VALIDATION_PROMPT = "Reply with exactly OK."131AUTH_FAILURE_TEXT_PATTERNS = (132	"are you signed in",133	"sign in",134	"signed in",135	"log in",136	"logged in",137)138DEFAULT_USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/144.0.0.0 Safari/537.36 Edg/144.0.0.0"139 140os.environ.setdefault("GEMINI_COOKIE_PATH", GEMINI_COOKIE_PATH)141 142 143async def background_delete_chat(client: GeminiClient, cid: str):144	"""Deletes a chat conversation in the background to avoid blocking the main thread."""145	if not cid:146		return147	try:148		await client.delete_chat(cid)149	except Exception as e:150		logger.error(f"Failed to auto-delete chat {cid}: {e}")151 152 153def response_indicates_auth_failure(text: str) -> bool:154	"""Return True if the response text looks like a signed-out or degraded session."""155	normalized = (text or "").strip().lower()156	if not normalized:157		return True158	return any(pattern in normalized for pattern in AUTH_FAILURE_TEXT_PATTERNS)159 160 161async def fetch_readable_chat_response(client: GeminiClient, cid: str, retry_delays: List[int]) -> Optional[object]:162	"""Poll Gemini history until the chat becomes readable or retries are exhausted."""163	for attempt, delay in enumerate(retry_delays, start=1):164		try:165			if delay:166				await asyncio.sleep(delay)167 168			recovered = await client.fetch_latest_chat_response(cid)169			if recovered and getattr(recovered, "text", ""):170				return recovered171		except Exception as e:172			logger.exception(173				"Gemini history read failed for cid=%s on retry %s/%s after %ss delay: %s",174				cid,175				attempt,176				len(retry_delays),177				delay,178				e,179			)180			continue181 182	return None183 184 185async def background_verify_chat_persistence(client: GeminiClient, cid: str, source: str):186	"""Best-effort verification that a returned cid is readable from Gemini history."""187	if not cid:188		return189 190	retry_delays = [1, 3, 8]191	recovered = await fetch_readable_chat_response(client, cid, retry_delays)192	if recovered:193		logger.debug(194			"Gemini history verification succeeded: source=%s cid=%s text_len=%s metadata=%s",195			source,196			cid,197			len(recovered.text),198			getattr(recovered, "metadata", None),199		)200		return201 202	logger.warning(203		"Gemini history verification exhausted retries for cid=%s source=%s",204		cid,205		source,206	)207 208 209async def validate_gemini_client_session(client: GeminiClient, source: str):210	"""Verify that an initialized client can create and read back a normal persistent Gemini chat."""211	validation_cid = None212	try:213		response = await client.generate_content(SESSION_VALIDATION_PROMPT, temporary=False)214		response_text = getattr(response, "text", "") or ""215		metadata = getattr(response, "metadata", None) or []216		validation_cid = metadata[0] if metadata else None217 218		if response_indicates_auth_failure(response_text):219			raise ValueError("validation probe returned signed-out or empty content")220 221		if not validation_cid:222			raise ValueError("validation probe returned no persistent chat metadata")223 224		recovered = await fetch_readable_chat_response(client, validation_cid, [1, 3, 8])225		if not recovered or response_indicates_auth_failure(getattr(recovered, "text", "") or ""):226			raise ValueError("validation probe chat was not readable from Gemini history")227 228		logger.info("Gemini session validation succeeded using %s credentials", source)229	finally:230		if validation_cid:231			try:232				await client.delete_chat(validation_cid)233			except Exception:234				logger.debug("Failed to delete Gemini validation chat %s", validation_cid)235 236 237def load_or_generate_secret() -> str:238	"""239	Load the signature secret from file, or generate a new one if not found.240	"""241	if os.path.exists(SECRET_FILE_PATH):242		try:243			with open(SECRET_FILE_PATH, "r") as f:244				secret = f.read().strip()245				if secret:246					logger.info(f"Loaded proxy secret from {SECRET_FILE_PATH}")247					return secret248		except Exception as e:249			logger.warning(f"Failed to read secret file, trying to generate a new one: {e}")250 251	# Generate new secret if not found or error occurred252	new_secret = secrets.token_hex(32)253	try:254		# Ensure directory exists255		os.makedirs(os.path.dirname(SECRET_FILE_PATH), exist_ok=True)256		with open(SECRET_FILE_PATH, "w") as f:257			f.write(new_secret)258 259		# Set restrictive permissions (user-only readable/writable)260		try:261			os.chmod(SECRET_FILE_PATH, 0o600)262		except Exception as e:263			logger.warning(f"Failed to set restrictive permissions on {SECRET_FILE_PATH}: {e}")264 265		logger.info(f"Generated new proxy secret and saved to {SECRET_FILE_PATH}")266		return new_secret267	except Exception as e:268		logger.error(f"Error writing secret file: {e}")269		# if unable to save, return an in-memory ephemeral secret instead of using API_KEY or SECURE_1PSID270		ephemeral_secret = secrets.token_urlsafe(32)271		logger.warning("Using an in-memory secret to proxy images for this session.")272		return ephemeral_secret273 274 275SIGNATURE_SECRET = load_or_generate_secret()276 277# Watermark removal constants278ASSETS_DIR = os.path.join(os.path.dirname(__file__), "assets")279ALPHA_MAP_CACHE = {}280 281 282def get_alpha_map(size: int) -> np.ndarray:283	"""Load and cache the alpha map from the background capture image."""284	if size in ALPHA_MAP_CACHE:285		return ALPHA_MAP_CACHE[size]286 287	bg_path = os.path.join(ASSETS_DIR, f"bg_{size}.png")288	if not os.path.exists(bg_path):289		logger.warning(f"Watermark asset not found: {bg_path}")290		return None291 292	try:293		with Image.open(bg_path) as img:294			img_data = np.array(img.convert("RGB"))295			alpha_map = np.max(img_data, axis=2) / 255.0296			ALPHA_MAP_CACHE[size] = alpha_map297			return alpha_map298	except Exception as e:299		logger.error(f"Error loading alpha map {size}: {e}")300		return None301 302 303def remove_gemini_watermark(image_bytes: bytes) -> bytes:304	"""Remove Gemini watermark using Reverse Alpha Blending."""305	try:306		with Image.open(io.BytesIO(image_bytes)) as img:307			width, height = img.size308			orig_format = img.format309 310			if width > 1024 and height > 1024:311				logo_size, margin = 96, 64312			else:313				logo_size, margin = 48, 32314 315			alpha_map = get_alpha_map(logo_size)316			if alpha_map is None:317				return image_bytes318 319			x = width - margin - logo_size320			y = height - margin - logo_size321			if x < 0 or y < 0:322				logger.warning(f"Image too small for watermark removal: {width}x{height}")323				return image_bytes324 325			# Reverse Alpha Blending: original = (watermarked - α × 255) / (1 - α)326			img_array = np.array(img.convert("RGB")).astype(np.float64)327			roi = img_array[y : y + logo_size, x : x + logo_size].copy()328 329			alpha = np.clip(alpha_map, 0.002, 0.99)330			alpha_expanded = np.expand_dims(alpha, axis=2)331			cleaned_roi = (roi - alpha_expanded * 255.0) / (1.0 - alpha_expanded)332			cleaned_roi = np.clip(np.round(cleaned_roi), 0, 255).astype(np.uint8)333 334			img_array_uint8 = np.array(img.convert("RGB"))335			img_array_uint8[y : y + logo_size, x : x + logo_size] = cleaned_roi336 337			out_io = io.BytesIO()338			save_format = orig_format or "PNG"339			if save_format.upper() == "JPEG":340				Image.fromarray(img_array_uint8).save(out_io, format="JPEG", quality=95)341			else:342				Image.fromarray(img_array_uint8).save(out_io, format=save_format)343			return out_io.getvalue()344 345	except Exception as e:346		logger.error(f"Error removing watermark: {e}")347		return image_bytes348 349 350if not SECURE_1PSID or not SECURE_1PSIDTS:351	logger.warning("Gemini credentials are missing; set SECURE_1PSID and SECURE_1PSIDTS before serving requests.")352else:353	logger.info(354		"Startup config: thinking=%s temporary_chat=%s auto_delete_chat=%s public_base_url=%s gemini_webapi=%s",355		ENABLE_THINKING,356		TEMPORARY_CHAT,357		AUTO_DELETE_CHAT,358		bool(PUBLIC_BASE_URL),359		get_gemini_webapi_version(),360	)361	if not re.match("^[\\w\\-\\.]+$", SECURE_1PSID):362		logger.warning(363			"SECURE_1PSID contains characters outside the safe cache filename pattern. This may be valid for auth, but cached 1PSIDTS lookup will fall back to the env value."364		)365 366if not API_KEY:367	logger.info("API key authentication is disabled.")368else:369	logger.info("API key authentication is enabled.")370 371 372def correct_markdown(md_text: str) -> str:373	"""374	修正Markdown文本，移除Google搜索链接包装器，并根据显示文本简化目标URL。375	"""376 377	def simplify_link_target(text_content: str) -> str:378		match_colon_num = re.match(r"([^:]+:\d+)", text_content)379		if match_colon_num:380			return match_colon_num.group(1)381		return text_content382 383	def replacer(match: re.Match) -> str:384		outer_open_paren = match.group(1)385		display_text = match.group(2)386 387		new_target_url = simplify_link_target(display_text)388		new_link_segment = f"[`{display_text}`]({new_target_url})"389 390		if outer_open_paren:391			return f"{outer_open_paren}{new_link_segment})"392		else:393			return new_link_segment394 395	pattern = r"(\()?\[`([^`]+?)`\]\((https://www.google.com/search\?q=)(.*?)(?<!\\)\)\)*(\))?"396 397	fixed_google_links = re.sub(pattern, replacer, md_text)398	# fix wrapped markdownlink399	pattern = r"`(\[[^\]]+\]\([^\)]+\))`"400	return re.sub(pattern, r"\1", fixed_google_links)401 402 403# Pydantic models for API requests and responses404class ContentItem(BaseModel):405	type: str406	text: Optional[str] = None407	image_url: Optional[Dict[str, str]] = None408 409 410class Message(BaseModel):411	role: str412	content: Union[str, List[ContentItem]]413	name: Optional[str] = None414 415 416class ChatCompletionRequest(BaseModel):417	model: str418	messages: List[Message]419	temperature: Optional[float] = 0.7420	top_p: Optional[float] = 1.0421	n: Optional[int] = 1422	stream: Optional[bool] = False423	max_tokens: Optional[int] = None424	presence_penalty: Optional[float] = 0425	frequency_penalty: Optional[float] = 0426	user: Optional[str] = None427 428 429class Choice(BaseModel):430	index: int431	message: Message432	finish_reason: str433 434 435class Usage(BaseModel):436	prompt_tokens: int437	completion_tokens: int438	total_tokens: int439 440 441class ChatCompletionResponse(BaseModel):442	id: str443	object: str = "chat.completion"444	created: int445	model: str446	choices: List[Choice]447	usage: Usage448 449 450class ModelData(BaseModel):451	id: str452	object: str = "model"453	created: int454	owned_by: str = "google"455 456 457class ModelList(BaseModel):458	object: str = "list"459	data: List[ModelData]460 461 462# Authentication dependency463async def verify_api_key(authorization: str = Header(None)):464	"""465	Verify the API key extracted from the Authorization header.466 467	Raises:468		HTTPException: If the authorization header is missing, incorrectly formatted, or the token is invalid.469	"""470	if not API_KEY:471		# If API_KEY is not set in environment, skip validation (for development)472		return473 474	if not authorization:475		raise HTTPException(status_code=401, detail="Missing Authorization header")476 477	try:478		scheme, token = authorization.split()479		if scheme.lower() != "bearer":480			raise HTTPException(481				status_code=401,482				detail="Invalid authentication scheme. Use Bearer token",483			)484 485		if token != API_KEY:486			raise HTTPException(status_code=401, detail="Invalid API key")487	except ValueError:488		raise HTTPException(489			status_code=401,490			detail="Invalid authorization format. Use 'Bearer YOUR_API_KEY'",491		)492 493	return token494 495 496# Simple error handler middleware497@app.middleware("http")498async def error_handling(request: Request, call_next):499	"""500	Global middleware to catch unhandled exceptions, log the error,501	and return a standardized HTTP 500 response.502	"""503	try:504		return await call_next(request)505	except Exception:506		logger.exception("Request failed")507		return JSONResponse(508			status_code=500,509			content={510				"error": {511					"message": "Internal server error",512					"type": "internal_server_error",513				}514			},515		)516 517 518# Get list of available models519@app.get("/v1/models")520async def list_models():521	"""返回 gemini_webapi 中声明的模型列表"""522	now = int(datetime.now(tz=timezone.utc).timestamp())523	data = [524		{525			"id": m.model_name,  # 如 "gemini-2.0-flash"526			"object": "model",527			"created": now,528			"owned_by": "google-gemini-web",529		}530		for m in Model531	]532	return {"object": "list", "data": data}533 534 535# Helper to convert between Gemini and OpenAI model names536def map_model_name(openai_model_name: str) -> Model:537	"""根据模型名称字符串查找匹配的 Model 枚举值"""538	normalized_openai_model_name = openai_model_name.lower()539 540	# 首先尝试直接查找匹配的模型名称541	for m in Model:542		model_name = m.model_name if hasattr(m, "model_name") else str(m)543		if normalized_openai_model_name in model_name.lower():544			return m545 546	# 如果找不到匹配项，使用默认映射547	model_keywords = {548		"gemini-pro": ["pro", "2.0"],549		"gemini-pro-vision": ["vision", "pro"],550		"gemini-flash": ["flash", "2.0"],551		"gemini-1.5-pro": ["1.5", "pro"],552		"gemini-1.5-flash": ["1.5", "flash"],553	}554 555	# 根据关键词模糊匹配556	keywords = None557	for key, candidate_keywords in model_keywords.items():558		normalized_key = key.lower()559		matches_key = normalized_key in normalized_openai_model_name560		matches_any_kw = any(kw.lower() in normalized_openai_model_name for kw in candidate_keywords)561		if matches_key or matches_any_kw:562			keywords = candidate_keywords563			break564 565	if keywords is None:566		if "flash" in normalized_openai_model_name:567			keywords = ["flash"]568		elif "vision" in normalized_openai_model_name:569			keywords = ["vision"]570		else:571			keywords = ["pro"]572 573	for m in Model:574		model_name = m.model_name if hasattr(m, "model_name") else str(m)575		if all(kw.lower() in model_name.lower() for kw in keywords):576			return m577 578	# 如果还是找不到，返回第一个模型579	return next(iter(Model))580 581 582# Prepare conversation history from OpenAI messages format583def prepare_conversation(messages: List[Message]) -> tuple:584	"""585	Convert a list of OpenAI-formatted message objects into a586	flat string conversation format suitable for the Gemini API.587	Also extracts and saves base64 images to temporary files.588 589	Returns:590		A tuple containing the constructed conversation string and a list of paths to temporary image files.591	"""592	conversation = ""593	temp_files = []594 595	for msg in messages:596		if isinstance(msg.content, str):597			# String content handling598			if msg.role == "system":599				conversation += f"System: {msg.content}\n\n"600			elif msg.role == "user":601				conversation += f"Human: {msg.content}\n\n"602			elif msg.role == "assistant":603				conversation += f"Assistant: {msg.content}\n\n"604		else:605			# Mixed content handling606			if msg.role == "user":607				conversation += "Human: "608			elif msg.role == "system":609				conversation += "System: "610			elif msg.role == "assistant":611				conversation += "Assistant: "612 613			for item in msg.content:614				if item.type == "text":615					conversation += item.text or ""616				elif item.type == "image_url" and item.image_url:617					# Handle image618					image_url = item.image_url.get("url", "")619					if image_url.startswith("data:image/"):620						# Process base64 encoded image621						try:622							# Extract the base64 part623							base64_data = image_url.split(",")[1]624							image_data = base64.b64decode(base64_data)625 626							# Create temporary file to hold the image627							with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:628								tmp.write(image_data)629								temp_files.append(tmp.name)630						except Exception as e:631							logger.error(f"Error processing base64 image: {str(e)}")632 633			conversation += "\n\n"634 635	# Add a final prompt for the assistant to respond to636	conversation += "Assistant: "637 638	return conversation, temp_files639 640 641# Dependency to get the initialized Gemini client642async def get_gemini_client():643	"""644	Get or initialize the global GeminiClient instance.645 646	Raises:647		HTTPException: If initialization fails due to invalid parameters or connection issues.648	"""649	global gemini_client650	if gemini_client is not None:651		return gemini_client652 653	async with gemini_client_lock:654		if gemini_client is not None:655			return gemini_client656 657		try:658			psid = SECURE_1PSID659			cached_psidts = load_cached_1psidts(psid)660			attempts = []661 662			if cached_psidts:663				attempts.append(("cache", cached_psidts))664			if SECURE_1PSIDTS:665				attempts.append(("environment", SECURE_1PSIDTS))666 667			seen_psidts = set()668			new_attempts = []669			for source, psidts in attempts:670				if not psidts or psidts in seen_psidts:671					continue672				seen_psidts.add(psidts)673				new_attempts.append((source, psidts))674			attempts = new_attempts675 676			if not attempts:677				raise HTTPException(678					status_code=500,679					detail="Missing SECURE_1PSIDTS and no cached rotated 1PSIDTS is available",680				)681 682			last_error = None683			for source, psidts in attempts:684				tmp_client = None685				try:686					logger.info("Initializing Gemini client using %s credentials", source)687 688					tmp_client = GeminiClient(psid, psidts)689					await tmp_client.init(timeout=300)690					await validate_gemini_client_session(tmp_client, source)691 692					gemini_client = tmp_client693					break694				except Exception as e:695					last_error = e696					logger.warning(f"Gemini session setup failed using {source} 1PSIDTS: {e}")697					if tmp_client is not None:698						try:699							await tmp_client.close()700						except Exception:701							pass702 703			if gemini_client is None:704				raise last_error705 706		except Exception as e:707			logger.error(f"Failed to initialize Gemini client: {str(e)}")708			raise HTTPException(status_code=500, detail=f"Failed to initialize Gemini client: {str(e)}")709	return gemini_client710 711 712def get_image_signature(url: str) -> str:713	"""714	Generate a HMAC-SHA256 signature for the image URL using the persistent SIGNATURE_SECRET.715	"""716	secret = SIGNATURE_SECRET.encode()717	return hmac.new(secret, url.encode(), hashlib.sha256).hexdigest()718 719 720def postprocess_text(text: str) -> str:721	"""Apply text cleanup and markdown corrections to response text."""722	text = text.replace("&lt;", "<").replace("\\<", "<").replace("\\_", "_").replace("\\>", ">")723	return correct_markdown(text)724 725 726def extract_image_markdown(response, base_url: str) -> str:727	"""Extract images from a response and return markdown image links."""728	result = ""729	if hasattr(response, "images") and response.images:730		for img in response.images:731			img_url = getattr(img, "url", None)732			if img_url:733				sig = get_image_signature(img_url)734				proxy_url = f"{base_url}/gemini-proxy/image?url={quote(img_url)}&sig={sig}"735				result += f"\n\n![🎨 Loading image...]({proxy_url})"736	return result737 738 739@app.post("/v1/chat/completions")740async def create_chat_completion(741	request: ChatCompletionRequest,742	raw_request: Request,743	api_key: str = Depends(verify_api_key),744):745	"""746	Handle chat completion requests, translating from OpenAI API format to Gemini API format.747	Supports both streaming and non-streaming responses, caching, thinking features,748	and background conversation cleanup based on configuration.749	"""750	try:751		# 确保客户端已初始化752		global gemini_client753		gemini_client = await get_gemini_client()754 755		# 转换消息为对话格式756		conversation, temp_files = prepare_conversation(request.messages)757		logger.info(758			"Chat completion request: stream=%s requested_model=%s messages=%s temp_files=%s",759			request.stream,760			request.model,761			len(request.messages),762			len(temp_files),763		)764 765		# 获取适当的模型766		model = map_model_name(request.model)767 768		# 创建响应对象769		completion_id = f"chatcmpl-{uuid.uuid4()}"770		created_time = int(time.time())771		base_url = PUBLIC_BASE_URL or str(raw_request.base_url).rstrip("/")772 773		# Prepare generate_content arguments774		gen_kwargs = {"model": model}775		if TEMPORARY_CHAT:776			gen_kwargs["temporary"] = True777		if temp_files:778			gen_kwargs["files"] = temp_files779 780		if request.stream:781			# Real streaming using upstream generate_content_stream782			async def generate_stream():783				try:784 785					def make_chunk(delta: dict, finish_reason=None):786						return (787							"data: "788							+ json.dumps(789								{790									"id": completion_id,791									"object": "chat.completion.chunk",792									"created": created_time,793									"model": request.model,794									"choices": [795										{796											"index": 0,797											"delta": delta,798											"finish_reason": finish_reason,799										}800									],801								}802							)803							+ "\n\n"804						)805 806					# Send initial role chunk807					yield make_chunk({"role": "assistant"})808 809					thinking_started = False810					thinking_ended = False811					yielded_images = 0812					text_buffer = ""813					captured_cid = None814					chunk_count = 0815					last_metadata = None816 817					async for chunk in gemini_client.generate_content_stream(conversation, **gen_kwargs):818						chunk_count += 1819						if hasattr(chunk, "metadata") and chunk.metadata:820							last_metadata = chunk.metadata821						# Capture conversation ID for auto-deletion822						if AUTO_DELETE_CHAT and captured_cid is None and hasattr(chunk, "metadata") and chunk.metadata and len(chunk.metadata) > 0:823							captured_cid = chunk.metadata[0]824 825						# Handle thinking/thoughts delta826						if ENABLE_THINKING and hasattr(chunk, "thoughts_delta") and chunk.thoughts_delta:827							if not thinking_started:828								yield make_chunk({"content": "<think>\n"})829								thinking_started = True830 831							# Also include reasoning_content for full Open WebUI native compatibility832							yield make_chunk(833								{834									"content": chunk.thoughts_delta,835									"reasoning_content": chunk.thoughts_delta,836								}837							)838 839						# Handle text delta840						if hasattr(chunk, "text_delta") and chunk.text_delta:841							# Close thinking tag before first text content842							if thinking_started and not thinking_ended:843								thinking_ended = True844								yield make_chunk({"content": "\n</think>\n\n"})845 846							text_buffer += chunk.text_delta847							safe_to_yield = False848 849							# Yield if buffer ends with whitespace and looks like it's outside a markdown link850							if (851								text_buffer[-1].isspace()852								and text_buffer.count("[") == text_buffer.count("]")853								and text_buffer.count("(") == text_buffer.count(")")854							):855								safe_to_yield = True856							elif len(text_buffer) > 500:857								safe_to_yield = True858 859							if safe_to_yield:860								yield make_chunk({"content": postprocess_text(text_buffer)})861								text_buffer = ""862 863						# Handle inline images as they arrive864						if hasattr(chunk, "images") and chunk.images and len(chunk.images) > yielded_images:865							# Close thinking tag if an image arrives before any text866							if thinking_started and not thinking_ended:867								thinking_ended = True868								yield make_chunk({"content": "\n</think>\n\n"})869 870							new_images = chunk.images[yielded_images:]871							for img in new_images:872								img_url = getattr(img, "url", None)873								if img_url:874									sig = get_image_signature(img_url)875									proxy_url = f"{base_url}/gemini-proxy/image?url={quote(img_url)}&sig={sig}"876									img_md = f"\n\n![🎨 Loading image...]({proxy_url})\n\n"877									yield make_chunk({"content": img_md})878							yielded_images = len(chunk.images)879 880					# Flush any remaining text881					if text_buffer:882						yield make_chunk({"content": postprocess_text(text_buffer)})883 884					# Close thinking tag if it was never closed885					if thinking_started and not thinking_ended:886						yield make_chunk({"content": "\n</think>\n\n"})887 888					# Send finish chunk889					yield make_chunk({}, finish_reason="stop")890					yield "data: [DONE]\n\n"891 892					logger.info(893						"Streaming response completed: chunks=%s images=%s",894						chunk_count,895						yielded_images,896					)897					if last_metadata and len(last_metadata) > 0 and not AUTO_DELETE_CHAT:898						asyncio.create_task(background_verify_chat_persistence(gemini_client, last_metadata[0], "stream"))899				except Exception as e:900					logger.error(f"Error during streaming: {str(e)}", exc_info=True)901					# Send error as a content chunk so the client sees it902					error_msg = "\n\n[An internal error occurred while streaming]"903					yield make_chunk({"content": error_msg})904					yield make_chunk({}, finish_reason="stop")905					yield "data: [DONE]\n\n"906				finally:907					# Create background task to delete the chat if AUTO_DELETE_CHAT is enabled908					if AUTO_DELETE_CHAT and captured_cid:909						asyncio.create_task(background_delete_chat(gemini_client, captured_cid))910 911					# 清理临时文件912					for temp_file in temp_files:913						try:914							os.unlink(temp_file)915						except Exception as e:916							logger.warning(f"Failed to delete temp file {temp_file}: {str(e)}")917 918			return StreamingResponse(generate_stream(), media_type="text/event-stream")919		else:920			# Non-streaming response921			try:922				response = await gemini_client.generate_content(conversation, **gen_kwargs)923 924				if AUTO_DELETE_CHAT and hasattr(response, "metadata") and response.metadata and len(response.metadata) > 0:925					cid = response.metadata[0]926					asyncio.create_task(background_delete_chat(gemini_client, cid))927				elif hasattr(response, "metadata") and response.metadata and len(response.metadata) > 0:928					asyncio.create_task(background_verify_chat_persistence(gemini_client, response.metadata[0], "non-stream"))929				elif not getattr(response, "metadata", None):930					logger.warning("Non-stream response returned no Gemini metadata. This request may not map to a persistent Gemini chat.")931 932			finally:933				# 清理临时文件934				for temp_file in temp_files:935					try:936						os.unlink(temp_file)937					except Exception as e:938						logger.warning(f"Failed to delete temp file {temp_file}: {str(e)}")939 940			# 提取文本响应941			reply_text = ""942			if ENABLE_THINKING and hasattr(response, "thoughts") and response.thoughts:943				reply_text += f"<think>\n{response.thoughts}\n</think>\n\n"944			if hasattr(response, "text"):945				reply_text += response.text946			else:947				reply_text += str(response)948 949			# 提取并追加图片响应950			reply_text += extract_image_markdown(response, base_url)951			reply_text = postprocess_text(reply_text)952 953			if not reply_text or reply_text.strip() == "":954				logger.warning("Empty response received from Gemini")955				reply_text = "Server returned an empty response. Please check that Gemini API credentials are valid."956 957			result = {958				"id": completion_id,959				"object": "chat.completion",960				"created": created_time,961				"model": request.model,962				"choices": [963					{964						"index": 0,965						"message": {"role": "assistant", "content": reply_text},966						"finish_reason": "stop",967					}968				],969				"usage": {970					"prompt_tokens": len(conversation.split()),971					"completion_tokens": len(reply_text.split()),972					"total_tokens": len(conversation.split()) + len(reply_text.split()),973				},974			}975 976			logger.info("Non-streaming response completed")977			return result978 979	except Exception as e:980		logger.error(f"Error generating completion: {str(e)}", exc_info=True)981		raise HTTPException(status_code=500, detail=f"Error generating completion: {str(e)}")982 983 984@app.get("/gemini-proxy/image")985async def proxy_image(url: str, sig: str):986	"""987	Proxy images from Google domains to bypass browser security policies.988	Requires a valid HMAC signature.989	"""990	# Verify signature991	expected_sig = get_image_signature(url)992	if not hmac.compare_digest(sig, expected_sig):993		logger.warning(f"Invalid signature for proxy request: {url}")994		raise HTTPException(status_code=403, detail="Invalid signature")995 996	# Prevent open proxying997	allowed_domains = ["google.com", "googleusercontent.com", "gstatic.com"]998 999	try:1000		parsed = urlparse(url)1001		if parsed.scheme not in ["http", "https"]:1002			logger.warning(f"Invalid scheme in proxy request: {parsed.scheme}")1003			raise HTTPException(status_code=400, detail="Invalid URL scheme")1004 1005		hostname = parsed.hostname1006		if not hostname:1007			logger.warning(f"No hostname in proxy request: {url}")1008			raise HTTPException(status_code=400, detail="Invalid URL")1009 1010		hostname = hostname.lower()1011		is_allowed = any(hostname == d or hostname.endswith("." + d) for d in allowed_domains)1012 1013		if not is_allowed:1014			logger.warning(f"Blocked proxy request for domain: {hostname}")1015			raise HTTPException(status_code=403, detail="Domain not allowed")1016	except ValueError:1017		logger.warning(f"Malformed URL in proxy request: {url}")1018		raise HTTPException(status_code=400, detail="Invalid URL")1019 1020	# Minimal browser-like headers1021	headers = {1022		"User-Agent": DEFAULT_USER_AGENT,1023		"Accept": "image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8",1024		"Accept-Language": "en-US,en;q=0.9",1025		"Referer": "https://gemini.google.com/",1026	}1027 1028	# 10MB limit1029	MAX_BYTES = 10 * 1024 * 10241030 1031	# Use scoped cookies to prevent leakage during redirects1032	jar = httpx.Cookies()1033 1034	# Use the freshest available 1PSIDTS without overriding env cookies up front.1035	psid = SECURE_1PSID1036	psidts = get_cookie_value(getattr(gemini_client, "cookies", None), "__Secure-1PSIDTS") or load_cached_1psidts(psid) or SECURE_1PSIDTS1037 1038	jar.set("__Secure-1PSID", psid, domain=".google.com")1039	jar.set("__Secure-1PSIDTS", psidts, domain=".google.com")1040	jar.set("__Secure-1PSID", psid, domain=".googleusercontent.com")1041	jar.set("__Secure-1PSIDTS", psidts, domain=".googleusercontent.com")1042 1043	async with httpx.AsyncClient(http2=True, cookies=jar, follow_redirects=True) as client:1044		try:1045			# Fetch original resolution to keep watermark at expected size/position1046			fetch_url = re.sub(r"=s\d+$", "=s0", url) if re.search(r"=s\d+$", url) else url + "=s0"1047 1048			async with client.stream("GET", fetch_url, timeout=15.0, headers=headers) as resp:1049				if resp.status_code != 200:1050					logger.error(f"Google returned {resp.status_code} for image: {url}")1051 1052				resp.raise_for_status()1053 1054				content = bytearray()1055				async for chunk in resp.aiter_bytes():1056					content.extend(chunk)1057					if len(content) > MAX_BYTES:1058						logger.warning(f"Image too large: {url} (exceeded {MAX_BYTES} bytes)")1059						raise HTTPException(status_code=413, detail="Image too large")1060				# Validate Content-Type to prevent XSS/MIME sniffing1061				upstream_content_type = resp.headers.get("content-type", "image/png").lower()1062				if not upstream_content_type.startswith("image/"):1063					logger.warning(f"Rejected non-image Content-Type: {upstream_content_type} for {url}")1064					media_type = "image/png"1065				else:1066					media_type = upstream_content_type1067 1068				# Process watermark removal1069				if media_type in ["image/png", "image/jpeg", "image/webp"]:1070					processed_content = remove_gemini_watermark(bytes(content))1071				else:1072					processed_content = bytes(content)1073 1074				return Response(1075					content=processed_content,1076					media_type=media_type,1077					headers={1078						"Cross-Origin-Resource-Policy": "cross-origin",1079						"Access-Control-Allow-Origin": "*",1080						"Cache-Control": "public, max-age=86400",  # Cache for 24 hours1081						"X-Content-Type-Options": "nosniff",1082					},1083				)1084		except httpx.HTTPStatusError as e:1085			logger.error(f"Failed to fetch image: {e.response.status_code} for {url}")1086			raise HTTPException(1087				status_code=e.response.status_code,1088				detail=f"Failed to fetch image: Google returned {e.response.status_code}",1089			)1090		except HTTPException:1091			raise1092		except Exception as e:1093			logger.error(f"Proxy error: {str(e)}")1094			raise HTTPException(status_code=500, detail="Internal proxy error")1095 1096 1097@app.get("/")1098async def root():1099	"""1100	Health check endpoint to verify the API server is currently running.1101	"""1102	return {"status": "online", "message": "Gemini API FastAPI Server is running"}1103 1104 1105if __name__ == "__main__":1106	import uvicorn1107 1108	uvicorn.run("main:app", host="0.0.0.0", port=8000, log_level="info")1109