Rename adapter.ts to main.py

adapter.ts

@@ -1,255 +0,0 @@
-// deno run --allow-net --allow-env adapter.ts
-
-import { serve } from "https://deno.land/std@0.203.0/http/server.ts";
-
-// --- Configuration from Environment Variables (Safer for deployment) ---
-
-function getKeysFromEnv(envVarName: string): Set<string> {
-  const keysString = Deno.env.get(envVarName);
-  if (!keysString) {
-    console.warn(`Environment variable ${envVarName} is not set.`);
-    return new Set();
-  }
-  // Split by comma and trim whitespace, filter out empty strings
-  return new Set(keysString.split(',').map(k => k.trim()).filter(Boolean));
-}
-
-// Client keys will be read from Hugging Face Secrets
-const CLIENT_API_KEYS = getKeysFromEnv("CLIENT_KEYS"); 
-
-// CodeGeeX tokens will also be read from Hugging Face Secrets
-const codegeeXTokensRaw = Array.from(getKeysFromEnv("CODEGEEX_KEYS"));
-
-const CODEGEEX_TOKENS: {
-  token: string;
-  isValid: boolean;
-  lastUsed: number;
-  errorCount: number;
-}[] = codegeeXTokensRaw.map(token => ({
-    token: token,
-    isValid: true,
-    lastUsed: 0,
-    errorCount: 0
-}));
-
-
-const MAX_ERROR_COUNT = 3;
-const ERROR_COOLDOWN = 300 * 1000; // ms
-
-// --- Utilities ---
-function now(): number {
-  return Date.now();
-}
-
-function rotateToken(): typeof CODEGEEX_TOKENS[0] | null {
-  if (CODEGEEX_TOKENS.length === 0) {
-      console.error("CODEGEEX_TOKENS array is empty. Check your CODEGEEX_KEYS secret.");
-      return null;
-  }
-    
-  const available = CODEGEEX_TOKENS.filter(t => {
-    if (!t.isValid) return false;
-    if (t.errorCount >= MAX_ERROR_COUNT && now() - t.lastUsed < ERROR_COOLDOWN) return false;
-    return true;
-  });
-  if (available.length === 0) return null;
-
-  // reset cooled-down tokens
-  for (const t of available) {
-    if (t.errorCount >= MAX_ERROR_COUNT && now() - t.lastUsed >= ERROR_COOLDOWN) {
-      t.errorCount = 0;
-    }
-  }
-
-  // pick the one least recently used, then lowest errorCount
-  available.sort((a, b) => a.lastUsed - b.lastUsed || a.errorCount - b.errorCount);
-  const tok = available[0];
-  tok.lastUsed = now();
-  return tok;
-}
-
-// This function translates the OpenAI format to CodeGeeX format
-function convertToCodeGeeXPayload(params: { model: string; messages: any[] }) {
-    // CodeGeeX seems to use the last message's content as the main prompt.
-    // The history part is more complex, here we simplify it.
-    const lastMessage = params.messages.slice(-1)[0];
-    const history = params.messages.slice(0, -1)
-        .filter(msg => msg.role === 'user' || msg.role === 'assistant')
-        .map(msg => ({
-            role: msg.role,
-            content: msg.content
-        }));
-
-    return {
-        user_role: 0, // This seems to be a fixed value
-        ide: "HuggingFace", // Let's identify the source
-        prompt: lastMessage?.content || "",
-        history: history, // Passing a simplified history
-        model: params.model,
-    };
-}
-
-
-async function proxyChat(req: Request, params: { stream: boolean; model: string; messages: any[] }) {
-  const tokenObj = rotateToken();
-  if (!tokenObj) {
-    return new Response(JSON.stringify({ error: { message: "No valid CodeGeeX tokens available", type: "server_error" } }), { status: 503, headers: { "Content-Type": "application/json" }});
-  }
-
-  const payload = convertToCodeGeeXPayload(params);
-
-  try {
-    const response = await fetch("https://codegeex.cn/prod/code/chatCodeSseV3/chat", {
-      method: "POST",
-      headers: {
-        "Content-Type": "application/json",
-        "Accept": "text/event-stream",
-        "code-token": tokenObj.token,
-      },
-      body: JSON.stringify(payload),
-    });
-
-    if (!response.ok) {
-      console.error(`Upstream error from CodeGeeX: ${response.status}`);
-      if (response.status === 401 || response.status === 403) {
-        tokenObj.isValid = false;
-        console.warn(`Token ${tokenObj.token.substring(0, 15)}... marked as invalid due to 401/403 error.`);
-      } else {
-        tokenObj.errorCount++;
-        console.warn(`Token ${tokenObj.token.substring(0, 15)}... error count increased to ${tokenObj.errorCount}.`);
-      }
-      const errorBody = await response.text();
-      return new Response(JSON.stringify({ error: { message: `Upstream error ${response.status}: ${errorBody}`, type: "upstream_error" } }), { status: 502, headers: { "Content-Type": "application/json" }});
-    }
-
-    // For stream, we must transform the raw CodeGeeX SSE to OpenAI format
-    if (params.stream) {
-        const { readable, writable } = new TransformStream();
-        const writer = writable.getWriter();
-        const encoder = new TextEncoder();
-
-        // This function processes the stream from CodeGeeX and sends OpenAI compatible chunks
-        (async () => {
-            const reader = response.body?.getReader();
-            if (!reader) {
-                await writer.close();
-                return;
-            }
-            const decoder = new TextDecoder();
-            const completionId = `chatcmpl-${crypto.randomUUID()}`;
-            const creationTime = Math.floor(now() / 1000);
-
-            try {
-                while(true) {
-                    const { done, value } = await reader.read();
-                    if (done) break;
-
-                    const chunkText = decoder.decode(value);
-                    // A simple transformation: assume the raw chunk is the content delta
-                    const openAIChunk = {
-                        id: completionId,
-                        object: "chat.completion.chunk",
-                        created: creationTime,
-                        model: params.model,
-                        choices: [{ delta: { content: chunkText }, index: 0, finish_reason: null }]
-                    };
-                    await writer.write(encoder.encode(`data: ${JSON.stringify(openAIChunk)}\n\n`));
-                }
-                // Send the final DONE chunk
-                await writer.write(encoder.encode(`data: [DONE]\n\n`));
-            } catch (e) {
-                console.error("Error while transforming stream:", e);
-            } finally {
-                await writer.close();
-            }
-        })();
-        
-        return new Response(readable, {
-            status: 200,
-            headers: { "Content-Type": "text/event-stream", "Cache-Control": "no-cache", "Connection": "keep-alive" },
-        });
-    } else {
-      // accumulate and return JSON
-      const text = await response.text();
-      return new Response(JSON.stringify({
-        id: `chatcmpl-${crypto.randomUUID()}`,
-        object: "chat.completion",
-        created: Math.floor(now() / 1000),
-        model: params.model,
-        choices: [{ message: { role: "assistant", content: text }, index: 0, finish_reason: "stop" }],
-        usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 } // Placeholder usage
-      }), {
-        status: 200,
-        headers: { "Content-Type": "application/json" },
-      });
-    }
-  } catch (err) {
-    tokenObj.errorCount++;
-    console.error("Fetch to CodeGeeX failed:", err);
-    return new Response(JSON.stringify({ error: { message: err.message, type: "server_error" } }), { status: 500, headers: { "Content-Type": "application/json" }});
-  }
-}
-
-// --- Main Handler ---
-async function handler(req: Request): Promise<Response> {
-  const url = new URL(req.url);
-  console.log(`Received request: ${req.method} ${url.pathname}`);
-
-  // CORS preflight request handler for web clients
-  if (req.method === 'OPTIONS') {
-    return new Response(null, {
-      status: 204,
-      headers: {
-        'Access-Control-Allow-Origin': '*',
-        'Access-Control-Allow-Methods': 'GET, POST, OPTIONS',
-        'Access-Control-Allow-Headers': 'Content-Type, Authorization',
-      },
-    });
-  }
-  
-  // Authentication middleware
-  const auth = req.headers.get("Authorization")?.replace(/^Bearer\s+/, "");
-  if (CLIENT_API_KEYS.size === 0) {
-    console.error("Server misconfigured: CLIENT_KEYS secret is not set or empty.");
-    return new Response(JSON.stringify({ error: { message: "Server misconfigured: no client keys", type: "server_error" }}), { status: 503, headers: { "Content-Type": "application/json" }});
-  }
-  if (!auth || !CLIENT_API_KEYS.has(auth)) {
-    return new Response(JSON.stringify({ error: { message: "Invalid or missing API key", type: "auth_error" }}), {
-      status: 401,
-      headers: { "WWW-Authenticate": "Bearer", "Content-Type": "application/json" },
-    });
-  }
-
-  // GET /v1/models
-  if (url.pathname === "/v1/models" && req.method === "GET") {
-    const modelData = [
-      { id: "codegeex-4", object: "model", created: Math.floor(now() / 1000), owned_by: "codegeex" },
-      { id: "codegeex-pro", object: "model", created: Math.floor(now() / 1000), owned_by: "codegeex" }
-    ];
-    return new Response(JSON.stringify({ object: "list", data: modelData }), {
-      headers: { "Content-Type": "application/json" },
-    });
-  }
-
-  // POST /v1/chat/completions
-  if (url.pathname === "/v1/chat/completions" && req.method === "POST") {
-    try {
-        const body = await req.json();
-        const { model, messages, stream = true } = body;
-        if (!model || !Array.isArray(messages) || messages.length === 0) {
-          return new Response(JSON.stringify({ error: { message: "Bad Request: 'model' and 'messages' are required.", type: "invalid_request_error" } }), { status: 400, headers: { "Content-Type": "application/json" }});
-        }
-        return proxyChat(req, { model, messages, stream });
-    } catch (e) {
-        return new Response(JSON.stringify({ error: { message: "Invalid JSON body.", type: "invalid_request_error" } }), { status: 400, headers: { "Content-Type": "application/json" }});
-    }
-  }
-
-  // Not found
-  return new Response(JSON.stringify({ error: "Not Found" }), { status: 404, headers: { "Content-Type": "application/json" }});
-}
-
-// --- Start Server ---
-const PORT = 7860; // Use the standard port for Hugging Face Spaces
-console.log(`Starting Deno CodeGeeX Adapter on http://0.0.0.0:${PORT}`);
-serve(handler, { port: PORT });

main.py

@@ -0,0 +1,260 @@
+# ===================================================================
+# main.py (已修改以适配 Hugging Face Secrets)
+# ===================================================================
+
+import json
+import os
+import time
+import uuid
+import threading
+from typing import Any, Dict, List, Optional, TypedDict, Union
+
+import requests
+from fastapi import FastAPI, HTTPException, Depends, Query
+from fastapi.responses import StreamingResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from pydantic import BaseModel, Field
+
+# --- 类型定义和全局变量 (与原始文件一致) ---
+class CodeGeeXToken(TypedDict):
+    token: str
+    is_valid: bool
+    last_used: float
+    error_count: int
+
+VALID_CLIENT_KEYS: set = set()
+CODEGEEX_TOKENS: List[CodeGeeXToken] = []
+CODEGEEX_MODELS: List[str] = ["claude-3-7-sonnet", "claude-sonnet-4"]
+token_rotation_lock = threading.Lock()
+MAX_ERROR_COUNT = 3
+ERROR_COOLDOWN = 300
+DEBUG_MODE = os.environ.get("DEBUG_MODE", "false").lower() == "true"
+
+# --- Pydantic 模型 (与原始文件一致) ---
+class ChatMessage(BaseModel):
+    role: str
+    content: Union[str, List[Dict[str, Any]]]
+    reasoning_content: Optional[str] = None
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[ChatMessage]
+    stream: bool = True
+    temperature: Optional[float] = None
+    max_tokens: Optional[int] = None
+    top_p: Optional[float] = None
+
+class ModelInfo(BaseModel):
+    id: str
+    object: str = "model"
+    created: int
+    owned_by: str
+
+class ModelList(BaseModel):
+    object: str = "list"
+    data: List[ModelInfo]
+# ... (其他 Pydantic 模型与原始文件一致)
+class ChatCompletionChoice(BaseModel):
+    message: ChatMessage
+    index: int = 0
+    finish_reason: str = "stop"
+class ChatCompletionResponse(BaseModel):
+    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex}")
+    object: str = "chat.completion"
+    created: int = Field(default_factory=lambda: int(time.time()))
+    model: str
+    choices: List[ChatCompletionChoice]
+    usage: Dict[str, int] = Field(default_factory=lambda: {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0})
+class StreamChoice(BaseModel):
+    delta: Dict[str, Any] = Field(default_factory=dict)
+    index: int = 0
+    finish_reason: Optional[str] = None
+class StreamResponse(BaseModel):
+    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex}")
+    object: str = "chat.completion.chunk"
+    created: int = Field(default_factory=lambda: int(time.time()))
+    model: str
+    choices: List[StreamChoice]
+
+# --- FastAPI App ---
+app = FastAPI(title="CodeGeeX OpenAI API Adapter")
+security = HTTPBearer(auto_error=False)
+
+def log_debug(message: str):
+    if DEBUG_MODE:
+        print(f"[DEBUG] {message}")
+
+# --- 配置加载函数 (已修改为从 Secrets 读取) ---
+def load_client_api_keys_from_secrets():
+    """从环境变量加载客户端 API Keys"""
+    global VALID_CLIENT_KEYS
+    try:
+        keys_str = os.environ.get("CLIENT_API_KEYS")
+        if not keys_str:
+            raise ValueError("Secret 'CLIENT_API_KEYS' not found.")
+        keys = json.loads(keys_str)
+        VALID_CLIENT_KEYS = set(keys) if isinstance(keys, list) else set()
+        print(f"Successfully loaded {len(VALID_CLIENT_KEYS)} client API keys from secrets.")
+    except Exception as e:
+        print(f"FATAL: Error loading client API keys from secrets: {e}")
+        VALID_CLIENT_KEYS = set()
+
+def load_codegeex_tokens_from_secrets():
+    """从环境变量加载 CodeGeeX Tokens"""
+    global CODEGEEX_TOKENS
+    CODEGEEX_TOKENS = []
+    try:
+        tokens_str = os.environ.get("CODEGEEX_TOKENS")
+        if not tokens_str:
+            raise ValueError("Secret 'CODEGEEX_TOKENS' not found.")
+        tokens = json.loads(tokens_str) # 假设Secret是一个JSON数组
+        if not isinstance(tokens, list):
+             raise TypeError("Secret 'CODEGEEX_TOKENS' must be a JSON list of strings.")
+
+        for token in tokens:
+            if isinstance(token, str) and token:
+                CODEGEEX_TOKENS.append({
+                    "token": token, "is_valid": True, "last_used": 0, "error_count": 0
+                })
+        print(f"Successfully loaded {len(CODEGEEX_TOKENS)} CodeGeeX tokens from secrets.")
+    except Exception as e:
+        print(f"FATAL: Error loading CodeGeeX tokens from secrets: {e}")
+
+# --- 核心逻辑 (与原始文件保持一致) ---
+def get_best_codegeex_token() -> Optional[CodeGeeXToken]:
+    with token_rotation_lock:
+        now = time.time()
+        valid_tokens = [t for t in CODEGEEX_TOKENS if t["is_valid"] and (t["error_count"] < MAX_ERROR_COUNT or now - t["last_used"] > ERROR_COOLDOWN)]
+        if not valid_tokens: return None
+        for token in valid_tokens:
+            if token["error_count"] >= MAX_ERROR_COUNT and now - token["last_used"] > ERROR_COOLDOWN: token["error_count"] = 0
+        valid_tokens.sort(key=lambda x: (x["last_used"], x["error_count"]))
+        token = valid_tokens[0]
+        token["last_used"] = now
+        return token
+
+def _convert_messages_to_codegeex_format(messages: List[ChatMessage]):
+    if not messages: return "", []
+    last_user_msg = next((msg for msg in reversed(messages) if msg.role == "user"), None)
+    if not last_user_msg: raise HTTPException(status_code=400, detail="No user message found.")
+    prompt = last_user_msg.content if isinstance(last_user_msg.content, str) else ""
+    history, user_content, assistant_content = [], "", ""
+    for msg in messages:
+        if msg == last_user_msg: break
+        if msg.role == "user":
+            if user_content and assistant_content: history.append({"query": user_content, "answer": assistant_content, "id": f"{uuid.uuid4()}"}); user_content, assistant_content = "", ""
+            user_content = msg.content if isinstance(msg.content, str) else ""
+        elif msg.role == "assistant":
+            assistant_content = msg.content if isinstance(msg.content, str) else ""
+            if user_content: history.append({"query": user_content, "answer": assistant_content, "id": f"{uuid.uuid4()}"}); user_content, assistant_content = "", ""
+    if user_content and not assistant_content: prompt = user_content + "\n" + prompt
+    return prompt, history
+
+async def authenticate_client(auth: Optional[HTTPAuthorizationCredentials] = Depends(security)):
+    if not VALID_CLIENT_KEYS: raise HTTPException(status_code=503, detail="Service unavailable: Client API keys not configured.")
+    if not auth or not auth.credentials: raise HTTPException(status_code=401, detail="API key required.", headers={"WWW-Authenticate": "Bearer"})
+    if auth.credentials not in VALID_CLIENT_KEYS: raise HTTPException(status_code=403, detail="Invalid client API key.")
+
+@app.on_event("startup")
+async def startup():
+    print("Starting CodeGeeX OpenAI API Adapter server...")
+    load_client_api_keys_from_secrets()
+    load_codegeex_tokens_from_secrets()
+    print("Server initialization completed.")
+
+def get_models_list_response() -> ModelList:
+    return ModelList(data=[ModelInfo(id=model, created=int(time.time()), owned_by="anthropic") for model in CODEGEEX_MODELS])
+
+@app.get("/v1/models", response_model=ModelList)
+async def list_v1_models(_: None = Depends(authenticate_client)):
+    return get_models_list_response()
+
+@app.get("/models", response_model=ModelList)
+async def list_models_no_auth():
+    return get_models_list_response()
+
+# ... (所有路由和核心函数都与原始文件一致, 此处省略以保持简洁, 但它们都在上面的完整代码块中)
+# --- The rest of the original code follows ---
+# This includes _codegeex_stream_generator, _build_codegeex_non_stream_response, chat_completions, etc.
+# They are included in the full code block above.
+def _codegeex_stream_generator(response, model: str):
+    stream_id = f"chatcmpl-{uuid.uuid4().hex}"
+    created_time = int(time.time())
+    yield f"data: {StreamResponse(id=stream_id, created=created_time, model=model, choices=[StreamChoice(delta={'role': 'assistant'})]).json()}\n\n"
+    buffer = ""
+    try:
+        for chunk in response.iter_content(chunk_size=1024):
+            if not chunk: continue
+            buffer += chunk.decode("utf-8", errors='ignore')
+            while "\n\n" in buffer:
+                event_data, buffer = buffer.split("\n\n", 1)
+                event_data = event_data.strip()
+                if not event_data: continue
+                event_type, data_json = None, None
+                for line in event_data.split("\n"):
+                    if line.startswith("event:"): event_type = line[6:].strip()
+                    elif line.startswith("data:"):
+                        try: data_json = json.loads(line[5:].strip())
+                        except: continue
+                if not event_type or not data_json: continue
+                if event_type == "add":
+                    delta = data_json.get("text", "")
+                    if delta: yield f"data: {StreamResponse(id=stream_id, created=created_time, model=model, choices=[StreamChoice(delta={'content': delta})]).json()}\n\n"
+                elif event_type == "finish":
+                    yield f"data: {StreamResponse(id=stream_id, created=created_time, model=model, choices=[StreamChoice(delta={}, finish_reason='stop')]).json()}\n\n"
+                    yield "data: [DONE]\n\n"
+                    return
+    except Exception as e:
+        log_debug(f"Stream processing error: {e}")
+        yield f"data: {json.dumps({'error': str(e)})}\n\n"
+    yield f"data: {StreamResponse(id=stream_id, created=created_time, model=model, choices=[StreamChoice(delta={}, finish_reason='stop')]).json()}\n\n"
+    yield "data: [DONE]\n\n"
+
+def _build_codegeex_non_stream_response(response, model: str) -> ChatCompletionResponse:
+    full_content = ""
+    buffer = ""
+    for chunk in response.iter_content(chunk_size=1024):
+        if not chunk: continue
+        buffer += chunk.decode("utf-8", errors='ignore')
+        while "\n\n" in buffer:
+            event_data, buffer = buffer.split("\n\n", 1)
+            event_data = event_data.strip()
+            if not event_data: continue
+            event_type, data_json = None, None
+            for line in event_data.split("\n"):
+                if line.startswith("event:"): event_type = line[6:].strip()
+                elif line.startswith("data:"):
+                    try: data_json = json.loads(line[5:].strip())
+                    except: continue
+                if not event_type or not data_json: continue
+                if event_type == "add": full_content += data_json.get("text", "")
+                elif event_type == "finish":
+                    finish_text = data_json.get("text", "")
+                    if finish_text: full_content = finish_text
+                    return ChatCompletionResponse(model=model, choices=[ChatCompletionChoice(message=ChatMessage(role="assistant", content=full_content))])
+    return ChatCompletionResponse(model=model, choices=[ChatCompletionChoice(message=ChatMessage(role="assistant", content=full_content))])
+
+@app.post("/v1/chat/completions")
+async def chat_completions(request: ChatCompletionRequest, _: None = Depends(authenticate_client)):
+    if request.model not in CODEGEEX_MODELS: raise HTTPException(status_code=404, detail=f"Model '{request.model}' not found.")
+    if not request.messages: raise HTTPException(status_code=400, detail="No messages provided.")
+    try: prompt, history = _convert_messages_to_codegeex_format(request.messages)
+    except Exception as e: raise HTTPException(status_code=400, detail=f"Failed to process messages: {e}")
+    for attempt in range(len(CODEGEEX_TOKENS) + 1):
+        if attempt == len(CODEGEEX_TOKENS): raise HTTPException(status_code=503, detail="All attempts to contact CodeGeeX API failed.")
+        token = get_best_codegeex_token()
+        if not token: raise HTTPException(status_code=503, detail="No valid CodeGeeX tokens available.")
+        try:
+            payload = {"user_role": 0, "ide": "VSCode", "prompt": prompt, "model": request.model, "history": history, "talkId": f"{uuid.uuid4()}", "plugin_version": "", "locale": "", "agent": None, "candidates": {"candidate_msg_id": "", "candidate_type": "", "selected_candidate": ""}, "ide_version": "", "machineId": ""}
+            headers = {"User-Agent": "Mozilla/5.0", "Accept": "text/event-stream", "Content-Type": "application/json", "code-token": token["token"]}
+            response = requests.post("https://codegeex.cn/prod/code/chatCodeSseV3/chat", data=json.dumps(payload), headers=headers, stream=True, timeout=300.0)
+            response.raise_for_status()
+            if request.stream: return StreamingResponse(_codegeex_stream_generator(response, request.model), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"})
+            else: return _build_codegeex_non_stream_response(response, request.model)
+        except requests.HTTPError as e:
+            status_code = getattr(e.response, "status_code", 500)
+            with token_rotation_lock:
+                if status_code in [401, 403]: token["is_valid"] = False
+                elif status_code in [429, 500, 502, 503, 504]: token["error_count"] += 1
+        except Exception as e:
+            with token_rotation_lock: token["error_count"] += 1