File size: 2,249 Bytes
6acbc8e 70629ed 3164c55 6acbc8e 3164c55 e9b2bc1 3164c55 e9b2bc1 3164c55 6acbc8e e9b2bc1 6acbc8e 3164c55 70629ed |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
from fastapi import APIRouter, HTTPException
from .utils.PoeBot import SendMessage, GenerateImage
from .Schemas import BotRequest
from aiohttp import ClientSession
from pydantic import BaseModel
from ballyregan.models import Protocols, Anonymities
from ballyregan import ProxyFetcher
# Setting the debug mode to True, defaults to False
fetcher = ProxyFetcher()
proxies = fetcher.get(
limit=10,
protocols=[Protocols.HTTP],
anonymities=[Anonymities.ELITE],
)
chat_router = APIRouter(tags=["Chat"])
proxy = ""
class InputData(BaseModel):
input: dict
version: str = "727e49a643e999d602a896c774a0658ffefea21465756a6ce24b7ea4165eba6a"
async def fetch_predictions(data):
async with ClientSession() as session:
for p in proxies:
if proxy != "":
if p != proxy:
continue
try:
async with session.post(
"https://replicate.com/api/predictions",
json=data,
timeout=5,
proxy=str(p),
) as response:
if response.status == 403:
continue
proxy = str(p)
return await response.json(), response.status
except:
pass
async def fetch_result(id):
url = f"https://replicate.com/api/predictions/{id}"
async with ClientSession() as session:
async with session.get(url) as response:
return await response.json(), response.status
@chat_router.post("/predictions")
async def get_predictions(input_data: InputData):
data = {
"input": input_data.input,
"is_training": False,
"create_model": "0",
"stream": False,
"version": input_data.version,
}
try:
predictions, status_code = await fetch_predictions(data)
return predictions, status_code
except Exception as e:
raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}")
@chat_router.post("/chat")
async def chat(req: BotRequest):
return await SendMessage(req)
@chat_router.post("/generate_image")
async def chat(req: BotRequest):
return await GenerateImage(req)
|