perplexity/client.py

import re
import sys
import json
import random
import mimetypes
from uuid import uuid4
from curl_cffi import CurlMime, requests

from .config import (
    DEFAULT_HEADERS,
    ENDPOINT_AUTH_SESSION,
    ENDPOINT_AUTH_SIGNIN,
    ENDPOINT_SSE_ASK,
    ENDPOINT_UPLOAD_URL,
    MODEL_MAPPINGS,
)
from .emailnator import Emailnator
from .exceptions import (
    AuthenticationError,
    FileUploadError,
    InvalidModeError,
    InvalidModelError,
    InvalidSourceError,
    QueryLimitExceededError,
)
from .utils import parse_nested_json_response, validate_search_params, validate_query_limits
from .logger import get_logger

logger = get_logger("client")


class Client:
    def __init__(self, cookies={}):
        self.session = requests.Session(
            headers=DEFAULT_HEADERS.copy(),
            cookies=cookies,
            impersonate="chrome",
        )
        self.own = bool(cookies)
        self.copilot = 0 if not cookies else 5
        self.file_upload = 0 if not cookies else 10
        self.signin_regex = re.compile(
            r'"(https://www\.perplexity\.ai/api/auth/callback/email\?callbackUrl=.*?)"'
        )
        self.timestamp = format(random.getrandbits(32), "08x")
        self.session.get(ENDPOINT_AUTH_SESSION, timeout=15)

    def _get_csrf(self):
        csrf = self.session.cookies.get_dict().get("next-auth.csrf-token", "")
        if csrf:
            return csrf.split("%")[0]
        try:
            resp = self.session.get("https://www.perplexity.ai/api/auth/csrf", timeout=15)
            return resp.json().get("csrfToken", "")
        except Exception:
            return ""

    def create_account(self, cookies):
        while True:
            try:
                emailnator_cli = Emailnator(cookies)
                csrf = self._get_csrf()
                resp = self.session.post(
                    ENDPOINT_AUTH_SIGNIN,
                    json={
                        "email": emailnator_cli.email,
                        "csrfToken": csrf,
                        "callbackUrl": "https://www.perplexity.ai/",
                    },
                    timeout=15,
                )
                if resp.ok:
                    new_msgs = emailnator_cli.reload(
                        wait_for=lambda x: x["subject"] == "Sign in to Perplexity",
                        timeout=20,
                    )
                    if new_msgs:
                        break
                else:
                    logger.error(f"Account creation failed: {resp.status_code} {resp.text[:200]}")
            except Exception as e:
                logger.warning(f"Account creation attempt failed: {e}")

        msg = emailnator_cli.get(func=lambda x: x["subject"] == "Sign in to Perplexity")
        new_account_link = self.signin_regex.search(emailnator_cli.open(msg["messageID"])).group(1)
        self.session.get(new_account_link, timeout=15)
        self.copilot = 5
        self.file_upload = 10
        return True

    def search(
        self,
        query,
        mode="auto",
        model=None,
        sources=["web"],
        files={},
        stream=False,
        language="en-US",
        follow_up=None,
        incognito=False,
    ):
        try:
            validate_search_params(mode, model, sources, self.own)
            validate_query_limits(self.copilot, self.file_upload, mode, len(files))
        except Exception as e:
            raise

        self.copilot = (
            self.copilot - 1 if mode in ["pro", "reasoning", "deep research"] else self.copilot
        )
        self.file_upload = self.file_upload - len(files) if files else self.file_upload

        uploaded_files = []
        for filename, file in files.items():
            file_type = mimetypes.guess_type(filename)[0]
            file_upload_info = (
                self.session.post(
                    ENDPOINT_UPLOAD_URL,
                    params={"version": "2.18", "source": "default"},
                    json={
                        "content_type": file_type,
                        "file_size": sys.getsizeof(file),
                        "filename": filename,
                        "force_image": False,
                        "source": "default",
                    },
                    timeout=30,
                )
            ).json()

            mp = CurlMime()
            for key, value in file_upload_info["fields"].items():
                mp.addpart(name=key, data=value)
            mp.addpart(
                name="file",
                content_type=file_type,
                filename=filename,
                data=file,
            )

            upload_resp = self.session.post(file_upload_info["s3_bucket_url"], multipart=mp, timeout=30)

            if not upload_resp.ok:
                raise FileUploadError(f"File upload failed: {upload_resp.status_code}")

            if "image/upload" in file_upload_info["s3_object_url"]:
                uploaded_url = re.sub(
                    r"/private/s--.*?--/v\d+/user_uploads/",
                    "/private/user_uploads/",
                    upload_resp.json()["secure_url"],
                )
            else:
                uploaded_url = file_upload_info["s3_object_url"]

            uploaded_files.append(uploaded_url)

        json_data = {
            "query_str": query,
            "params": {
            "attachments": (
                uploaded_files + follow_up.get("attachments", []) if follow_up else uploaded_files
            ),
                "frontend_context_uuid": str(uuid4()),
                "frontend_uuid": str(uuid4()),
                "is_incognito": incognito,
                "language": language,
                "last_backend_uuid": (follow_up["backend_uuid"] if follow_up else None),
                "mode": "concise" if mode == "auto" else "copilot",
                "model_preference": MODEL_MAPPINGS[mode][model],
                "source": "default",
                "sources": sources,
                "version": "2.18",
            },
        }

        resp = self.session.post(ENDPOINT_SSE_ASK, json=json_data, stream=True, timeout=120)
        chunks = []

        def stream_response(resp):
            for chunk in resp.iter_lines(delimiter=b"\r\n\r\n"):
                content = chunk.decode("utf-8")
                if content.startswith("event: message\r\n"):
                    try:
                        content_json = json.loads(content[len("event: message\r\ndata: "):])
                        content_json = parse_nested_json_response(content_json)
                        chunks.append(content_json)
                        yield chunks[-1]
                    except (json.JSONDecodeError, KeyError):
                        continue
                elif content.startswith("event: end_of_stream\r\n"):
                    return

        if stream:
            return stream_response(resp)

        last_complete = {}
        for chunk in resp.iter_lines(delimiter=b"\r\n\r\n"):
            content = chunk.decode("utf-8")
            if content.startswith("event: message\r\n"):
                try:
                    content_json = json.loads(content[len("event: message\r\ndata: "):])
                    content_json = parse_nested_json_response(content_json)
                    chunks.append(content_json)
                    # Keep track of chunks that have an answer field
                    if content_json.get("answer"):
                        last_complete = content_json
                except (json.JSONDecodeError, KeyError):
                    continue
            elif content.startswith("event: end_of_stream\r\n"):
                # Return the last chunk that had an answer, or the last chunk
                return last_complete if last_complete else (chunks[-1] if chunks else {})