novelai-python 0.4.6__tar.gz → 0.4.7__tar.gz

{novelai_python-0.4.6 → novelai_python-0.4.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: novelai-python
-Version: 0.4.6
+Version: 0.4.7
 Summary: NovelAI Python Binding With Pydantic
 Author-Email: sudoskys <coldlando@hotmail.com>
 License: Apache License 2.0
@@ -24,6 +24,9 @@ Requires-Dist: fake-useragent>=1.4.0
 Requires-Dist: tenacity>=8.2.3
 Requires-Dist: sentencepiece>=0.2.0
 Requires-Dist: pynacl>=1.5.0
+Requires-Dist: ftfy>=6.2.0
+Requires-Dist: regex>=2023.12.25
+Requires-Dist: tokenizers>=0.15.2
 Requires-Dist: pytest>=7.4.4; extra == "testing"
 Requires-Dist: pytest-asyncio>=0.23.4; extra == "testing"
 Provides-Extra: testing
@@ -36,7 +39,7 @@ Description-Content-Type: text/markdown
 [![PyPI version](https://badge.fury.io/py/novelai-python.svg)](https://badge.fury.io/py/novelai-python)
 [![Downloads](https://pepy.tech/badge/novelai_python)](https://pepy.tech/project/novelai_python)
 
-✨ NovelAI api python sdk with Pydantic.
+✨ NovelAI api python sdk with Pydantic, modern and user-friendly.
 
 The goal of this repository is to use Pydantic to build legitimate requests to access the NovelAI API service.
 
@@ -52,10 +55,11 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 - [x] /ai/upscale
 - [x] /ai/generate-image/suggest-tags
 - [x] /ai/generate-voice
+- [x] /ai/generate-stream
+- [x] /ai/generate
 - [ ] /ai/annotate-image
 - [ ] /ai/classify
 - [ ] /ai/generate-prompt
-- [ ] /ai/generate
 
 > GenerateImageInfer.calculate_cost is correct in most cases, but please request account information to get accurate
 > consumption information.
@@ -78,34 +82,20 @@ import os
 from dotenv import load_dotenv
 from pydantic import SecretStr
 
-from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential, LoginCredential, ApiCredential
+from novelai_python import GenerateImageInfer, ImageGenerateResp, ApiCredential
 
 load_dotenv()
-
 enhance = "year 2023,dynamic angle,  best quality, amazing quality, very aesthetic, absurdres"
+session = ApiCredential(api_token=SecretStr(os.getenv("NOVELAI_JWT")))  # pst-***
 
 
 async def main():
-    globe_s = JwtCredential(
-        jwt_token=SecretStr(os.getenv("NOVELAI_JWT"))  # ey****
-    )
-    globe_s1 = ApiCredential(
-        api_token=SecretStr(os.getenv("NOVELAI_JWT"))  # pst-***
-    )
-    globe_s2 = LoginCredential(
-        username=os.getenv("NOVELAI_USERNAME"),
-        password=SecretStr(os.getenv("NOVELAI_PASSWORD"))
-    )
-
-    gen = await GenerateImageInfer.build(
-        prompt=f"1girl,{enhance}")
+    gen = await GenerateImageInfer.build(prompt=f"1girl,{enhance}")
     cost = gen.calculate_cost(is_opus=True)
     print(f"charge: {cost} if you are vip3")
-
-    resp = gen.request(session=globe_s)
+    resp = gen.request(session=session)
     resp: ImageGenerateResp
     print(resp.meta)
-
     file = resp.files[0]
     with open(file[0], "wb") as f:
         f.write(file[1])
@@ -116,13 +106,47 @@ loop.run_until_complete(main())
 
 ```
 
+#### LLM
+
+```python
+import asyncio
+import os
+
+from dotenv import load_dotenv
+
+from novelai_python import APIError, Login
+from novelai_python.sdk.ai.generate import TextLLMModel, LLM
+
+load_dotenv()
+username = os.getenv("NOVELAI_USER", None)
+assert username is not None
+# credential = JwtCredential(jwt_token=SecretStr(jwt))
+login_credential = Login.build(
+    user_name=os.getenv("NOVELAI_USER"),
+    password=os.getenv("NOVELAI_PASS")
+)
+
+
+async def chat(prompt: str):
+    try:
+        agent = LLM.build(prompt=prompt, model=TextLLMModel.Kayra)
+        result = await agent.request(session=login_credential)
+    except APIError as e:
+        raise Exception(f"Error: {e.message}")
+    print(f"Result: \n{result.text}")
+
+
+loop = asyncio.get_event_loop()
+loop.run_until_complete(chat("Hello"))
+```
+
 #### Random Prompt
 
 ```python
 from novelai_python.tool.random_prompt import RandomPromptGenerator
 
-s = RandomPromptGenerator(nsfw_enabled=False).random_prompt()
-print(s)
+prompt = RandomPromptGenerator(nsfw_enabled=False).random_prompt()
+print(prompt)
 ```
 
 #### Run A Server

{novelai_python-0.4.6 → novelai_python-0.4.7}/README.md

@@ -5,7 +5,7 @@
 [![PyPI version](https://badge.fury.io/py/novelai-python.svg)](https://badge.fury.io/py/novelai-python)
 [![Downloads](https://pepy.tech/badge/novelai_python)](https://pepy.tech/project/novelai_python)
 
-✨ NovelAI api python sdk with Pydantic.
+✨ NovelAI api python sdk with Pydantic, modern and user-friendly.
 
 The goal of this repository is to use Pydantic to build legitimate requests to access the NovelAI API service.
 
@@ -21,10 +21,11 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 - [x] /ai/upscale
 - [x] /ai/generate-image/suggest-tags
 - [x] /ai/generate-voice
+- [x] /ai/generate-stream
+- [x] /ai/generate
 - [ ] /ai/annotate-image
 - [ ] /ai/classify
 - [ ] /ai/generate-prompt
-- [ ] /ai/generate
 
 > GenerateImageInfer.calculate_cost is correct in most cases, but please request account information to get accurate
 > consumption information.
@@ -47,34 +48,20 @@ import os
 from dotenv import load_dotenv
 from pydantic import SecretStr
 
-from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential, LoginCredential, ApiCredential
+from novelai_python import GenerateImageInfer, ImageGenerateResp, ApiCredential
 
 load_dotenv()
-
 enhance = "year 2023,dynamic angle,  best quality, amazing quality, very aesthetic, absurdres"
+session = ApiCredential(api_token=SecretStr(os.getenv("NOVELAI_JWT")))  # pst-***
 
 
 async def main():
-    globe_s = JwtCredential(
-        jwt_token=SecretStr(os.getenv("NOVELAI_JWT"))  # ey****
-    )
-    globe_s1 = ApiCredential(
-        api_token=SecretStr(os.getenv("NOVELAI_JWT"))  # pst-***
-    )
-    globe_s2 = LoginCredential(
-        username=os.getenv("NOVELAI_USERNAME"),
-        password=SecretStr(os.getenv("NOVELAI_PASSWORD"))
-    )
-
-    gen = await GenerateImageInfer.build(
-        prompt=f"1girl,{enhance}")
+    gen = await GenerateImageInfer.build(prompt=f"1girl,{enhance}")
     cost = gen.calculate_cost(is_opus=True)
     print(f"charge: {cost} if you are vip3")
-
-    resp = gen.request(session=globe_s)
+    resp = gen.request(session=session)
     resp: ImageGenerateResp
     print(resp.meta)
-
     file = resp.files[0]
     with open(file[0], "wb") as f:
         f.write(file[1])
@@ -85,13 +72,47 @@ loop.run_until_complete(main())
 
 ```
 
+#### LLM
+
+```python
+import asyncio
+import os
+
+from dotenv import load_dotenv
+
+from novelai_python import APIError, Login
+from novelai_python.sdk.ai.generate import TextLLMModel, LLM
+
+load_dotenv()
+username = os.getenv("NOVELAI_USER", None)
+assert username is not None
+# credential = JwtCredential(jwt_token=SecretStr(jwt))
+login_credential = Login.build(
+    user_name=os.getenv("NOVELAI_USER"),
+    password=os.getenv("NOVELAI_PASS")
+)
+
+
+async def chat(prompt: str):
+    try:
+        agent = LLM.build(prompt=prompt, model=TextLLMModel.Kayra)
+        result = await agent.request(session=login_credential)
+    except APIError as e:
+        raise Exception(f"Error: {e.message}")
+    print(f"Result: \n{result.text}")
+
+
+loop = asyncio.get_event_loop()
+loop.run_until_complete(chat("Hello"))
+```
+
 #### Random Prompt
 
 ```python
 from novelai_python.tool.random_prompt import RandomPromptGenerator
 
-s = RandomPromptGenerator(nsfw_enabled=False).random_prompt()
-print(s)
+prompt = RandomPromptGenerator(nsfw_enabled=False).random_prompt()
+print(prompt)
 ```
 
 #### Run A Server

{novelai_python-0.4.6 → novelai_python-0.4.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.4.6"
+version = "0.4.7"
 description = "NovelAI Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "coldlando@hotmail.com" },
@@ -23,6 +23,9 @@ dependencies = [
     "tenacity>=8.2.3",
     "sentencepiece>=0.2.0",
     "pynacl>=1.5.0",
+    "ftfy>=6.2.0",
+    "regex>=2023.12.25",
+    "tokenizers>=0.15.2",
 ]
 requires-python = ">=3.8"
 readme = "README.md"

{novelai_python-0.4.6 → novelai_python-0.4.7}/src/novelai_python/__init__.py

@@ -12,14 +12,22 @@ from ._exceptions import (
 )
 from .credential import JwtCredential, LoginCredential, ApiCredential
 from .sdk import GenerateImageInfer, ImageGenerateResp
-from .sdk import VoiceGenerate, VoiceResponse
-from .sdk import SuggestTags, SuggestTagsResp
 from .sdk import Information, InformationResp
+from .sdk import LLM, LLMResp
+from .sdk import LLMStream, LLMStreamResp
 from .sdk import Login, LoginResp
 from .sdk import Subscription, SubscriptionResp
+from .sdk import SuggestTags, SuggestTagsResp
 from .sdk import Upscale, UpscaleResp
+from .sdk import VoiceGenerate, VoiceResponse
 
 __all__ = [
+    "LLM",
+    "LLMResp",
+
+    "LLMStream",
+    "LLMStreamResp",
+
     "GenerateImageInfer",
     "ImageGenerateResp",
 

novelai_python-0.4.7/src/novelai_python/_response/ai/generate.py

@@ -0,0 +1,23 @@
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, ConfigDict
+
+from novelai_python.sdk.ai.generate._enum import TOKENIZER, TextLLMModel  # noqa
+from novelai_python.tokenizer import LLMTokenizer
+from novelai_python.utils.encode import b64_to_tokens
+
+if TYPE_CHECKING:
+    pass
+
+
+class LLMResp(BaseModel):
+    """
+    response.json().get("output",None)
+    """
+    output: str
+    text: str
+    model_config = ConfigDict(extra="ignore", arbitrary_types_allowed=False)
+
+    @staticmethod
+    def decode_token(token_str, model: TextLLMModel) -> str:
+        return LLMTokenizer.decode(b64_to_tokens(token_str), tokenizer_name=TOKENIZER.get(model))

novelai_python-0.4.7/src/novelai_python/_response/ai/generate_stream.py

@@ -0,0 +1,23 @@
+from typing import Optional, Any
+
+from pydantic import BaseModel, ConfigDict
+
+from novelai_python.sdk.ai.generate._enum import TOKENIZER, TextLLMModel  # noqa
+from novelai_python.tokenizer import LLMTokenizer
+from novelai_python.utils.encode import b64_to_tokens
+
+
+class LLMStreamResp(BaseModel):
+    """
+    {"token":"LRI=","ptr":0,"final":false,"logprobs":null}
+    """
+    token: str
+    ptr: int
+    final: bool
+    logprobs: Optional[Any]
+    text: Optional[str] = None
+    model_config = ConfigDict(extra="ignore", arbitrary_types_allowed=False)
+
+    @staticmethod
+    def decode(token_str, model: TextLLMModel) -> str:
+        return LLMTokenizer.decode(b64_to_tokens(token_str), tokenizer_name=TOKENIZER.get(model))

{novelai_python-0.4.6 → novelai_python-0.4.7}/src/novelai_python/sdk/__init__.py

@@ -4,8 +4,10 @@
 # @File    : __init__.py.py
 # @Software: PyCharm
 
+from .ai.generate import LLM, LLMResp  # noqa 401
 from .ai.generate_image import GenerateImageInfer, ImageGenerateResp  # noqa 401
 from .ai.generate_image.suggest_tags import SuggestTags, SuggestTagsResp  # noqa 401
+from .ai.generate_stream import LLMStream, LLMStreamResp  # noqa 401
 from .ai.generate_voice import VoiceGenerate, VoiceResponse  # noqa 401
 from .ai.upscale import Upscale, UpscaleResp  # noqa 401
 from .user.information import Information, InformationResp  # noqa 401

novelai_python-0.4.7/src/novelai_python/sdk/ai/generate/__init__.py

@@ -0,0 +1,234 @@
+# -*- coding: utf-8 -*-
+# @Author  : sudoskys
+from typing import Optional, Union, List
+from urllib.parse import urlparse
+
+import curl_cffi
+import httpx
+import pydantic
+from curl_cffi.requests import AsyncSession
+from loguru import logger
+from pydantic import ConfigDict, PrivateAttr, model_validator, Field
+from tenacity import wait_random, retry, stop_after_attempt, retry_if_exception
+
+from ._const import d_repetition_penalty_whitelist, d_bad_words_ids
+from ._enum import TextLLMModel, PenStyle, generate_order, TOKENIZER
+from ._preset import Params, PRESET
+from ...schema import ApiBaseModel
+from ...._exceptions import APIError, SessionHttpError
+from ...._response.ai.generate import LLMResp
+from ....credential import CredentialBase
+from ....tokenizer import LLMTokenizer
+from ....utils.encode import tokens_to_b64
+
+
+class LLM(ApiBaseModel):
+    """
+    LLM Stream for /ai/generate-stream
+    """
+    _endpoint: str = PrivateAttr("https://api.novelai.net")
+    input: str = Field(..., description="Base64 encoded token text")
+    model: Union[TextLLMModel, str] = TextLLMModel.Kayra
+    """Model to use"""
+    parameters: Params = Params()
+
+    model_config = ConfigDict(extra="ignore")
+
+    @property
+    def endpoint(self):
+        return self._endpoint
+
+    @endpoint.setter
+    def endpoint(self, value):
+        self._endpoint = value
+
+    @property
+    def base_url(self):
+        return f"{self.endpoint.strip('/')}/ai/generate"
+
+    async def necessary_headers(self, request_data) -> dict:
+        """
+        :param request_data: dict
+        :return: dict
+        """
+        return {
+            "Host": urlparse(self.endpoint).netloc,
+            "accept": "*/*",
+            "accept-language": "zh-CN,zh;q=0.9",
+            "cache-control": "no-cache",
+            "content-type": "application/json",
+            "pragma": "no-cache",
+            "Referer": "https://novelai.net/",
+            "Referrer-Policy": "strict-origin-when-cross-origin"
+        }
+
+    @model_validator(mode="after")
+    def validate_model(self):
+        if isinstance(self.input, str):
+            prompt = LLMTokenizer.encode(self.input, tokenizer_name=TOKENIZER.get(self.model))
+            self.input = tokens_to_b64(prompt)
+        if self.parameters.stop_sequences:
+            if not isinstance(self.parameters.stop_sequences, list):
+                raise pydantic.ValidationError("stop_sequences must be a list")
+            stop_sequences = []
+            for i, obj in enumerate(self.parameters.stop_sequences):
+                if isinstance(obj, str):
+                    stop_sequences.append(LLMTokenizer.encode(obj, tokenizer_name=TOKENIZER.get(self.model)))
+                elif not isinstance(obj, list):
+                    raise ValueError(
+                        f"Expected type 'str' or 'list' for item #{i} of 'stop_sequences', " f"but got '{type(obj)}'"
+                    )
+            self.parameters.stop_sequences = stop_sequences
+        if self.parameters.repetition_penalty:
+            # adjust repetition penalty value for Sigurd and Euterpe
+            if self.model in (TextLLMModel.Sigurd, TextLLMModel.Euterpe):
+                rep_pen = self.parameters.repetition_penalty
+                self.parameters.repetition_penalty = (0.525 * (rep_pen - 1) / 7) + 1
+
+        return self
+
+    @classmethod
+    def build(cls,
+              prompt: str,
+              model: TextLLMModel = TextLLMModel.Kayra,
+              min_length: Optional[int] = None,
+              max_length: Optional[int] = None,
+              temperature: Optional[float] = None,
+              top_k: Optional[int] = None,
+              top_p: Optional[float] = None,
+              frequency_penalty: Optional[float] = None,
+              presence_penalty: Optional[float] = None,
+              stop_sequences: Optional[Union[List[str]]] = None,
+              repetition_penalty: Optional[float] = None,
+              pen_style: Optional[PenStyle] = None,
+              **kwargs
+              ) -> "LLM":
+        """
+        Generate Prompt
+        """
+        assert isinstance(prompt, str), "Prompt must be a string"
+        # 清理空值
+        kwarg = {
+            "min_length": min_length,
+            "max_length": max_length,
+            "temperature": temperature,
+            "top_k": top_k,
+            "top_p": top_p,
+            "frequency_penalty": frequency_penalty,
+            "presence_penalty": presence_penalty,
+            "stop_sequences": stop_sequences,
+            "repetition_penalty": repetition_penalty,
+            "pen_style": pen_style
+        }
+        kwargs = {**kwarg, **kwargs}
+        param = {k: v for k, v in kwargs.items() if v is not None}
+        _preset_model = PRESET.get(model.name, {})
+        if not _preset_model:
+            _build_prop = Params.model_validate(param)
+        else:
+            _build_prop = _preset_model.model_copy(deep=True, update=param)
+        assert _build_prop is not None, "Params validate failed"
+        return cls(
+            input=prompt,
+            model=model,
+            parameters=_build_prop
+        )
+
+    @retry(
+        wait=wait_random(min=1, max=3),
+        stop=stop_after_attempt(3),
+        retry=retry_if_exception(lambda e: hasattr(e, "code") and str(e.code) == "500"),
+        reraise=True
+    )
+    async def request(self,
+                      session: Union[AsyncSession, "CredentialBase"],
+                      *,
+                      override_headers: Optional[dict] = None
+                      ) -> LLMResp:
+        """
+        生成图片
+        :param override_headers:
+        :param session:  session
+        :return: LLMStreamResp
+        """
+        # Data Build
+        request_data = self.model_dump(mode="json", exclude_none=True)
+        request_data["order"] = generate_order(request_data) or []
+        assert request_data.get("parameters"), "Parameters is required"
+        # Header
+        if isinstance(session, AsyncSession):
+            session.headers.update(await self.necessary_headers(request_data))
+        elif isinstance(session, CredentialBase):
+            update_header = await self.necessary_headers(request_data)
+            session = await session.get_session(update_headers=update_header)
+        if override_headers:
+            session.headers.clear()
+            session.headers.update(override_headers)
+        # Log
+        logger.debug(f"LLM request data: {request_data}")
+        # Request
+        try:
+            assert hasattr(session, "post"), "session must have post method."
+            response = await session.post(
+                self.base_url,
+                json=request_data,
+                stream=False
+            )
+            header_type = response.headers.get('Content-Type')
+            if "application/json" not in header_type:
+                try:
+                    _msg = response.json()
+                except Exception as e:
+                    _msg = {"statusCode": response.status_code, "message": response.content}
+                status_code = _msg.get("statusCode", response.status_code)
+                message = _msg.get("message", "Unknown error")
+                if status_code == 400:
+                    raise APIError(
+                        "A validation error occured.",
+                        request=request_data, code=status_code, response=_msg
+                    )
+
+                elif status_code == 401:
+                    raise APIError(
+                        "Access Token is incorrect.",
+                        request=request_data, code=status_code, response=_msg
+                    )
+                elif status_code == 402:
+                    raise APIError(
+                        "An active subscription is required to access this endpoint.",
+                        request=request_data, code=status_code, response=_msg
+                    )
+                elif status_code == 409:
+                    raise APIError(
+                        "A conflict error occured.",
+                        request=request_data, code=status_code, response=_msg
+                    )
+                else:
+                    raise APIError(
+                        f"An unknown error occured.{response.status_code}",
+                        request=request_data, code=status_code, response=_msg
+                    )
+            else:
+                output = response.json().get("output", None)
+                assert output, APIError("No Content Returned",
+                                        request=request_data, code=response.status_code, response=response.content
+                                        )
+                return LLMResp(
+                    output=output,
+                    text=LLMResp.decode_token(model=self.model, token_str=output)
+                )
+        except curl_cffi.requests.errors.RequestsError as exc:
+            logger.exception(exc)
+            raise SessionHttpError(
+                "An AsyncSession RequestsError occurred, perhaps it's an SSL error. Please try again later!"
+            )
+        except httpx.HTTPError as exc:
+            logger.exception(exc)
+            raise SessionHttpError(
+                "An HTTPError occurred, perhaps it's an SSL error. Please try again later!"
+            )
+        except APIError as e:
+            raise e
+        except Exception as e:
+            logger.opt(exception=e).exception("An unexpected error occurred")
+            raise e