feat: creates stubs

Author: flatypus

flowchat/__init__.pyi

@@ -0,0 +1,3 @@
+from .types import *
+from .autodedent import autodedent as autodedent
+from .chain import Chain as Chain

flowchat/autodedent.pyi

@@ -0,0 +1,3 @@
+from typing import Any
+
+def autodedent(*text_lines: Any) -> str: ...

flowchat/chain.py

@@ -2,8 +2,10 @@
 from .private._private_helpers import encode_image, wrap_stream_and_count, async_wrap_stream_and_count
 from .types import *
 from datetime import datetime
+from openai.types.chat.chat_completion import ChatCompletion
 from typing import List, Optional, Union, Callable, Any, Generator, AsyncGenerator
 from typing_extensions import Unpack
+from openai import Stream, AsyncStream
 import json
 import logging
 import openai
@@ -125,7 +127,8 @@ def _ask(
             return None
 
         if stream and isinstance(completion, Stream):
-            return wrap_stream_and_count(completion, model, self._add_token_count, plain_text_stream)
+            return wrap_stream_and_count(
+                completion, model, self._add_token_count, plain_text_stream)  # type: ignore
 
         return self._post_completion(completion, model, json_schema)
 
@@ -152,7 +155,8 @@ async def async_ask(
             return None
 
         if stream and isinstance(completion, AsyncStream):
-            return async_wrap_stream_and_count(completion, model, self._add_token_count, plain_text_stream)
+            return async_wrap_stream_and_count(
+                completion, model, self._add_token_count, plain_text_stream)  # type: ignore
 
         return self._post_completion(completion, model, json_schema)
 
@@ -253,7 +257,7 @@ def link(self, modifier: Union[Callable[[str | Any | None], str], str], model: s
             )
         return self
 
-    def setup_pull(self, asynchronous: bool = False, json_schema: Optional[dict[Any, Any]] = None, **params: Unpack[RequestParams]):
+    def setup_pull(self, asynchronous: bool = False, json_schema: Optional[dict[Any, Any]] = None, **params: Unpack[RequestParams]) -> Any:
         params['model'] = params.get('model', self.model)
 
         if len(self.user_prompt) == 0:

flowchat/chain.pyi

@@ -0,0 +1,37 @@
+from .types import *
+from .autodedent import autodedent as autodedent
+from .private._private_helpers import async_wrap_stream_and_count as async_wrap_stream_and_count, encode_image as encode_image, wrap_stream_and_count as wrap_stream_and_count
+from _typeshed import Incomplete
+from typing import Any, AsyncGenerator, Callable, Generator
+from typing_extensions import Unpack
+
+class Chain:
+    client: Incomplete
+    asyncClient: Incomplete
+    model: Incomplete
+    system: Incomplete
+    user_prompt: Incomplete
+    model_response: Incomplete
+    raw_model_response: Incomplete
+    usage: Incomplete
+    detailed_usage: Incomplete
+    def __init__(self, model: str, api_key: str = '', environ_key: str = 'OPENAI_API_KEY') -> None: ...
+    async def async_ask(self, system: Message | None, user_messages: list[Message], json_schema: dict[Any, Any] | None = None, stream: bool = False, plain_text_stream: bool = False, **params: Unpack[RequestParams]) -> Any: ...
+    def unhook(self) -> Chain: ...
+    def anchor(self, system_prompt: str) -> Chain: ...
+    def transform(self, function: Callable[[Any], Any]) -> Chain: ...
+    def link(self, modifier: Callable[[str | Any | None], str] | str, model: str | None = None, assistant: bool = False, images: str | Any | list[str | Any] | ImageFormat = None) -> Chain: ...
+    def setup_pull(self, asynchronous: bool = False, json_schema: dict[Any, Any] | None = None, **params: Unpack[RequestParams]) -> Any: ...
+    def pull(self, json_schema: dict[Any, Any] | None = None, **params: Unpack[RequestParams]) -> Chain: ...
+    async def async_pull(self, json_schema: dict[Any, Any] | None = None, **params: Unpack[RequestParams]) -> Chain: ...
+    def stream(self, plain_text_stream: bool = False, **params: Unpack[RequestParams]) -> Generator[str | Any | None, None, None]: ...
+    async def async_stream(self, plain_text_stream: bool = False, **params: Unpack[RequestParams]) -> AsyncGenerator[str | Any | None, None]: ...
+    def last(self) -> Any: ...
+    def token_usage(self) -> Usage: ...
+    def detailed_token_usage(self) -> list[DetailedUsage]: ...
+    def reset_token_usage(self) -> Chain: ...
+    def subscribe_token_usage(self, usage_object: Usage) -> Chain: ...
+    def subscribe_detailed_token_usage(self, detailed_usage_object: list[DetailedUsage]) -> Chain: ...
+    def log(self) -> Chain: ...
+    def log_tokens(self) -> Chain: ...
+    def log_detailed_tokens(self) -> Chain: ...

flowchat/private/__init__.pyi

@@ -1,18 +1,19 @@
-from typing import Callable
+from typing import Callable, Generator, AsyncGenerator
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
 from ..types import *
 from io import BytesIO
 from PIL.Image import Image as PILImage
 import base64
 
 
-def encode_image(image: PILImage, format_type: str = "PNG"):
+def encode_image(image: PILImage, format_type: str = "PNG") -> str:
     buffered = BytesIO()
     image.save(buffered, format=format_type)
     img_str = base64.b64encode(buffered.getvalue())
     return f"data:image/png;base64,{img_str.decode('utf-8')}"
 
 
-def wrap_stream_and_count(generator: StreamCompletion, model: str, callback: Callable[[int, int, str], None], plain_text_stream: bool = False):
+def wrap_stream_and_count(generator: StreamCompletion, model: str, callback: Callable[[int, int, str], None], plain_text_stream: bool = False) -> Generator[str | ChatCompletionChunk | None, None, None]:
     for response in generator:
         if len(response.choices) == 0:
             usage = response.usage
@@ -23,7 +24,7 @@ def wrap_stream_and_count(generator: StreamCompletion, model: str, callback: Cal
         yield content if plain_text_stream else response
 
 
-async def async_wrap_stream_and_count(generator: AsyncStreamCompletion, model: str, callback: Callable[[int, int, str], None], plain_text_stream: bool = False):
+async def async_wrap_stream_and_count(generator: AsyncStreamCompletion, model: str, callback: Callable[[int, int, str], None], plain_text_stream: bool = False) -> AsyncGenerator[str | ChatCompletionChunk | None, None]:
     async for response in generator:
         if len(response.choices) == 0:
             usage = response.usage

flowchat/private/_private_helpers.py

@@ -0,0 +1,8 @@
+from ..types import *
+from PIL.Image import Image as PILImage
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk as ChatCompletionChunk
+from typing import AsyncGenerator, Callable, Generator
+
+def encode_image(image: PILImage, format_type: str = 'PNG') -> str: ...
+def wrap_stream_and_count(generator: StreamCompletion, model: str, callback: Callable[[int, int, str], None], plain_text_stream: bool = False) -> Generator[str | ChatCompletionChunk | None, None, None]: ...
+async def async_wrap_stream_and_count(generator: AsyncStreamCompletion, model: str, callback: Callable[[int, int, str], None], plain_text_stream: bool = False) -> AsyncGenerator[str | ChatCompletionChunk | None, None]: ...

flowchat/private/_private_helpers.pyi

@@ -0,0 +1,70 @@
+from PIL.Image import Image as PILImage
+from _typeshed import Incomplete
+from datetime import datetime
+from openai import AsyncStream, Stream
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+from typing import Any, Literal, NotRequired
+from typing_extensions import TypedDict
+
+StreamCompletion = Stream[ChatCompletionChunk]
+AsyncStreamCompletion = AsyncStream[ChatCompletionChunk]
+CreateResponse: Incomplete
+
+class Message(TypedDict):
+    role: str
+    content: str | list[Any]
+
+class ResponseFormat(TypedDict):
+    type: Literal['text', 'json_object']
+
+class ImageFormat(TypedDict):
+    url: str | PILImage
+    format_type: str
+    detail: Literal['low', 'high']
+
+class StreamOptions(TypedDict):
+    include_usage: bool
+
+class Function(TypedDict):
+    name: str
+    description: NotRequired[str]
+    parameters: NotRequired[Any]
+
+class Tool(TypedDict):
+    type: str
+    function: Function
+
+class ToolChoiceFunction(TypedDict):
+    name: str
+
+class ToolChoice(TypedDict):
+    type: str
+    function: ToolChoiceFunction
+
+class RequestParams(TypedDict, total=False):
+    model: NotRequired[str]
+    frequency_penalty: NotRequired[float | int]
+    logit_bias: NotRequired[dict[Any, Any]]
+    logprobs: NotRequired[bool]
+    top_logprobs: NotRequired[int]
+    max_tokens: NotRequired[float | int]
+    n: NotRequired[float | int]
+    presence_penalty: NotRequired[float | int]
+    response_format: NotRequired[ResponseFormat]
+    seed: NotRequired[int]
+    stop: NotRequired[str | list[str]]
+    stream_options: NotRequired[StreamOptions]
+    temperature: NotRequired[float | int]
+    top_p: NotRequired[float | int]
+    tools: NotRequired[list[Tool]]
+    tool_choice: NotRequired[str | ToolChoice]
+    user: NotRequired[str]
+
+class Usage(TypedDict):
+    prompt_tokens: int
+    completion_tokens: int
+
+class DetailedUsage(TypedDict):
+    model: str
+    usage: Usage
+    time: datetime

flowchat/types.pyi

@@ -4,7 +4,7 @@
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()
 
-VERSION = '1.3.3'
+VERSION = '1.4.0'
 DESCRIPTION = 'Streamlining the process of multi-prompting LLMs with chains'
 
 setup(

setup.py

@@ -0,0 +1,3 @@
+import subprocess
+
+subprocess.run(["stubgen", "-p", "flowchat", "-o", "."])