Update main.py

main.py

@@ -8,103 +8,60 @@ from typing import Any, Dict, List, Optional
 import httpx
 import uvicorn
 from dotenv import load_dotenv
-from fastapi import FastAPI, HTTPException, Depends
+from fastapi import FastAPI, HTTPException, Depends, Request
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 from starlette.middleware.cors import CORSMiddleware
 from starlette.responses import StreamingResponse, Response
 
+# Configure Logging
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 logger = logging.getLogger(__name__)
 
+# Load Environment Variables
 load_dotenv()
+APP_SECRET = os.getenv("APP_SECRET", "786")
+ACCESS_TOKEN = os.getenv("SD_ACCESS_TOKEN", "")
+
+# Initialize FastAPI
 app = FastAPI()
-BASE_URL = "https://aichatonlineorg.erweima.ai/aichatonline"
-APP_SECRET = os.getenv("APP_SECRET","786")
-ACCESS_TOKEN = os.getenv("SD_ACCESS_TOKEN","")
-headers = {
-    'accept': '*/*',
-    'accept-language': 'en-US,en;q=0.9',
-    'authorization': f'Bearer {ACCESS_TOKEN}',
-    'cache-control': 'no-cache',
-    'origin': 'chrome-extension://difoiogjjojoaoomphldepapgpbgkhkb',
-    'pragma': 'no-cache',
-    'priority': 'u=1, i',
-    'cookie': 'lang=en; source=gg; p1=pricing; p2=search; _clck=ih7kjx%7C2%7Cfqq%7C0%7C1774; _gcl_gs=2.1.k1$i1731168118$u7016880; _ga=GA1.1.1378917294.1731168124; _gcl_au=1.1.436410932.1731168125; _fbp=fb.1.1731168125483.909164969447382788; _gcl_aw=GCL.1731168126.EAIaIQobChMI_aOK5c_PiQMVAqSDBx2OJh5bEAAYASAAEgJr9PD_BwE; __stripe_mid=b2de6d00-4ac3-4c15-9ba8-355dd63bbdf0cbeb53; _uetsid=f73638b09eb311efa0cf6b0f13512e16; _uetvid=f73635709eb311efbd8d55293947782a; token=Bearer%20eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyX2lkIjoxNTYzMzI1NSwicmVnaXN0ZXJfdHlwZSI6Im9hdXRoMiIsImFwcF9uYW1lIjoiQ2hpdENoYXRfV2ViIiwidG9rZW5faWQiOiI2MGRmZGMwMi1lY2QzLTQzNzktYjMzNy1kYmNlNGY5M2M1N2IiLCJpc3MiOiJzaWRlci5haSIsImF1ZCI6WyIiXSwiZXhwIjoxNzYyMjcyNTk5LCJuYmYiOjE3MzExNjg1OTksImlhdCI6MTczMTE2ODU5OX0.cT9hXiCFm-4Or2RV5Lf3IY-blu40UOvjqD01BGoDwh8; refresh_token=discard; userinfo-avatar=https://chitchat-avatar.s3.amazonaws.com/eb0e281f19c442d681cf552f3b8896dd-1731167278.png; userinfo-name=Chat%20GPT%20free; userinfo-type=oauth2; CloudFront-Policy=eyJTdGF0ZW1lbnQiOlt7IlJlc291cmNlIjoiaHR0cHM6Ly9maWxlLWNkbi5zaWRlci5haS8qL1UwMUFIRUc4NkVBLyoiLCJDb25kaXRpb24iOnsiRGF0ZUxlc3NUaGFuIjp7IkFXUzpFcG9jaFRpbWUiOjE3MzM3NjA2MDF9fX1dfQ__; CloudFront-Signature=fuXgBZW3E5TCNvvi8mb8DKpNTneaPIGunH~zxWPoVK~OaG-n6PRVz2qlxOFCY3HLTZyA72mQ4T2OtKCZGoWCTL4QL6DP5LnRw7HPt50K~D8Wte5M3GQjuSeuBe~mK44Sk~Xqn1tOa2dmOtri84vlLdo7ud4ZGuMML60AUve13l7eS8uAJO88cQ9rhdZMJ26opJhayI0MSJgCFZymzMm1iEwwSF3ufV5c2elpTuKqVTP2HqZGHzgJQvRl7U~gU4qFiH7KxCG0MF5X0HDsz3Pn0qghNle~FLAtCumI-mzoGNCXEM9~qpYAXJLO3a--ThhnbTrWsPLl2l~zCRiKQ5wB3g__; CloudFront-Key-Pair-Id=K344F5VVSSM536; _rdt_uuid=1731168120149.9c981ceb-cb94-490d-99d8-0f1530348c3c; _clsk=r6ig1v%7C1731170609000%7C8%7C1%7Cu.clarity.ms%2Fcollect; _ga_0PRFKME4HP=GS1.1.1731171204.2.1.1731171205.59.0.0',
-    'sec-fetch-dest': 'empty',
-    'sec-fetch-mode': 'cors',
-    'sec-fetch-site': 'none',
-    'sec-fetch-site': 'none',
-    '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',
-}
 
+# Define Allowed Models
 ALLOWED_MODELS = [
-    {"id": "claude-3.5-sonnet", "name": "claude-3.5-sonnet"},
-    {"id": "claude-3-opus", "name": "claude-3-opus"},
-    {"id": "gemini-1.5-pro", "name": "gemini-1.5-pro"},
-    {"id": "gpt-4o", "name": "gpt-4o"},
-    {"id": "o1-preview", "name": "o1-preview"},
-    {"id": "o1-mini", "name": "o1-mini"},
-    {"id": "gpt-4o-mini", "name": "gpt-4o-mini"},
+    {"id": "claude-3.5-sonnet", "name": "Claude 3.5 Sonnet"},
+    {"id": "claude-3-opus", "name": "Claude 3 Opus"},
+    {"id": "gemini-1.5-pro", "name": "Gemini 1.5 Pro"},
+    {"id": "gpt-4o", "name": "GPT-4o"},
+    {"id": "o1-preview", "name": "O1 Preview"},
+    {"id": "o1-mini", "name": "O1 Mini"},
+    {"id": "gpt-4o-mini", "name": "GPT-4o Mini"},
 ]
-# Configure CORS
+
+# Configure CORS Middleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Allow all sources, you can restrict specific sources if needed
+    allow_origins=["https://yourdomain.com"],  # Replace with your trusted domains
     allow_credentials=True,
-    allow_methods=["*"],  # All methods allowed
-    allow_headers=["*"],  # Allow all headers
+    allow_methods=["POST", "OPTIONS"],
+    allow_headers=["Content-Type", "Authorization"],
 )
-security = HTTPBearer()
 
+# Security Scheme
+security = HTTPBearer()
 
+# Pydantic Models
 class Message(BaseModel):
     role: str
     content: str
 
-
 class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
     stream: Optional[bool] = False
 
-
-def simulate_data(content, model):
-    return {
-        "id": f"chatcmpl-{uuid.uuid4()}",
-        "object": "chat.completion.chunk",
-        "created": int(datetime.now().timestamp()),
-        "model": model,
-        "choices": [
-            {
-                "index": 0,
-                "delta": {"content": content, "role": "assistant"},
-                "finish_reason": None,
-            }
-        ],
-        "usage": None,
-    }
-
-
-def stop_data(content, model):
-    return {
-        "id": f"chatcmpl-{uuid.uuid4()}",
-        "object": "chat.completion.chunk",
-        "created": int(datetime.now().timestamp()),
-        "model": model,
-        "choices": [
-            {
-                "index": 0,
-                "delta": {"content": content, "role": "assistant"},
-                "finish_reason": "stop",
-            }
-        ],
-        "usage": None,
-    }
-    
-    
+# Helper Functions
 def create_chat_completion_data(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
@@ -121,13 +78,12 @@ 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:
         raise HTTPException(status_code=403, detail="Invalid APP_SECRET")
     return credentials.credentials
 
-
+# Routes
 @app.options("/hf/v1/chat/completions")
 async def chat_completions_options():
     return Response(
@@ -139,16 +95,10 @@ async def chat_completions_options():
         },
     )
 
-
-def replace_escaped_newlines(input_string: str) -> str:
-    return input_string.replace("\\n", "\n")
-
-
 @app.get("/hf/v1/models")
 async def list_models():
     return {"object": "list", "data": ALLOWED_MODELS}
 
-
 @app.post("/hf/v1/chat/completions")
 async def chat_completions(
     request: ChatRequest, app_secret: str = Depends(verify_app_secret)
@@ -156,65 +106,91 @@ async def chat_completions(
     logger.info(f"Received chat completion request for model: {request.model}")
 
     if request.model not in [model['id'] for model in ALLOWED_MODELS]:
+        allowed = ', '.join(model['id'] for model in ALLOWED_MODELS)
         raise HTTPException(
             status_code=400,
-            detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(model['id'] for model in ALLOWED_MODELS)}",
+            detail=f"Model '{request.model}' is not allowed. Allowed models are: {allowed}",
         )
+
     # Generate a UUID
-    original_uuid = uuid.uuid4()
-    uuid_str = str(original_uuid).replace("-", "")
+    uuid_str = str(uuid.uuid4()).replace("-", "")
+
+    # Prepare Headers (Move sensitive data to environment variables or secure storage)
+    headers = {
+        'accept': '*/*',
+        'accept-language': 'en-US,en;q=0.9',
+        'authorization': f'Bearer {ACCESS_TOKEN}',
+        'cache-control': 'no-cache',
+        'origin': 'chrome-extension://difoiogjjojoaoomphldepapgpbgkhkb',
+        'pragma': 'no-cache',
+        '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/130.0.0.0 Safari/537.36',
+    }
 
-    # Using the OpenAI API
+    # Prepare JSON Payload
     json_data = {
         'prompt': "\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.28.0',
         'tz_name': 'Asia/Karachi',
         'cid': 'C092SEMXM9BJ',
         'model': request.model,
-        'search': False,          # Explicitly turn off search
-        'auto_search': False,      # Turn off auto-search
+        'search': False,
+        'auto_search': False,
         'from': 'chat',
         'group_id': 'default',
         'prompt_template': {
-        'key': 'artifacts',
-        'attributes': {
-            'lang': 'original',
+            'key': 'artifacts',
+            'attributes': {
+                'lang': 'original',
+            },
         },
-    },
         'tools': {
-            'auto': ['text_to_image', 'data_analysis'],  # Remove 'search' from tools if not needed
-    },
+            'auto': ['text_to_image', 'data_analysis'],
+        },
         'extra_info': {
             'origin_url': 'chrome-extension://difoiogjjojoaoomphldepapgpbgkhkb/standalone.html?from=sidebar',
             'origin_title': 'Sider',
-    },
-}
-
-
+        },
+    }
 
     async def generate():
         async with httpx.AsyncClient() as client:
             try:
-                async with client.stream('POST', 'https://sider.ai/api/v2/completion/text', headers=headers, json=json_data, timeout=120.0) as response:
+                async with client.stream(
+                    'POST',
+                    'https://sider.ai/api/v2/completion/text',
+                    headers=headers,
+                    json=json_data,
+                    timeout=120.0
+                ) as response:
                     response.raise_for_status()
                     async for line in response.aiter_lines():
                         if line and ("[DONE]" not in line):
-                            content = json.loads(line[5:])["data"]
-                            yield f"data: {json.dumps(create_chat_completion_data(content.get('text',''), request.model))}\n\n"
+                            # Assuming the response starts with some prefix, adjust accordingly
+                            try:
+                                data = json.loads(line[5:]).get("data", {})
+                                content = data.get("text", "")
+                                yield f"data: {json.dumps(create_chat_completion_data(content, request.model))}\n\n"
+                            except json.JSONDecodeError as e:
+                                logger.error(f"JSON decode error: {e}")
+                    # Indicate the end of the stream
                     yield f"data: {json.dumps(create_chat_completion_data('', request.model, 'stop'))}\n\n"
                     yield "data: [DONE]\n\n"
             except httpx.HTTPStatusError as e:
                 logger.error(f"HTTP error occurred: {e}")
                 raise HTTPException(status_code=e.response.status_code, detail=str(e))
             except httpx.RequestError as e:
-                logger.error(f"An error occurred while requesting: {e}")
+                logger.error(f"Request error occurred: {e}")
                 raise HTTPException(status_code=500, detail=str(e))
 
     if request.stream:
@@ -225,11 +201,10 @@ async def chat_completions(
         full_response = ""
         async for chunk in generate():
             if chunk.startswith("data: ") and not chunk[6:].startswith("[DONE]"):
-                # print(chunk)
                 data = json.loads(chunk[6:])
                 if data["choices"][0]["delta"].get("content"):
                     full_response += data["choices"][0]["delta"]["content"]
-        
+
         return {
             "id": f"chatcmpl-{uuid.uuid4()}",
             "object": "chat.completion",
@@ -245,7 +220,6 @@ async def chat_completions(
             "usage": None,
         }
 
-
-
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+# Remove the Uvicorn run block for production deployment
+# if __name__ == "__main__":
+#     uvicorn.run(app, host="0.0.0.0", port=7860)