app.py

import json
import time
import uuid
import threading
from typing import Any, AsyncGenerator, Dict, List, Optional

import httpx
import uvicorn
from fastapi import FastAPI, HTTPException, Depends, Header
from fastapi.responses import StreamingResponse
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel, Field

# Configuration
CONVERSATION_CACHE_MAX_SIZE = 100
DEFAULT_REQUEST_TIMEOUT = 30.0

# Global variables
VALID_CLIENT_KEYS: set = set()
JETBRAINS_JWTS: list = []
current_jwt_index: int = 0
jwt_rotation_lock = threading.Lock()
models_data: Dict[str, Any] = {}
http_client: Optional[httpx.AsyncClient] = None

# Pydantic Models
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

class ModelInfo(BaseModel):
    id: str
    object: str = "model"
    created: int
    owned_by: str

class ModelList(BaseModel):
    object: str = "list"
    data: List[ModelInfo]

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="JetBrains AI OpenAI Compatible API")
security = HTTPBearer(auto_error=False)

# Helper functions
def load_models():
    """加载模型配置"""
    try:
        with open("models.json", "r", encoding="utf-8") as f:
            model_ids = json.load(f)
        
        processed_models = []
        if isinstance(model_ids, list):
            for model_id in model_ids:
                if isinstance(model_id, str):
                    processed_models.append({
                        "id": model_id,
                        "object": "model",
                        "created": int(time.time()),
                        "owned_by": "jetbrains-ai"
                    })
        
        return {"data": processed_models}
    except Exception as e:
        print(f"加载 models.json 时出错: {e}")
        return {"data": []}

def load_client_api_keys():
    """加载客户端 API 密钥"""
    global VALID_CLIENT_KEYS
    try:
        with open("client_api_keys.json", "r", encoding="utf-8") as f:
            keys = json.load(f)
            if not isinstance(keys, list):
                print("警告: client_api_keys.json 应包含密钥列表")
                VALID_CLIENT_KEYS = set()
                return
            VALID_CLIENT_KEYS = set(keys)
            if not VALID_CLIENT_KEYS:
                print("警告: client_api_keys.json 为空")
            else:
                print(f"成功加载 {len(VALID_CLIENT_KEYS)} 个客户端 API 密钥")
    except FileNotFoundError:
        print("错误: 未找到 client_api_keys.json")
        VALID_CLIENT_KEYS = set()
    except Exception as e:
        print(f"加载 client_api_keys.json 时出错: {e}")
        VALID_CLIENT_KEYS = set()

def load_jetbrains_jwts():
    """加载 JetBrains AI 认证 JWT"""
    global JETBRAINS_JWTS
    try:
        with open("jetbrainsai.json", "r", encoding="utf-8") as f:
            # 假设 jetbrainsai.json 包含一个对象列表，每个对象都有 'jwt' 键
            jwt_data = json.load(f)
            if isinstance(jwt_data, list):
                JETBRAINS_JWTS = [item.get("jwt") for item in jwt_data if "jwt" in item]
        
        if not JETBRAINS_JWTS:
            print("警告: jetbrainsai.json 中未找到有效的 JWT")
        else:
            print(f"成功加载 {len(JETBRAINS_JWTS)} 个 JetBrains AI JWT")
            
    except FileNotFoundError:
        print("错误: 未找到 jetbrainsai.json 文件")
        JETBRAINS_JWTS = []
    except Exception as e:
        print(f"加载 jetbrainsai.json 时出错: {e}")
        JETBRAINS_JWTS = []

def get_model_item(model_id: str) -> Optional[Dict]:
    """根据模型ID获取模型配置"""
    for model in models_data.get("data", []):
        if model.get("id") == model_id:
            return model
    return None

async def authenticate_client(auth: Optional[HTTPAuthorizationCredentials] = Depends(security)):
    """客户端认证"""
    if not VALID_CLIENT_KEYS:
        raise HTTPException(status_code=503, detail="服务不可用: 未配置客户端 API 密钥")
    
    if not auth or not auth.credentials:
        raise HTTPException(
            status_code=401,
            detail="需要在 Authorization header 中提供 API 密钥",
            headers={"WWW-Authenticate": "Bearer"},
        )
    
    if auth.credentials not in VALID_CLIENT_KEYS:
        raise HTTPException(status_code=403, detail="无效的客户端 API 密钥")

def get_next_jetbrains_jwt() -> str:
    """轮询获取下一个 JetBrains JWT"""
    global current_jwt_index
    
    if not JETBRAINS_JWTS:
        raise HTTPException(status_code=503, detail="服务不可用: 未配置 JetBrains JWT")
    
    with jwt_rotation_lock:
        if not JETBRAINS_JWTS:
             raise HTTPException(status_code=503, detail="服务不可用: JetBrains JWT 不可用")
        token_to_use = JETBRAINS_JWTS[current_jwt_index]
        current_jwt_index = (current_jwt_index + 1) % len(JETBRAINS_JWTS)
    return token_to_use

# FastAPI 生命周期事件
@app.on_event("startup")
async def startup():
    global models_data, http_client
    models_data = load_models()
    load_client_api_keys()
    load_jetbrains_jwts()
    http_client = httpx.AsyncClient(timeout=None)
    print("JetBrains AI OpenAI Compatible API 服务器已启动")

@app.on_event("shutdown")
async def shutdown():
    global http_client
    if http_client:
        await http_client.aclose()

# API 端点
@app.get("/v1/models", response_model=ModelList)
async def list_models(_: None = Depends(authenticate_client)):
    """列出可用模型"""
    model_list = []
    for model in models_data.get("data", []):
        model_list.append(ModelInfo(
            id=model.get("id", ""),
            created=model.get("created", int(time.time())),
            owned_by=model.get("owned_by", "jetbrains-ai")
        ))
    return ModelList(data=model_list)

async def openai_stream_adapter(
    api_stream_generator: AsyncGenerator[str, None],
    model_name: str
) -> AsyncGenerator[str, None]:
    """将 JetBrains API 的流转换为 OpenAI 格式的 SSE"""
    stream_id = f"chatcmpl-{uuid.uuid4().hex}"
    first_chunk_sent = False
    
    try:
        async for line in api_stream_generator:
            if not line or line == "data: end":
                continue

            if line.startswith('data: '):
                try:
                    data = json.loads(line[6:])
                    event_type = data.get("type")

                    if event_type == "Content":
                        content = data.get("content", "")
                        if not content:
                            continue
                        
                        delta_payload = {}
                        if not first_chunk_sent:
                            delta_payload = {"role": "assistant", "content": content}
                            first_chunk_sent = True
                        else:
                            delta_payload = {"content": content}
                        
                        stream_resp = StreamResponse(id=stream_id, model=model_name, choices=[StreamChoice(delta=delta_payload)])
                        yield f"data: {stream_resp.json()}\n\n"

                    elif event_type == "FinishMetadata":
                        final_resp = StreamResponse(id=stream_id, model=model_name, choices=[StreamChoice(delta={}, finish_reason="stop")])
                        yield f"data: {final_resp.json()}\n\n"
                        break
                except json.JSONDecodeError:
                    print(f"警告: 无法解析的 JSON 行: {line}")
                    continue
        
        yield "data: [DONE]\n\n"
        
    except Exception as e:
        print(f"流式适配器错误: {e}")
        error_resp = StreamResponse(
            id=stream_id,
            model=model_name,
            choices=[StreamChoice(
                delta={"role": "assistant", "content": f"内部错误: {str(e)}"},
                index=0,
                finish_reason="stop"
            )]
        )
        yield f"data: {error_resp.json()}\n\n"
        yield "data: [DONE]\n\n"

async def aggregate_stream_for_non_stream_response(
    openai_sse_stream: AsyncGenerator[str, None],
    model_name: str
) -> ChatCompletionResponse:
    """聚合流式响应为完整响应"""
    content_parts = []
    
    async for sse_line in openai_sse_stream:
        if sse_line.startswith("data: ") and sse_line.strip() != "data: [DONE]":
            try:
                data = json.loads(sse_line[6:].strip())
                if data.get("choices") and len(data["choices"]) > 0:
                    delta = data["choices"][0].get("delta", {})
                    if "content" in delta:
                        content_parts.append(delta["content"])
            except:
                pass
    
    full_content = "".join(content_parts)
    
    return ChatCompletionResponse(
        model=model_name,
        choices=[ChatCompletionChoice(
            message=ChatMessage(role="assistant", content=full_content),
            finish_reason="stop"
        )]
    )

@app.post("/v1/chat/completions")
async def chat_completions(
    request: ChatCompletionRequest,
    _: None = Depends(authenticate_client)
):
    """创建聊天完成"""
    model_config = get_model_item(request.model)
    if not model_config:
        raise HTTPException(status_code=404, detail=f"模型 {request.model} 未找到")

    auth_token = get_next_jetbrains_jwt()

    # 将 OpenAI 格式的消息转换为 JetBrains 格式
    jetbrains_messages = []
    for msg in request.messages:
        # JetBrains API 需要一个特定的交替格式，这里我们简化处理
        # 实际可能需要更复杂的逻辑来确保用户/助手消息交替
        jetbrains_messages.append({"type": f"{msg.role}_message", "content": msg.content})

    # 创建 API 请求的 payload
    payload = {
        "prompt": "ij.chat.request.new-chat-on-start", # or other relevant prompt
        "profile": request.model,
        "chat": {
            "messages": jetbrains_messages
        },
        "parameters": {"data": []},
    }

    headers = {
        "User-Agent": "ktor-client",
        "Accept": "text/event-stream",
        "Content-Type": "application/json",
        "Accept-Charset": "UTF-8",
        "Cache-Control": "no-cache",
        "grazie-agent": '{"name":"aia:pycharm","version":"251.26094.80.13:251.26094.141"}', # 可根据需要更新
        "grazie-authenticate-jwt": auth_token,
    }

    async def api_stream_generator():
        """一个包装 httpx 请求的异步生成器"""
        async with http_client.stream("POST", "https://api.jetbrains.ai/user/v5/llm/chat/stream/v7", 
                                       json=payload, headers=headers, timeout=300) as response:
            response.raise_for_status()
            async for line in response.aiter_lines():
                yield line

    # 创建 OpenAI 格式的流
    openai_sse_stream = openai_stream_adapter(
        api_stream_generator(),
        request.model
    )

    # 返回流式或非流式响应
    if request.stream:
        return StreamingResponse(
            openai_sse_stream,
            media_type="text/event-stream"
        )
    else:
        return await aggregate_stream_for_non_stream_response(
            openai_sse_stream,
            request.model
        )

# 主程序入口
if __name__ == "__main__":
    import os
    
    # 创建示例配置文件（如果不存在）
    if not os.path.exists("client_api_keys.json"):
        with open("client_api_keys.json", "w", encoding="utf-8") as f:
            json.dump(["sk-your-custom-key-here"], f, indent=2)
        print("已创建示例 client_api_keys.json 文件")

    if not os.path.exists("jetbrainsai.json"):
        with open("jetbrainsai.json", "w", encoding="utf-8") as f:
            json.dump([{"jwt": "your-jwt-here"}], f, indent=2)
        print("已创建示例 jetbrainsai.json 文件")

    if not os.path.exists("models.json"):
        with open("models.json", "w", encoding="utf-8") as f:
            json.dump(["anthropic-claude-3.5-sonnet"], f, indent=2)
        print("已创建示例 models.json 文件")
    
    print("正在启动 JetBrains AI OpenAI Compatible API 服务器...")
    print("端点:")
    print("  GET  /v1/models")
    print("  POST /v1/chat/completions")
    print("\n在 Authorization header 中使用客户端 API 密钥 (Bearer sk-xxx)")
    
    uvicorn.run(app, host="0.0.0.0", port=8000)