poe router

App/Chat/PoeChatrouter.py

@@ -12,49 +12,53 @@ from ballyregan import ProxyFetcher
 
 chat_router = APIRouter(tags=["Chat"])
 proxy = ""
-proxies = [
-    "http://51.89.14.70:80",
-    "http://52.151.210.204:9002",
-    "http://38.180.36.19:80",
-    "http://38.54.79.150:80",
-    "https://80.91.26.137:3128",
-    "http://82.223.102.92:9443",
-    "http://189.240.60.166:9090",
-    "https://189.240.60.168:9090",
-    "http://189.240.60.171:9090",
-]
+proxies = []
 
 
 class InputData(BaseModel):
     input: dict
     version: str = "727e49a643e999d602a896c774a0658ffefea21465756a6ce24b7ea4165eba6a"
     proxies: list[str] = []
+    is_proxied: bool = False
 
 
-async def fetch_predictions(data):
+async def fetch_predictions(data, is_proxied=False):
     global proxy, proxies
-    proxy_set = proxy != ""
-    async with ClientSession() as session:
-        for p in proxies:
-            if proxy_set:
-                if p != proxy:
-                    continue
-            try:
-                async with session.post(
-                    "https://replicate.com/api/predictions",
-                    json=data,
-                    timeout=5,
-                ) as response:
-                    if str(response.status).startswith("4"):
+    if is_proxied:
+        proxy_set = proxy != ""
+        async with ClientSession() as session:
+            for p in proxies:
+                if proxy_set:
+                    if p != proxy:
                         continue
-                    proxy = str(p)
-                    temp = await response.json()
-                    print(temp)
-                    return temp
-            except Exception as e:
-                print("Error fetching", e)
-                pass
-        proxy = ""
+                try:
+                    async with session.post(
+                        "https://replicate.com/api/predictions",
+                        json=data,
+                        timeout=5,
+                    ) as response:
+                        if str(response.status).startswith("4"):
+                            continue
+                        proxy = str(p)
+                        temp = await response.json()
+                        print(temp)
+                        return temp
+                except Exception as e:
+                    print("Error fetching", e)
+                    pass
+            proxy = ""
+    else:
+        try:
+            async with session.post(
+                "https://replicate.com/api/predictions",
+                json=data,
+                timeout=5,
+            ) as response:
+                temp = await response.json()
+                return temp
+        except Exception as e:
+            print("Error fetching", e)
+            pass
 
 
 @chat_router.post("/predictions")
@@ -62,6 +66,13 @@ async def get_predictions(input_data: InputData):
     global proxies
     if input_data.proxies != []:
         proxies = input_data.proxies
+    else:
+
+        proxies = [
+            "http://51.89.14.70:80",
+            "http://52.151.210.204:9002",
+            "http://38.180.36.19:80",
+        ]
     data = {
         "input": input_data.input,
         "is_training": False,
@@ -70,7 +81,7 @@ async def get_predictions(input_data: InputData):
         "version": input_data.version,
     }
     try:
-        predictions = await fetch_predictions(data)
+        predictions = await fetch_predictions(data, input_data.is_proxied)
         return predictions
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}")