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 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

# Configure 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
from dotenv import load_dotenv

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 = {
    'accept': '*/*',
    'accept-encoding': 'gzip, deflate, br, zstd',
    'accept-language': 'en-US,en;q=0.9',
    'authorization': f'Bearer {ACCESS_TOKEN}',
    '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/130.0.0.0 Safari/537.36',
}

# 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"},
]

# Configure CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # ⚠️ IMPORTANT: Restrict this to specific origins in production for security
    allow_credentials=True,
    allow_methods=["*"],  # Allow all HTTP methods
    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 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:
        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


# CORS Preflight Options Endpoint
@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",
        },
    )


# 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(
    request: ChatRequest, app_secret: str = Depends(verify_app_secret)
):
    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}",
        )

    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(
            [
                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,
        'model': request.model,
        'search': False,
        'auto_search': False,
        'filter_search_history': False,
        'from': 'chat',
        'group_id': 'default',
        'chat_models': [],  # As per the sample payload
        'files': [],
        'prompt_template': {
            'key': PROMPT_TEMPLATE_KEY,
            'attributes': {
                'lang': PROMPT_TEMPLATE_LANG,
            },
        },
        'tools': {
            'auto': TOOLS_AUTO,
        },
        'extra_info': {
            'origin_url': ORIGIN_URL,
            'origin_title': ORIGIN_TITLE,
        },
    }

    logger.debug(f"JSON Data Sent to External API: {json.dumps(json_data, indent=2)}")

    async def generate():
        async with httpx.AsyncClient() as client:
            try:
                async with client.stream(
                    'POST',
                    'https://sider.ai/api/v3/completion/text',  # Updated endpoint
                    headers=OUTGOING_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):
                            # 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"
            except httpx.HTTPStatusError as e:
                logger.error(f"HTTP error occurred: {e} - Response: {e.response.text}")
                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.")
        return StreamingResponse(generate(), media_type="text/event-stream")
    else:
        logger.info("Non-streaming response initiated.")
        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
        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,
        }


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