main.py

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)