Update main.py

main.py

@@ -7,7 +7,6 @@ from typing import Any, Dict, List, Optional
 
 import httpx
 import uvicorn
-from dotenv import load_dotenv
 from fastapi import FastAPI, HTTPException, Depends
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
@@ -20,26 +19,43 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# Load Environment Variables
+# Load Environment Variables for Sensitive Information
+from dotenv import load_dotenv
+
 load_dotenv()
+
 app = FastAPI()
 
-# Configuration Constants
-BASE_URL = "https://aichatonlineorg.erweima.ai/aichatonline"
+# Configuration Constants (Hardcoded as per requirements)
+APP_NAME = "ChitChat_Chrome_Ext"
+APP_VERSION = "4.28.0"
+ORIGIN_URL = "chrome-extension://difoiogjjojoaoomphldepapgpbgkhkb/standalone.html?from=sidebar"
+ORIGIN_TITLE = "Sider"
+TZ_NAME = "Asia/Karachi"
+PROMPT_TEMPLATE_KEY = "artifacts"
+PROMPT_TEMPLATE_LANG = "original"
+TOOLS_AUTO = ["search", "text_to_image", "data_analysis"]
+
+# Security Configuration
 APP_SECRET = os.getenv("APP_SECRET", "666")
 ACCESS_TOKEN = os.getenv("SD_ACCESS_TOKEN", "")
-headers = {
+if not ACCESS_TOKEN:
+    logger.error("SD_ACCESS_TOKEN is not set in the environment variables.")
+    raise RuntimeError("SD_ACCESS_TOKEN is required but not set.")
+
+# Outgoing Request Headers (As per provided sample)
+OUTGOING_HEADERS = {
     'accept': '*/*',
-    'accept-language': 'zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6',
+    'accept-encoding': 'gzip, deflate, br, zstd',
+    'accept-language': 'en-US,en;q=0.9',
     'authorization': f'Bearer {ACCESS_TOKEN}',
-    'cache-control': 'no-cache',
-    'origin': 'chrome-extension://dhoenijjpgpeimemopealfcbiecgceod',
-    'pragma': 'no-cache',
+    'content-type': 'application/json',
+    'origin': 'chrome-extension://difoiogjjojoaoomphldepapgpbgkhkb',
     'priority': 'u=1, i',
     'sec-fetch-dest': 'empty',
     'sec-fetch-mode': 'cors',
     'sec-fetch-site': 'none',
-    'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36 Edg/129.0.0.0',
+    'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36',
 }
 
 # Updated ALLOWED_MODELS List (No Duplicates)
@@ -59,25 +75,28 @@ ALLOWED_MODELS = [
 # Configure CORS
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Restrict this to specific origins in production
+    allow_origins=["*"],  # ⚠️ IMPORTANT: Restrict this to specific origins in production for security
     allow_credentials=True,
-    allow_methods=["*"],  # All methods allowed
+    allow_methods=["*"],  # Allow all HTTP methods
     allow_headers=["*"],  # Allow all headers
 )
 
-# Security Configuration
+# Security Dependency
 security = HTTPBearer()
 
+
 # Pydantic Models
 class Message(BaseModel):
     role: str
     content: str
 
+
 class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
     stream: Optional[bool] = False
 
+
 # Utility Functions
 def create_chat_completion_data(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
     return {
@@ -95,6 +114,7 @@ def create_chat_completion_data(content: str, model: str, finish_reason: Optiona
         "usage": None,
     }
 
+
 def verify_app_secret(credentials: HTTPAuthorizationCredentials = Depends(security)):
     if credentials.credentials != APP_SECRET:
         logger.warning(f"Invalid APP_SECRET provided: {credentials.credentials}")
@@ -102,6 +122,7 @@ def verify_app_secret(credentials: HTTPAuthorizationCredentials = Depends(securi
     logger.info("APP_SECRET verified successfully.")
     return credentials.credentials
 
+
 # CORS Preflight Options Endpoint
 @app.options("/hf/v1/chat/completions")
 async def chat_completions_options():
@@ -114,15 +135,13 @@ async def chat_completions_options():
         },
     )
 
-# Replace Escaped Newlines
-def replace_escaped_newlines(input_string: str) -> str:
-    return input_string.replace("\\n", "\n")
 
 # List Available Models
 @app.get("/hf/v1/models")
 async def list_models():
     return {"object": "list", "data": ALLOWED_MODELS}
 
+
 # Chat Completions Endpoint
 @app.post("/hf/v1/chat/completions")
 async def chat_completions(
@@ -133,7 +152,7 @@ async def chat_completions(
     # Validate Selected Model
     if request.model not in [model['id'] for model in ALLOWED_MODELS]:
         allowed = ', '.join(model['id'] for model in ALLOWED_MODELS)
-        logger.error(f"Model {request.model} is not allowed.")
+        logger.error(f"Model '{request.model}' is not allowed.")
         raise HTTPException(
             status_code=400,
             detail=f"Model '{request.model}' is not allowed. Allowed models are: {allowed}",
@@ -141,48 +160,48 @@ async def chat_completions(
 
     logger.info(f"Using model: {request.model}")
 
-    # Generate a UUID
-    original_uuid = uuid.uuid4()
-    uuid_str = str(original_uuid).replace("-", "")
-    logger.debug(f"Generated UUID: {uuid_str}")
+    # Generate a unique CID for each request
+    cid = str(uuid.uuid4()).replace("-", "").upper()[:12]  # Example: C0MES13070J1
+    logger.debug(f"Generated CID: {cid}")
 
     # Prepare JSON Payload for External API
-    json_data = {
-        'prompt': "\n".join(
+    if not request.messages:
+        prompt_text = "make a dog"
+    else:
+        prompt_text = "\n".join(
             [
                 f"{'User' if msg.role == 'user' else 'Assistant'}: {msg.content}"
                 for msg in request.messages
             ]
-        ),
-        'stream': True,
-        'app_name': 'ChitChat_Edge_Ext',
-        'app_version': '4.26.1',
-        'tz_name': 'Asia/Karachi',
-        'cid': '',
-        'model': request.model,  # Use the selected model directly
+        )
+
+    json_data = {
+        'prompt': prompt_text,
+        'stream': request.stream,
+        'app_name': APP_NAME,
+        'app_version': APP_VERSION,
+        'tz_name': TZ_NAME,
+        'cid': cid,
+        'model': request.model,
         'search': False,
         'auto_search': False,
         'filter_search_history': False,
         'from': 'chat',
         'group_id': 'default',
-        'chat_models': [request.model],  # Include the model in chat_models
+        'chat_models': [],  # As per the sample payload
         'files': [],
         'prompt_template': {
-            'key': '',
+            'key': PROMPT_TEMPLATE_KEY,
             'attributes': {
-                'lang': 'original',
+                'lang': PROMPT_TEMPLATE_LANG,
             },
         },
         'tools': {
-            'auto': [
-                'search',
-                'text_to_image',
-                'data_analysis',
-            ],
+            'auto': TOOLS_AUTO,
         },
         'extra_info': {
-            'origin_url': '',
-            'origin_title': '',
+            'origin_url': ORIGIN_URL,
+            'origin_title': ORIGIN_TITLE,
         },
     }
 
@@ -193,8 +212,8 @@ async def chat_completions(
             try:
                 async with client.stream(
                     'POST',
-                    'https://sider.ai/api/v2/completion/text',
-                    headers=headers,
+                    'https://sider.ai/api/v3/completion/text',  # Updated endpoint
+                    headers=OUTGOING_HEADERS,
                     json=json_data,
                     timeout=120.0
                 ) as response:
@@ -237,7 +256,7 @@ async def chat_completions(
                         full_response += data["choices"][0]["delta"]["content"]
                 except json.JSONDecodeError:
                     logger.warning(f"Failed to decode JSON from chunk: {chunk}")
-        
+
         # Final Response Structure
         return {
             "id": f"chatcmpl-{uuid.uuid4()}",
@@ -254,6 +273,7 @@ async def chat_completions(
             "usage": None,
         }
 
+
 # Entry Point
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)