Calmlo/SambaNova2API

1"""2SambaNova OpenAI 接口代理 (支持模型列表透传和自动登录)3"""4 5import os6import uuid7import json8import time9import asyncio10import httpx11import secrets12import urllib.parse13from typing import Optional, Dict, Any14from fastapi import FastAPI, Request, HTTPException, Depends, Header15from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse16from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials17from fake_useragent import UserAgent18# 修复 Pydantic 导入19try:20    # 尝试从 pydantic-settings 导入 (Pydantic v2)21    from pydantic_settings import BaseSettings22except ImportError:23    # 回退到旧版本 (Pydantic v1)24    from pydantic import BaseSettings25 26# ================ 配置 ================27class Settings(BaseSettings):28    # SambaNova 配置29    SAMBA_EMAIL: str = os.getenv("SAMBA_EMAIL", "")30    SAMBA_PASSWORD: str = os.getenv("SAMBA_PASSWORD", "")31    SAMBA_COMPLETION_URL: str = os.getenv("SAMBA_COMPLETION_URL", "https://cloud.sambanova.ai/api/completion")32    SAMBA_MODELS_URL: str = os.getenv("SAMBA_MODELS_URL", "https://api.sambanova.ai/v1/models")33    34    # 本地API密钥配置35    LOCAL_API_KEY: str = os.getenv("LOCAL_API_KEY", secrets.token_urlsafe(32))36    37    # 其他配置38    TOKEN_CACHE_TIME: int = int(os.getenv("TOKEN_CACHE_TIME", 604800))  # 默认缓存7天 (7*24*60*60=604800秒)39    FINGERPRINT_PREFIX: str = os.getenv("FINGERPRINT_PREFIX", "anon_")40    41    class Config:42        env_file = ".env"43 44settings = Settings()45# =====================================46 47app = FastAPI(title="SambaNova OpenAI Proxy with Auto-Login")48security = HTTPBearer()49 50# 全局变量存储访问令牌和过期时间51access_token = None52token_expiry = 053token_lock = asyncio.Lock()54 55def generate_fingerprint() -> str:56    """生成符合格式要求的随机指纹"""57    return f"{settings.FINGERPRINT_PREFIX}{uuid.uuid4().hex[:20]}"58 59async def validate_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)) -> str:60    """验证本地API密钥并返回SambaNova访问令牌"""61    api_key = credentials.credentials62    63    # 如果未配置本地API密钥或为空，则跳过验证64    if settings.LOCAL_API_KEY and settings.LOCAL_API_KEY.strip():65        # 验证本地API密钥66        if api_key != settings.LOCAL_API_KEY:67            raise HTTPException(68                status_code=401,69                detail="Invalid API key",70                headers={"WWW-Authenticate": "Bearer"},71            )72    else:73        print("[警告] LOCAL_API_KEY未配置或为空，跳过API密钥验证")74    75    # 获取或刷新SambaNova访问令牌76    token = await get_samba_token()77    if not token:78        raise HTTPException(79            status_code=500,80            detail="Failed to obtain SambaNova access token. Check server logs for details."81        )82    83    return token84 85async def get_samba_token() -> Optional[str]:86    """获取或刷新SambaNova访问令牌"""87    global access_token, token_expiry88    89    # 使用锁防止并发请求同时刷新令牌90    async with token_lock:91        current_time = time.time()92        93        # 如果令牌有效，直接返回94        if access_token and current_time < token_expiry:95            print(f"[令牌] 使用缓存令牌: {access_token}")96            return access_token97        98        # 否则获取新令牌99        try:100            # 检查凭据是否已配置101            if not settings.SAMBA_EMAIL or not settings.SAMBA_PASSWORD:102                print("[错误] 未配置SambaNova凭据，请设置SAMBA_EMAIL和SAMBA_PASSWORD环境变量")103                return None104                105            print(f"[令牌] 开始获取新令牌... 邮箱: {settings.SAMBA_EMAIL}")106            auth = SambaAuthAsync(settings.SAMBA_EMAIL, settings.SAMBA_PASSWORD)107            new_token = await auth.login()108            109            if new_token:110                access_token = new_token111                token_expiry = current_time + settings.TOKEN_CACHE_TIME112                print(f"[令牌更新成功] 完整令牌: {new_token}")113                print(f"[令牌更新成功] 令牌将在 {settings.TOKEN_CACHE_TIME} 秒后过期")114                return access_token115            else:116                print("[令牌获取失败] 请检查SambaNova凭据是否正确")117                return None118        except Exception as e:119            print(f"[令牌获取异常] {str(e)}")120            return None121 122def reset_token_expiry():123    """重置令牌过期时间，强制下次请求重新获取令牌"""124    global token_expiry125    token_expiry = 0126    print("[令牌] 令牌已过期，将在下次请求时重新获取")127 128async def forward_get_request(url: str, token: str) -> httpx.Response:129    """转发 GET 请求到目标接口"""130    headers = {131        "accept": "application/json",132        "user-agent": "SambaNova-Proxy/1.0",133        "origin": "https://cloud.sambanova.ai",134        "referer": "https://cloud.sambanova.ai/"135    }136    137    cookies = {138        "access_token": token139    }140    141    async with httpx.AsyncClient() as client:142        try:143            resp = await client.get(144                url,145                headers=headers,146                cookies=cookies,147                timeout=10.0148            )149            150            # 检查是否需要刷新令牌151            if resp.status_code == 401:152                # 令牌已过期，需要刷新153                reset_token_expiry()154                raise HTTPException(401, "Token expired, please retry")155                156            resp.raise_for_status()157            return resp158        except httpx.HTTPStatusError as e:159            if e.response.status_code == 401:160                # 令牌已过期，需要刷新161                reset_token_expiry()162                raise HTTPException(401, "Token expired, please retry")163            raise HTTPException(e.response.status_code, f"Upstream error: {e.response.text}")164 165async def forward_post_request(url: str, payload: dict, token: str) -> httpx.Response:166    """转发 POST 请求到目标接口"""167    headers = {168        "content-type": "application/json",169        "user-agent": "SambaNova-Proxy/1.0",170        "origin": "https://cloud.sambanova.ai",171        "referer": "https://cloud.sambanova.ai/"172    }173    174    cookies = {175        "access_token": token176    }177    178    async with httpx.AsyncClient() as client:179        try:180            resp = await client.post(181                url,182                json=payload,183                headers=headers,184                cookies=cookies,185                timeout=30.0186            )187            188            # 检查是否需要刷新令牌189            if resp.status_code == 401:190                # 令牌已过期，需要刷新191                reset_token_expiry()192                raise HTTPException(401, "Token expired, please retry")193                194            resp.raise_for_status()195            return resp196        except httpx.HTTPStatusError as e:197            if e.response.status_code == 401:198                # 令牌已过期，需要刷新199                reset_token_expiry()200                raise HTTPException(401, "Token expired, please retry")201            raise HTTPException(e.response.status_code, f"Upstream error: {e.response.text}")202 203@app.get("/v1/models")204async def list_models(token: str = Depends(validate_api_key)):205    """透传模型列表接口"""206    try:207        resp = await forward_get_request(settings.SAMBA_MODELS_URL, token)208        content = resp.json()209        json_str = json.dumps(content, separators=(',', ':'), ensure_ascii=False)210        json_bytes = json_str.encode('utf-8')211        return JSONResponse(212            content=content,213            headers={214                "Content-Type": "application/json",215                "Content-Length": str(len(json_bytes)),216                "Cache-Control": "public, max-age=300"217            }218        )219    except httpx.RequestError as e:220        raise HTTPException(504, f"Gateway timeout: {str(e)}")221    except Exception as e:222        raise HTTPException(500, f"Internal server error: {str(e)}")223 224@app.post("/v1/chat/completions")225async def chat_completions(226    request: Request,227    token: str = Depends(validate_api_key)228):229    """处理对话请求"""230    try:231        openai_payload = await request.json()232        print(f"[请求] 收到聊天请求，模型: {openai_payload.get('model', 'DeepSeek-R1')}")233        234        samba_payload = {235            "body": {236                "model": openai_payload.get("model", "DeepSeek-R1"),237                "messages": openai_payload["messages"],238                "stream": True,239                "stop": openai_payload.get("stop", ["<|eot_id|>"]),240                "temperature": openai_payload.get("temperature", 0),241                "max_tokens": openai_payload.get("max_tokens", 2048),242                "do_sample": openai_payload.get("temperature", 0) > 0243            },244            "env_type": "text",245            "fingerprint": generate_fingerprint()246        }247        248        print(f"[转发] 使用令牌 {token[:10]}... 转发请求到 SambaNova")249        resp = await forward_post_request(settings.SAMBA_COMPLETION_URL, samba_payload, token)250        print(f"[响应] 成功获取响应，开始流式传输")251        252        return StreamingResponse(253            resp.aiter_bytes(),254            media_type="text/event-stream",255            headers={256                "X-Proxy-Version": "1.0",257                "X-Request-ID": str(uuid.uuid4())258            }259        )260    except HTTPException as e:261        print(f"[错误] HTTP异常: {e.detail}")262        raise263    except httpx.RequestError as e:264        print(f"[错误] 请求错误: {str(e)}")265        raise HTTPException(504, f"Gateway timeout: {str(e)}")266    except Exception as e:267        print(f"[错误] 未处理异常: {str(e)}")268        raise HTTPException(500, f"Internal server error: {str(e)}")269 270@app.get("/info")271async def get_info():272    """获取服务信息"""273    return {274        "status": "running",275        "api_key_configured": bool(settings.LOCAL_API_KEY),276        "samba_credentials_configured": bool(settings.SAMBA_EMAIL and settings.SAMBA_PASSWORD),277        "token_status": "active" if access_token and time.time() < token_expiry else "not_available",278        "token_expires_in": max(0, int(token_expiry - time.time())) if access_token else 0279    }280 281@app.get("/debug/token", include_in_schema=False)282async def debug_token():283    """调试端点：检查当前令牌状态"""284    global access_token, token_expiry285    current_time = time.time()286    287    return {288        "token_exists": access_token is not None,289        "token_prefix": access_token[:10] + "..." if access_token else None,290        "token_valid": access_token is not None and current_time < token_expiry,291        "expires_in_seconds": max(0, int(token_expiry - current_time)) if access_token else 0,292        "current_time": current_time,293        "expiry_time": token_expiry,294    }295 296@app.get("/", response_class=HTMLResponse)297async def root():298    """根路由健康检查，返回HTML界面"""299    current_time = time.time()300    token_valid = access_token is not None and current_time < token_expiry301    expires_in = max(0, int(token_expiry - current_time)) if access_token else 0302    303    # 计算过期时间的可读格式304    if expires_in > 0:305        days = expires_in // 86400306        hours = (expires_in % 86400) // 3600307        minutes = (expires_in % 3600) // 60308        expiry_readable = f"{days}天 {hours}小时 {minutes}分钟"309    else:310        expiry_readable = "已过期"311    312    # 使用东八区时间（中国标准时间）313    import datetime314    from datetime import timezone, timedelta315    316    # 创建东八区时区对象317    china_tz = timezone(timedelta(hours=8))318    # 获取当前UTC时间并转换为东八区时间319    current_time_china = datetime.datetime.now(china_tz)320    formatted_time = current_time_china.strftime('%Y-%m-%d %H:%M:%S')321    322    html_content = f"""323    <!DOCTYPE html>324    <html>325    <head>326        <title>SambaNova OpenAI 代理服务</title>327        <meta charset="UTF-8">328        <meta name="viewport" content="width=device-width, initial-scale=1.0">329        <style>330            body {{331                font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;332                line-height: 1.6;333                color: #333;334                max-width: 800px;335                margin: 0 auto;336                padding: 20px;337            }}338            h1 {{339                color: #2c3e50;340                border-bottom: 1px solid #eee;341                padding-bottom: 10px;342            }}343            .status-card {{344                background-color: #f8f9fa;345                border-radius: 8px;346                padding: 20px;347                margin-bottom: 20px;348                box-shadow: 0 2px 4px rgba(0,0,0,0.1);349            }}350            .status-item {{351                margin-bottom: 10px;352                display: flex;353                justify-content: space-between;354            }}355            .status-label {{356                font-weight: bold;357                color: #555;358            }}359            .status-value {{360                text-align: right;361            }}362            .status-healthy {{363                color: #28a745;364                font-weight: bold;365            }}366            .status-warning {{367                color: #ffc107;368                font-weight: bold;369            }}370            .status-error {{371                color: #dc3545;372                font-weight: bold;373            }}374            .code-block {{375                background-color: #f1f1f1;376                padding: 15px;377                border-radius: 5px;378                font-family: monospace;379                overflow-x: auto;380            }}381            .footer {{382                margin-top: 30px;383                font-size: 0.9em;384                color: #6c757d;385                text-align: center;386            }}387        </style>388    </head>389    <body>390        <h1>SambaNova OpenAI 代理服务</h1>391        392        <div class="status-card">393            <h2>服务状态</h2>394            <div class="status-item">395                <span class="status-label">状态:</span>396                <span class="status-value status-healthy">运行中</span>397            </div>398            <div class="status-item">399                <span class="status-label">版本:</span>400                <span class="status-value">1.0.0</span>401            </div>402            <div class="status-item">403                <span class="status-label">令牌状态:</span>404                <span class="status-value {('status-healthy' if token_valid else 'status-error')}">405                    {('有效' if token_valid else '无效')}406                </span>407            </div>408            <div class="status-item">409                <span class="status-label">令牌过期时间:</span>410                <span class="status-value">{expiry_readable}</span>411            </div>412            <div class="status-item">413                <span class="status-label">SambaNova 凭据:</span>414                <span class="status-value {('status-healthy' if settings.SAMBA_EMAIL and settings.SAMBA_PASSWORD else 'status-error')}">415                    {('已配置' if settings.SAMBA_EMAIL and settings.SAMBA_PASSWORD else '未配置')}416                </span>417            </div>418            <div class="status-item">419                <span class="status-label">本地API密钥:</span>420                <span class="status-value {('status-healthy' if settings.LOCAL_API_KEY else 'status-warning')}">421                    {('已配置' if settings.LOCAL_API_KEY else '未配置')}422                </span>423            </div>424        </div>425                426        <div class="footer">427            <p>当前时间: {formatted_time} (中国标准时间)</p>428        </div>429    </body>430    </html>431    """432    433    return html_content434 435class SambaAuthAsync:436    def __init__(self, email, password):437        self.email = email438        self.password = password439        self.client = httpx.AsyncClient()440        self.ua = UserAgent()441        self.base_headers = {442            "accept": "*/*",443            "accept-language": "zh-CN,zh;q=0.9,en-US;q=0.8,en;q=0.7",444            "origin": "https://cloud.sambanova.ai",445            "referer": "https://cloud.sambanova.ai/",446            "user-agent": self.ua.random447        }448        self.config = None449        self.nonce = None  # 确保nonce属性存在450 451    async def _get_config(self):452        """获取动态配置信息"""453        config_url = "https://cloud.sambanova.ai/api/config"454        response = await self.client.get(config_url, headers=self.base_headers)455        response.raise_for_status()456        self.config = response.json()457        print(f"[配置获取成功] ClientID: {self.config['clientId']}")458 459    async def _get_login_ticket(self):460        """获取登录票据"""461        auth_url = f"https://{self.config['issuerBaseUrl']}/co/authenticate"462        payload = {463            "client_id": self.config["clientId"],464            "username": self.email,465            "password": self.password,466            "realm": "Username-Password-Authentication",467            "credential_type": "http://auth0.com/oauth/grant-type/password-realm"468        }469        470        headers = {**self.base_headers, "content-type": "application/json"}471 472        response = await self.client.post(auth_url, headers=headers, json=payload)473        response.raise_for_status()474        return response.json()["login_ticket"]475 476    async def _get_auth_code(self, login_ticket: str):477        """获取授权码"""478        state = secrets.token_urlsafe(32)479        self.nonce = secrets.token_urlsafe(32)  # 保存nonce到实例变量480        481        params = {482            "client_id": self.config["clientId"],483            "response_type": "code",484            "redirect_uri": self.config["redirectURL"],485            "scope": "openid profile email",486            "nonce": self.nonce,487            "state": state,488            "login_ticket": login_ticket,489            "realm": "Username-Password-Authentication",490            "auth0Client": "eyJuYW1lIjoibG9jay5qcyIsInZlcnNpb24iOiIxMi4zLjAiLCJlbnYiOnsiYXV0aDAuanMiOiI5LjIyLjEiLCJhdXRoMC5qcy11bHAiOiI5LjIyLjEifX0="491        }492 493        auth_url = f"https://{self.config['issuerBaseUrl']}/authorize"494        response = await self.client.get(495            auth_url,496            params=params,497            follow_redirects=False498        )499        500        if response.status_code == 302:501            location = response.headers["location"]502            parsed = urllib.parse.urlparse(location)503            query = urllib.parse.parse_qs(parsed.query)504            return query.get("code", [None])[0], state505        raise Exception(f"未收到302重定向，实际状态码：{response.status_code}")506 507    async def _exchange_token(self, code: str, state: str):508        """交换访问令牌"""509        # 设置必要的cookies510        self.client.cookies.set("nonce", self.nonce, domain="cloud.sambanova.ai")511        512        callback_url = f"{self.config['redirectURL']}?code={code}&state={state}"513        response = await self.client.get(514            callback_url,515            headers={516                **self.base_headers,517                "sec-fetch-site": "same-site",518                "sec-fetch-mode": "navigate",519                "sec-fetch-user": "?1",520                "sec-fetch-dest": "document"521            },522            follow_redirects=True523        )524        525        # 从cookies中提取access_token526        for cookie in self.client.cookies.jar:527            if cookie.name == "access_token" and "sambanova.ai" in cookie.domain:528                return cookie.value529        raise Exception("未找到access_token")530 531    async def login(self):532        """完整登录流程"""533        try:534            await self._get_config()535            login_ticket = await self._get_login_ticket()536            print(f"[登录票据获取成功] 完整票据: {login_ticket}")537            538            auth_code, state = await self._get_auth_code(login_ticket)539            if not auth_code:540                raise Exception("授权码获取失败")541            print(f"[授权码获取成功] 完整授权码: {auth_code}")542            print(f"[授权状态] state: {state}")543            544            token = await self._exchange_token(auth_code, state)545            print(f"[令牌获取成功] 完整令牌: {token}")546            return token547            548        except Exception as e:549            print(f"[登录失败] 详细错误: {str(e)}")550            return None551        finally:552            await self.client.aclose()553 554@app.on_event("startup")555async def startup_event():556    """应用启动时预获取令牌"""557    print("\n" + "="*50)558    print("[启动] SambaNova OpenAI 代理服务启动")559    print("="*50)560    561    # 检查环境变量562    print(f"[环境] SAMBA_EMAIL: {'已设置' if settings.SAMBA_EMAIL else '未设置'}")563    print(f"[环境] SAMBA_PASSWORD: {'已设置' if settings.SAMBA_PASSWORD else '未设置'}")564    print(f"[环境] LOCAL_API_KEY: {'已设置' if settings.LOCAL_API_KEY else '未设置'}")565    566    # 尝试直接登录567    print("[登录] 开始尝试登录...")568    try:569        auth = SambaAuthAsync(settings.SAMBA_EMAIL, settings.SAMBA_PASSWORD)570        token = await auth.login()571        572        if token:573            global access_token, token_expiry574            access_token = token575            token_expiry = time.time() + settings.TOKEN_CACHE_TIME576            print(f"[登录] 登录成功! 令牌: {token}")577            print(f"[登录] 令牌将在 {settings.TOKEN_CACHE_TIME} 秒后过期")578        else:579            print("[登录] 登录失败，未获取到令牌")580    except Exception as e:581        print(f"[登录] 登录过程发生异常: {str(e)}")582    583    print("="*50 + "\n")584 585if __name__ == "__main__":586    import uvicorn587    uvicorn.run(app, host="0.0.0.0", port=7860)