Update app.py

app.py

@@ -4,7 +4,6 @@ import uuid
 import json
 from typing import List, Optional
 from pydantic import BaseModel, ValidationError
-import logging
 from API_provider import API_Inference
 from core_logic import (
     check_api_key_validity,
@@ -16,7 +15,6 @@ from core_logic import (
 )
 
 app = Flask(__name__)
-logging.basicConfig(level=logging.DEBUG)
 
 class Message(BaseModel):
     role: str
@@ -53,94 +51,92 @@ def index():
 @app.route('/chat/completions', methods=['POST', 'GET'])
 @requires_api_key
 def chat_completions(api_key):
-    return jsonify({'detail': "YOUUUUUUUUUU"}), 500
-    # logging.info("Received request for chat completions")
-    # print("requess received")
-    # try:
-    #     logging.info("Received request for chat completions")
-    #     # Parse and validate request data
-    #     try:
-    #         data = request.get_json()
-    #         chat_request = ChatCompletionRequest(**data)
-    #     except ValidationError as e:
-    #         return jsonify({'detail': e.errors()}), 400
-
-    #     # Check API key validity and rate limit
-    #     is_valid, error_message = check_api_key_validity(api_key)
-    #     if not is_valid:
-    #         return jsonify({'detail': error_message}), 401
-
-    #     messages = [{"role": msg.role, "content": msg.content} for msg in chat_request.messages]
-
-    #     # Get model info
-    #     model_info = get_model_info(chat_request.model)
-    #     if not model_info:
-    #         return jsonify({'detail': 'Invalid model specified'}), 400
-
-    #     # Model mapping
-    #     model_mapping = {
-    #         "meta-llama-405b-turbo": "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo",
-    #         "claude-3.5-sonnet": "claude-3-sonnet-20240229",
-    #     }
-    #     model_name = model_mapping.get(chat_request.model, chat_request.model)
-    #     credits_reduction = {
-    #         "gpt-4o": 1,
-    #         "claude-3-sonnet-20240229": 1,
-    #         "gemini-1.5-pro": 1,
-    #         "gemini-1-5-flash": 1,
-    #         "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": 1,
-    #         "o1-mini": 2,
-    #         "o1-preview": 3,
-    #     }.get(model_name, 0)
-
-    #     if chat_request.stream:
-    #         def generate():
-    #             try:
-    #                 for chunk in API_Inference(messages, model=model_name, stream=True,
-    #                                            max_tokens=chat_request.max_tokens,
-    #                                            temperature=chat_request.temperature,
-    #                                            top_p=chat_request.top_p):
-    #                     data = json.dumps({'choices': [{'delta': {'content': chunk}}]})
-    #                     yield f"data: {data}\n\n"
-    #                 yield f"data: [DONE]\n\nCredits used: {credits_reduction}\n\n"
-    #                 update_request_count(api_key, credits_reduction)
-    #             except Exception as e:
-    #                 yield f"data: [ERROR] {str(e)}\n\n"
-
-    #         return Response(generate(), mimetype='text/event-stream')
-    #     else:
-    #         response = API_Inference(messages, model=model_name, stream=False,
-    #                                  max_tokens=chat_request.max_tokens,
-    #                                  temperature=chat_request.temperature,
-    #                                  top_p=chat_request.top_p)
-    #         update_request_count(api_key, credits_reduction)
-    #         prompt_tokens = sum(len(msg['content'].split()) for msg in messages)
-    #         completion_tokens = len(response.split())
-    #         total_tokens = prompt_tokens + completion_tokens
-    #         return jsonify({
-    #             "id": f"chatcmpl-{str(uuid.uuid4())}",
-    #             "object": "chat.completion",
-    #             "created": int(uuid.uuid1().time // 1e7),
-    #             "model": model_name,
-    #             "choices": [
-    #                 {
-    #                     "index": 0,
-    #                     "message": {
-    #                         "role": "assistant",
-    #                         "content": response
-    #                     },
-    #                     "finish_reason": "stop"
-    #                 }
-    #             ],
-    #             "usage": {
-    #                 "prompt_tokens": prompt_tokens,
-    #                 "completion_tokens": completion_tokens,
-    #                 "total_tokens": total_tokens
-    #             },
-    #             "credits_used": credits_reduction
-    #         })
-    # except Exception as e:
-    #     return jsonify({'detail': str(e)}), 500
+    print("requess received")
+    try:
+        logging.info("Received request for chat completions")
+        # Parse and validate request data
+        try:
+            data = request.get_json()
+            chat_request = ChatCompletionRequest(**data)
+        except ValidationError as e:
+            return jsonify({'detail': e.errors()}), 400
+
+        # Check API key validity and rate limit
+        is_valid, error_message = check_api_key_validity(api_key)
+        if not is_valid:
+            return jsonify({'detail': error_message}), 401
+
+        messages = [{"role": msg.role, "content": msg.content} for msg in chat_request.messages]
+
+        # Get model info
+        model_info = get_model_info(chat_request.model)
+        if not model_info:
+            return jsonify({'detail': 'Invalid model specified'}), 400
+
+        # Model mapping
+        model_mapping = {
+            "meta-llama-405b-turbo": "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo",
+            "claude-3.5-sonnet": "claude-3-sonnet-20240229",
+        }
+        model_name = model_mapping.get(chat_request.model, chat_request.model)
+        credits_reduction = {
+            "gpt-4o": 1,
+            "claude-3-sonnet-20240229": 1,
+            "gemini-1.5-pro": 1,
+            "gemini-1-5-flash": 1,
+            "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": 1,
+            "o1-mini": 2,
+            "o1-preview": 3,
+        }.get(model_name, 0)
+
+        if chat_request.stream:
+            def generate():
+                try:
+                    for chunk in API_Inference(messages, model=model_name, stream=True,
+                                               max_tokens=chat_request.max_tokens,
+                                               temperature=chat_request.temperature,
+                                               top_p=chat_request.top_p):
+                        data = json.dumps({'choices': [{'delta': {'content': chunk}}]})
+                        yield f"data: {data}\n\n"
+                    yield f"data: [DONE]\n\nCredits used: {credits_reduction}\n\n"
+                    update_request_count(api_key, credits_reduction)
+                except Exception as e:
+                    yield f"data: [ERROR] {str(e)}\n\n"
+
+            return Response(generate(), mimetype='text/event-stream')
+        else:
+            response = API_Inference(messages, model=model_name, stream=False,
+                                     max_tokens=chat_request.max_tokens,
+                                     temperature=chat_request.temperature,
+                                     top_p=chat_request.top_p)
+            update_request_count(api_key, credits_reduction)
+            prompt_tokens = sum(len(msg['content'].split()) for msg in messages)
+            completion_tokens = len(response.split())
+            total_tokens = prompt_tokens + completion_tokens
+            return jsonify({
+                "id": f"chatcmpl-{str(uuid.uuid4())}",
+                "object": "chat.completion",
+                "created": int(uuid.uuid1().time // 1e7),
+                "model": model_name,
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": response
+                        },
+                        "finish_reason": "stop"
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": total_tokens
+                },
+                "credits_used": credits_reduction
+            })
+    except Exception as e:
+        return jsonify({'detail': str(e)}), 500
 
 @app.route('/rate_limit/status', methods=['GET'])
 @requires_api_key