Spaces:
Runtime error
Runtime error
File size: 4,292 Bytes
007bea3 |
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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
from os import listdir
from uuid import uuid4
from time import sleep, time
from threading import Thread
from json import loads, dumps
from random import getrandbits
from websocket import WebSocketApp
from requests import Session, get, post
class Labs:
def __init__(self) -> None:
self.history: list = []
self.session: Session = Session()
self.user_agent: dict = { "User-Agent": "Ask/2.2.1/334 (iOS; iPhone) isiOSOnMac/false", "X-Client-Name": "Perplexity-iOS" }
self.session.headers.update(self.user_agent)
self._init_session_without_login()
self.t: str = self._get_t()
self.sid: str = self._get_sid()
self.queue: list = []
self.finished: bool = True
assert self._ask_anonymous_user(), "failed to ask anonymous user"
self.ws: WebSocketApp = self._init_websocket()
self.ws_thread: Thread = Thread(target=self.ws.run_forever).start()
self._auth_session()
while not (self.ws.sock and self.ws.sock.connected):
sleep(0.01)
def _init_session_without_login(self) -> None:
self.session.get(url=f"https://www.perplexity.ai/search/{str(uuid4())}")
self.session.headers.update(self.user_agent)
def _auth_session(self) -> None:
self.session.get(url="https://www.perplexity.ai/api/auth/session")
def _get_t(self) -> str:
return format(getrandbits(32), "08x")
def _get_sid(self) -> str:
return loads(self.session.get(
url=f"https://labs-api.perplexity.ai/socket.io/?transport=polling&EIO=4"
).text[1:])["sid"]
def _ask_anonymous_user(self) -> bool:
response = self.session.post(
url=f"https://labs-api.perplexity.ai/socket.io/?EIO=4&transport=polling&t={self.t}&sid={self.sid}",
data="40{\"jwt\":\"anonymous-ask-user\"}"
).text
return response == "OK"
def _get_cookies_str(self) -> str:
cookies = ""
for key, value in self.session.cookies.get_dict().items():
cookies += f"{key}={value}; "
return cookies[:-2]
def _init_websocket(self) -> WebSocketApp:
def on_open(ws: WebSocketApp) -> None:
ws.send("2probe")
ws.send("5")
def on_message(ws: WebSocketApp, message: str) -> None:
if message == "2":
ws.send("3")
elif message.startswith("42"):
message = loads(message[2:])[1]
if "status" not in message:
self.queue.append(message)
elif message["status"] == "completed":
self.finished = True
self.history.append({"role": "assistant", "content": message["output"], "priority": 0})
elif message["status"] == "failed":
self.finished = True
headers: dict = self.user_agent
headers["Cookie"] = self._get_cookies_str()
return WebSocketApp(
url=f"wss://labs-api.perplexity.ai/socket.io/?EIO=4&transport=websocket&sid={self.sid}",
header=headers,
on_open=on_open,
on_message=on_message,
on_error=lambda ws, err: print(f"websocket error: {err}")
)
def _c(self, prompt: str, model: str) -> dict:
assert self.finished, "already searching"
assert model in ["codellama-34b-instruct", "llama-2-7b-chat", "llama-2-13b-chat", "llama-2-70b-chat", "mistral-7b-instruct", "pplx-70b-chat"]
self.finished = False
self.history.append({"role": "user", "content": prompt, "priority": 0})
self.ws.send("42[\"perplexity_playground\",{\"version\":\"2.1\",\"source\":\"default\",\"model\":\"" + model + "\",\"messages\":" + dumps(self.history) + "}]")
def chat(self, prompt: str, model: str = "llama-2-7b-chat") -> dict:
self._c(prompt, model)
while (not self.finished) or (len(self.queue) != 0):
if len(self.queue) > 0:
yield self.queue.pop(0)
def chat_sync(self, prompt: str, model: str = "llama-2-7b-chat") -> dict:
self._c(prompt, model)
while not self.finished:
pass
return self.queue.pop(-1)
def close(self) -> None:
self.ws.close()
|