"""Gemini-compatible application service handlers.""" from __future__ import annotations import json from fastapi import HTTPException from fastapi.responses import StreamingResponse from aistudio_api.api.response_models import ( GeminiCandidateResponse, GeminiContentResponse, GeminiGenerateContentResponse, ) from aistudio_api.api.responses import to_gemini_parts, to_gemini_usage_metadata from aistudio_api.api.schemas import GeminiGenerateContentRequest from aistudio_api.api.state import runtime_state from aistudio_api.application.api_service_common import ( MAX_RETRIES, ensure_active_account, logger, record_rotator_event, require_busy_lock, try_switch_account, ) from aistudio_api.application.chat_service import ( cleanup_files, normalize_gemini_request, ) from aistudio_api.domain.errors import ( AistudioError, AuthError, RequestError, UsageLimitExceeded, ) from aistudio_api.infrastructure.gateway.client import AIStudioClient async def handle_gemini_generate_content( model_path: str, req: GeminiGenerateContentRequest, client: AIStudioClient, *, stream: bool, ): busy_lock = require_busy_lock() if stream: return _build_gemini_streaming_response( client=client, req=req, model_path=model_path ) last_error = None for attempt in range(MAX_RETRIES): async with busy_lock: await ensure_active_account(attempt) normalized = None try: normalized = normalize_gemini_request(req, model_path) logger.info( "Gemini: model=%s, contents=%s, stream=False, attempt=%d", normalized.model, len(req.contents), attempt + 1, ) output = await client.generate_content( model=normalized.model, capture_prompt=normalized.capture_prompt, capture_images=normalized.capture_images, contents=normalized.contents, system_instruction_content=normalized.system_instruction, tools=normalized.tools, safety_settings=normalized.safety_settings, temperature=normalized.temperature, top_p=normalized.top_p, top_k=normalized.top_k, max_tokens=normalized.max_tokens, generation_config_overrides=normalized.generation_config_overrides, sanitize_plain_text=False, ) record_rotator_event("success") runtime_state.record(normalized.model, "success", output.usage) return GeminiGenerateContentResponse( candidates=[ GeminiCandidateResponse( content=GeminiContentResponse( parts=to_gemini_parts( output.text, function_calls=output.function_calls, function_responses=output.function_responses, thinking=output.thinking, images=output.images, reasoning_images=output.reasoning_images, ), ), finishReason="STOP" if not output.function_calls else "FUNCTION_CALL", ) ], usageMetadata=to_gemini_usage_metadata(output.usage), ) except ValueError as exc: raise HTTPException( 400, detail={"message": str(exc), "type": "bad_request"} ) from exc except UsageLimitExceeded as exc: runtime_state.record(model_path, "rate_limited") last_error = exc record_rotator_event("rate_limited") if await try_switch_account(): logger.info( "Gemini 429 限流，已切换账号，重试 %d/%d", attempt + 1, MAX_RETRIES, ) continue logger.warning("Gemini 429 限流，无法切换账号") raise HTTPException( 429, detail={"message": str(exc), "type": "rate_limit_exceeded"} ) from exc except AistudioError as exc: runtime_state.record(model_path, "errors") record_rotator_event("error") raise HTTPException( 500, detail={"message": str(exc), "type": "server_error"} ) from exc except Exception as exc: runtime_state.record(model_path, "errors") record_rotator_event("error") logger.error("Gemini error: %s", exc, exc_info=True) raise HTTPException( 500, detail={"message": str(exc), "type": "server_error"} ) from exc finally: if normalized is not None and not stream: cleanup_files(normalized.cleanup_paths) raise HTTPException( 429, detail={"message": str(last_error), "type": "rate_limit_exceeded"} ) from last_error def _build_gemini_streaming_response( *, client: AIStudioClient, req: GeminiGenerateContentRequest, model_path: str, ) -> StreamingResponse: async def stream_response(): busy_lock = runtime_state.busy_lock if busy_lock is None: yield ( "data: " + json.dumps( {"error": {"message": "Server not ready"}}, ensure_ascii=False ) + "\n\n" ) return async with busy_lock: normalized = None try: await ensure_active_account(0) normalized = normalize_gemini_request(req, model_path) logger.info( "Gemini stream: model=%s, contents=%s", normalized.model, len(req.contents), ) final_usage = None has_yielded_data = False for stream_attempt in range(MAX_RETRIES): try: has_yielded_data = False async for event_type, text in client.stream_generate_content( model=normalized.model, capture_prompt=normalized.capture_prompt, capture_images=normalized.capture_images, contents=normalized.contents, system_instruction_content=normalized.system_instruction, tools=normalized.tools, safety_settings=normalized.safety_settings, temperature=normalized.temperature, top_p=normalized.top_p, top_k=normalized.top_k, max_tokens=normalized.max_tokens, generation_config_overrides=normalized.generation_config_overrides, sanitize_plain_text=False, force_refresh_capture=stream_attempt > 0, ): has_yielded_data = True if event_type == "body" and text: yield ( "data: " + json.dumps( { "candidates": [ { "content": { "role": "model", "parts": [{"text": text}], }, "finishReason": None, } ] }, ensure_ascii=False, ) + "\n\n" ) elif event_type == "images" and text: img_list = text if isinstance(text, list) else [] yield ( "data: " + json.dumps( { "candidates": [ { "content": { "role": "model", "parts": [ part.model_dump( mode="json", exclude_none=True, ) for part in to_gemini_parts( "", images=img_list ) ], }, "finishReason": None, } ] }, ensure_ascii=False, ) + "\n\n" ) elif event_type == "reasoning_images" and text: r_img_list = text if isinstance(text, list) else [] yield ( "data: " + json.dumps( { "candidates": [ { "content": { "role": "model", "parts": [ part.model_dump( mode="json", exclude_none=True, ) for part in to_gemini_parts( "", reasoning_images=r_img_list, ) ], }, "finishReason": None, } ] }, ensure_ascii=False, ) + "\n\n" ) elif event_type == "thinking" and text: yield ( "data: " + json.dumps( { "candidates": [ { "content": { "role": "model", "parts": [ { "text": text, "thought": True, } ], }, "finishReason": None, } ] }, ensure_ascii=False, ) + "\n\n" ) elif event_type == "usage": final_usage = text if isinstance(text, dict) else None break except UsageLimitExceeded: if normalized is not None: runtime_state.record(normalized.model, "rate_limited") record_rotator_event("rate_limited") if not has_yielded_data and await try_switch_account(): logger.warning( "Gemini stream 429 限流，已切换账号，重试 %d/%d", stream_attempt + 1, MAX_RETRIES, ) continue raise except RequestError as exc: if exc.status == 204 and stream_attempt == 0: logger.warning( "Gemini stream 收到 204，清理 snapshot 缓存后重试一次" ) client.clear_snapshot_cache() continue raise except AuthError as exc: if stream_attempt == 0: logger.warning( "Gemini stream 鉴权异常，清理 snapshot 缓存后重试一次: %s", exc, ) client.clear_snapshot_cache() continue raise record_rotator_event("success") if normalized is not None: runtime_state.record(normalized.model, "success", final_usage) if final_usage: yield ( "data: " + json.dumps( { "candidates": [], "usageMetadata": to_gemini_usage_metadata( final_usage ).model_dump(mode="json"), }, ensure_ascii=False, ) + "\n\n" ) yield "data: [DONE]\n\n" except Exception as exc: logger.error("Gemini stream error: %s", exc, exc_info=True) if not isinstance(exc, UsageLimitExceeded): record_rotator_event("error") if normalized is not None: runtime_state.record(normalized.model, "errors") yield ( "data: " + json.dumps({"error": {"message": str(exc)}}, ensure_ascii=False) + "\n\n" ) finally: if normalized is not None: cleanup_files(normalized.cleanup_paths) return StreamingResponse( stream_response(), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, )