gpt4free

from __future__ import annotations

import os
from typing import AsyncIterator

from ..base_provider import AsyncAuthedProvider
from ..Copilot import Copilot, readHAR, has_nodriver, get_access_token_and_cookies
from ...providers.response import AuthResult, RequestLogin
from ...typing import AsyncResult, Messages
from ...errors import NoValidHarFileError
from ... import debug

class CopilotAccount(Copilot, AsyncAuthedProvider):
    needs_auth = True
    use_nodriver = True
    parent = "Copilot"
    default_model = "Copilot"
    default_vision_model = default_model

    @classmethod
    async def on_auth_async(cls, proxy: str = None, **kwargs) -> AsyncIterator:
        try:
            cls._access_token, cls._cookies = readHAR(cls.url)
        except NoValidHarFileError as h:
            debug.log(f"Copilot: {h}")
            if has_nodriver:
                yield RequestLogin(cls.label, os.environ.get("G4F_LOGIN_URL", ""))
                cls._access_token, cls._cookies = await get_access_token_and_cookies(cls.url, proxy)
            else:
                raise h
        yield AuthResult(
            api_key=cls._access_token,
            cookies=cls.cookies_to_dict()
        )

    @classmethod
    async def create_authed(
        cls,
        model: str,
        messages: Messages,
        auth_result: AuthResult,
        **kwargs
    ) -> AsyncResult:
        cls._access_token = getattr(auth_result, "api_key")
        cls._cookies = getattr(auth_result, "cookies")
        async for chunk in cls.create_async_generator(model, messages, **kwargs):
            yield chunk
        auth_result.cookies = cls.cookies_to_dict()

    @classmethod
    def cookies_to_dict(cls):
        return cls._cookies if isinstance(cls._cookies, dict) else {c.name: c.value for c in cls._cookies}