Update main.py

main.py

@@ -1,6 +1,8 @@
 import base64
 import json
 import os
+import secrets  # <-- Import secrets
+import string   # <-- Import string
 import time
 from typing import List, Optional
 
@@ -15,25 +17,20 @@ from pydantic import BaseModel
 load_dotenv()
 
 # Env variables for external services
-IMAGE_API_URL = os.environ.get("IMAGE_API_URL", "https://image.api.example.com") # Add a default for safety
+IMAGE_API_URL = os.environ.get("IMAGE_API_URL", "https://image.api.example.com")
 SNAPZION_UPLOAD_URL = "https://upload.snapzion.com/api/public-upload"
-SNAPZION_API_KEY = os.environ.get("SNAP", "") # Add a default for safety
+SNAPZION_API_KEY = os.environ.get("SNAP", "")
 
 # --- Dummy Model Definitions ---
 # In a real application, these would be defined properly.
-# For this example, we define them here.
-
 AVAILABLE_MODELS = [
     {"id": "gpt-4-turbo", "object": "model", "created": int(time.time()), "owned_by": "system"},
     {"id": "gpt-4o", "object": "model", "created": int(time.time()), "owned_by": "system"},
     {"id": "gpt-3.5-turbo", "object": "model", "created": int(time.time()), "owned_by": "system"},
     {"id": "dall-e-3", "object": "model", "created": int(time.time()), "owned_by": "system"},
-    # Add any other models you support
 ]
 
-MODEL_ALIASES = {
-    # Example: "gpt-4": "gpt-4-turbo"
-}
+MODEL_ALIASES = {}
 
 # --- FastAPI Application ---
 
@@ -43,9 +40,19 @@ app = FastAPI(
     version="1.0.0"
 )
 
-def unix_id():
-    """Generates a Unix timestamp-based ID."""
-    return str(int(time.time() * 1000))
+
+# --- Helper Function for Random ID Generation ---
+def generate_random_id(prefix: str, length: int = 29) -> str:
+    """
+    Generates a cryptographically secure, random alphanumeric ID.
+    The default length of 29 characters is common for OpenAI IDs.
+    The example 'bwvaLjbI0KEKMadGmFbSsjYNLgaI' is 30 characters.
+    You can adjust the length as needed.
+    """
+    population = string.ascii_letters + string.digits
+    random_part = "".join(secrets.choice(population) for _ in range(length))
+    return f"{prefix}{random_part}"
+
 
 # === API Endpoints ===
 
@@ -54,6 +61,7 @@ async def list_models():
     """Lists the available models."""
     return {"object": "list", "data": AVAILABLE_MODELS}
 
+
 # === Chat Completion ===
 
 class Message(BaseModel):
@@ -64,19 +72,17 @@ class ChatRequest(BaseModel):
     messages: List[Message]
     model: str
     stream: Optional[bool] = False
-    # Add other common parameters for compatibility if needed
-    # max_tokens: Optional[int] = None
-    # temperature: Optional[float] = None
-    # user: Optional[str] = None
-
 
 @app.post("/v1/chat/completions")
 async def chat_completion(request: ChatRequest):
     """
     Handles chat completion requests, supporting both streaming and non-streaming responses.
-    This endpoint is updated to match the OpenAI streaming chunk format precisely.
+    This endpoint now uses a long, random ID for completions.
     """
     model_id = MODEL_ALIASES.get(request.model, request.model)
+    
+    # Generate the ID once for the entire request
+    chat_id = generate_random_id("chatcmpl-")
 
     headers = {
         'accept': 'text/event-stream',
@@ -93,7 +99,6 @@ async def chat_completion(request: ChatRequest):
 
     if request.stream:
         async def event_stream():
-            chat_id = f"chatcmpl-{unix_id()}"
             created = int(time.time())
             
             is_first_chunk = True
@@ -121,7 +126,7 @@ async def chat_completion(request: ChatRequest):
                                         is_first_chunk = False
 
                                     chunk_data = {
-                                        "id": chat_id,
+                                        "id": chat_id, # Use the pre-generated ID
                                         "object": "chat.completion.chunk",
                                         "created": created,
                                         "model": model_id,
@@ -140,7 +145,6 @@ async def chat_completion(request: ChatRequest):
                                     pass
                                 break
                 
-                # After the loop, send the final chunk with finish_reason and usage
                 final_usage = None
                 if usage_info:
                     prompt_tokens = usage_info.get("promptTokens", 0)
@@ -152,7 +156,7 @@ async def chat_completion(request: ChatRequest):
                     }
 
                 done_chunk = {
-                    "id": chat_id,
+                    "id": chat_id, # Use the pre-generated ID
                     "object": "chat.completion.chunk",
                     "created": created,
                     "model": model_id,
@@ -180,7 +184,7 @@ async def chat_completion(request: ChatRequest):
 
         return StreamingResponse(event_stream(), media_type="text/event-stream")
     else:
-        # Non-streaming logic remains the same
+        # Non-streaming logic
         assistant_response = ""
         usage_info = {}
 
@@ -203,7 +207,7 @@ async def chat_completion(request: ChatRequest):
                                 continue
             
             return JSONResponse(content={
-                "id": f"chatcmpl-{unix_id()}",
+                "id": chat_id, # Use the pre-generated ID
                 "object": "chat.completion",
                 "created": int(time.time()),
                 "model": model_id,
@@ -291,8 +295,10 @@ async def generate_images(request: ImageGenerationRequest):
 
     return {"created": int(time.time()), "data": results}
 
+
 if __name__ == "__main__":
     import uvicorn
-    # Make sure you have a .env file with SNAP and IMAGE_API_URL
-    # Example: uvicorn your_script_name:app --reload --port 8000
+    # To run this file:
+    # 1. Make sure you have a .env file with your SNAP key.
+    # 2. Run in your terminal: uvicorn your_script_name:app --reload --port 8000
     uvicorn.run(app, host="0.0.0.0", port=8000)