Create main.py

main.py

@@ -0,0 +1,128 @@
+import os
+import time
+import random
+import json
+import asyncio
+import requests
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+from typing import List, Optional, Union
+
+app = FastAPI()
+
+class ChatCompletionMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[ChatCompletionMessage]
+    temperature: Optional[float] = 1.0
+    max_tokens: Optional[int] = None
+    stream: Optional[bool] = False
+
+class ChatCompletionResponse(BaseModel):
+    id: str
+    object: str
+    created: int
+    model: str
+    choices: List[dict]
+    usage: dict
+
+def generate_random_ip():
+    return f"{random.randint(1,255)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(0,255)}"
+
+async def fetch_response(messages: List[ChatCompletionMessage], model: str):
+    your_api_url = "https://chatpro.ai-pro.org/api/ask/openAI"
+    headers = {
+        "content-type": "application/json",
+        "X-Forwarded-For": generate_random_ip(),
+        "origin": "https://chatpro.ai-pro.org",
+        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36"
+    }
+    
+    conversation = "\n".join([f"{msg.role}: {msg.content}" for msg in messages])
+    conversation += "\n请关注并回复user最近的消息并避免总结对话历史的回答"
+    
+    data = {
+        "text": conversation,
+        "endpoint": "openAI",
+        "model": model
+    }
+
+    response = requests.post(your_api_url, headers=headers, json=data)
+
+    if response.status_code != 200:
+        raise HTTPException(status_code=response.status_code, detail="Error from upstream API")
+
+    return response.json()
+
+async def stream_response(content: str):
+    chunk_size = len(content)  # 将整个内容作为一个块发送
+    chat_id = f"chatcmpl-{os.urandom(12).hex()}"
+    
+    # 发送开始的块
+    yield f"data: {json.dumps({
+        'id': chat_id,
+        'object': 'chat.completion.chunk',
+        'created': int(time.time()),
+        'model': 'gpt-3.5-turbo-0613',
+        'choices': [{
+            'index': 0,
+            'delta': {
+                'content': content
+            },
+            'finish_reason': None
+        }]
+    })}\n\n"
+
+    # 发送结束的块
+    yield f"data: {json.dumps({
+        'id': chat_id,
+        'object': 'chat.completion.chunk',
+        'created': int(time.time()),
+        'model': 'gpt-3.5-turbo-0613',
+        'choices': [{
+            'index': 0,
+            'delta': {},
+            'finish_reason': 'stop'
+        }]
+    })}\n\n"
+
+    yield 'data: [DONE]\n\n'
+
+@app.post("/hf/v1/chat/completions")
+async def chat_completions(request: Request):
+    body = await request.json()
+    chat_request = ChatCompletionRequest(**body)
+
+    api_response = await fetch_response(chat_request.messages, chat_request.model)
+
+    content = api_response.get("response", "")
+
+    if chat_request.stream:
+        return StreamingResponse(stream_response(content), media_type="text/event-stream")
+    else:
+        openai_response = ChatCompletionResponse(
+            id="chatcmpl-" + os.urandom(12).hex(),
+            object="chat.completion",
+            created=int(time.time()),
+            model=chat_request.model,
+            choices=[
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": content
+                    },
+                    "finish_reason": "stop"
+                }
+            ],
+            usage={
+                "prompt_tokens": sum(len(msg.content) for msg in chat_request.messages),
+                "completion_tokens": len(content),
+                "total_tokens": sum(len(msg.content) for msg in chat_request.messages) + len(content)
+            }
+        )
+        return openai_response