Update main.py

main.py

@@ -7,99 +7,102 @@ from typing import Any, Dict, List, Optional
 
 import httpx
 import uvicorn
-from fastapi import FastAPI, HTTPException, Depends, Request, status
+from dotenv import load_dotenv
+from fastapi import FastAPI, HTTPException, Depends
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 from starlette.middleware.cors import CORSMiddleware
 from starlette.responses import StreamingResponse, Response
 
-from dotenv import load_dotenv
-
-# Retry Mechanism Libraries
-from tenacity import retry, wait_exponential, stop_after_attempt, retry_if_exception_type
-
-# Initialize Logging
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 logger = logging.getLogger(__name__)
 
-# Load Environment Variables for Sensitive Information
 load_dotenv()
-
 app = FastAPI()
-
-# 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", "")
-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 = {
+BASE_URL = "https://aichatonlineorg.erweima.ai/aichatonline"
+APP_SECRET = os.getenv("APP_SECRET","666")
+ACCESS_TOKEN = os.getenv("SD_ACCESS_TOKEN","")
+headers = {
     'accept': '*/*',
-    'accept-encoding': 'gzip, deflate, br, zstd',
-    'accept-language': 'en-US,en;q=0.9',
+    'accept-language': 'zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6',
     'authorization': f'Bearer {ACCESS_TOKEN}',
-    'content-type': 'application/json',
-    'origin': 'chrome-extension://difoiogjjojoaoomphldepapgpbgkhkb',
+    'cache-control': 'no-cache',
+    'origin': 'chrome-extension://dhoenijjpgpeimemopealfcbiecgceod',
+    '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',
+    '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',
 }
 
-# Updated ALLOWED_MODELS List (No Duplicates)
 ALLOWED_MODELS = [
-    {"id": "claude-3.5-sonnet", "name": "Claude 3.5 Sonnet"},
-    {"id": "sider", "name": "Sider"},
-    {"id": "gpt-4o-mini", "name": "GPT-4o Mini"},
-    {"id": "claude-3-haiku", "name": "Claude 3 Haiku"},
-    {"id": "claude-3.5-haiku", "name": "Claude 3.5 Haiku"},
-    {"id": "gemini-1.5-flash", "name": "Gemini 1.5 Flash"},
-    {"id": "llama-3", "name": "Llama 3"},
-    {"id": "gpt-4o", "name": "GPT-4o"},
-    {"id": "gemini-1.5-pro", "name": "Gemini 1.5 Pro"},
-    {"id": "llama-3.1-405b", "name": "Llama 3.1 405b"},
+    {"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
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # ⚠️ IMPORTANT: Restrict this to specific origins in production for security
+    allow_origins=["*"],  # Allow all sources, you can restrict specific sources if needed
     allow_credentials=True,
-    allow_methods=["*"],  # Allow all HTTP methods
+    allow_methods=["*"],  # All methods allowed
     allow_headers=["*"],  # Allow all headers
 )
-
-# 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 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,
+    }
+    
+    
 def create_chat_completion_data(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
@@ -116,35 +119,13 @@ 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}")
         raise HTTPException(status_code=403, detail="Invalid APP_SECRET")
-    logger.info("APP_SECRET verified successfully.")
     return credentials.credentials
 
-# Retry Configuration using Tenacity
-def is_retryable_exception(exception):
-    return isinstance(exception, httpx.HTTPStatusError) and exception.response.status_code == 429
-
-@retry(
-    retry=retry_if_exception_type(httpx.HTTPStatusError) & is_retryable_exception,
-    wait=wait_exponential(multiplier=1, min=2, max=10),
-    stop=stop_after_attempt(5),
-    reraise=True
-)
-async def send_request_with_retry(json_data: Dict[str, Any]) -> httpx.Response:
-    async with httpx.AsyncClient() as client:
-        response = await client.post(
-            'https://sider.ai/api/v3/completion/text',  # Updated endpoint
-            headers=OUTGOING_HEADERS,
-            json=json_data,
-            timeout=120.0
-        )
-        response.raise_for_status()
-        return response
 
-# CORS Preflight Options Endpoint
 @app.options("/hf/v1/chat/completions")
 async def chat_completions_options():
     return Response(
@@ -156,130 +137,102 @@ async def chat_completions_options():
         },
     )
 
-# List Available Models
+
+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}
 
-# Chat Completions Endpoint
+
 @app.post("/hf/v1/chat/completions")
 async def chat_completions(
-    request: ChatRequest, app_secret: str = Depends(verify_app_secret), req: Request = None
+    request: ChatRequest, app_secret: str = Depends(verify_app_secret)
 ):
-    client_ip = req.client.host if req else "unknown"
-    logger.info(f"Received chat completion request from {client_ip} for model: {request.model}")
+    logger.info(f"Received chat completion request for model: {request.model}")
 
-    # 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.")
         raise HTTPException(
             status_code=400,
-            detail=f"Model '{request.model}' is not allowed. Allowed models are: {allowed}",
+            detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(model['id'] for model in ALLOWED_MODELS)}",
         )
+    # Generate a UUID
+    original_uuid = uuid.uuid4()
+    uuid_str = str(original_uuid).replace("-", "")
 
-    logger.info(f"Using model: {request.model}")
-
-    # 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
-    if not request.messages:
-        prompt_text = "make a dog"
-    else:
-        prompt_text = "\n".join(
+    # Using the OpenAI API
+    json_data = {
+        'prompt': "\n".join(
             [
                 f"{'User' if msg.role == 'user' else 'Assistant'}: {msg.content}"
                 for msg in request.messages
             ]
-        )
-
-    json_data = {
-        'prompt': prompt_text,
-        'stream': request.stream,
-        'app_name': APP_NAME,
-        'app_version': APP_VERSION,
-        'tz_name': TZ_NAME,
-        'cid': cid,
+        ),
+        'stream': True,
+        'app_name': 'ChitChat_Edge_Ext',
+        'app_version': '4.26.1',
+        'tz_name': 'Asia/Karachi',
+        'cid': '',
         'model': request.model,
         'search': False,
         'auto_search': False,
         'filter_search_history': False,
         'from': 'chat',
         'group_id': 'default',
-        'chat_models': [],  # As per the sample payload
+        'chat_models': [],
         'files': [],
         'prompt_template': {
-            'key': PROMPT_TEMPLATE_KEY,
+            'key': '',
             'attributes': {
-                'lang': PROMPT_TEMPLATE_LANG,
+                'lang': 'original',
             },
         },
         'tools': {
-            'auto': TOOLS_AUTO,
+            'auto': [
+                'search',
+                'text_to_image',
+                'data_analysis',
+            ],
         },
         'extra_info': {
-            'origin_url': ORIGIN_URL,
-            'origin_title': ORIGIN_TITLE,
+            'origin_url': '',
+            'origin_title': '',
         },
     }
 
-    logger.debug(f"JSON Data Sent to External API: {json.dumps(json_data, indent=2)}")
-
-    try:
-        response = await send_request_with_retry(json_data)
-    except httpx.HTTPStatusError as e:
-        status_code = e.response.status_code
-        if status_code == 429:
-            retry_after = e.response.headers.get("Retry-After", "60")
-            logger.warning(f"Rate limited by Sider AI. Retry after {retry_after} seconds.")
-            raise HTTPException(
-                status_code=429,
-                detail=f"Rate limited by external service. Please retry after {retry_after} seconds."
-            )
-        else:
-            logger.error(f"HTTP error occurred: {e} - Response: {e.response.text}")
-            raise HTTPException(status_code=status_code, detail=str(e))
-    except httpx.RequestError as e:
-        logger.error(f"An error occurred while requesting: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-
     async def generate():
-        async for line in response.aiter_lines():
-            if line and ("[DONE]" not in line):
-                # Assuming the line starts with 'data: ' followed by JSON
-                if line.startswith("data: "):
-                    json_line = line[6:]
-                    if json_line.startswith("{"):
-                        try:
-                            data = json.loads(json_line)
-                            content = data.get("data", {}).get("text", "")
-                            logger.debug(f"Received content: {content}")
-                            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} - Line: {json_line}")
-        # Send the stop signal
-        yield f"data: {json.dumps(create_chat_completion_data('', request.model, 'stop'))}\n\n"
-        yield "data: [DONE]\n\n"
+        async with httpx.AsyncClient() as client:
+            try:
+                async with client.stream('POST', 'https://sider.ai/api/v3/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"
+                    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}")
+                raise HTTPException(status_code=500, detail=str(e))
 
     if request.stream:
-        logger.info("Streaming response initiated.")
+        logger.info("Streaming response")
         return StreamingResponse(generate(), media_type="text/event-stream")
     else:
-        logger.info("Non-streaming response initiated.")
+        logger.info("Non-streaming response")
         full_response = ""
         async for chunk in generate():
             if chunk.startswith("data: ") and not chunk[6:].startswith("[DONE]"):
-                # Parse the JSON part after 'data: '
-                try:
-                    data = json.loads(chunk[6:])
-                    if data["choices"][0]["delta"].get("content"):
-                        full_response += data["choices"][0]["delta"]["content"]
-                except json.JSONDecodeError:
-                    logger.warning(f"Failed to decode JSON from chunk: {chunk}")
-
-        # Final Response Structure
+                # 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",
@@ -295,6 +248,7 @@ async def chat_completions(
             "usage": None,
         }
 
-# Entry Point
+
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)