File size: 3,361 Bytes
2a81db1 c4f87f9 2a81db1 c4f87f9 2a81db1 c4f87f9 2a81db1 dad78f4 2a81db1 |
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 |
import logging
from datetime import datetime, timedelta
from flask import request, Response, current_app as app
from app.config import IGNORED_MODEL_NAMES, IMAGE_MODEL_NAMES, AUTH_TOKEN, HISTORY_MSG_LIMIT
from app.config import configure_logging
from app.utils import send_chat_message, fetch_channel_id, map_model_name, process_content, get_user_contents, \
generate_hash, get_next_auth_token, handle_error, get_request_parameters
configure_logging()
storage_map = {}
@app.route("/hf/v1/chat/completions", methods=["GET", "POST", "OPTIONS"])
def onRequest():
try:
return fetch(request)
except Exception as e:
logging.error("An error occurred with chat : %s", e)
return handle_error(e)
@app.route('/hf/v1/models')
def list_models():
return {
"object": "list",
"data": [{
"id": m,
"object": "model",
"created": int(datetime.now().timestamp()),
"owned_by": "popai"
} for m in IGNORED_MODEL_NAMES]
}
@app.route('/hf/v1/images/generations', methods= ["post"])
def image():
try:
request.get_json()["model"] = IMAGE_MODEL_NAMES[0]
return fetch(request)
except Exception as e:
logging.error("An error occurred with image : %s", e)
return handle_error(e)
def get_channel_id(hash_value, token, model_name, content, template_id):
if hash_value in storage_map:
channel_id, expiry_time = storage_map[hash_value]
if expiry_time > datetime.now() and channel_id:
logging.info("Returning channel id from cache")
return channel_id
channel_id = fetch_channel_id(token, model_name, content, template_id)
expiry_time = datetime.now() + timedelta(days=1)
storage_map[hash_value] = (channel_id, expiry_time)
return channel_id
def fetch(req):
if req.method == "OPTIONS":
return handle_options_request()
token = req.headers.get("Authorization").replace("Bearer ", "")
messages, model_name, prompt, user_stream = get_request_parameters(req.get_json())
model_to_use = map_model_name(model_name)
template_id = 2000000 if model_name in IMAGE_MODEL_NAMES else ''
if not messages and prompt:
final_user_content = prompt
first_user_message = final_user_content
image_url = None
elif messages:
last_message = messages[-1]
first_user_message, end_user_message, concatenated_messages = get_user_contents(messages, HISTORY_MSG_LIMIT)
final_user_content, image_url = process_content(last_message.get('content'))
final_user_content = concatenated_messages + '\n' + final_user_content if concatenated_messages else final_user_content
# channel_id = get_channel_id(hash_value, token, model_to_use, final_user_content, template_id)
hash_value = generate_hash(first_user_message, model_to_use, token)
channel_id = get_channel_id(hash_value, token, model_to_use, final_user_content, template_id)
if final_user_content is None:
return Response("No user message found", status=400)
return send_chat_message(req, token, channel_id, final_user_content, model_to_use, user_stream, image_url, model_name)
def handle_options_request():
return Response(status=204, headers={'Access-Control-Allow-Origin': '*', 'Access-Control-Allow-Headers': '*'})
|