app.py

from fastapi import FastAPI, Request, HTTPException, Response
from fastapi.responses import StreamingResponse
import httpx
import logging
import os, json
import re  # 用于URL路径处理
from pprint import pprint
from key_selector import KeySelector

def get_target_url(url: str) -> str:
    """将url参数变了转换为合法的目标url；from http/ or https/ to http:// or https://"""
    url = re.sub(r"^http/", "http://", url)
    url = re.sub(r"^https/", "https://", url)
    return url

app = FastAPI()

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("uvicorn.error")

# 从环境变量获取配置
# UPSTREAM_HOST = os.getenv("UPSTREAM_HOST", "http://127.0.0.1")  # 必须包含端口号
# AUTH_TOKEN = os.getenv("AUTH_TOKEN", "YIG8ANC8q2QxFV_Gf8qwkPdBj2EpsqGqlfc3qvSdg7ksVkZcokOUtQn43XGK0NK3UUdBlIkYzNWefco_Wu4RcKnB0kpNgtZ2nTeqNum0i3fTUEhEcWSlJtT8FQgRK7bi")  # 安全认证令牌
X_Goog_Api_Key = os.getenv("X_Goog_Api_Key", "")

@app.get("/")
async def read_root():
    return {"message": "FastAPI Proxy is running"}

@app.post("/v25/{path:path}")
async def proxy(request: Request, path: str):
    # 添加流式请求判断逻辑
    is_streaming = ":streamGenerateContent" in path.lower()
    print(f"path: {path}")
    # path = "https/generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"
    # target_url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"
    target_url = get_target_url(path)
    # target_url = f"{target_url}?key={X_Goog_Api_Key}"
    print(f"target_url: {target_url}")
    method = request.method
    headers = {k: v for k, v in request.headers.items()
        # if k.lower() not in ["host", "connection", "Postman-Token", "content-length"]}
        if k.lower() not in ["host", "content-length"]}
    
    key_selector = KeySelector()
    # print('apikey',key_selector.get_api_key()['key_value'])

    headers["X-Goog-Api-Key"] = key_selector.get_api_key()['key_value']
    # print(key_selector.api_key_info.key_value)
    # headers["host"] = "generativelanguage.googleapis.com"

    try:
        # 关键修复：禁用KeepAlive防止连接冲突
        transport = httpx.AsyncHTTPTransport(retries=3, http1=True)

        async with httpx.AsyncClient(
            transport=transport,
            timeout=httpx.Timeout(300.0, connect=30.0)
        ) as client:
            # 处理请求体
            req_content = await request.body()
            print(f"Request headers: {headers}")
            # print(f"req_content: {req_content}")
            # req_content_str = req_content.decode('utf-8')  # 假设内容是 UTF-8 编码
            # print(f"req_content_str: {req_content_str}")
            print(f'target_url: {target_url}')
            print(f'method: {method}')
            # 发送请求到上游服务
            response = await client.request(
                method=method,
                url=target_url,
                headers=headers,
                content=req_content,
                follow_redirects=True  # 自动处理重定向
            )

            if is_streaming:
                # 流式响应处理
                async def stream_generator():
                    try:
                        async for chunk in response.aiter_bytes():
                            yield chunk
                    except Exception as e:
                        logger.error(f"Stream interrupted: {str(e)}")
                        yield json.dumps({"error": "流中断"}).encode()
            
                # 移除冲突头部
                headers = dict(response.headers)
                headers.pop("Content-Length", None)
            
                return StreamingResponse(
                    content=stream_generator(),
                    status_code=response.status_code,
                    headers=headers,
                    media_type="application/x-ndjson"  # Gemini流式格式
                )
            else:

                print(f"response.status_code: {response.status_code}")
                print(f"response.text: {response.text}")
                # 解析 JSON 字符串
                try:
                    data = json.loads(response.text)
                    # 格式化输出 JSON 数据
                    formatted_json = json.dumps(data, ensure_ascii=False, indent=4)
                    print('formatted_json')
                    print(formatted_json)
                    # return formatted_json
                    return Response(
                        content=formatted_json,
                        media_type="application/json"
                    )
                except json.JSONDecodeError as e:
                    print(f"Error decoding JSON: {e}")


        # # 详细记录错误响应
        # if response.status_code >= 400:
        #     error_detail = response.text[:1000]  # 增加错误详情长度
        #     logger.error(f"Upstream error {response.status_code}: {error_detail}")
        #     # 可选：记录完整响应到文件
        #     # with open("error.log", "a") as f:
        #     #     f.write(f"\n--- {target_url} ---\n{response.text}\n")

        # # 流式传输响应
        # return StreamingResponse(
        #     content=response.aiter_bytes(),
        #     status_code=response.status_code,
        #     headers=dict(response.headers)
        # )

    except httpx.ConnectError as e:
        logger.error(f"Connection failed to {target_url}: {e}")
        raise HTTPException(502, f"无法连接到上游服务: {target_url}") # Modified error message
    except httpx.ReadTimeout as e:
        logger.error(f"Timeout: {e}")
        raise HTTPException(504, "上游服务响应超时")
    except httpx.HTTPError as e:  # 捕获所有HTTP异常
        print('111111111111111111111')

        try:
            # 安全地获取异常信息
            error_type = type(e).__name__

            # 尝试获取状态码（如果存在）
            status_code = getattr(e, 'response', None) and e.response.status_code

            # 安全地获取错误详情
            error_detail = ""
            try:
                # 尝试获取文本响应（限制长度）
                if hasattr(e, 'response') and e.response:
                    error_detail = e.response.text[:500]  # 只取前500个字符
            except Exception as ex:
                error_detail = f"无法获取错误详情: {type(ex).__name__}"

            # 安全地记录日志
            logger.error(
                "HTTP代理错误 | "
                f"类型: {error_type} | "
                f"状态码: {status_code or 'N/A'} | "
                f"目标URL: {target_url} | "
                f"详情: {error_detail[:200]}"  # 日志中只记录前200字符
            )

            # 打印到控制台以便调试
            print(f"111111111111111111111 - HTTP代理错误: {error_type}")
            print(f"目标URL: {target_url}")
            print(f"状态码: {status_code or 'N/A'}")
            print(f"错误详情: {error_detail[:500]}")

        except Exception as ex:
            # 如果记录日志本身出错，使用最安全的方式记录
            logger.error(f"记录HTTP错误时发生异常: {type(ex).__name__}")
            print(f"严重错误: 记录HTTP错误时发生异常: {ex}")

        # 返回用户友好的错误响应
        raise HTTPException(
            status_code=502,
            detail=f"网关服务错误: {error_type} (上游状态: {status_code or '未知'})"
        )


        print('111111111111111111111')
        # logger.error(f"HTTP error: {str(e)}")

        raise HTTPException(502, f"网关错误: {str(e)}")
    except Exception as e:
        logger.exception("Unexpected proxy error")
        raise HTTPException(500, f"内部服务器错误: {str(e)}")