Update main.py

main.py

@@ -9,10 +9,8 @@ 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
+from starlette.responses import StreamingResponse
 
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
@@ -21,41 +19,17 @@ 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","")
+APP_SECRET = os.getenv("APP_SECRET", "666")
+ACCESS_TOKEN = os.getenv("SD_ACCESS_TOKEN", "")
 headers = {
-    'accept': '*/*',
-    '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}',
-    '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/129.0.0.0 Safari/537.36 Edg/129.0.0.0',
 }
 
 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"},
+    "claude-3.5-sonnet", "sider", "gpt-4o-mini", "claude-3-haiku", "claude-3.5-haiku",
+    "gemini-1.5-flash", "llama-3", "gpt-4o", "gemini-1.5-pro", "llama-3.1-405b"
 ]
-# 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):
@@ -69,40 +43,6 @@ class ChatRequest(BaseModel):
     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()}",
@@ -120,109 +60,36 @@ def create_chat_completion_data(content: str, model: str, finish_reason: Optiona
     }
 
 
-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)
-):
-    # Log the model requested by the client
+async def chat_completions(request: ChatRequest):
     logger.info(f"Received chat completion request for model: {request.model}")
 
-    # Log allowed models for verification purposes
-    logger.info(f"Allowed models: {[model['id'] for model in ALLOWED_MODELS]}")
-
-    if request.model not in [model['id'] for model in ALLOWED_MODELS]:
+    if request.model not 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: {', '.join(model['id'] for model in ALLOWED_MODELS)}",
+            detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(ALLOWED_MODELS)}",
         )
 
-    # Log the JSON payload for the request to the external API
-    logger.info(f"Sending request to external API with model: {request.model}")
-    
-    # Generate a UUID
-    original_uuid = uuid.uuid4()
-    uuid_str = str(original_uuid).replace("-", "")
-
-    # Prepare json_data for the external request
     json_data = {
         'prompt': "\n".join(
-            [
-                f"{'User' if msg.role == 'user' else 'Assistant'}: {msg.content}"
-                for msg in request.messages
-            ]
+            [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.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': [],
-        'files': [],
-        'prompt_template': {
-            'key': '',
-            'attributes': {
-                'lang': 'original',
-            },
-        },
-        'tools': {
-            'auto': [
-                'search',
-                'text_to_image',
-                'data_analysis',
-            ],
-        },
-        'extra_info': {
-            'origin_url': '',
-            'origin_title': '',
-        },
     }
+    
+    logger.info(f"Sending request to external API with data: {json_data}")
 
     async def generate():
         async with httpx.AsyncClient() as client:
             try:
-                # Log request details before making the API call
-                logger.info(f"External API request json_data: {json.dumps(json_data, indent=2)}")
-                
                 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(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:
@@ -240,14 +107,11 @@ async def chat_completions(
         full_response = ""
         async for chunk in generate():
             if chunk.startswith("data: ") and not chunk[6:].startswith("[DONE]"):
-                # Parse the chunk data and log it for debugging
                 data = json.loads(chunk[6:])
-                logger.info(f"Chunk data received: {data}")
-                
                 if data["choices"][0]["delta"].get("content"):
                     full_response += data["choices"][0]["delta"]["content"]
         
-        logger.info(f"Full response generated for model {request.model}: {full_response}")
+        logger.info(f"Full response generated: {full_response}")
         return {
             "id": f"chatcmpl-{uuid.uuid4()}",
             "object": "chat.completion",
@@ -262,3 +126,7 @@ async def chat_completions(
             ],
             "usage": None,
         }
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)