gemini-balance/app/middleware/smart_routing_middleware.py

from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from app.config.config import settings
from app.log.logger import get_main_logger
import re

logger = get_main_logger()

class SmartRoutingMiddleware(BaseHTTPMiddleware):
    def __init__(self, app):
        super().__init__(app)
        # 简化的路由规则 - 直接根据检测结果路由
        pass

    async def dispatch(self, request: Request, call_next):
        if not settings.URL_NORMALIZATION_ENABLED:
            return await call_next(request)
        logger.debug(f"request: {request}")
        original_path = str(request.url.path)
        method = request.method
        
        # 尝试修复URL
        fixed_path, fix_info = self.fix_request_url(original_path, method, request)

        if fixed_path != original_path:
            logger.info(f"URL fixed: {method} {original_path} → {fixed_path}")
            if fix_info:
                logger.debug(f"Fix details: {fix_info}")

            # 重写请求路径
            request.scope["path"] = fixed_path
            request.scope["raw_path"] = fixed_path.encode()
        
        return await call_next(request)

    def fix_request_url(self, path: str, method: str, request: Request) -> tuple:
        """简化的URL修复逻辑"""

        # 首先检查是否已经是正确的格式，如果是则不处理
        if self.is_already_correct_format(path):
            return path, None

        # 1. 最高优先级：包含generateContent → Gemini格式
        if "generatecontent" in path.lower() or "v1beta/models" in path.lower():
            return self.fix_gemini_by_operation(path, method, request)

        # 2. 第二优先级：包含/openai/ → OpenAI格式
        if "/openai/" in path.lower():
            return self.fix_openai_by_operation(path, method)

        # 3. 第三优先级：包含/v1/ → v1格式
        if "/v1/" in path.lower():
            return self.fix_v1_by_operation(path, method)

        # 4. 第四优先级：包含/chat/completions → chat功能
        if "/chat/completions" in path.lower():
            return "/v1/chat/completions", {"type": "v1_chat"}

        # 5. 默认：原样传递
        return path, None

    def is_already_correct_format(self, path: str) -> bool:
        """检查是否已经是正确的API格式"""
        # 检查是否已经是正确的端点格式
        correct_patterns = [
            r"^/v1beta/models/[^/:]+:(generate|streamGenerate)Content$",  # Gemini原生
            r"^/gemini/v1beta/models/[^/:]+:(generate|streamGenerate)Content$",  # Gemini带前缀
            r"^/v1beta/models$",  # Gemini模型列表
            r"^/gemini/v1beta/models$",  # Gemini带前缀的模型列表
            r"^/v1/(chat/completions|models|embeddings|images/generations|audio/speech)$",  # v1格式
            r"^/openai/v1/(chat/completions|models|embeddings|images/generations|audio/speech)$",  # OpenAI格式
            r"^/hf/v1/(chat/completions|models|embeddings|images/generations|audio/speech)$",  # HF格式
            r"^/vertex-express/v1beta/models/[^/:]+:(generate|streamGenerate)Content$",  # Vertex Express Gemini格式
            r"^/vertex-express/v1beta/models$",  # Vertex Express模型列表
            r"^/vertex-express/v1/(chat/completions|models|embeddings|images/generations)$",  # Vertex Express OpenAI格式
        ]

        for pattern in correct_patterns:
            if re.match(pattern, path):
                return True

        return False

    def fix_gemini_by_operation(
        self, path: str, method: str, request: Request
    ) -> tuple:
        """根据Gemini操作修复，考虑端点偏好"""
        if method == "GET":
            return "/v1beta/models", {
                "role": "gemini_models",
            }

        # 提取模型名称
        try:
            model_name = self.extract_model_name(path, request)
        except ValueError:
            # 无法提取模型名称，返回原路径不做处理
            return path, None

        # 检测是否为流式请求
        is_stream = self.detect_stream_request(path, request)

        # 检查是否有vertex-express偏好
        if "/vertex-express/" in path.lower():
            if is_stream:
                target_url = (
                    f"/vertex-express/v1beta/models/{model_name}:streamGenerateContent"
                )
            else:
                target_url = (
                    f"/vertex-express/v1beta/models/{model_name}:generateContent"
                )

            fix_info = {
                "rule": (
                    "vertex_express_generate"
                    if not is_stream
                    else "vertex_express_stream"
                ),
                "preference": "vertex_express_format",
                "is_stream": is_stream,
                "model": model_name,
            }
        else:
            # 标准Gemini端点
            if is_stream:
                target_url = f"/v1beta/models/{model_name}:streamGenerateContent"
            else:
                target_url = f"/v1beta/models/{model_name}:generateContent"

            fix_info = {
                "rule": "gemini_generate" if not is_stream else "gemini_stream",
                "preference": "gemini_format",
                "is_stream": is_stream,
                "model": model_name,
            }

        return target_url, fix_info

    def fix_openai_by_operation(self, path: str, method: str) -> tuple:
        """根据操作类型修复OpenAI格式"""
        if method == "POST":
            if "chat" in path.lower() or "completion" in path.lower():
                return "/openai/v1/chat/completions", {"type": "openai_chat"}
            elif "embedding" in path.lower():
                return "/openai/v1/embeddings", {"type": "openai_embeddings"}
            elif "image" in path.lower():
                return "/openai/v1/images/generations", {"type": "openai_images"}
            elif "audio" in path.lower():
                return "/openai/v1/audio/speech", {"type": "openai_audio"}
        elif method == "GET":
            if "model" in path.lower():
                return "/openai/v1/models", {"type": "openai_models"}

        return path, None

    def fix_v1_by_operation(self, path: str, method: str) -> tuple:
        """根据操作类型修复v1格式"""
        if method == "POST":
            if "chat" in path.lower() or "completion" in path.lower():
                return "/v1/chat/completions", {"type": "v1_chat"}
            elif "embedding" in path.lower():
                return "/v1/embeddings", {"type": "v1_embeddings"}
            elif "image" in path.lower():
                return "/v1/images/generations", {"type": "v1_images"}
            elif "audio" in path.lower():
                return "/v1/audio/speech", {"type": "v1_audio"}
        elif method == "GET":
            if "model" in path.lower():
                return "/v1/models", {"type": "v1_models"}

        return path, None

    def detect_stream_request(self, path: str, request: Request) -> bool:
        """检测是否为流式请求"""
        # 1. 路径中包含stream关键词
        if "stream" in path.lower():
            return True

        # 2. 查询参数
        if request.query_params.get("stream") == "true":
            return True

        return False

    def extract_model_name(self, path: str, request: Request) -> str:
        """从请求中提取模型名称，用于构建Gemini API URL"""
        # 1. 从请求体中提取
        try:
            if hasattr(request, "_body") and request._body:
                import json

                body = json.loads(request._body.decode())
                if "model" in body and body["model"]:
                    return body["model"]
        except Exception:
            pass

        # 2. 从查询参数中提取
        model_param = request.query_params.get("model")
        if model_param:
            return model_param

        # 3. 从路径中提取（用于已包含模型名称的路径）
        match = re.search(r"/models/([^/:]+)", path, re.IGNORECASE)
        if match:
            return match.group(1)

        # 4. 如果无法提取模型名称，抛出异常
        raise ValueError("Unable to extract model name from request")