Merge pull request #8 from LlmKira/dev

feat(test):/user/information

playground/information.py

@@ -0,0 +1,39 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/8 下午4:34
+# @Author  : sudoskys
+# @File    : information.py
+# @Software: PyCharm
+import asyncio
+import os
+
+from dotenv import load_dotenv
+from loguru import logger
+from pydantic import SecretStr
+
+from novelai_python import APIError
+from novelai_python import Information, InformationResp, JwtCredential
+
+load_dotenv()
+
+enhance = "year 2023,dynamic angle, best quality, amazing quality, very aesthetic, absurdres"
+token = None
+jwt = os.getenv("NOVELAI_JWT") or token
+
+
+async def main():
+    globe_s = JwtCredential(jwt_token=SecretStr(jwt))
+    try:
+        _res = await Information().request(
+            session=globe_s
+        )
+        _res: InformationResp
+        print(f"Information: {_res}")
+        print(_res.model_dump())
+    except APIError as e:
+        logger.exception(e)
+        print(e.__dict__)
+        return
+
+
+loop = asyncio.get_event_loop()
+loop.run_until_complete(main())

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.2.1"
+version = "0.2.2"
 description = "Novelai Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "[email protected]" },
@@ -30,6 +30,7 @@ Issues = "https://github.com/LlmKira/novelai-python/issues"
 [project.optional-dependencies]
 testing = [
     "pytest>=7.4.4",
+    "pytest-asyncio>=0.23.4",
 ]
 [build-system]
 requires = ["pdm-backend"]

src/novelai_python/__init__.py

@@ -10,6 +10,7 @@
 )
 from .credential import JwtCredential, LoginCredential
 from .sdk import GenerateImageInfer, ImageGenerateResp
+from .sdk import Information, InformationResp
 from .sdk import Login, LoginResp
 from .sdk import Subscription, SubscriptionResp
 
@@ -23,6 +24,9 @@
     "Login",
     "LoginResp",
 
+    "Information",
+    "InformationResp",
+
     "JwtCredential",
     "LoginCredential",
 

src/novelai_python/_response/__init__.py

@@ -4,11 +4,13 @@
 # @File    : __init__.py.py
 # @Software: PyCharm
 from .ai.generate_image import ImageGenerateResp
+from .user.information import InformationResp
 from .user.login import LoginResp
 from .user.subscription import SubscriptionResp
 
 __all__ = [
     "ImageGenerateResp",
     "SubscriptionResp",
-    "LoginResp"
+    "LoginResp",
+    "InformationResp"
 ]

src/novelai_python/_response/user/information.py

@@ -0,0 +1,16 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/8 下午3:10
+# @Author  : sudoskys
+# @File    : information.py
+# @Software: PyCharm
+
+from pydantic import BaseModel, Field
+
+
+class InformationResp(BaseModel):
+    emailVerified: bool = Field(..., description="Email verification status")
+    emailVerificationLetterSent: bool = Field(..., description="Email verification letter sent status")
+    trialActivated: bool = Field(..., description="Trial activation status")
+    trialActionsLeft: int = Field(..., description="Number of trial actions left")
+    trialImagesLeft: int = Field(..., description="Number of trial images left")
+    accountCreatedAt: int = Field(..., description="Account creation time")

src/novelai_python/credential/ApiToken.py

@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/8 下午3:05
+# @Author  : sudoskys
+# @File    : ApiToken.py
+# @Software: PyCharm
+from curl_cffi.requests import AsyncSession
+from loguru import logger
+from pydantic import SecretStr, Field, field_validator
+
+from ._base import CredentialBase
+
+
+class ApiCredential(CredentialBase):
+    """
+    JwtCredential is the base class for all credential.
+    """
+    api_token: SecretStr = Field(None, description="api token")
+    _session: AsyncSession = None
+
+    async def get_session(self, timeout: int = 180):
+        if not self._session:
+            self._session = AsyncSession(timeout=timeout, headers={
+                "Authorization": f"Bearer {self.api_token.get_secret_value()}",
+                "Content-Type": "application/json",
+                "Origin": "https://novelai.net",
+                "Referer": "https://novelai.net/",
+            }, impersonate="chrome110")
+        return self._session
+
+    @field_validator('api_token')
+    def check_api_token(cls, v: SecretStr):
+        if not v.get_secret_value().startswith("pst"):
+            logger.warning("api token should start with pst-")
+        return v

src/novelai_python/credential/JwtToken.py

@@ -4,7 +4,8 @@
 # @File    : JwtToken.py
 # @Software: PyCharm
 from curl_cffi.requests import AsyncSession
-from pydantic import SecretStr, Field
+from loguru import logger
+from pydantic import SecretStr, Field, field_validator
 
 from ._base import CredentialBase
 
@@ -25,3 +26,9 @@ async def get_session(self, timeout: int = 180):
                 "Referer": "https://novelai.net/",
             }, impersonate="chrome110")
         return self._session
+
+    @field_validator('jwt_token')
+    def check_jwt_token(cls, v: SecretStr):
+        if not v.get_secret_value().startswith("ey"):
+            logger.warning("jwt_token should start with ey")
+        return v

src/novelai_python/credential/_base.py

@@ -9,6 +9,7 @@
 
 class CredentialBase(BaseModel):
     _session: AsyncSession = None
+    """会话"""
 
     async def get_session(self, timeout: int = 180):
         raise NotImplementedError

src/novelai_python/sdk/__init__.py

@@ -5,5 +5,6 @@
 # @Software: PyCharm
 
 from .ai.generate_image import GenerateImageInfer, ImageGenerateResp  # noqa 401
+from .user.information import Information, InformationResp  # noqa 401
 from .user.login import Login, LoginResp  # noqa 401
 from .user.subscription import Subscription, SubscriptionResp  # noqa 401

src/novelai_python/sdk/ai/__init__.py

@@ -3,5 +3,3 @@
 # @Author  : sudoskys
 # @File    : __init__.py.py
 # @Software: PyCharm
-
-from .generate_image import *  # noqa: F403

src/novelai_python/sdk/ai/generate_image.py

@@ -18,8 +18,8 @@
 
 from ..._exceptions import APIError, AuthError
 from ..._response import ImageGenerateResp
-from ...utils import try_jsonfy, NovelAiMetadata
 from ...credential import CredentialBase
+from ...utils import try_jsonfy, NovelAiMetadata
 
 
 class GenerateImageInfer(BaseModel):
@@ -297,6 +297,9 @@ async def generate(self, session: Union[AsyncSession, "CredentialBase"],
                 data=json.dumps(request_data).encode("utf-8")
             )
             if response.headers.get('Content-Type') not in ['binary/octet-stream', 'application/x-zip-compressed']:
+                logger.error(
+                    f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
+                )
                 try:
                     _msg = response.json()
                 except Exception:

src/novelai_python/sdk/user/__init__.py

@@ -3,5 +3,3 @@
 # @Author  : sudoskys
 # @File    : __init__.py.py
 # @Software: PyCharm
-
-from .subscription import *  # noqa: F403

src/novelai_python/sdk/user/information.py

@@ -0,0 +1,86 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/8 下午3:09
+# @Author  : sudoskys
+# @File    : information.py
+# @Software: PyCharm
+from typing import Optional, Union
+
+import httpx
+from curl_cffi.requests import AsyncSession, RequestsError
+from loguru import logger
+from pydantic import BaseModel, PrivateAttr
+
+from ..._exceptions import APIError, AuthError
+from ..._response.user.information import InformationResp
+from ...credential import CredentialBase
+from ...utils import try_jsonfy
+
+
+class Information(BaseModel):
+    _endpoint: Optional[str] = PrivateAttr("https://api.novelai.net")
+
+    @property
+    def base_url(self):
+        return f"{self.endpoint.strip('/')}/user/information"
+
+    @property
+    def endpoint(self):
+        return self._endpoint
+
+    @endpoint.setter
+    def endpoint(self, value):
+        self._endpoint = value
+
+    async def request(self,
+                      session: Union[AsyncSession, CredentialBase],
+                      ) -> InformationResp:
+        """
+        Request to get user subscription information
+        :param session: 
+        :return: 
+        """
+        if isinstance(session, CredentialBase):
+            session = await session.get_session()
+        request_data = {}
+        logger.debug("Information")
+        try:
+            assert hasattr(session, "get"), "session must have get method."
+            response = await session.get(
+                self.base_url,
+            )
+            if "application/json" not in response.headers.get('Content-Type') or response.status_code != 200:
+                logger.error(
+                    f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
+                )
+                try:
+                    _msg = response.json()
+                except Exception:
+                    raise APIError(
+                        message=f"Unexpected content type: {response.headers.get('Content-Type')}",
+                        request=request_data,
+                        status_code=response.status_code,
+                        response=try_jsonfy(response.content)
+                    )
+                status_code = _msg.get("statusCode", response.status_code)
+                message = _msg.get("message", "Unknown error")
+                if status_code in [400, 401, 402]:
+                    # 400 : validation error
+                    # 401 : unauthorized
+                    # 402 : payment required
+                    # 409 : conflict
+                    raise AuthError(message, request=request_data, status_code=status_code, response=_msg)
+                if status_code in [500]:
+                    # An unknown error occured.
+                    raise APIError(message, request=request_data, status_code=status_code, response=_msg)
+                raise APIError(message, request=request_data, status_code=status_code, response=_msg)
+            return InformationResp.model_validate(response.json())
+        except RequestsError as exc:
+            logger.exception(exc)
+            raise RuntimeError(f"An AsyncSession error occurred: {exc}")
+        except httpx.HTTPError as exc:
+            raise RuntimeError(f"An HTTP error occurred: {exc}")
+        except APIError as e:
+            raise e
+        except Exception as e:
+            logger.opt(exception=e).exception("An Unexpected error occurred")
+            raise e

src/novelai_python/sdk/user/login.py

@@ -18,6 +18,7 @@
 
 class Login(BaseModel):
     _endpoint: Optional[str] = PrivateAttr("https://api.novelai.net")
+    _session: Optional[AsyncSession] = PrivateAttr(None)
     key: str = Field(..., description="User's key")
 
     @property
@@ -34,11 +35,19 @@ def base_url(self):
 
     @property
     def session(self):
-        return AsyncSession(timeout=180, headers={
-            "Content-Type": "application/json",
-            "Origin": "https://novelai.net",
-            "Referer": "https://novelai.net/",
-        }, impersonate="chrome110")
+        if self._session is None:
+            self._session = AsyncSession(timeout=180, headers={
+                "Content-Type": "application/json",
+                "Origin": "https://novelai.net",
+                "Referer": "https://novelai.net/",
+            }, impersonate="chrome110")
+        return self._session
+
+    @session.setter
+    def session(self, value):
+        if not isinstance(value, AsyncSession):
+            raise ValueError("session must be an instance of AsyncSession")
+        self._session = value
 
     @classmethod
     def build(cls, *, user_name: str, password: str):
@@ -65,7 +74,9 @@ async def request(self,
                 data=json.dumps(request_data).encode("utf-8")
             )
             if "application/json" not in response.headers.get('Content-Type') or response.status_code != 201:
-                logger.error(f"Unexpected content type: {response.headers.get('Content-Type')}")
+                logger.error(
+                    f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
+                )
                 try:
                     _msg = response.json()
                 except Exception:

src/novelai_python/sdk/user/subscription.py

@@ -3,7 +3,7 @@
 # @Author  : sudoskys
 # @File    : subscription.py.py
 # @Software: PyCharm
-from typing import Optional, Union, Type
+from typing import Optional, Union
 
 import httpx
 from curl_cffi.requests import AsyncSession, RequestsError
@@ -49,7 +49,9 @@ async def request(self,
                 self.base_url,
             )
             if "application/json" not in response.headers.get('Content-Type') or response.status_code != 200:
-                logger.error(f"Unexpected content type: {response.headers.get('Content-Type')}")
+                logger.error(
+                    f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
+                )
                 try:
                     _msg = response.json()
                 except Exception: