Create main.py

main.py

@@ -0,0 +1,386 @@
+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)