#python学霸from flask import Flask, request, jsonify, Responseimport osimport requestsimport jsonfrom typing import Dict, Anyapp = Flask(__name__)NIM_MODEL = os.getenv("ANTHROPIC_MODEL")NIM_API_KEY = os.getenv("ANTHROPIC_AUTH_TOKEN")NIM_BASE_URL = "https://integrate.api.nvidia.com/v1"PROXY_PORT = 5000def convert_claude_to_nim(claude_request: Dict[str, Any]) -> Dict[str, Any]: nim_request = { "model": NIM_MODEL, "messages": claude_request.get("messages", []), "temperature": claude_request.get("temperature", 1.0), "max_tokens": claude_request.get("max_tokens", claude_request.get("max_completion_tokens", 1024)), "top_p": claude_request.get("top_p", 1.0) } for k in ["frequency_penalty", "presence_penalty", "stop", "stream", "tools", "tool_choice"]: if k in claude_request: nim_request[k] = claude_request[k] return nim_requestdef convert_nim_to_claude(nim_response: Dict[str, Any]) -> Dict[str, Any]: claude_response = { "id": nim_response.get("id", "chatcmpl-unknown"), "object": "chat.completion", "created": nim_response.get("created", 0), "model": NIM_MODEL } if "choices" in nim_response: claude_response["choices"] = [] for choice in nim_response["choices"]: claude_choice = { "index": choice.get("index", 0), "message": { "role": choice.get("message", {}).get("role", "assistant"), "content": choice.get("message", {}).get("content", "") }, "finish_reason": choice.get("finish_reason", "stop") } if "tool_calls" in choice.get("message", {}): claude_choice["message"]["tool_calls"] = choice["message"]["tool_calls"] claude_response["choices"].append(claude_choice) if "usage" in nim_response: claude_response["usage"] = { "prompt_tokens": nim_response["usage"].get("prompt_tokens", 0), "completion_tokens": nim_response["usage"].get("completion_tokens", 0), "total_tokens": nim_response["usage"].get("total_tokens", 0) } return claude_response@app.route('/health', methods=['GET'])def health_check(): return jsonify({"status": "healthy", "port": PROXY_PORT})@app.route('/v1/models', methods=['GET'])def list_models(): try: headers = {"Authorization": f"Bearer {NIM_API_KEY}", "Content-Type": "application/json"} res = requests.get(f"{NIM_BASE_URL}/models", headers=headers) res.raise_for_status() return jsonify(res.json()) except Exception as e: return jsonify({"error": {"message": str(e)}}), 500@app.route('/v1/messages', methods=['POST'])def messages(): try: claude_req = request.get_json() if not claude_req: return jsonify({"error": {"message": "缺少请求体"}}), 400 nim_req = { "model": NIM_MODEL, "messages": claude_req.get("messages", []), "temperature": claude_req.get("temperature", 1.0), "max_tokens": claude_req.get("max_tokens", 1024), "stream": claude_req.get("stream", False) } headers = {"Authorization": f"Bearer {NIM_API_KEY}", "Content-Type": "application/json"} nim_url = f"{NIM_BASE_URL}/chat/completions" if nim_req["stream"]: return handle_stream(nim_url, headers, nim_req) res = requests.post(nim_url, json=nim_req, headers=headers) res.raise_for_status() nim_res = res.json() claude_res = { "id": nim_res.get("id", "msg-unknown"), "type": "message", "role": "assistant", "content": [{"type": "text", "text": nim_res.get("choices", [{}])[0].get("message", {}).get("content", "")}], "model": NIM_MODEL, "stop_reason": nim_res.get("choices", [{}])[0].get("finish_reason", "end_turn"), "usage": { "input_tokens": nim_res.get("usage", {}).get("prompt_tokens", 0), "output_tokens": nim_res.get("usage", {}).get("completion_tokens", 0) } } return jsonify(claude_res) except Exception as e: return jsonify({"error": {"message": str(e)}}), 500@app.route('/v1/chat/completions', methods=['POST'])def chat_completions(): try: claude_req = request.get_json() if not claude_req: return jsonify({"error": {"message": "缺少请求体"}}), 400 nim_req = convert_claude_to_nim(claude_req) headers = {"Authorization": f"Bearer {NIM_API_KEY}", "Content-Type": "application/json"} nim_url = f"{NIM_BASE_URL}/chat/completions" if nim_req.get("stream", False): return handle_stream(nim_url, headers, nim_req) res = requests.post(nim_url, json=nim_req, headers=headers) res.raise_for_status() return jsonify(convert_nim_to_claude(res.json())) except requests.exceptions.HTTPError as e: err = e.response.json() if e.response.headers.get("content-type", "").startswith("application/json") else {"message": str(e)} return jsonify({"error": err.get("error", err)}), e.response.status_code except Exception as e: return jsonify({"error": {"message": str(e)}}), 500def handle_stream(nim_url: str, headers: Dict[str, str], nim_req: Dict[str, Any]) -> Response: def generate(): try: res = requests.post(nim_url, json=nim_req, headers=headers, stream=True) res.raise_for_status() for line in res.iter_lines(): if line: yield line.decode("utf-8") + "\n" except Exception as e: yield f"data: {json.dumps({'error': {'message': str(e)}})}\n\n" return Response(generate(), mimetype="text/event-stream")@app.errorhandler(404)def not_found(_): return jsonify({"error": {"message": f"接口不存在: {request.path}"}}), 404@app.errorhandler(500)def server_error(_): return jsonify({"error": {"message": "服务器内部错误"}}), 500if __name__ == '__main__': app.run(host='0.0.0.0', port=PROXY_PORT, debug=False)