Update app.py

app.py

@@ -6,18 +6,27 @@ import random
 import logging
 import sys
 import re
+from logging.handlers import RotatingFileHandler
 
 app = Flask(__name__)
 
 # 配置日志
-logging.basicConfig(
-    level=logging.DEBUG,
-    format='%(asctime)s [%(levelname)s] %(message)s',
-    handlers=[
-        logging.StreamHandler(sys.stdout)
-    ]
+class RequestFormatter(logging.Formatter):
+    def format(self, record):
+        record.url = request.url
+        record.remote_addr = request.remote_addr
+        return super().format(record)
+
+formatter = RequestFormatter(
+    '%(asctime)s - %(remote_addr)s - %(url)s - %(levelname)s - %(message)s'
 )
-logger = logging.getLogger(__name__)
+
+handler = RotatingFileHandler('app.log', maxBytes=10000, backupCount=3)
+handler.setFormatter(formatter)
+handler.setLevel(logging.INFO)
+
+app.logger.addHandler(handler)
+app.logger.setLevel(logging.INFO)
 
 SYSTEM_ASSISTANT = """作为 Stable Diffusion Prompt 提示词专家，您将从关键词中创建提示，通常来自 Danbooru 等数据库。
 提示通常描述图像，使用常见词汇，按重要性排列，并用逗号分隔。避免使用"-"或"."，但可以接受空格和自然语言。避免词汇重复。
@@ -59,7 +68,7 @@ def get_random_token(auth_header):
     tokens = [token.strip() for token in auth_header.split(',') if token.strip()]
     if not tokens:
         return None
-    return f"Bearer {random.choice(tokens)}"
+    return random.choice(tokens)
 
 def translate_and_enhance_prompt(prompt, auth_token):
     translate_url = 'https://api.siliconflow.cn/v1/chat/completions'
@@ -72,24 +81,13 @@ def translate_and_enhance_prompt(prompt, auth_token):
     }
     headers = {
         'Content-Type': 'application/json',
-        'Authorization': auth_token
+        'Authorization': f'Bearer {auth_token}'
     }
     
-    logger.info(f"Sending request to {translate_url}")
-    logger.info(f"Request headers: {headers}")
-    logger.info(f"Request body: {json.dumps(translate_body, ensure_ascii=False)}")
-    
-    try:
-        response = requests.post(translate_url, headers=headers, json=translate_body, timeout=30)
-        logger.info(f"Response status code: {response.status_code}")
-        logger.info(f"Response content: {response.text}")
-        
-        response.raise_for_status()
-        result = response.json()
-        return result['choices'][0]['message']['content']
-    except requests.exceptions.RequestException as e:
-        logger.error(f"Error in translate_and_enhance_prompt: {str(e)}")
-        raise
+    response = requests.post(translate_url, headers=headers, json=translate_body, timeout=30)
+    response.raise_for_status()
+    result = response.json()
+    return result['choices'][0]['message']['content']
 
 def extract_params_from_prompt(prompt):
     size_match = re.search(r'-s\s+(\S+)', prompt)
@@ -107,11 +105,34 @@ def extract_params_from_prompt(prompt):
         clean_prompt = re.sub(r'-o', '', clean_prompt).strip()
     
     image_size = RATIO_MAP.get(size, RATIO_MAP["16:9"])
-    return image_size, clean_prompt, use_original
+    return image_size, clean_prompt, use_original, size
 
 @app.route('/')
 def index():
-    return "text-to-image with siliconflow", 200
+    usage = """
+    Welcome to the text-to-image API with siliconflow!
+
+    Usage:
+    1. Send a POST request to /ai/v1/chat/completions
+    2. Include your prompt in the 'content' field of the last message
+    3. Optional parameters:
+       -s <ratio>: Set image size ratio (e.g., -s 1:1, -s 16:9)
+       -o: Use original prompt without enhancement
+
+    Example:
+    {
+        "model": "flux",
+        "messages": [
+            {
+                "role": "user",
+                "content": "A beautiful landscape -s 16:9"
+            }
+        ]
+    }
+
+    For more details, please refer to the API documentation.
+    """
+    return usage, 200
 
 @app.route('/ai/v1/chat/completions', methods=['POST'])
 def handle_request():
@@ -124,20 +145,17 @@ def handle_request():
             return jsonify({"error": "Bad Request: Missing required fields"}), 400
         
         prompt = messages[-1]['content']
-        image_size, clean_prompt, use_original = extract_params_from_prompt(prompt)
+        image_size, clean_prompt, use_original, size_param = extract_params_from_prompt(prompt)
         
-        random_token = get_random_token(request.headers.get('Authorization'))
+        auth_header = request.headers.get('Authorization')
+        random_token = get_random_token(auth_header)
         if not random_token:
             return jsonify({"error": "Unauthorized: Invalid or missing Authorization header"}), 401
         
         if use_original:
             enhanced_prompt = clean_prompt
         else:
-            try:
-                enhanced_prompt = translate_and_enhance_prompt(clean_prompt, random_token)
-            except Exception as e:
-                logger.error(f"Error in translate_and_enhance_prompt: {str(e)}")
-                return jsonify({"error": "Failed to enhance prompt"}), 500
+            enhanced_prompt = translate_and_enhance_prompt(clean_prompt, random_token)
         
         new_url = f'https://api.siliconflow.cn/v1/{model}/text-to-image'
         new_request_body = {
@@ -151,50 +169,38 @@ def handle_request():
         headers = {
             'accept': 'application/json',
             'content-type': 'application/json',
-            'Authorization': random_token
+            'Authorization': f'Bearer {random_token}'
         }
         
-        logger.info(f"Sending request to {new_url}")
-        logger.info(f"Request headers: {headers}")
-        logger.info(f"Request body: {json.dumps(new_request_body, ensure_ascii=False)}")
+        response = requests.post(new_url, headers=headers, json=new_request_body, timeout=60)
+        response.raise_for_status()
+        response_body = response.json()
         
-        try:
-            response = requests.post(new_url, headers=headers, json=new_request_body, timeout=60)
-            logger.info(f"Response status code: {response.status_code}")
-            logger.info(f"Response content: {response.text}")
-            
-            response.raise_for_status()
-            response_body = response.json()
-            
-            if 'images' in response_body and response_body['images'] and 'url' in response_body['images'][0]:
-                image_url = response_body['images'][0]['url']
-                logger.info(f"Successfully retrieved image URL: {image_url}")
-            else:
-                logger.error(f"Unexpected response structure: {response_body}")
-                return jsonify({"error": "Unexpected response structure from image generation API"}), 500
-        except requests.exceptions.RequestException as e:
-            logger.error(f"Error in image generation request: {str(e)}")
-            return jsonify({"error": "Failed to generate image"}), 500
-        except (KeyError, IndexError, ValueError) as e:
-            logger.error(f"Error parsing image generation response: {str(e)}")
-            return jsonify({"error": "Failed to parse image generation response"}), 500
+        if 'images' in response_body and response_body['images'] and 'url' in response_body['images'][0]:
+            image_url = response_body['images'][0]['url']
+        else:
+            raise ValueError("Unexpected response structure from image generation API")
         
-        unique_id = str(int(time.time() * 1000))  # 生成字符串类型的 unique_id
+        unique_id = str(int(time.time() * 1000))
         current_timestamp = int(time.time())
         system_fingerprint = "fp_" + ''.join(random.choices('abcdefghijklmnopqrstuvwxyz0123456789', k=9))
         
         image_data = {'data': [{'url': image_url}]}
         
+        # Log the key information
+        params = f"-s {size_param}" if size_param != "16:9" else ""
+        params += " -o" if use_original else ""
+        app.logger.info(f"Status: 200 - Token: {random_token[:6]}... - Params: {params} - Image URL: {image_url}")
+        
         if stream:
             return stream_response(unique_id, image_data, clean_prompt, enhanced_prompt, image_size, current_timestamp, model, system_fingerprint, use_original)
         else:
             return non_stream_response(unique_id, image_data, clean_prompt, enhanced_prompt, image_size, current_timestamp, model, system_fingerprint, use_original)
     except Exception as e:
-        logger.error(f"Unexpected error in handle_request: {str(e)}")
+        app.logger.error(f"Error: {str(e)}")
         return jsonify({"error": f"Internal Server Error: {str(e)}"}), 500
 
 def stream_response(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint, use_original):
-    logger.debug("Starting stream response")
     return Response(stream_with_context(generate_stream(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint, use_original)), content_type='text/event-stream')
 
 def generate_stream(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint, use_original):