File size: 5,989 Bytes
48f8b4b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
@app.route('/hf/v1/chat/completions', methods=['POST'])
def chat_completions():
global current_api_key
is_authenticated, auth_error, status_code = authenticate_request(request)
if not is_authenticated:
return auth_error if auth_error else jsonify({'error': 'Unauthorized'}), status_code if status_code else 401
try:
request_data = request.get_json()
messages = request_data.get('messages', [])
model = request_data.get('model', 'gemini-exp-1206')
temperature = request_data.get('temperature', 1)
max_tokens = request_data.get('max_tokens', 8192)
stream = request_data.get('stream', False)
logger.info(colored(f"\n{model} [r] -> {current_api_key[:11]}...", 'yellow'))
# 将 OpenAI 格式的消息转换为 Gemini 格式
gemini_history = []
for message in messages:
role = message.get('role')
content = message.get('content')
if role == 'system':
gemini_history.append({"role": "user", "parts": [content]})
elif role == 'user':
gemini_history.append({"role": "user", "parts": [content]})
elif role == 'assistant':
gemini_history.append({"role": "model", "parts": [content]})
user_message = gemini_history[-1]['parts'][0] if gemini_history else messages[-1]['content']
gemini_history = gemini_history[:-1] if gemini_history else []
genai.configure(api_key=current_api_key)
# 调用 Gemini API
generation_config = {
"temperature": temperature,
"max_output_tokens": max_tokens
}
gen_model = genai.GenerativeModel(
model_name=model,
generation_config=generation_config,
safety_settings=safety_settings
)
if stream:
# 流式响应
if gemini_history:
chat_session = gen_model.start_chat(history=gemini_history)
response = chat_session.send_message(user_message, stream=True)
else:
response = gen_model.generate_content(user_message, stream=True)
def generate():
try:
for chunk in response:
if chunk.text:
data = {
'choices': [
{
'delta': {
'content': chunk.text
},
'finish_reason': None,
'index': 0
}
],
'object': 'chat.completion.chunk'
}
yield f"data: {json.dumps(data)}\n\n"
data = {
'choices': [
{
'delta': {},
'finish_reason': 'stop',
'index': 0
}
],
'object': 'chat.completion.chunk'
}
yield f"data: {json.dumps(data)}\n\n"
except Exception as e:
logger.error(f"Error during streaming: {str(e)}")
data = {
'error': {
'message': str(e),
'type': 'internal_server_error'
}
}
yield f"data: {json.dumps(data)}\n\n"
return Response(stream_with_context(generate()), mimetype='text/event-stream')
else:
# 调用 API
if gemini_history:
chat_session = gen_model.start_chat(history=gemini_history)
response = chat_session.send_message(user_message)
else:
response = gen_model.generate_content(user_message)
try:
text_content = response.candidates[0].content.parts[0].text
except (AttributeError, IndexError, TypeError) as e:
logger.error(colored(f"Error getting text content: {str(e)}",'red'))
text_content = "Error: Unable to get text content."
response_data = {
'id': 'chatcmpl-xxxxxxxxxxxx',
'object': 'chat.completion',
'created': int(datetime.now().timestamp()),
'model': model,
'choices': [{
'index': 0,
'message': {
'role': 'assistant',
'content': text_content
},
'finish_reason': 'stop'
}],
'usage':{
'prompt_tokens': 0,
'completion_tokens': 0,
'total_tokens': 0
}
}
logger.info(colored(f"Generation Success", 'green'))
return jsonify(response_data)
except Exception as e:
logger.error(f"Error in chat completions: {str(e)}")
return jsonify({
'error': {
'message': str(e),
'type': 'invalid_request_error'
}
}), 500
finally:
current_api_key = key_manager.get_available_key()
logger.info(colored(f"API KEY Switched -> {current_api_key[:11]}...", 'aqua')) |