Update main.py

main.py

@@ -14,28 +14,17 @@ from pydantic import BaseModel
 from starlette.middleware.cors import CORSMiddleware
 from starlette.responses import StreamingResponse, Response
 
-# ==============================
-# Configuration and Setup
-# ==============================
-
-# Configure logging
 logging.basicConfig(
-    level=logging.DEBUG,  # Set to DEBUG for detailed logs
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 logger = logging.getLogger(__name__)
 
-# Load environment variables from .env file
 load_dotenv()
-
-# Initialize FastAPI app
 app = FastAPI()
-
-# Constants and configurations
 BASE_URL = "https://aichatonlineorg.erweima.ai/aichatonline"
-APP_SECRET = os.getenv("APP_SECRET", "666")
-ACCESS_TOKEN = os.getenv("SD_ACCESS_TOKEN", "")
-headers = {    
+APP_SECRET = os.getenv("APP_SECRET","666")
+ACCESS_TOKEN = os.getenv("SD_ACCESS_TOKEN","")
+headers = {
     'accept': '*/*',
     'accept-language': 'en-US,en;q=0.9',
     'authorization': f'Bearer {ACCESS_TOKEN}',
@@ -43,13 +32,14 @@ headers = {
     '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',
-    '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', 
+    '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"},
@@ -59,38 +49,29 @@ ALLOWED_MODELS = [
     {"id": "o1-mini", "name": "o1-mini"},
     {"id": "gpt-4o-mini", "name": "gpt-4o-mini"},
 ]
-
-# Configure CORS middleware
+# Configure CORS
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Allow all origins; restrict if necessary
+    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 configuration
 security = HTTPBearer()
 
-# ==============================
-# Pydantic Models
-# ==============================
 
 class Message(BaseModel):
     role: str
     content: str
 
+
 class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
     stream: Optional[bool] = False
 
-# ==============================
-# Helper Functions
-# ==============================
 
-def simulate_data(content: str, model: str) -> Dict[str, Any]:
-    """Simulate chunked response data."""
+def simulate_data(content, model):
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion.chunk",
@@ -106,8 +87,8 @@ def simulate_data(content: str, model: str) -> Dict[str, Any]:
         "usage": None,
     }
 
-def stop_data(content: str, model: str) -> Dict[str, Any]:
-    """Simulate the final chunk indicating the end of the response."""
+
+def stop_data(content, model):
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion.chunk",
@@ -122,9 +103,9 @@ def stop_data(content: str, model: str) -> Dict[str, Any]:
         ],
         "usage": None,
     }
-
+    
+    
 def create_chat_completion_data(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
-    """Create a structured chat completion data chunk."""
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion.chunk",
@@ -140,23 +121,15 @@ def create_chat_completion_data(content: str, model: str, finish_reason: Optiona
         "usage": None,
     }
 
-def verify_app_secret(credentials: HTTPAuthorizationCredentials = Depends(security)) -> str:
-    """Verify the provided APP_SECRET."""
+
+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
 
-def replace_escaped_newlines(input_string: str) -> str:
-    """Replace escaped newline characters with actual newlines."""
-    return input_string.replace("\\n", "\n")
-
-# ==============================
-# API Endpoints
-# ==============================
 
 @app.options("/hf/v1/chat/completions")
 async def chat_completions_options():
-    """Handle CORS preflight requests."""
     return Response(
         status_code=200,
         headers={
@@ -166,31 +139,32 @@ 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():
-    """List all allowed 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)
 ):
-    """Handle chat completion requests."""
     logger.info(f"Received chat completion request for model: {request.model}")
 
-    # Validate model
-    allowed_model_ids = [model['id'] for model in ALLOWED_MODELS]
-    if request.model not in allowed_model_ids:
+    if request.model not in [model['id'] for model in ALLOWED_MODELS]:
         raise HTTPException(
             status_code=400,
-            detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(allowed_model_ids)}",
+            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("-", "")
 
-    # Construct the payload to send to the external API
+    # Using the OpenAI API
     json_data = {
         'prompt': "\n".join(
             [
@@ -200,12 +174,12 @@ async def chat_completions(
         ),
         'stream': True,
         'app_name': 'ChitChat_Edge_Ext',
-        'app_version': '4.28.0',
+        'app_version': '4.26.1',
         'tz_name': 'Asia/Karachi',
-        'cid': 'C092SEMXM9BJ',
+        'cid': '',
         'model': request.model,
-        'search': False,  # Disable search
-        'auto_search': False,  # Disable auto_search
+        'search': False,
+        'auto_search': False,
         'filter_search_history': False,
         'from': 'chat',
         'group_id': 'default',
@@ -219,7 +193,7 @@ async def chat_completions(
         },
         'tools': {
             'auto': [
-                'search',  # Re-add search to maintain API expectations
+                'search',
                 'text_to_image',
                 'data_analysis',
             ],
@@ -230,53 +204,17 @@ async def chat_completions(
         },
     }
 
-    # Define the asynchronous generator for streaming responses
     async def generate():
         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:
+                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:
-                            logger.debug(f"Raw line received: {line}")  # Log raw line
                         if line and ("[DONE]" not in line):
-                            try:
-                                # Remove 'data: ' prefix if present
-                                if line.startswith("data: "):
-                                    line_content = line[6:]
-                                else:
-                                    line_content = line
-                                
-                                # Log the content before parsing
-                                logger.debug(f"Line content to parse: {line_content}")
-                                
-                                # Check if the line is not empty after stripping
-                                if not line_content.strip():
-                                    logger.warning("Received an empty line, skipping.")
-                                    continue
-                                
-                                # Attempt to parse JSON
-                                parsed_json = json.loads(line_content)
-                                
-                                # Ensure 'data' key exists
-                                if "data" not in parsed_json:
-                                    logger.error(f"'data' key not found in the response: {parsed_json}")
-                                    continue
-                                
-                                content_data = parsed_json["data"]
-                                
-                                # Extract text content if available
-                                text_content = content_data.get("text", "")
-                                
-                                # Yield the formatted data
-                                yield f"data: {json.dumps(create_chat_completion_data(text_content, request.model))}\n\n"
-                            except json.JSONDecodeError as e:
-                                logger.error(f"JSON decode error: {e} | Line: {line_content}")
-                                continue
-                        else:
-                            if line and "[DONE]" in line:
-                                yield f"data: {json.dumps(create_chat_completion_data('', request.model, 'stop'))}\n\n"
-                                yield "data: [DONE]\n\n"
+                            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))
@@ -292,14 +230,11 @@ async def chat_completions(
         full_response = ""
         async for chunk in generate():
             if chunk.startswith("data: ") and not chunk[6:].startswith("[DONE]"):
-                try:
-                    data = json.loads(chunk[6:])
-                    if data["choices"][0]["delta"].get("content"):
-                        full_response += data["choices"][0]["delta"]["content"]
-                except json.JSONDecodeError as e:
-                    logger.error(f"JSON decode error in non-streaming response: {e}")
-                    continue
-
+                # 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",
@@ -315,9 +250,7 @@ async def chat_completions(
             "usage": None,
         }
 
-# ==============================
-# Entry Point
-# ==============================
+
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)