Spaces:
Paused
Paused
File size: 3,763 Bytes
6a0eaad |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
import asyncio
import multiprocessing
import random
import time
import json
from typing import List, Optional, Dict, Any, Union, AsyncGenerator, Tuple, Callable, Set
import os
import traceback
from contextlib import asynccontextmanager
import sys
import platform
import logging
import logging.handlers
import socket # 保留 socket 以便在 __main__ 中进行简单的直接运行提示
from asyncio import Queue, Lock, Future, Task, Event
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse, StreamingResponse, FileResponse
from fastapi import WebSocket, WebSocketDisconnect
from pydantic import BaseModel
from playwright.async_api import Page as AsyncPage, Browser as AsyncBrowser, Playwright as AsyncPlaywright, Error as PlaywrightAsyncError, expect as expect_async, BrowserContext as AsyncBrowserContext, Locator, TimeoutError
from playwright.async_api import async_playwright
from urllib.parse import urljoin, urlparse
import uuid
import datetime
import aiohttp
import stream
import queue
# --- 配置模块导入 ---
from config import *
# --- models模块导入 ---
from models import (
FunctionCall,
ToolCall,
MessageContentItem,
Message,
ChatCompletionRequest,
ClientDisconnectedError,
StreamToLogger,
WebSocketConnectionManager,
WebSocketLogHandler
)
# --- logging_utils模块导入 ---
from logging_utils import setup_server_logging, restore_original_streams
# --- browser_utils模块导入 ---
from browser_utils import (
_initialize_page_logic,
_close_page_logic,
signal_camoufox_shutdown,
_handle_model_list_response,
detect_and_extract_page_error,
save_error_snapshot,
get_response_via_edit_button,
get_response_via_copy_button,
_wait_for_response_completion,
_get_final_response_content,
get_raw_text_content,
switch_ai_studio_model,
load_excluded_models,
_handle_initial_model_state_and_storage,
_set_model_from_page_display
)
# --- api_utils模块导入 ---
from api_utils import (
generate_sse_chunk,
generate_sse_stop_chunk,
generate_sse_error_chunk,
use_helper_get_response,
use_stream_response,
clear_stream_queue,
prepare_combined_prompt,
validate_chat_request,
_process_request_refactored,
create_app,
queue_worker
)
# --- stream queue ---
STREAM_QUEUE:Optional[multiprocessing.Queue] = None
STREAM_PROCESS = None
# --- Global State ---
playwright_manager: Optional[AsyncPlaywright] = None
browser_instance: Optional[AsyncBrowser] = None
page_instance: Optional[AsyncPage] = None
is_playwright_ready = False
is_browser_connected = False
is_page_ready = False
is_initializing = False
# --- 全局代理配置 ---
PLAYWRIGHT_PROXY_SETTINGS: Optional[Dict[str, str]] = None
global_model_list_raw_json: Optional[List[Any]] = None
parsed_model_list: List[Dict[str, Any]] = []
model_list_fetch_event = asyncio.Event()
current_ai_studio_model_id: Optional[str] = None
model_switching_lock: Optional[Lock] = None
excluded_model_ids: Set[str] = set()
request_queue: Optional[Queue] = None
processing_lock: Optional[Lock] = None
worker_task: Optional[Task] = None
page_params_cache: Dict[str, Any] = {}
params_cache_lock: Optional[Lock] = None
logger = logging.getLogger("AIStudioProxyServer")
log_ws_manager = None
# --- FastAPI App 定义 ---
app = create_app()
# --- Main Guard ---
if __name__ == "__main__":
import uvicorn
port = int(os.environ.get("PORT", 2048))
uvicorn.run(
"server:app",
host="0.0.0.0",
port=port,
log_level="info",
access_log=False
) |