novelai-python 0.1.7__tar.gz → 0.1.8__tar.gz

{novelai_python-0.1.7 → novelai_python-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: novelai-python
-Version: 0.1.7
+Version: 0.1.8
 Summary: Novelai Python Binding With Pydantic
 Author-Email: sudoskys <coldlando@hotmail.com>
 License: Apache 2.0
@@ -32,6 +32,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 ### Roadmap 🚧
 
 - [x] /ai/generate-image
+- [x] /user/subscription
 - [ ] /ai/generate-image/suggest-tags
 - [ ] /ai/annotate-image
 - [ ] /ai/classify

{novelai_python-0.1.7 → novelai_python-0.1.8}/README.md

@@ -8,6 +8,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 ### Roadmap 🚧
 
 - [x] /ai/generate-image
+- [x] /user/subscription
 - [ ] /ai/generate-image/suggest-tags
 - [ ] /ai/annotate-image
 - [ ] /ai/classify

{novelai_python-0.1.7 → novelai_python-0.1.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.1.7"
+version = "0.1.8"
 description = "Novelai Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "coldlando@hotmail.com" },

{novelai_python-0.1.7 → novelai_python-0.1.8}/src/novelai_python/__init__.py

@@ -9,14 +9,19 @@ from ._exceptions import (
     AuthError,
     NovelAiError,
 )  # noqa: F401, F403
-from ._response import ImageGenerateResp  # noqa: F401, F403
 from .credential import JwtCredential  # noqa: F401, F403
-from .sdk.ai import GenerateImageInfer  # noqa: F401, F403
+from .sdk.ai import GenerateImageInfer, ImageGenerateResp  # noqa: F401, F403
+from .sdk.user import Subscription, SubscriptionResp  # noqa: F401, F403
 
 __all__ = [
     "GenerateImageInfer",
     "ImageGenerateResp",
+
+    "Subscription",
+    "SubscriptionResp",
+
     "JwtCredential",
+
     "APIError",
     "AuthError",
     "NovelAiError",

{novelai_python-0.1.7 → novelai_python-0.1.8}/src/novelai_python/_response/__init__.py

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

novelai_python-0.1.8/src/novelai_python/_response/user/__init__.py

@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/7 上午9:57
+# @Author  : sudoskys
+# @File    : __init__.py.py
+# @Software: PyCharm

novelai_python-0.1.8/src/novelai_python/_response/user/subscription.py

@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/7 上午9:57
+# @Author  : sudoskys
+# @File    : subscription.py
+# @Software: PyCharm
+from typing import Optional, Dict, Any, List
+
+from pydantic import BaseModel, Field
+
+
+class TrainingSteps(BaseModel):
+    fixedTrainingStepsLeft: int
+    purchasedTrainingSteps: int
+
+
+class ImageGenerationLimit(BaseModel):
+    resolution: int
+    maxPrompts: int
+
+
+class Perks(BaseModel):
+    maxPriorityActions: int
+    startPriority: int
+    moduleTrainingSteps: int
+    unlimitedMaxPriority: bool
+    voiceGeneration: bool
+    imageGeneration: bool
+    unlimitedImageGeneration: bool
+    unlimitedImageGenerationLimits: List[ImageGenerationLimit]
+    contextTokens: int
+
+
+class SubscriptionResp(BaseModel):
+    tier: int = Field(..., description="Subscription tier")
+    active: bool = Field(..., description="Subscription status")
+    expiresAt: int = Field(..., description="Subscription expiration time")
+    perks: Perks = Field(..., description="Subscription perks")
+    paymentProcessorData: Optional[Dict[Any, Any]]
+    trainingStepsLeft: TrainingSteps = Field(..., description="Training steps left")
+    accountType: int = Field(..., description="Account type")
+
+    @property
+    def is_active(self):
+        return self.active
+
+    @property
+    def anlas_left(self):
+        return self.trainingStepsLeft.fixedTrainingStepsLeft
+
+    @property
+    def is_unlimited_image_generation(self):
+        return self.perks.unlimitedImageGeneration

{novelai_python-0.1.7 → novelai_python-0.1.8}/src/novelai_python/sdk/__init__.py

@@ -5,4 +5,4 @@
 # @Software: PyCharm
 
 from .ai.generate_image import GenerateImageInfer, ImageGenerateResp  # noqa 401
-
+from .user.subscription import Subscription, SubscriptionResp  # noqa 401

novelai_python-0.1.8/src/novelai_python/sdk/user/__init__.py

@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/7 上午10:03
+# @Author  : sudoskys
+# @File    : __init__.py.py
+# @Software: PyCharm
+
+from .subscription import *  # noqa: F403

novelai_python-0.1.8/src/novelai_python/sdk/user/subscription.py

@@ -0,0 +1,81 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/7 上午10:04
+# @Author  : sudoskys
+# @File    : subscription.py.py
+# @Software: PyCharm
+from typing import Optional, Union
+
+import httpx
+from curl_cffi.requests import AsyncSession
+from loguru import logger
+from pydantic import BaseModel, PrivateAttr
+
+from ... import APIError, AuthError
+from ..._response.user.subscription import SubscriptionResp
+from ...credential import JwtCredential
+from ...utils import try_jsonfy
+
+
+class Subscription(BaseModel):
+    _endpoint: Optional[str] = PrivateAttr("https://api.novelai.net")
+
+    @property
+    def base_url(self):
+        return f"{self.endpoint.strip('/')}/user/subscription"
+
+    @property
+    def endpoint(self):
+        return self._endpoint
+
+    @endpoint.setter
+    def endpoint(self, value):
+        self._endpoint = value
+
+    async def request(self,
+                      session: Union[AsyncSession, JwtCredential],
+                      ) -> SubscriptionResp:
+        """
+        Request to get user subscription information
+        :param session: 
+        :return: 
+        """
+        if isinstance(session, JwtCredential):
+            session = session.session
+        request_data = {}
+        logger.debug("Subscription")
+        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'):
+                logger.error(f"Unexpected content type: {response.headers.get('Content-Type')}")
+                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 SubscriptionResp.model_validate(response.json())
+        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

{novelai_python-0.1.7 → novelai_python-0.1.8}/src/novelai_python/server.py

@@ -10,10 +10,12 @@ import zipfile
 import uvicorn
 from fastapi import FastAPI, Depends, Security
 from fastapi.security import APIKeyHeader
+from loguru import logger
 from starlette.responses import JSONResponse, StreamingResponse
 
 from .credential import JwtCredential, SecretStr
 from .sdk.ai.generate_image import GenerateImageInfer
+from .sdk.user.subscription import Subscription
 
 app = FastAPI()
 token_key = APIKeyHeader(name="Authorization")
@@ -35,6 +37,24 @@ async def health():
     return {"status": "ok"}
 
 
+@app.get("/user/subscription")
+async def subscription(
+        current_token: str = Depends(get_current_token)
+):
+    """
+    订阅
+    :param current_token: Authorization
+    :param req: Subscription
+    :return:
+    """
+    try:
+        _result = await Subscription().request(session=get_session(current_token))
+        return _result.model_dump()
+    except Exception as e:
+        logger.exception(e)
+        return JSONResponse(status_code=500, content=e.__dict__)
+
+
 @app.post("/ai/generate_image")
 async def generate_image(
         req: GenerateImageInfer,
@@ -58,6 +78,7 @@ async def generate_image(
             'Content-Disposition': 'attachment;filename=multiple_files.zip'
         })
     except Exception as e:
+        logger.exception(e)
         return JSONResponse(status_code=500, content=e.__dict__)
 
 

novelai_python-0.1.8/tests/test_server_run.py

@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/30 下午11:52
+# @Author  : sudoskys
+# @File    : test_server_run.py
+# @Software: PyCharm
+
+from unittest.mock import patch
+
+from fastapi.testclient import TestClient
+
+from src.novelai_python.server import app, get_session
+
+client = TestClient(app)
+
+
+def test_health_check():
+    response = client.get("/health")
+    assert response.status_code == 200
+    assert response.json() == {"status": "ok"}
+
+
+@patch('src.novelai_python.server.Subscription')
+def test_subscription_without_api_key(mock_subscription):
+    mock_subscription.return_value.request.return_value.model_dump.return_value = {"status": "subscribed"}
+    response = client.get("/user/subscription")
+    assert response.status_code == 403
+
+
+@patch('src.novelai_python.server.GenerateImageInfer')
+def test_generate_image_without_api_key(mock_generate_image):
+    mock_generate_image.return_value.generate.return_value = {"status": "image generated"}
+    response = client.post("/ai/generate_image")
+    assert response.status_code == 403
+
+
+def test_get_session_new_token():
+    token = "new_token"
+    session = get_session(token)
+    assert session.jwt_token.get_secret_value() == token
+
+
+def test_get_session_existing_token():
+    token = "existing_token"
+    get_session(token)
+    session = get_session(token)
+    assert session.jwt_token.get_secret_value() == token

novelai_python-0.1.7/tests/test_server_run.py

@@ -1,28 +0,0 @@
-# -*- coding: utf-8 -*-
-# @Time    : 2024/1/30 下午11:52
-# @Author  : sudoskys
-# @File    : test_server_run.py
-# @Software: PyCharm
-from fastapi.testclient import TestClient
-
-from src.novelai_python.sdk.ai.generate_image import GenerateImageInfer
-from src.novelai_python.server import app, get_session
-
-client = TestClient(app)
-
-
-def test_health_check():
-    response = client.get("/health")
-    assert response.status_code == 200
-    assert response.json() == {"status": "ok"}
-
-
-def test_generate_image_with_valid_token():
-    valid_token = "valid_token"
-    get_session(valid_token)  # to simulate a valid session
-    response = client.post(
-        "/ai/generate_image",
-        headers={"Authorization": valid_token},
-        json=GenerateImageInfer(input="1girl").model_dump()
-    )
-    assert response.status_code == 500