aipro / main.py
tianlong12's picture
Create main.py
e100b79 verified
raw
history blame
3.95 kB
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