main.py

import json
import logging
import os
import uuid
from datetime import datetime
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
from starlette.middleware.cors import CORSMiddleware
from starlette.responses import StreamingResponse, Response

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

load_dotenv()
app = FastAPI()
BASE_URL = "https://aichatonlineorg.erweima.ai/aichatonline"
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}',
    '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',
}

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"},
]
# Configure CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allow all sources, you can restrict specific sources if needed
    allow_credentials=True,
    allow_methods=["*"],  # All methods allowed
    allow_headers=["*"],  # Allow all headers
)
security = HTTPBearer()


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,
    }
    
    
def create_chat_completion_data(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
    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": finish_reason,
            }
        ],
        "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


@app.options("/hf/v1/chat/completions")
async def chat_completions_options():
    return Response(
        status_code=200,
        headers={
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Methods": "POST, OPTIONS",
            "Access-Control-Allow-Headers": "Content-Type, Authorization",
        },
    )


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)
):
    logger.info(f"Received chat completion request for model: {request.model}")

    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(model['id'] for model in ALLOWED_MODELS)}",
        )
    # Generate a UUID
    original_uuid = uuid.uuid4()
    uuid_str = str(original_uuid).replace("-", "")

    # Using the OpenAI API
    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',
        'model': request.model,
        'search': False,          # Explicitly turn off search
        'auto_search': False,      # Turn off auto-search
        'from': 'chat',
        'group_id': 'default',
        'prompt_template': {
        'key': 'artifacts',
        'attributes': {
            'lang': 'original',
        },
    },
        'tools': {
            'auto': ['text_to_image', 'data_analysis'],  # Remove 'search' from tools if not needed
    },
        '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:
                    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")
        return StreamingResponse(generate(), media_type="text/event-stream")
    else:
        logger.info("Non-streaming response")
        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",
            "created": int(datetime.now().timestamp()),
            "model": request.model,
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": full_response},
                    "finish_reason": "stop",
                }
            ],
            "usage": None,
        }



if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)