command2283/jetbrainsai2api
0
1import json2import time3import uuid4import threading5import asyncio6import os7from typing import Any, AsyncGenerator, Dict, List, Optional, Union8 9import httpx10import uvicorn11import aiofiles12from fastapi import FastAPI, HTTPException, Depends, Header13from fastapi.middleware.cors import CORSMiddleware14from fastapi.responses import StreamingResponse15from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials16from pydantic import BaseModel, Field17 18# Configuration19DEFAULT_REQUEST_TIMEOUT = 30.020 21# Global variables22VALID_CLIENT_KEYS: set = set()23JETBRAINS_ACCOUNTS: list = []24current_account_index: int = 025account_rotation_lock = asyncio.Lock()26file_write_lock = asyncio.Lock()27models_data: Dict[str, Any] = {}28anthropic_model_mappings: Dict[str, str] = {}29http_client: Optional[httpx.AsyncClient] = None30 31# Pydantic Models32class ChatMessage(BaseModel):33 role: str34 content: Optional[Union[str, List[Dict[str, Any]]]] = None35 tool_calls: Optional[List[Dict[str, Any]]] = None36 tool_call_id: Optional[str] = None37 38 39class ChatCompletionRequest(BaseModel):40 model: str41 messages: List[ChatMessage]42 stream: bool = False43 temperature: Optional[float] = None44 max_tokens: Optional[int] = None45 top_p: Optional[float] = None46 tools: Optional[List[Dict[str, Any]]] = None47 stop: Optional[Union[str, List[str]]] = None48 49 50# --- Anthropic-Compatible Models ---51 52 53class AnthropicContentBlock(BaseModel):54 type: str55 text: Optional[str] = None56 tool_use_id: Optional[str] = None57 content: Optional[Union[str, List[Dict[str, Any]]]] = None58 id: Optional[str] = None59 name: Optional[str] = None60 input: Optional[Dict[str, Any]] = None61 62 63class AnthropicMessage(BaseModel):64 role: str65 content: Union[str, List[AnthropicContentBlock]]66 67 68class AnthropicTool(BaseModel):69 name: str70 description: Optional[str] = None71 input_schema: Dict[str, Any]72 73 74class AnthropicMessageRequest(BaseModel):75 model: str76 messages: List[AnthropicMessage]77 system: Optional[Union[str, List[Dict[str, Any]]]] = None78 max_tokens: int79 stream: bool = False80 temperature: Optional[float] = None81 top_p: Optional[float] = None82 tools: Optional[List[AnthropicTool]] = None83 stop_sequences: Optional[List[str]] = None84 85 86# --- Anthropic-Compatible Response Models ---87 88 89class AnthropicUsage(BaseModel):90 input_tokens: int91 output_tokens: int92 93 94class AnthropicResponseContent(BaseModel):95 type: str96 id: Optional[str] = None97 name: Optional[str] = None98 input: Optional[Dict[str, Any]] = None99 text: Optional[str] = None100 101 102class AnthropicResponseMessage(BaseModel):103 id: str104 type: str = "message"105 role: str = "assistant"106 model: str107 content: List[AnthropicResponseContent]108 stop_reason: Optional[str]109 stop_sequence: Optional[str] = None110 usage: AnthropicUsage111 112 113# --- End Anthropic Models ---114 115 116class ModelInfo(BaseModel):117 id: str118 object: str = "model"119 created: int120 owned_by: str121 122 123class ModelList(BaseModel):124 object: str = "list"125 data: List[ModelInfo]126 127 128class ChatCompletionChoice(BaseModel):129 message: ChatMessage130 index: int = 0131 finish_reason: str = "stop"132 133 134class ChatCompletionResponse(BaseModel):135 id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex}")136 object: str = "chat.completion"137 created: int = Field(default_factory=lambda: int(time.time()))138 model: str139 choices: List[ChatCompletionChoice]140 usage: Dict[str, int] = Field(141 default_factory=lambda: {142 "prompt_tokens": 0,143 "completion_tokens": 0,144 "total_tokens": 0,145 }146 )147 148 149class StreamChoice(BaseModel):150 delta: Dict[str, Any] = Field(default_factory=dict)151 index: int = 0152 finish_reason: Optional[str] = None153 154 155class StreamResponse(BaseModel):156 id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex}")157 object: str = "chat.completion.chunk"158 created: int = Field(default_factory=lambda: int(time.time()))159 model: str160 choices: List[StreamChoice]161 162 163# FastAPI App164app = FastAPI(title="JetBrains AI OpenAI Compatible API")165 166# Add CORS middleware167app.add_middleware(168 CORSMiddleware,169 allow_origins=["*"], # Allows all origins170 allow_credentials=True,171 allow_methods=["*"], # Allows all methods172 allow_headers=["*"], # Allows all headers173)174 175security = HTTPBearer(auto_error=False)176 177 178# Helper functions179def load_models():180 """加载模型配置和映射规则"""181 global anthropic_model_mappings182 try:183 with open("models.json", "r", encoding="utf-8") as f:184 config = json.load(f)185 186 # 支持新格式（包含 models 和 anthropic_model_mappings）187 if isinstance(config, dict):188 if "models" in config:189 model_ids = config["models"]190 # 加载模型映射配置191 anthropic_model_mappings = config.get("anthropic_model_mappings", {})192 print(f"从 models.json 加载了 {len(anthropic_model_mappings)} 个模型映射规则")193 else:194 # 处理旧格式的字典（如果有其他字段但没有 models）195 model_ids = []196 anthropic_model_mappings = {}197 print("警告: models.json 使用非标准格式，没有找到 models 字段")198 # 支持旧格式（仅包含模型列表）199 elif isinstance(config, list):200 model_ids = config201 anthropic_model_mappings = {}202 print("警告: models.json 使用旧格式，没有找到模型映射配置")203 else:204 print("错误: models.json 格式不正确")205 return {"data": []}206 207 processed_models = []208 if isinstance(model_ids, list):209 for model_id in model_ids:210 if isinstance(model_id, str):211 processed_models.append(212 {213 "id": model_id,214 "object": "model",215 "created": int(time.time()),216 "owned_by": "jetbrains-ai",217 }218 )219 220 return {"data": processed_models}221 except Exception as e:222 print(f"加载 models.json 时出错: {e}")223 # 设置默认映射规则224 anthropic_model_mappings = {}225 return {"data": []}226 227 228async def _save_accounts_to_file():229 """将当前账户状态异步保存到文件"""230 async with file_write_lock:231 try:232 async with aiofiles.open("jetbrainsai.json", "w", encoding="utf-8") as f:233 await f.write(json.dumps(JETBRAINS_ACCOUNTS, indent=2))234 except Exception as e:235 print(f"保存 jetbrainsai.json 文件时出错: {e}")236 237 238def load_client_api_keys():239 """加载客户端 API 密钥（优先从环境变量读取）"""240 global VALID_CLIENT_KEYS241 242 # 首先尝试从环境变量读取243 client_keys_env = os.getenv("CLIENT_API_KEYS")244 if client_keys_env:245 keys = [key.strip() for key in client_keys_env.split(",") if key.strip()]246 if keys:247 VALID_CLIENT_KEYS = set(keys)248 print(f"从环境变量成功加载 {len(VALID_CLIENT_KEYS)} 个客户端 API 密钥")249 return250 251 # 回退到从文件读取252 try:253 with open("client_api_keys.json", "r", encoding="utf-8") as f:254 keys = json.load(f)255 if not isinstance(keys, list):256 print("警告: client_api_keys.json 应包含密钥列表")257 VALID_CLIENT_KEYS = set()258 return259 VALID_CLIENT_KEYS = set(keys)260 if not VALID_CLIENT_KEYS:261 print("警告: client_api_keys.json 为空")262 else:263 print(f"从文件成功加载 {len(VALID_CLIENT_KEYS)} 个客户端 API 密钥")264 except FileNotFoundError:265 print("错误: 未找到 client_api_keys.json 且未设置 CLIENT_API_KEYS 环境变量")266 VALID_CLIENT_KEYS = set()267 except Exception as e:268 print(f"加载 client_api_keys.json 时出错: {e}")269 VALID_CLIENT_KEYS = set()270 271 272def load_jetbrains_accounts():273 """加载 JetBrains AI 认证信息（优先从环境变量读取）"""274 global JETBRAINS_ACCOUNTS275 processed_accounts = []276 277 # 首先尝试从环境变量读取（支持JWT_TOKEN, LICENSE_IDS, AUTHORIZATION_TOKENS）278 jwt_tokens_env = os.getenv("JWT_TOKEN")279 license_ids_env = os.getenv("LICENSE_IDS")280 authorization_tokens_env = os.getenv("AUTHORIZATION_TOKENS")281 282 if jwt_tokens_env:283 # 解析逗号分隔的值284 jwt_tokens = [token.strip() for token in jwt_tokens_env.split(",") if token.strip()]285 license_ids = []286 authorization_tokens = []287 288 # 解析LICENSE_IDS（如果提供）289 if license_ids_env:290 license_ids = [lid.strip() for lid in license_ids_env.split(",") if lid.strip()]291 292 # 解析AUTHORIZATION_TOKENS（如果提供）293 if authorization_tokens_env:294 authorization_tokens = [auth.strip() for auth in authorization_tokens_env.split(",") if auth.strip()]295 296 # 按位置配对创建账户297 max_len = max(len(jwt_tokens), len(license_ids), len(authorization_tokens))298 for i in range(max_len):299 account = {300 "jwt": jwt_tokens[i] if i < len(jwt_tokens) else None,301 "licenseId": license_ids[i] if i < len(license_ids) else None,302 "authorization": authorization_tokens[i] if i < len(authorization_tokens) else None,303 "last_updated": int(time.time()),304 "has_quota": True,305 "last_quota_check": 0,306 }307 # 至少需要有jwt或者同时有licenseId和authorization308 if account["jwt"] or (account["licenseId"] and account["authorization"]):309 processed_accounts.append(account)310 311 if processed_accounts:312 JETBRAINS_ACCOUNTS = processed_accounts313 print(f"从环境变量成功加载 {len(JETBRAINS_ACCOUNTS)} 个 JetBrains AI 账户")314 # 打印账户配置概况315 for i, account in enumerate(JETBRAINS_ACCOUNTS):316 jwt_status = "✓" if account.get("jwt") else "✗"317 license_status = "✓" if account.get("licenseId") else "✗"318 auth_status = "✓" if account.get("authorization") else "✗"319 print(f" 账户 {i+1}: JWT({jwt_status}) LicenseId({license_status}) Authorization({auth_status})")320 return321 322 # 其次尝试从分别的环境变量读取（JWT_TOKEN_1, JWT_TOKEN_2, ...）323 for i in range(1, 11): # 支持最多10个token324 jwt_token = os.getenv(f"JWT_TOKEN_{i}")325 license_id = os.getenv(f"LICENSE_ID_{i}")326 authorization = os.getenv(f"AUTHORIZATION_TOKEN_{i}")327 328 if jwt_token or (license_id and authorization):329 processed_accounts.append({330 "licenseId": license_id.strip() if license_id else None,331 "authorization": authorization.strip() if authorization else None,332 "jwt": jwt_token.strip() if jwt_token else None,333 "last_updated": int(time.time()),334 "has_quota": True,335 "last_quota_check": 0,336 })337 338 if processed_accounts:339 JETBRAINS_ACCOUNTS = processed_accounts340 print(f"从分别的环境变量成功加载 {len(JETBRAINS_ACCOUNTS)} 个 JetBrains AI 账户")341 return342 343 # 最后回退到从文件读取344 try:345 with open("jetbrainsai.json", "r", encoding="utf-8") as f:346 accounts_data = json.load(f)347 348 if not isinstance(accounts_data, list):349 print("警告: jetbrainsai.json 格式不正确，应为对象列表")350 JETBRAINS_ACCOUNTS = []351 return352 353 for account in accounts_data:354 processed_accounts.append(355 {356 "licenseId": account.get("licenseId"),357 "authorization": account.get("authorization"),358 "jwt": account.get("jwt"),359 "last_updated": account.get("last_updated", 0),360 "has_quota": account.get("has_quota", True),361 "last_quota_check": account.get("last_quota_check", 0),362 }363 )364 365 JETBRAINS_ACCOUNTS = processed_accounts366 if not JETBRAINS_ACCOUNTS:367 print("警告: jetbrainsai.json 中未找到有效的认证信息")368 else:369 print(f"从文件成功加载 {len(JETBRAINS_ACCOUNTS)} 个 JetBrains AI 账户")370 371 except FileNotFoundError:372 print("错误: 未找到 jetbrainsai.json 文件且未设置 JWT_TOKEN 环境变量")373 JETBRAINS_ACCOUNTS = []374 except Exception as e:375 print(f"加载 jetbrainsai.json 时出错: {e}")376 JETBRAINS_ACCOUNTS = []377 378 379def get_model_item(model_id: str) -> Optional[Dict]:380 """根据模型ID获取模型配置"""381 for model in models_data.get("data", []):382 if model.get("id") == model_id:383 return model384 return None385 386 387async def authenticate_client(388 auth: Optional[HTTPAuthorizationCredentials] = Depends(security),389):390 """客户端认证 (OpenAI-style)"""391 if not VALID_CLIENT_KEYS:392 raise HTTPException(status_code=503, detail="服务不可用: 未配置客户端 API 密钥")393 394 if not auth or not auth.credentials:395 raise HTTPException(396 status_code=401,397 detail="需要在 Authorization header 中提供 API 密钥",398 headers={"WWW-Authenticate": "Bearer"},399 )400 401 if auth.credentials not in VALID_CLIENT_KEYS:402 raise HTTPException(status_code=403, detail="无效的客户端 API 密钥")403 404 405async def authenticate_any_client(406 auth: Optional[HTTPAuthorizationCredentials] = Depends(security),407 api_key: Optional[str] = Header(None, alias="x-api-key"),408):409 """客户端认证 (支持 OpenAI 和 Anthropic 风格)"""410 if not VALID_CLIENT_KEYS:411 raise HTTPException(status_code=503, detail="服务不可用: 未配置客户端 API 密钥")412 413 # 优先检查 x-api-key414 if api_key:415 if api_key in VALID_CLIENT_KEYS:416 return417 else:418 raise HTTPException(419 status_code=403, detail="无效的客户端 API 密钥 (x-api-key)"420 )421 422 # 其次检查 Authorization header423 if auth and auth.credentials:424 if auth.credentials in VALID_CLIENT_KEYS:425 return426 else:427 raise HTTPException(428 status_code=403, detail="无效的客户端 API 密钥 (Bearer token)"429 )430 431 # 如果两者都未提供432 raise HTTPException(433 status_code=401,434 detail="需要在 Authorization header (Bearer) 或 x-api-key header 中提供 API 密钥",435 headers={"WWW-Authenticate": "Bearer"},436 )437 438 439async def authenticate_anthropic_client(440 api_key: Optional[str] = Header(None, alias="x-api-key"),441):442 """客户端认证 (Anthropic-style)"""443 if not VALID_CLIENT_KEYS:444 raise HTTPException(status_code=503, detail="服务不可用: 未配置客户端 API 密钥")445 446 if not api_key:447 raise HTTPException(448 status_code=401,449 detail="需要在 x-api-key header 中提供 API 密钥",450 )451 452 if api_key not in VALID_CLIENT_KEYS:453 raise HTTPException(status_code=403, detail="无效的客户端 API 密钥")454 455 456async def _check_quota(account: dict):457 """检查指定账户的配额"""458 if not http_client:459 raise HTTPException(status_code=500, detail="HTTP 客户端未初始化")460 461 # 对于基于许可证的账户，如果 JWT 不存在，则先刷新462 if not account.get("jwt") and account.get("licenseId"):463 await _refresh_jetbrains_jwt(account)464 465 if not account.get("jwt"):466 account["has_quota"] = False467 return468 469 try:470 headers = {471 "User-Agent": "ktor-client",472 "Content-Length": "0",473 "Accept-Charset": "UTF-8",474 "grazie-agent": '{"name":"aia:pycharm","version":"251.26094.80.13:251.26094.141"}',475 "grazie-authenticate-jwt": account["jwt"],476 }477 response = await http_client.post(478 "https://api.jetbrains.ai/user/v5/quota/get", headers=headers, timeout=10.0479 )480 481 if response.status_code == 401 and account.get("licenseId"):482 print(f"JWT for {account['licenseId']} expired, refreshing...")483 await _refresh_jetbrains_jwt(account)484 headers["grazie-authenticate-jwt"] = account["jwt"]485 response = await http_client.post(486 "https://api.jetbrains.ai/user/v5/quota/get", headers=headers, timeout=10.0487 )488 489 response.raise_for_status()490 quota_data = response.json()491 492 has_quota = quota_data.get("dailyUsed", 0) < quota_data.get("dailyTotal", 1)493 account["has_quota"] = has_quota494 if not has_quota:495 print(f"Account {account.get('licenseId') or 'with static JWT'} has no quota.")496 497 except Exception as e:498 print(f"Error checking quota for account: {e}")499 # On error, assume it has no quota to be safe500 account["has_quota"] = False501 finally:502 account["last_quota_check"] = time.time()503 await _save_accounts_to_file()504 505 506async def _refresh_jetbrains_jwt(account: dict):507 """使用 licenseId 和 authorization 刷新 JWT"""508 if not http_client:509 raise HTTPException(status_code=500, detail="HTTP 客户端未初始化")510 511 print(f"正在为 licenseId {account['licenseId']} 刷新 JWT...")512 try:513 headers = {514 "User-Agent": "ktor-client",515 "Content-Type": "application/json",516 "Accept-Charset": "UTF-8",517 "authorization": f"Bearer {account['authorization']}",518 }519 payload = {"licenseId": account["licenseId"]}520 521 response = await http_client.post(522 "https://api.jetbrains.ai/auth/jetbrains-jwt/provide-access/license/v2",523 json=payload,524 headers=headers,525 timeout=DEFAULT_REQUEST_TIMEOUT,526 )527 response.raise_for_status()528 529 data = response.json()530 if data.get("state") == "PAID" and "token" in data:531 account["jwt"] = data["token"]532 account["last_updated"] = time.time()533 print(f"成功刷新 licenseId {account['licenseId']} 的 JWT")534 await _save_accounts_to_file()535 else:536 print(f"刷新 JWT 失败: 无效的响应状态 {data.get('state')}")537 raise HTTPException(status_code=500, detail=f"刷新 JWT 失败: {data}")538 539 except httpx.HTTPStatusError as e:540 print(f"刷新 JWT 时 HTTP 错误: {e.response.status_code} {e.response.text}")541 raise HTTPException(542 status_code=e.response.status_code,543 detail=f"刷新 JWT 失败: {e.response.text}",544 )545 except Exception as e:546 print(f"刷新 JWT 时发生未知错误: {e}")547 raise HTTPException(status_code=500, detail=f"刷新 JWT 时发生未知错误: {e}")548 549 550async def get_next_jetbrains_account() -> dict:551 """轮询获取下一个有配额的 JetBrains 账户"""552 global current_account_index553 554 if not JETBRAINS_ACCOUNTS:555 raise HTTPException(status_code=503, detail="服务不可用: 未配置 JetBrains 账户")556 557 async with account_rotation_lock:558 start_index = current_account_index559 for _ in range(len(JETBRAINS_ACCOUNTS)):560 account = JETBRAINS_ACCOUNTS[current_account_index]561 current_account_index = (current_account_index + 1) % len(562 JETBRAINS_ACCOUNTS563 )564 565 # 如果状态是 stale，检查配额566 is_quota_stale = (567 time.time() - account.get("last_quota_check", 0) > 3600568 ) # 1 hour cache569 if account.get("has_quota") and is_quota_stale:570 await _check_quota(account)571 572 if account.get("has_quota"):573 # 如果是基于许可证的账户，检查 JWT 是否需要刷新574 if account.get("licenseId"):575 is_jwt_stale = (576 time.time() - account.get("last_updated", 0) > 12 * 3600577 )578 if not account.get("jwt") or is_jwt_stale:579 await _refresh_jetbrains_jwt(account)580 # 刷新 JWT 后可能需要重新检查配额581 if not account.get("has_quota"):582 await _check_quota(account)583 if not account.get("has_quota"):584 continue585 586 if account.get("jwt"):587 return account588 589 # 循环完成，没有找到可用的账户590 raise HTTPException(status_code=429, detail="所有 JetBrains 账户均已超出配额或无效")591 592 593# FastAPI 生命周期事件594@app.on_event("startup")595async def startup():596 global models_data, http_client597 models_data = load_models()598 load_client_api_keys()599 load_jetbrains_accounts()600 http_client = httpx.AsyncClient(timeout=None)601 print("JetBrains AI OpenAI Compatible API 服务器已启动")602 603 604@app.on_event("shutdown")605async def shutdown():606 global http_client607 if http_client:608 await http_client.aclose()609 610 611# API 端点612@app.get("/v1/models", response_model=ModelList)613async def list_models(_: None = Depends(authenticate_any_client)):614 """列出可用模型"""615 model_list = [616 ModelInfo(617 id=model.get("id", ""),618 created=model.get("created", int(time.time())),619 owned_by=model.get("owned_by", "jetbrains-ai"),620 )621 for model in models_data.get("data", [])622 ]623 return ModelList(data=model_list)624 625 626async def openai_stream_adapter(627 api_stream_generator: AsyncGenerator[str, None],628 model_name: str,629 tools: Optional[List[Dict[str, Any]]],630) -> AsyncGenerator[str, None]:631 """将 JetBrains API 的流转换为 OpenAI 格式的 SSE"""632 stream_id = f"chatcmpl-{uuid.uuid4().hex}"633 first_chunk_sent = False634 tool_id = 0635 636 try:637 async for line in api_stream_generator:638 if not line or line == "data: end":639 continue640 641 if line.startswith("data: "):642 try:643 data = json.loads(line[6:])644 event_type = data.get("type")645 646 if event_type == "Content":647 content = data.get("content", "")648 if not content:649 continue650 651 delta_payload = {}652 if not first_chunk_sent:653 delta_payload = {"role": "assistant", "content": content}654 first_chunk_sent = True655 else:656 delta_payload = {"content": content}657 658 stream_resp = StreamResponse(659 id=stream_id,660 model=model_name,661 choices=[StreamChoice(delta=delta_payload)],662 )663 yield f"data: {stream_resp.json()}\n\n"664 665 elif event_type == "FunctionCall":666 func_name = data.get("name", None)667 func_argu = data.get("content", None)668 if func_name and tools:669 for tool_id, tool in enumerate(tools):670 if tool["name"] == func_name:671 break672 673 delta_payload = {674 "tool_calls": [675 {676 "index": tool_id,677 "id": f"call_{uuid.uuid4().hex}",678 "function": {679 "arguments": func_argu,680 "name": func_name,681 },682 "type": "function" if func_name else None,683 }684 ]685 }686 stream_resp = StreamResponse(687 id=stream_id,688 model=model_name,689 choices=[StreamChoice(delta=delta_payload)],690 )691 yield f"data: {stream_resp.json()}\n\n"692 693 elif event_type == "FinishMetadata":694 final_resp = StreamResponse(695 id=stream_id,696 model=model_name,697 choices=[StreamChoice(delta={}, finish_reason="stop")],698 )699 yield f"data: {final_resp.json()}\n\n"700 break701 except json.JSONDecodeError:702 print(f"警告: 无法解析的 JSON 行: {line}")703 continue704 705 yield "data: [DONE]\n\n"706 707 except Exception as e:708 print(f"流式适配器错误: {e}")709 error_resp = StreamResponse(710 id=stream_id,711 model=model_name,712 choices=[713 StreamChoice(714 delta={"role": "assistant", "content": f"内部错误: {str(e)}"},715 index=0,716 finish_reason="stop",717 )718 ],719 )720 yield f"data: {error_resp.json()}\n\n"721 yield "data: [DONE]\n\n"722 723 724async def aggregate_stream_for_non_stream_response(725 openai_sse_stream: AsyncGenerator[str, None], model_name: str726) -> ChatCompletionResponse:727 """聚合流式响应为完整响应"""728 content_parts = []729 tool_calls_map = {}730 final_finish_reason = "stop"731 732 async for sse_line in openai_sse_stream:733 if sse_line.startswith("data: ") and sse_line.strip() != "data: [DONE]":734 try:735 data = json.loads(sse_line[6:].strip())736 if not data.get("choices"):737 continue738 739 choice = data["choices"][0]740 delta = choice.get("delta", {})741 742 if choice.get("finish_reason"):743 final_finish_reason = choice.get("finish_reason")744 745 if delta.get("content"):746 content_parts.append(delta["content"])747 748 if "tool_calls" in delta:749 for tc_chunk in delta["tool_calls"]:750 idx = tc_chunk["index"]751 if idx not in tool_calls_map:752 tool_calls_map[idx] = {753 "type": "function",754 "function": {"name": "", "arguments": ""},755 }756 757 if tc_chunk.get("id"):758 tool_calls_map[idx]["id"] = tc_chunk["id"]759 760 func_chunk = tc_chunk.get("function", {})761 if func_chunk.get("name"):762 tool_calls_map[idx]["function"]["name"] = func_chunk["name"]763 if func_chunk.get("arguments"):764 tool_calls_map[idx]["function"]["arguments"] += func_chunk[765 "arguments"766 ]767 except json.JSONDecodeError:768 print(f"警告: 聚合时无法解析的 JSON 行: {sse_line}")769 770 final_tool_calls = []771 for k, v in sorted(tool_calls_map.items()):772 if "id" not in v:773 v["id"] = f"call_{uuid.uuid4().hex}"774 final_tool_calls.append(v)775 776 full_content = "".join(content_parts) or None777 778 if final_tool_calls:779 message = ChatMessage(780 role="assistant", content=full_content, tool_calls=final_tool_calls781 )782 final_finish_reason = "tool_calls"783 else:784 message = ChatMessage(role="assistant", content=full_content)785 786 return ChatCompletionResponse(787 model=model_name,788 choices=[789 ChatCompletionChoice(790 message=message,791 finish_reason=final_finish_reason,792 )793 ],794 )795 796 797def extract_text_content(content: Optional[Union[str, List[Dict[str, Any]]]]) -> str:798 """从消息内容中提取文本内容"""799 if isinstance(content, str):800 return content801 elif isinstance(content, list):802 # 处理多模态消息格式，提取所有文本内容803 text_parts = []804 for item in content:805 if isinstance(item, dict) and item.get("type") == "text":806 text_parts.append(item.get("text", ""))807 return " ".join(text_parts)808 return ""809 810 811@app.post("/v1/chat/completions")812async def chat_completions(813 request: ChatCompletionRequest, _: None = Depends(authenticate_client)814):815 """创建聊天完成"""816 model_config = get_model_item(request.model)817 if not model_config:818 raise HTTPException(status_code=404, detail=f"模型 {request.model} 未找到")819 820 account = await get_next_jetbrains_account()821 auth_token = account["jwt"]822 823 # 从历史消息中创建 tool_call_id 到 function_name 的映射824 tool_id_to_func_name_map = {}825 for m in request.messages:826 if m.role == "assistant" and m.tool_calls:827 for tc in m.tool_calls:828 if tc.get("id") and tc.get("function", {}).get("name"):829 tool_id_to_func_name_map[tc["id"]] = tc["function"]["name"]830 831 # 将 OpenAI 格式的消息转换为 JetBrains 格式832 jetbrains_messages = []833 for msg in request.messages:834 # 提取文本内容，处理多模态消息格式835 text_content = extract_text_content(msg.content)836 837 if msg.role in ["user", "system"]:838 jetbrains_messages.append(839 {"type": f"{msg.role}_message", "content": text_content}840 )841 842 elif msg.role == "assistant":843 if msg.tool_calls:844 # 只处理第一个工具调用，以匹配 JetBrains API 的限制845 first_tool_call = msg.tool_calls[0]846 tool_id_to_func_name_map[first_tool_call["id"]] = first_tool_call[847 "function"848 ]["name"]849 jetbrains_messages.append(850 {851 "type": "assistant_message",852 "content": text_content,853 "functionCall": {854 "functionName": first_tool_call["function"]["name"],855 "content": first_tool_call["function"]["arguments"],856 },857 }858 )859 else:860 jetbrains_messages.append(861 {"type": "assistant_message", "content": text_content}862 )863 864 elif msg.role == "tool":865 function_name = tool_id_to_func_name_map.get(msg.tool_call_id)866 if function_name:867 jetbrains_messages.append(868 {869 "type": "function_message",870 "content": text_content,871 "functionName": function_name,872 }873 )874 else:875 print(876 f"警告: 无法为 tool_call_id {msg.tool_call_id} 找到对应的函数调用"877 )878 else:879 jetbrains_messages.append({"type": "user_message", "content": text_content})880 881 data = []882 tools = None883 if request.tools:884 data.append({"type": "json", "fqdn": "llm.parameters.functions"})885 tools = [t["function"] for t in request.tools]886 data.append({"type": "json", "value": json.dumps(tools)})887 888 # 创建 API 请求的 payload889 payload = {890 "prompt": "ij.chat.request.new-chat-on-start",891 "profile": request.model,892 "chat": {"messages": jetbrains_messages},893 "parameters": {"data": data},894 }895 896 headers = {897 "User-Agent": "ktor-client",898 "Accept": "text/event-stream",899 "Content-Type": "application/json",900 "Accept-Charset": "UTF-8",901 "Cache-Control": "no-cache",902 "grazie-agent": '{"name":"aia:pycharm","version":"251.26094.80.13:251.26094.141"}',903 "grazie-authenticate-jwt": auth_token,904 }905 906 async def api_stream_generator():907 """一个包装 httpx 请求的异步生成器"""908 try:909 async with http_client.stream(910 "POST",911 "https://api.jetbrains.ai/user/v5/llm/chat/stream/v8",912 json=payload,913 headers=headers,914 ) as response:915 if response.status_code == 477:916 print(f"Account {account.get('licenseId') or 'with static JWT'} has no quota (received 477).")917 account["has_quota"] = False918 account["last_quota_check"] = time.time()919 await _save_accounts_to_file()920 response.raise_for_status()921 async for line in response.aiter_lines():922 yield line923 except httpx.HTTPStatusError as e:924 if e.response.status_code == 477:925 print(f"Account {account.get('licenseId') or 'with static JWT'} has no quota (received 477).")926 account["has_quota"] = False927 account["last_quota_check"] = time.time()928 await _save_accounts_to_file()929 raise e930 931 # 创建 OpenAI 格式的流932 openai_sse_stream = openai_stream_adapter(933 api_stream_generator(), request.model, tools or []934 )935 936 # 返回流式或非流式响应937 if request.stream:938 return StreamingResponse(openai_sse_stream, media_type="text/event-stream")939 else:940 return await aggregate_stream_for_non_stream_response(941 openai_sse_stream, request.model942 )943 944 945def convert_anthropic_to_openai(946 anthropic_req: AnthropicMessageRequest,947) -> ChatCompletionRequest:948 openai_messages = []949 tool_id_to_func_name_map = {}950 951 if anthropic_req.system:952 system_prompt = ""953 if isinstance(anthropic_req.system, str):954 system_prompt = anthropic_req.system955 elif isinstance(anthropic_req.system, list):956 system_prompt = " ".join(957 [958 item.get("text", "")959 for item in anthropic_req.system960 if isinstance(item, dict) and item.get("type") == "text"961 ]962 )963 if system_prompt:964 openai_messages.append(ChatMessage(role="system", content=system_prompt))965 966 for msg in anthropic_req.messages:967 if msg.role == "user":968 text_parts = []969 if isinstance(msg.content, str):970 text_parts.append(msg.content)971 else:972 for block in msg.content:973 if block.type == "text":974 text_parts.append(block.text)975 elif block.type == "tool_result" and block.tool_use_id:976 content_str = (977 block.content978 if isinstance(block.content, str)979 else json.dumps(block.content)980 )981 openai_messages.append(982 ChatMessage(983 role="tool",984 tool_call_id=block.tool_use_id,985 content=content_str,986 )987 )988 989 if text_parts:990 openai_messages.append(991 ChatMessage(role="user", content=" ".join(text_parts))992 )993 994 elif msg.role == "assistant":995 text_parts = []996 tool_calls = []997 if isinstance(msg.content, list):998 for block in msg.content:999 if block.type == "text":1000 text_parts.append(block.text)1001 elif block.type == "tool_use" and block.id and block.name:1002 arguments = (1003 json.dumps(block.input) if block.input is not None else "{}"1004 )1005 tool_calls.append(1006 {1007 "id": block.id,1008 "type": "function",1009 "function": {1010 "name": block.name,1011 "arguments": arguments,1012 },1013 }1014 )1015 tool_id_to_func_name_map[block.id] = block.name1016 1017 content_text = " ".join(text_parts) if text_parts else None1018 openai_messages.append(1019 ChatMessage(1020 role="assistant",1021 content=content_text,1022 tool_calls=tool_calls if tool_calls else None,1023 )1024 )1025 1026 openai_tools = None1027 if anthropic_req.tools:1028 openai_tools = [1029 {1030 "type": "function",1031 "function": {1032 "name": t.name,1033 "description": t.description,1034 "parameters": t.input_schema,1035 },1036 }1037 for t in anthropic_req.tools1038 ]1039 1040 return ChatCompletionRequest(1041 model=anthropic_req.model,1042 messages=openai_messages,1043 stream=anthropic_req.stream,1044 temperature=anthropic_req.temperature,1045 max_tokens=anthropic_req.max_tokens,1046 top_p=anthropic_req.top_p,1047 tools=openai_tools,1048 stop=anthropic_req.stop_sequences,1049 )1050 1051 1052def map_finish_reason(finish_reason: Optional[str]) -> Optional[str]:1053 if finish_reason == "stop":1054 return "end_turn"1055 if finish_reason == "length":1056 return "max_tokens"1057 if finish_reason == "tool_calls":1058 return "tool_use"1059 return finish_reason1060 1061 1062async def openai_to_anthropic_stream_adapter(1063 openai_stream: AsyncGenerator[str, None], model_name: str1064) -> AsyncGenerator[str, None]:1065 message_id = f"msg_{uuid.uuid4().hex.replace('-', '')}"1066 yield f"event: message_start\ndata: {json.dumps({'type': 'message_start', 'message': {'id': message_id, 'type': 'message', 'role': 'assistant', 'model': model_name, 'content': [], 'stop_reason': None, 'stop_sequence': None, 'usage': {'input_tokens': 0, 'output_tokens': 0}}})}\n\n"1067 yield f"event: ping\ndata: {json.dumps({'type': 'ping'})}\n\n"1068 1069 content_block_index = 01070 text_block_open = False1071 tool_blocks = {} # index -> {id, name, args}1072 1073 async for sse_line in openai_stream:1074 if not sse_line.startswith("data:") or sse_line.strip() == "data: [DONE]":1075 continue1076 1077 data_str = sse_line[6:].strip()1078 try:1079 data = json.loads(data_str)1080 if not data.get("choices"):1081 continue1082 1083 delta = data["choices"][0].get("delta", {})1084 finish_reason = data["choices"][0].get("finish_reason")1085 1086 if delta.get("content"):1087 if not text_block_open:1088 yield f"event: content_block_start\ndata: {json.dumps({'type': 'content_block_start', 'index': content_block_index, 'content_block': {'type': 'text', 'text': ''}})}\n\n"1089 text_block_open = True1090 1091 yield f"event: content_block_delta\ndata: {json.dumps({'type': 'content_block_delta', 'index': content_block_index, 'delta': {'type': 'text_delta', 'text': delta['content']}})}\n\n"1092 1093 if delta.get("tool_calls"):1094 if text_block_open:1095 yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': content_block_index})}\n\n"1096 text_block_open = False1097 content_block_index += 11098 1099 for tc in delta["tool_calls"]:1100 idx = tc["index"]1101 if idx not in tool_blocks:1102 tool_blocks[idx] = {1103 "id": tc.get("id"),1104 "name": tc.get("function", {}).get("name"),1105 "args": "",1106 }1107 yield f"event: content_block_start\ndata: {json.dumps({'type': 'content_block_start', 'index': content_block_index + idx, 'content_block': {'type': 'tool_use', 'id': tc.get('id'), 'name': tc.get('function', {}).get('name'), 'input': {}}})}\n\n"1108 1109 if tc.get("function", {}).get("arguments"):1110 args_delta = tc["function"]["arguments"]1111 tool_blocks[idx]["args"] += args_delta1112 yield f"event: content_block_delta\ndata: {json.dumps({'type': 'content_block_delta', 'index': content_block_index + idx, 'delta': {'type': 'input_json_delta', 'partial_json': args_delta}})}\n\n"1113 1114 if finish_reason:1115 if text_block_open:1116 yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': content_block_index})}\n\n"1117 1118 for i in range(len(tool_blocks)):1119 yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': content_block_index + i})}\n\n"1120 1121 message_delta_data = {1122 "type": "message_delta",1123 "delta": {1124 "stop_reason": map_finish_reason(finish_reason),1125 "stop_sequence": None,1126 },1127 "usage": {"input_tokens": 0, "output_tokens": 0},1128 }1129 yield f"event: message_delta\ndata: {json.dumps(message_delta_data)}\n\n"1130 break1131 except json.JSONDecodeError:1132 print(f"Anthropic adapter JSON decode error: {data_str}")1133 continue1134 1135 yield f"event: message_stop\ndata: {json.dumps({'type': 'message_stop'})}\n\n"1136 1137 1138def convert_openai_to_anthropic_response(1139 resp: ChatCompletionResponse,1140) -> AnthropicResponseMessage:1141 message = resp.choices[0].message1142 content_blocks = []1143 1144 if message.content:1145 content_blocks.append(1146 AnthropicResponseContent(type="text", text=message.content)1147 )1148 1149 if message.tool_calls:1150 for tc in message.tool_calls:1151 try:1152 tool_input = json.loads(tc["function"]["arguments"])1153 except json.JSONDecodeError:1154 tool_input = {1155 "error": "invalid JSON in arguments",1156 "arguments": tc["function"]["arguments"],1157 }1158 content_blocks.append(1159 AnthropicResponseContent(1160 type="tool_use",1161 id=tc["id"],1162 name=tc["function"]["name"],1163 input=tool_input,1164 )1165 )1166 1167 return AnthropicResponseMessage(1168 id=resp.id.replace("chatcmpl-", "msg_"),1169 model=resp.model,1170 content=content_blocks,1171 stop_reason=map_finish_reason(resp.choices[0].finish_reason),1172 usage=AnthropicUsage(1173 input_tokens=resp.usage.get("prompt_tokens", 0),1174 output_tokens=resp.usage.get("completion_tokens", 0),1175 ),1176 )1177 1178 1179@app.post("/v1/messages", response_model=None)1180async def messages_completions(1181 request: AnthropicMessageRequest, _: None = Depends(authenticate_anthropic_client)1182):1183 """创建符合 Anthropic 规范的聊天完成"""1184 openai_request = convert_anthropic_to_openai(request)1185 1186 # Apply model mapping specifically for /v1/messages endpoint using config from models.json1187 if openai_request.model in anthropic_model_mappings:1188 original_model = openai_request.model1189 openai_request.model = anthropic_model_mappings[openai_request.model]1190 print(f"Model mapping applied: {original_model} -> {openai_request.model}")1191 1192 model_config = get_model_item(openai_request.model)1193 if not model_config:1194 raise HTTPException(1195 status_code=404, detail=f"模型 {openai_request.model} 未找到"1196 )1197 1198 account = await get_next_jetbrains_account()1199 auth_token = account["jwt"]1200