Merge pull request #5 from LlmKira/dev

feat:LoginCredential | user/login

NOTICE.MD

@@ -0,0 +1,7 @@
+# NOTICE OF THIS PROJECT
+
+## MIT License
+
+The MIT license applies to the files in:
+
+    file: "src/novelai_python/sdk/ai/generate_image.py" from https://github.com/HanaokaYuzu/NovelAI-API

README.md

@@ -9,6 +9,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 
 - [x] /ai/generate-image
 - [x] /user/subscription
+- [x] /user/login
 - [ ] /ai/generate-image/suggest-tags
 - [ ] /ai/annotate-image
 - [ ] /ai/classify
@@ -19,22 +20,30 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 
 ### Usage 🖥️
 
+More examples can be found in the [playground](/playground) directory.
+
 ```python
 import asyncio
 import os
 
 from dotenv import load_dotenv
 from pydantic import SecretStr
 
-from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential
+from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential, LoginCredential
 
 load_dotenv()
 
 enhance = "year 2023,dynamic angle,  best quality, amazing quality, very aesthetic, absurdres"
 
 
 async def main():
-    globe_s = JwtCredential(jwt_token=SecretStr(os.getenv("NOVELAI_JWT")))
+    globe_s = JwtCredential(
+        jwt_token=SecretStr(os.getenv("NOVELAI_JWT"))
+    )
+    globe_s2 = LoginCredential(
+        username=os.getenv("NOVELAI_USERNAME"),
+        password=SecretStr(os.getenv("NOVELAI_PASSWORD"))
+    )
     _res = await GenerateImageInfer.build(
         prompt=f"1girl,{enhance}").generate(
         session=globe_s)
@@ -69,5 +78,9 @@ python3 -m novelai_python.server -h '0.0.0.0' -p 7888
 ## Acknowledgements 🙏
 
 [BackEnd](https://api.novelai.net/docs)
+
 [novelai-api](https://github.com/Aedial/novelai-api)
 
+[NovelAI-API](https://github.com/HanaokaYuzu/NovelAI-API)
+
+

playground/generate_image.py

@@ -9,7 +9,7 @@
 from dotenv import load_dotenv
 from pydantic import SecretStr
 
-from novelai_python import APIError
+from novelai_python import APIError, Login
 from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential
 
 load_dotenv()
@@ -21,10 +21,15 @@
 
 async def main():
     globe_s = JwtCredential(jwt_token=SecretStr(jwt))
+    _res = await Login.build(user_name=os.getenv("NOVELAI_USER"), password=os.getenv("NOVELAI_PASS")
+                             ).request()
     try:
-        _res = await GenerateImageInfer.build(
-            prompt=f"1girl, winter, jacket, sfw, angel, flower,{enhance}").generate(
-            session=globe_s)
+        gen = GenerateImageInfer.build(
+            prompt=f"1girl, winter, jacket, sfw, angel, flower,{enhance}")
+        print(f"charge: {gen.calculate_cost(is_opus=True)}")
+        _res = await gen.generate(
+            session=globe_s, remove_sign=True
+        )
     except APIError as e:
         print(e.response)
         return

playground/login.py

@@ -0,0 +1,33 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/7 下午12:07
+# @Author  : sudoskys
+# @File    : login.py
+# @Software: PyCharm
+import asyncio
+import os
+
+from dotenv import load_dotenv
+from loguru import logger
+
+from novelai_python import APIError
+from novelai_python import Login, LoginResp
+
+load_dotenv()
+
+
+async def main():
+    try:
+        _res = await Login.build(user_name=os.getenv("NOVELAI_USER"), password=os.getenv("NOVELAI_PASS")
+                                 ).request()
+    except APIError as e:
+        logger.exception(e)
+        print(e.__dict__)
+        return
+
+    _res: LoginResp
+    print(_res)
+    print(f"Access Token: {_res.accessToken}")
+
+
+loop = asyncio.get_event_loop()
+loop.run_until_complete(main())

playground/subscription.py

@@ -10,7 +10,7 @@
 from loguru import logger
 from pydantic import SecretStr
 
-from novelai_python import APIError, SubscriptionResp
+from novelai_python import APIError, SubscriptionResp, LoginCredential
 from novelai_python import Subscription, JwtCredential
 
 load_dotenv()
@@ -26,15 +26,28 @@ async def main():
         _res = await Subscription().request(
             session=globe_s
         )
+        _res: SubscriptionResp
+        print(f"JwtCredential/Subscription: {_res}")
+        print(_res.is_active)
+        print(_res.anlas_left)
     except APIError as e:
         logger.exception(e)
-        print(e.response)
+        print(e.__dict__)
         return
 
-    _res: SubscriptionResp
-    print(_res)
-    print(_res.is_active)
-    print(_res.anlas_left)
+    try:
+        cre = LoginCredential(
+            username=os.getenv("NOVELAI_USER"),
+            password=SecretStr(os.getenv("NOVELAI_PASS"))
+        )
+        _res = await Subscription().request(
+            session=cre
+        )
+        print(f"LoginCredential/User subscription: {_res}")
+        print(_res.is_active)
+        print(_res.anlas_left)
+    except Exception as e:
+        logger.exception(e)
 
 
 loop = asyncio.get_event_loop()

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.1.8"
+version = "0.1.9"
 description = "Novelai Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "[email protected]" },
@@ -17,6 +17,7 @@ dependencies = [
     "fastapi>=0.109.0",
     "uvicorn[standard]>=0.27.0.post1",
     "numpy>=1.24.4",
+    "argon2-cffi>=23.1.0",
 ]
 requires-python = ">=3.8"
 readme = "README.md"

src/novelai_python/__init__.py

@@ -5,13 +5,13 @@
 # @Software: PyCharm
 
 from ._exceptions import (
-    APIError,
     AuthError,
-    NovelAiError,
-)  # noqa: F401, F403
-from .credential import JwtCredential  # noqa: F401, F403
-from .sdk.ai import GenerateImageInfer, ImageGenerateResp  # noqa: F401, F403
-from .sdk.user import Subscription, SubscriptionResp  # noqa: F401, F403
+    NovelAiError, APIError,
+)
+from .credential import JwtCredential, LoginCredential
+from .sdk import GenerateImageInfer, ImageGenerateResp
+from .sdk import Login, LoginResp
+from .sdk import Subscription, SubscriptionResp
 
 __all__ = [
     "GenerateImageInfer",
@@ -20,7 +20,11 @@
     "Subscription",
     "SubscriptionResp",
 
+    "Login",
+    "LoginResp",
+
     "JwtCredential",
+    "LoginCredential",
 
     "APIError",
     "AuthError",

src/novelai_python/_exceptions.py

@@ -32,6 +32,15 @@ def __init__(self, message: str, request: dict, response: Union[dict, str], stat
         self.code = status_code
         self.response = response
 
+    @property
+    def __dict__(self):
+        return {
+            "message": self.message,
+            "request": self.request,
+            "response": self.response,
+            "code": self.code
+        }
+
 
 class AuthError(APIError):
     """

src/novelai_python/_response/__init__.py

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

src/novelai_python/_response/user/login.py

@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/7 上午11:57
+# @Author  : sudoskys
+# @File    : login.py
+# @Software: PyCharm
+from pydantic import BaseModel
+
+
+class LoginResp(BaseModel):
+    accessToken: str

src/novelai_python/credential/JwtToken.py

@@ -4,18 +4,19 @@
 # @File    : JwtToken.py
 # @Software: PyCharm
 from curl_cffi.requests import AsyncSession
-from pydantic import BaseModel, SecretStr, Field
+from pydantic import SecretStr, Field
 
+from ._base import CredentialBase
 
-class JwtCredential(BaseModel):
+
+class JwtCredential(CredentialBase):
     """
     JwtCredential is the base class for all credential.
     """
     jwt_token: SecretStr = Field(None, description="jwt token")
     _session: AsyncSession = None
 
-    @property
-    def session(self, timeout: int = 180):
+    async def get_session(self, timeout: int = 180):
         if not self._session:
             self._session = AsyncSession(timeout=timeout, headers={
                 "Authorization": f"Bearer {self.jwt_token.get_secret_value()}",

src/novelai_python/credential/UserAuth.py

@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/7 下午12:14
+# @Author  : sudoskys
+# @File    : UserAuth.py
+# @Software: PyCharm
+import time
+from typing import Optional
+
+from curl_cffi.requests import AsyncSession
+from pydantic import SecretStr, Field
+
+from ._base import CredentialBase
+
+
+class LoginCredential(CredentialBase):
+    """
+    JwtCredential is the base class for all credential.
+    """
+    username: str = Field(None, description="username")
+    password: SecretStr = Field(None, description="password")
+    _session: Optional[AsyncSession] = None
+    _update_at: Optional[int] = None
+
+    async def get_session(self, timeout: int = 180):
+        # 30 天有效期
+        if not self._session or int(time.time()) - self._update_at > 29 * 24 * 60 * 60:
+            from ..sdk import Login
+            resp = await Login.build(user_name=self.username, password=self.password.get_secret_value()).request()
+            self._session = AsyncSession(timeout=timeout, headers={
+                "Authorization": f"Bearer {resp.accessToken}",
+                "Content-Type": "application/json",
+                "Origin": "https://novelai.net",
+                "Referer": "https://novelai.net/",
+            }, impersonate="chrome110")
+            self._update_at = int(time.time())
+        return self._session

src/novelai_python/credential/__init__.py

@@ -3,9 +3,15 @@
 # @Author  : sudoskys
 # @File    : __init__.py.py
 # @Software: PyCharm
-from .JwtToken import JwtCredential
 from pydantic import SecretStr
+
+from .JwtToken import JwtCredential
+from .UserAuth import LoginCredential
+from ._base import CredentialBase
+
 __all__ = [
     "JwtCredential",
+    "LoginCredential",
+    "CredentialBase",
     "SecretStr"
 ]