Update app.py

app.py

@@ -1,46 +1,64 @@
 import os
-import asyncio
-import uuid
 import json
+import uuid
+import asyncio
+import aiohttp
+from aiohttp import web
+from datetime import datetime
+import socketio
+
+# 从环境变量中获取API密钥
+API_KEY = os.getenv('PPLX_KEY')
+PROXY = os.getenv('HTTP_PROXY')  # 新增：从环境变量获取代理设置
+
+sio = socketio.AsyncClient()
+
+# 配置选项
+opts = {
+    "auth": {
+        "jwt": "anonymous-ask-user",
+    },
+    "reconnection": False,
+    "transports": ["websocket"],
+    "path": "/socket.io",
+    "hostname": "www.perplexity.ai",
+    "secure": True,
+    "port": "443",
+    "extraHeaders": {
+        "Cookie": os.getenv('PPLX_COOKIE'),
+        "User-Agent": os.getenv('USER_AGENT'),
+        "Accept": "*/*",
+        "priority": "u=1, i",
+        "Referer": "https://www.perplexity.ai/",
+    }
+}
 
-from fastapi import FastAPI, HTTPException, Request
-from fastapi.responses import StreamingResponse
-from socketio import AsyncClient
-
-# 从环境变量中获取 API 密钥
-API_KEY = os.environ.get("PPLX_KEY")
-PPLX_COOKIE = os.environ.get("PPLX_COOKIE")
-USER_AGENT = os.environ.get("USER_AGENT")
-# 从环境变量中获取代理地址
-PROXY_URL = os.environ.get("PROXY_URL")
-
-app = FastAPI()
-
-# 添加中间件来验证 API 密钥
-async def validate_api_key(request: Request, call_next):
-    api_key = request.headers.get("x-api-key")
+# API密钥验证中间件
+@web.middleware
+async def validate_api_key(request, handler):
+    api_key = request.headers.get('x-api-key')
     if api_key != API_KEY:
         log_request(request, 401)
-        raise HTTPException(status_code=401, detail="Invalid API key")
-    response = await call_next(request)
-    return response
-
-app.middleware("http")(validate_api_key)
-
+        return web.json_response({"error": "Invalid API key"}, status=401)
+    return await handler(request)
 
 # 日志记录函数
-def log_request(request: Request, status_code: int):
-    timestamp = datetime.datetime.now().isoformat()
-    ip = request.client.host
-    route = request.url.path
-    print(f"{timestamp} - {ip} - {route} - {status_code}")
-
+def log_request(request, status):
+    timestamp = datetime.now().isoformat()
+    ip = request.remote
+    route = request.path
+    print(f"{timestamp} - {ip} - {route} - {status}")
+
+# 创建事件流工具函数
+def create_event(event, data):
+    if isinstance(data, dict):
+        data = json.dumps(data)
+    return f"event: {event}\ndata: {data}\n\n"
 
-# 根路由处理
-@app.get("/")
-async def root(request: Request):
+# 路由处理函数
+async def handle_root(request):
     log_request(request, 200)
-    return {
+    return web.json_response({
         "message": "Welcome to the Perplexity AI Proxy API",
         "endpoints": {
             "/ai/v1/messages": {
@@ -48,183 +66,183 @@ async def root(request: Request):
                 "description": "Send a message to the AI",
                 "headers": {
                     "x-api-key": "Your API key (required)",
-                    "Content-Type": "application/json",
+                    "Content-Type": "application/json"
                 },
                 "body": {
                     "messages": "Array of message objects",
                     "stream": "Boolean (true for streaming response)",
-                    # 其他可能的参数...
-                },
+                }
             }
-        },
-    }
+        }
+    })
 
-
-#  AI 消息处理
-@app.post("/ai/v1/messages")
-async def handle_ai_message(request: Request):
+async def handle_messages(request):
     try:
         json_body = await request.json()
-        if not json_body.get("stream"):
+        if not json_body.get('stream', True):
             log_request(request, 200)
-            return {
+            return web.json_response({
                 "id": str(uuid.uuid4()),
                 "content": [
                     {"text": "Please turn on streaming."},
-                    {"id": "string", "name": "string", "input": {}},
+                    {"id": "string", "name": "string", "input": {}}
                 ],
                 "model": "string",
                 "stop_reason": "end_turn",
                 "stop_sequence": "string",
-                "usage": {"input_tokens": 0, "output_tokens": 0},
-            }
-        elif json_body.get("stream"):
-            async def event_stream(json_body):
-                # 计算用户消息长度
-                user_message = [{"question": "", "answer": ""}]
-                last_update = True
-                if json_body.get("system"):
-                    # 把系统消息加入 messages 的首条
-                    json_body["messages"].insert(0, {"role": "system", "content": json_body.get("system")})
-                for msg in json_body.get("messages", []):
-                    if msg["role"] in ("system", "user"):
-                        if last_update:
-                            user_message[-1]["question"] += msg["content"] + "\n"
-                        elif not user_message[-1]["question"]:
-                            user_message[-1]["question"] += msg["content"] + "\n"
-                        else:
-                            user_message.append({"question": msg["content"] + "\n", "answer": ""})
-                        last_update = True
-                    elif msg["role"] == "assistant":
-                        if not last_update:
-                            user_message[-1]["answer"] += msg["content"] + "\n"
-                        elif not user_message[-1]["answer"]:
-                            user_message[-1]["answer"] += msg["content"] + "\n"
-                        else:
-                            user_message.append({"question": "", "answer": msg["content"] + "\n"})
-                        last_update = False
-
-                # user message to plaintext
-                previous_messages = "\n\n".join([msg["content"] for msg in json_body.get("messages", [])])
-                msgid = str(uuid.uuid4())
-                # send message start
-                yield create_event(
-                    "message_start",
-                    {
-                        "type": "message_start",
-                        "message": {
-                            "id": msgid,
-                            "type": "message",
-                            "role": "assistant",
-                            "content": [],
-                            "model": "claude-3-opus-20240229",
-                            "stop_reason": None,
-                            "stop_sequence": None,
-                            "usage": {"input_tokens": 8, "output_tokens": 1},
-                        },
-                    },
-                )
-                yield create_event(
-                    "content_block_start",
-                    {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
-                )
-                yield create_event("ping", {"type": "ping"})
-
-                # 设置代理
-                engineio_kwargs = {}
-                if PROXY_URL:
-                    engineio_kwargs = {"http_proxy": PROXY_URL, "https_proxy": PROXY_URL}
-
-                # proxy response
-                async with AsyncClient(logger=True, engineio_logger=True, **engineio_kwargs) as socket:
+                "usage": {"input_tokens": 0, "output_tokens": 0}
+            })
+        else:
+            response = web.StreamResponse()
+            response.headers['Content-Type'] = 'text/event-stream;charset=utf-8'
+            await response.prepare(request)
+
+            # 处理消息
+            user_message = [{"question": "", "answer": ""}]
+            last_update = True
+            if 'system' in json_body:
+                json_body['messages'].insert(0, {"role": "system", "content": json_body['system']})
+
+            for msg in json_body['messages']:
+                if msg['role'] in ['system', 'user']:
+                    if last_update:
+                        user_message[-1]['question'] += msg['content'] + "\n"
+                    elif user_message[-1]['question'] == "":
+                        user_message[-1]['question'] += msg['content'] + "\n"
+                    else:
+                        user_message.append({"question": msg['content'] + "\n", "answer": ""})
+                    last_update = True
+                elif msg['role'] == 'assistant':
+                    if not last_update:
+                        user_message[-1]['answer'] += msg['content'] + "\n"
+                    elif user_message[-1]['answer'] == "":
+                        user_message[-1]['answer'] += msg['content'] + "\n"
+                    else:
+                        user_message.append({"question": "", "answer": msg['content'] + "\n"})
+                    last_update = False
+
+            previous_messages = "\n\n".join(msg['content'] for msg in json_body['messages'])
+            msg_id = str(uuid.uuid4())
+
+            await response.write(create_event("message_start", {
+                "type": "message_start",
+                "message": {
+                    "id": msg_id,
+                    "type": "message",
+                    "role": "assistant",
+                    "content": [],
+                    "model": "claude-3-opus-20240229",
+                    "stop_reason": None,
+                    "stop_sequence": None,
+                    "usage": {"input_tokens": 8, "output_tokens": 1},
+                }
+            }).encode())
+
+            await response.write(create_event("content_block_start", {
+                "type": "content_block_start",
+                "index": 0,
+                "content_block": {"type": "text", "text": ""}
+            }).encode())
+
+            await response.write(create_event("ping", {"type": "ping"}).encode())
+
+            # 设置代理
+            connector = None
+            if PROXY:
+                connector = aiohttp.TCPConnector(ssl=False, proxy=PROXY)
+
+            async with aiohttp.ClientSession(connector=connector) as session:
+                await sio.connect('wss://www.perplexity.ai/', **opts, http_session=session)
+
+                @sio.on('connect')
+                async def on_connect():
+                    print(" > [Connected]")
                     try:
-                        await socket.connect(
-                            "https://www.perplexity.ai/",
-                            headers={
-                                "Cookie": PPLX_COOKIE,
-                                "User-Agent": USER_AGENT,
-                                "Accept": "*/*",
-                                "priority": "u=1, i",
-                                "Referer": "https://www.perplexity.ai/",
-                            },
-                            transports=["websocket"],
-                        )
-                        print(" > [Connected]")
-                        await socket.emit(
-                            "perplexity_ask",
-                            previous_messages,
-                            {
-                                "version": "2.9",
-                                "source": "default",
-                                "attachments": [],
-                                "language": "en-GB",
-                                "timezone": "Europe/London",
-                                "search_focus": "writing",
-                                "frontend_uuid": str(uuid.uuid4()),
-                                "mode": "concise",
-                                "is_related_query": False,
-                                "is_default_related_query": False,
-                                "visitor_id": str(uuid.uuid4()),
-                                "frontend_context_uuid": str(uuid.uuid4()),
-                                "prompt_source": "user",
-                                "query_source": "home",
-                            },
-                        )
-                        response = await socket.wait()
+                        response = await sio.call('perplexity_ask', previous_messages, {
+                            "version": "2.9",
+                            "source": "default",
+                            "attachments": [],
+                            "language": "en-GB",
+                            "timezone": "Europe/London",
+                            "search_focus": "writing",
+                            "frontend_uuid": str(uuid.uuid4()),
+                            "mode": "concise",
+                            "is_related_query": False,
+                            "is_default_related_query": False,
+                            "visitor_id": str(uuid.uuid4()),
+                            "frontend_context_uuid": str(uuid.uuid4()),
+                            "prompt_source": "user",
+                            "query_source": "home"
+                        })
                         print(response)
-                        yield create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
-                        yield create_event(
-                            "message_delta",
-                            {
-                                "type": "message_delta",
-                                "delta": {"stop_reason": "end_turn", "stop_sequence": None},
-                                "usage": {"output_tokens": 12},
-                            },
-                        )
-                        yield create_event("message_stop", {"type": "message_stop"})
+                        await response.write(create_event("content_block_stop", {"type": "content_block_stop", "index": 0}).encode())
+                        await response.write(create_event("message_delta", {
+                            "type": "message_delta",
+                            "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+                            "usage": {"output_tokens": 12}
+                        }).encode())
+                        await response.write(create_event("message_stop", {"type": "message_stop"}).encode())
+                        await response.write_eof()
                         log_request(request, 200)
                     except Exception as e:
-                        print(e)
+                        if str(e) != "socket has been disconnected":
+                            print(e)
                         log_request(request, 500)
-                    finally:
-                        await socket.disconnect()
-
-                    @socket.on("query_progress")
-                    async def on_query_progress(data):
-                        if data.get("text"):
-                            text = json.loads(data["text"])
-                            chunk = text["chunks"][-1] if text.get("chunks") else None
-                            if chunk:
-                                yield create_event(
-                                    "content_block_delta",
-                                    {
-                                        "type": "content_block_delta",
-                                        "index": 0,
-                                        "delta": {"type": "text_delta", "text": chunk},
-                                    },
-                                )
-
-            return StreamingResponse(event_stream(json_body), media_type="text/event-stream;charset=utf-8")
-        else:
-            raise HTTPException(status_code=400, detail="Invalid request")
+
+                @sio.event
+                async def query_progress(data):
+                    if 'text' in data:
+                        text = json.loads(data['text'])
+                        chunk = text['chunks'][-1] if text['chunks'] else None
+                        if chunk:
+                            chunk_json = json.dumps({
+                                "type": "content_block_delta",
+                                "index": 0,
+                                "delta": {"type": "text_delta", "text": chunk},
+                            })
+                            await response.write(create_event("content_block_delta", chunk_json).encode())
+
+                @sio.event
+                async def disconnect():
+                    print(" > [Disconnected]")
+
+                @sio.event
+                async def connect_error(error):
+                    error_message = "Failed to connect to Perplexity.ai 连接到Perplexity失败\nPlease refer to the log for more information 请查看日志以获取更多信息"
+                    await response.write(create_event("content_block_delta", json.dumps({
+                        "type": "content_block_delta",
+                        "index": 0,
+                        "delta": {"type": "text_delta", "text": error_message},
+                    })).encode())
+                    await response.write(create_event("content_block_stop", {"type": "content_block_stop", "index": 0}).encode())
+                    await response.write(create_event("message_delta", {
+                        "type": "message_delta",
+                        "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+                        "usage": {"output_tokens": 12}
+                    }).encode())
+                    await response.write(create_event("message_stop", {"type": "message_stop"}).encode())
+                    await response.write_eof()
+                    print(error)
+                    log_request(request, 500)
+
+                await sio.wait()
+
+            return response
+
     except Exception as e:
         print(e)
         log_request(request, 400)
-        raise HTTPException(status_code=400, detail=str(e))
-
-
-# eventStream util
-def create_event(event: str, data: dict):
-    # if data is object, stringify it
-    if isinstance(data, dict):
-        data = json.dumps(data)
-    return f"event: {event}\ndata: {data}\n\n"
-
-
-if __name__ == "__main__":
-    import uvicorn
-
+        return web.json_response({"error": str(e)}, status=400)
+
+# 创建应用
+app = web.Application(middlewares=[validate_api_key])
+app.router.add_get('/', handle_root)
+app.router.add_post('/ai/v1/messages', handle_messages)
+
+# 启动服务器
+if __name__ == '__main__':
+    port = int(os.getenv('PORT', 8081))
+    web.run_app(app, port=port)
+    print(f"Perplexity proxy listening on port {port}")
     if not API_KEY:
-        print("Warning: PPLX_KEY environment variable is not set. API key validation will fail.")
-    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 8081)))
-
+        print("警告：PPLX_KEY 环境变量未设置。API 密钥验证将失败。")