| 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 |
| |
|
|
| 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': '*'}) |
|
|
