Create main.py

main.py

@@ -0,0 +1,211 @@
+# ===================================================================
+# main.py (已修改以适配 Hugging Face Secrets)
+# ===================================================================
+
+import httpx
+import time
+import json
+import os
+from fastapi import FastAPI, Request, HTTPException, Depends
+from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from pydantic import BaseModel, Field
+from typing import List, Optional, Dict, Any, Union
+
+# --- Pydantic 模型定义 ---
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[ChatMessage]
+    stream: bool = False
+    temperature: Optional[float] = None
+    max_tokens: Optional[int] = None
+    top_p: Optional[float] = None
+
+# --- 全局变量 ---
+config: Dict[str, Any] = {}
+app = FastAPI(title="Kodu2API Adapter", description="将 Kodu AI API 转换为 OpenAI 格式")
+security = HTTPBearer()
+
+# --- 核心函数 (已修改) ---
+
+def load_config():
+    """
+    从环境变量加载配置 (适配 Hugging Face Secrets).
+    """
+    global config
+    print("Loading configuration from environment variables...")
+    
+    try:
+        # 必需的 Secret: Kodu Refresh Token
+        kodu_refresh_token = os.environ.get("KODU_REFRESH_TOKEN")
+        if not kodu_refresh_token:
+            raise ValueError("Secret 'KODU_REFRESH_TOKEN' is required but not found.")
+
+        # 必需的 Secret: 你的服务密码 (API Keys)
+        service_api_keys_str = os.environ.get("SERVICE_API_KEYS")
+        if not service_api_keys_str:
+            raise ValueError("Secret 'SERVICE_API_KEYS' is required but not found.")
+        service_api_keys = json.loads(service_api_keys_str)
+
+        # 可选的 Secret: Kodu 模型映射
+        kodu_models_str = os.environ.get("KODU_MODELS")
+        kodu_models = json.loads(kodu_models_str) if kodu_models_str else {}
+
+        # 可选的 Secret: 代理 URL
+        proxy_url = os.environ.get("PROXY_URL") or None
+        
+        # 可选的 Secret: Base URL
+        base_url = os.environ.get("BASE_URL", "https://api.kodu.ai")
+
+        # 组合成 config 字典
+        config = {
+            "kodu_refresh_token": kodu_refresh_token,
+            "kodu_access_token": "",  # 初始为空，后续会自动刷新
+            "kodu_models": kodu_models,
+            "service_api_keys": service_api_keys,
+            "proxy_url": proxy_url,
+            "base_url": base_url
+        }
+        
+        # 验证 service_api_keys 的格式
+        if not isinstance(config["service_api_keys"], list) or not all(isinstance(i, str) for i in config["service_api_keys"]):
+            raise TypeError("Secret 'SERVICE_API_KEYS' must be a JSON list of strings (e.g., [\"sk-key1\", \"sk-key2\"]).")
+            
+        print("Configuration loaded successfully from secrets.")
+        print(f"Loaded {len(config['service_api_keys'])} service API key(s).")
+        print(f"Loaded {len(config['kodu_models'])} Kodu model mapping(s).")
+        if config["proxy_url"]:
+            print(f"Using proxy: {config['proxy_url']}")
+
+    except Exception as e:
+        print(f"FATAL: Could not load configuration from secrets. Error: {e}")
+        # 在配置失败时，将 config 设置为空字典以阻止应用不完整地运行
+        config = {}
+        raise
+
+async def refresh_access_token():
+    """
+    使用 refresh_token 获取新的 access_token.
+    """
+    if not config:
+        print("Cannot refresh token, config not loaded.")
+        return False
+        
+    print("Attempting to refresh Kodu access token...")
+    async with httpx.AsyncClient(proxies=config.get("proxy_url")) as client:
+        try:
+            response = await client.post(
+                f"{config['base_url']}/auth/refresh-token",
+                headers={"Authorization": f"Bearer {config['kodu_refresh_token']}"}
+            )
+            response.raise_for_status()
+            config['kodu_access_token'] = response.json().get('accessToken')
+            print("Successfully refreshed Kodu access token.")
+            return True
+        except httpx.HTTPStatusError as e:
+            print(f"Error refreshing access token: {e.response.status_code} - {e.response.text}")
+            return False
+        except Exception as e:
+            print(f"An unexpected error occurred during token refresh: {e}")
+            return False
+
+async def authenticate_client(credentials: HTTPAuthorizationCredentials = Depends(security)):
+    """
+    验证客户端提供的 API Key.
+    """
+    if not config:
+         raise HTTPException(status_code=503, detail="Service Unavailable: Server configuration is not loaded.")
+    
+    if credentials.credentials not in config.get("service_api_keys", []):
+        raise HTTPException(status_code=401, detail="Invalid API Key")
+
+# --- FastAPI 事件和路由 ---
+
+@app.on_event("startup")
+async def startup_event():
+    """
+    应用启动时执行的事件.
+    """
+    load_config()
+    if config:
+        await refresh_access_token()
+
+@app.get("/v1/models", dependencies=[Depends(authenticate_client)])
+async def list_models():
+    """
+    列出可用的模型.
+    """
+    model_data = [
+        {"id": model_name, "object": "model", "owned_by": "kodu-ai", "created": int(time.time())}
+        for model_name in config.get("kodu_models", {})
+    ]
+    return {"object": "list", "data": model_data}
+
+@app.post("/v1/chat/completions", dependencies=[Depends(authenticate_client)])
+async def create_chat_completion(request: ChatCompletionRequest):
+    """
+    处理聊天请求.
+    """
+    if not config.get('kodu_access_token'):
+        # 尝试再次刷新 token
+        if not await refresh_access_token():
+            raise HTTPException(status_code=503, detail="Kodu AI service unavailable, could not refresh token.")
+
+    kodu_model = config.get("kodu_models", {}).get(request.model)
+    if not kodu_model:
+        raise HTTPException(status_code=404, detail=f"Model '{request.model}' not found or mapped.")
+
+    payload = {
+        "model": kodu_model,
+        "messages": [{"role": msg.role, "content": msg.content} for msg in request.messages],
+        "stream": request.stream,
+    }
+
+    headers = {
+        "Authorization": f"Bearer {config['kodu_access_token']}",
+        "Content-Type": "application/json",
+    }
+
+    async def stream_generator():
+        async with httpx.AsyncClient(proxies=config.get("proxy_url"), timeout=300) as client:
+            try:
+                async with client.stream("POST", f"{config['base_url']}/chat/completions", headers=headers, json=payload) as response:
+                    if response.status_code != 200:
+                        error_content = await response.aread()
+                        print(f"Kodu API Error: {response.status_code} - {error_content.decode()}")
+                        # 尝试刷新 token 并重试一次
+                        if response.status_code == 401:
+                            print("Received 401, attempting to refresh token and retry...")
+                            await refresh_access_token()
+                            # 需要重新构建流请求，这里为了简化，直接返回错误
+                        yield f"data: {json.dumps({'error': {'message': f'Kodu API Error: {error_content.decode()}'}})}\n\n"
+                        yield "data: [DONE]\n\n"
+                        return
+
+                    async for chunk in response.aiter_bytes():
+                        yield chunk.decode('utf-8')
+            except Exception as e:
+                print(f"An error occurred during streaming: {e}")
+                yield f"data: {json.dumps({'error': {'message': 'An internal error occurred.'}})}\n\n"
+                yield "data: [DONE]\n\n"
+
+    if request.stream:
+        return StreamingResponse(stream_generator(), media_type="text/event-stream")
+    else:
+        async with httpx.AsyncClient(proxies=config.get("proxy_url"), timeout=300) as client:
+            try:
+                response = await client.post(f"{config['base_url']}/chat/completions", headers=headers, json=payload)
+                if response.status_code == 401:
+                    print("Received 401, refreshing token and retrying...")
+                    await refresh_access_token()
+                    headers["Authorization"] = f"Bearer {config['kodu_access_token']}"
+                    response = await client.post(f"{config['base_url']}/chat/completions", headers=headers, json=payload)
+                
+                response.raise_for_status()
+                return JSONResponse(content=response.json())
+            except httpx.HTTPStatusError as e:
+                raise HTTPException(status_code=e.response.status_code, detail=e.response.text)