Upload 3 files

Dockerfile

@@ -0,0 +1,15 @@
+# 第一阶段：构建阶段
+FROM python:3.9-slim as builder
+
+WORKDIR /app
+COPY requirements.txt .
+RUN pip install --user -r requirements.txt
+
+# 第二阶段：运行时镜像
+FROM python:3.9-slim
+WORKDIR /app
+COPY --from=builder /root/.local /root/.local
+COPY . .
+ENV PATH=/root/.local/bin:$PATH
+EXPOSE 8000
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

main.py

@@ -0,0 +1,247 @@
+from fastapi import FastAPI, Request, Response, HTTPException
+
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+import httpx
+import json
+import uuid
+from typing import Optional, List, Dict, Any
+from pydantic import BaseModel
+import asyncio
+
+# 创建FastAPI应用
+app = FastAPI()
+
+# 配置CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# 定义数据模型
+class Message(BaseModel):
+    role: str
+    content: str
+
+class ChatRequest(BaseModel):
+    messages: List[Message]
+    model: str
+    stream: Optional[bool] = True
+
+class ChatResponse(BaseModel):
+    id: str
+    object: str = "chat.completion"
+    created: int
+    model: str
+    choices: List[Dict[str, Any]]
+    usage: Optional[Dict[str, int]] = None
+
+# 模型映射
+MODEL_MAPPING = {
+    "gpt-4o-mini-abacus": "OPENAI_GPT4O_MINI",
+    "claude-3.5-sonnet-abacus": "CLAUDE_V3_5_SONNET",
+    "claude-3.7-sonnet-abacus": "CLAUDE_V3_7_SONNET", 
+    "claude-3.7-sonnet-thinking-abacus": "CLAUDE_V3_7_SONNET_THINKING", 
+    "o3-mini-abacus": "OPENAI_O3_MINI",
+    "o3-mini-high-abacus": "OPENAI_O3_MINI_HIGH",
+    "o1-mini-abacus": "OPENAI_O1_MINI",
+    "deepseek-r1-abacus": "DEEPSEEK_R1",
+    "gemini-2-pro-abacus": "GEMINI_2_PRO",
+    "gemini-2-flash-thinking-abacus": "GEMINI_2_FLASH_THINKING",
+    "gemini-2-flash-abacus": "GEMINI_2_FLASH",
+    "gemini-1.5-pro-abacus": "GEMINI_1_5_PRO",
+    "xai-grok-abacus": "XAI_GROK",
+    "deepseek-v3-abacus": "DEEPSEEK_V3",
+    "llama3-1-405b-abacus": "LLAMA3_1_405B",
+    "gpt-4o-abacus": "OPENAI_GPT4O",
+    "gpt-4o-2024-08-06-abacus": "OPENAI_GPT4O", 
+    "gpt-3.5-turbo-abacus": "OPENAI_O3_MINI",  
+    "gpt-3.5-turbo-16k-abacus": "OPENAI_O3_MINI_HIGH" 
+}
+
+BASE_URL = "https://pa002.abacus.ai"
+
+TIMEOUT = 30.0  # 请求超时时间（秒）
+MAX_RETRIES = 3  # 最大重试次数
+RETRY_DELAY = 1  # 重试延迟（秒）
+
+@app.get("/v1/models")
+async def list_models():
+    """返回支持的模型列表"""
+    models = [
+        {
+            "id": model_id,
+            "object": "model",
+            "created": 1677610602,
+            "owned_by": "system",
+        }
+        for model_id in MODEL_MAPPING.keys()
+    ]
+    return {
+        "object": "list",
+        "data": models
+    }
+
+# 工具函数：获取请求头
+def get_headers(auth_token: str) -> Dict[str, str]:
+    """生成请求头"""
+    return {
+        "sec-ch-ua-platform": "Windows",
+        "sec-ch-ua": '"Not(A:Brand";v="99", "Microsoft Edge";v="133", "Chromium";v="133"',
+        "sec-ch-ua-mobile": "?0",
+        "X-Abacus-Org-Host": "apps",
+        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36 Edg/133.0.0.0",
+        "Sec-Fetch-Site": "same-site",
+        "Sec-Fetch-Mode": "cors",
+        "Sec-Fetch-Dest": "empty",
+        "host": "pa002.abacus.ai",
+        "Cookie": auth_token,
+        "Accept": "text/event-stream",
+        "Content-Type": "text/plain;charset=UTF-8"
+    }
+
+def process_messages(messages: List[Message]) -> str:
+    """处理消息列表，合并成单个消息"""
+    system_message = next((msg.content for msg in messages if msg.role == "system"), None)
+    context_messages = [msg for msg in messages if msg.role != "system"][:-1]
+    current_message = messages[-1].content
+    
+    full_message = current_message
+    
+    if system_message:
+        full_message = f"System: {system_message}\n\n{full_message}"
+        
+    if context_messages:
+        context_str = "\n".join(f"{msg.role}: {msg.content}" for msg in context_messages)
+        full_message = f"Previous conversation:\n{context_str}\nCurrent message: {full_message}"
+        
+    return full_message
+
+@app.post("/v1/chat/completions")
+async def chat_completions(request: Request, chat_request: ChatRequest):
+    """处理聊天完成请求"""
+    # 获取认证token
+    auth_header = request.headers.get("Authorization", "")
+    if not auth_header.startswith("Bearer "):
+        return Response(
+            content=json.dumps({"error": "未提供有效的Authorization header"}),
+            status_code=401
+        )
+    
+    auth_token = auth_header.replace("Bearer ", "")
+    
+    # 创建会话ID
+    conversation_id = str(uuid.uuid4())
+    
+    # 处理消息
+    full_message = process_messages(chat_request.messages)
+    
+    # 准备请求数据
+    request_data = {
+        "requestId": str(uuid.uuid4()),
+        "deploymentConversationId": conversation_id,
+        "message": full_message,
+        "isDesktop": True,
+        "chatConfig": {
+            "timezone": "Asia/Shanghai",
+            "language": "zh-CN"
+        },
+        "llmName": MODEL_MAPPING.get(chat_request.model, chat_request.model),
+        "externalApplicationId": str(uuid.uuid4())
+    }
+
+    # 流式请求处理
+    async def generate_stream():
+        headers = get_headers(auth_token)
+        
+        for retry in range(MAX_RETRIES):
+            try:
+                async with httpx.AsyncClient() as client:
+                    async with client.stream(
+                        "POST",
+                        f"{BASE_URL}/api/_chatLLMSendMessageSSE",
+                        headers=headers,
+                        content=json.dumps(request_data),
+                        timeout=TIMEOUT
+                    ) as response:
+                        async for line in response.aiter_lines():
+                            if not line.strip():
+                                continue
+                                
+                            try:
+                                data = json.loads(line)
+                                
+                                if data.get("type") == "text" and data.get("title") != "Thinking...":
+                                    chunk = {
+                                        "id": str(uuid.uuid4()),
+                                        "object": "chat.completion.chunk",
+                                        "created": int(uuid.uuid1().time_low),
+                                        "model": chat_request.model,
+                                        "choices": [{
+                                            "delta": {
+                                                "role": "assistant",
+                                                "content": data.get("segment", "")
+                                            },
+                                            "index": 0
+                                        }]
+                                    }
+                                    yield f"data: {json.dumps(chunk)}\n\n"
+                                    
+                                if data.get("end"):
+                                    # 发送结束标记
+                                    chunk = {
+                                        "id": str(uuid.uuid4()),
+                                        "object": "chat.completion.chunk",
+                                        "created": int(uuid.uuid1().time_low),
+                                        "model": chat_request.model,
+                                        "choices": [{
+                                            "delta": {"content": ""},
+                                            "index": 0,
+                                            "finish_reason": "stop"
+                                        }]
+                                    }
+                                    yield f"data: {json.dumps(chunk)}\n\n"
+                                    yield "data: [DONE]\n\n"
+                                    break  # 成功完成，退出重试循环
+                                    
+                            except json.JSONDecodeError:
+                                continue
+            except (httpx.TimeoutException, httpx.RequestError) as e:
+                if retry == MAX_RETRIES - 1:  # 最后一次重试
+                    yield f"data: {json.dumps({'error': str(e)})}\n\n"
+                    yield "data: [DONE]\n\n"
+                    return
+                await asyncio.sleep(RETRY_DELAY)
+    
+    return StreamingResponse(
+        generate_stream(),
+        media_type="text/event-stream"
+    )
+
+@app.get("/")
+async def health_check():
+    """健康检查"""
+    return {"status": "ok", "version": "1.0.0"}
+
+@app.exception_handler(Exception)
+async def global_exception_handler(request: Request, exc: Exception):
+    """全局异常处理"""
+    error_message = str(exc)
+    return Response(
+        content=json.dumps({
+            "error": {
+                "message": error_message,
+                "type": exc.__class__.__name__,
+                "code": 500
+            }
+        }),
+        status_code=500,
+        media_type="application/json"
+    )
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)

requirements.txt

@@ -0,0 +1,7 @@
+fastapi>=0.68.0,<0.69.0
+uvicorn>=0.15.0,<0.16.0
+httpx>=0.24.0
+pydantic>=1.8.0,<2.0.0
+python-multipart>=0.0.5
+typing-extensions>=4.0.0
+openai>=1.0.0,<2.0.0