jmux 0.0.1__py3-none-any.whl

jmux/error.py

@@ -0,0 +1,141 @@
+from enum import Enum
+from typing import Sequence
+
+
+class MissingAttributeError(Exception):
+    def __init__(self, object_name: str, attribute: str) -> None:
+        super().__init__(f"'{object_name}' is missing required attribute '{attribute}'")
+
+
+class UnexpectedAttributeTypeError(Exception):
+    def __init__(self, object_name: str, attribute: str, expected_type: str) -> None:
+        super().__init__(
+            f"'{object_name}' has attribute '{attribute}' with unexpected type. "
+            f"Attribute must conform to {expected_type}."
+        )
+
+
+class EmptyKeyError(Exception):
+    def __init__(self, message: str | None = None) -> None:
+        super().__init__("Key cannot be empty" + (f": {message}" if message else ""))
+        self.message = message
+
+
+class TypeEmitError(Exception):
+    def __init__(self, expected_type: str, actual_type: str) -> None:
+        super().__init__(
+            f"Cannot emit to current sink. Type mismatch: expected {expected_type}, "
+            f"got {actual_type}"
+        )
+        self.expected_type = expected_type
+        self.actual_type = actual_type
+
+
+class NoCurrentSinkError(Exception):
+    def __init__(self, message: str | None = None) -> None:
+        super().__init__(
+            "No current sink available" + (f": {message}" if message else "")
+        )
+
+
+class ParsePrimitiveError(Exception):
+    def __init__(self, message: str | None = None) -> None:
+        super().__init__(
+            "Failed to parse primitive value" + (f": {message}" if message else "")
+        )
+
+
+class NothingEmittedError(Exception):
+    def __init__(self, message: str | None = None) -> None:
+        super().__init__(
+            "Nothing was emitted to the sink" + (f": {message}" if message else "")
+        )
+        self.message = message
+
+
+class SinkClosedError(Exception):
+    def __init__(self, message: str | None = None) -> None:
+        super().__init__(
+            "Sink is closed and cannot accept new items"
+            + (f": {message}" if message else "")
+        )
+        self.message = message
+
+
+class NotAllObjectPropertiesSetError(Exception):
+    def __init__(self, message: str | None = None) -> None:
+        super().__init__(
+            "Not all properties were set before closing the sink respective stream"
+            + (f": {message}" if message else "")
+        )
+
+
+class ObjectAlreadyClosedError(Exception):
+    def __init__(self, object_name: str, message: str | None = None) -> None:
+        super().__init__(
+            f"Object '{object_name}' is already closed"
+            + (f": {message}" if message else "")
+        )
+
+
+class ObjectMissmatchedError(Exception):
+    def __init__(
+        self,
+        jmux_model: str,
+        pydantic_model: str,
+        attribute: str,
+        message: str | None = None,
+    ) -> None:
+        super().__init__(
+            f"JMux object '{jmux_model}' and '{pydantic_model}' are missmatched on "
+            f"attribute '{attribute}'" + (f": {message}" if message else "")
+        )
+
+
+class ForbiddenTypeHintsError(Exception):
+    def __init__(self, message: str | None = None) -> None:
+        super().__init__(
+            "Forbidden type hints used in object definition"
+            + (f": {message}" if message else "")
+        )
+
+
+class UnexpectedCharacterError(Exception):
+    def __init__(
+        self,
+        character: str,
+        pda_stack: Sequence[Enum] | str,
+        pda_state: Enum | str,
+        message: str | None,
+    ) -> None:
+        pda_stack_str = [
+            item.value if isinstance(item, Enum) else item for item in pda_stack
+        ]
+        pda_state_str = pda_state.value if isinstance(pda_state, Enum) else pda_state
+        super().__init__(
+            f"Received unexpected character '{character}' in state '{pda_state_str}' "
+            f"with stack {pda_stack_str}" + (f": {message}" if message else "")
+        )
+
+
+class StreamParseError(Exception):
+    def __init__(self, message: str | None = None) -> None:
+        super().__init__("Failed to parse stream" + (f": {message}" if message else ""))
+        self.message = message
+
+
+class UnexpectedStateError(Exception):
+    def __init__(
+        self,
+        pda_stack: Sequence[Enum] | str,
+        pda_state: Enum | str,
+        message: str | None = None,
+    ) -> None:
+        pda_stack_str = [
+            item.value if isinstance(item, Enum) else item for item in pda_stack
+        ]
+        pda_state_str = pda_state.value if isinstance(pda_state, Enum) else pda_state
+        super().__init__(
+            f"Unexpected state '{pda_state_str}' with stack {pda_stack_str}"
+            + (f": {message}" if message else "")
+        )

jmux/helpers.py

@@ -0,0 +1,60 @@
+from types import NoneType, UnionType
+from typing import Set, Tuple, Type, Union, get_args, get_origin
+
+
+def is_json_whitespace(ch: str) -> bool:
+    return ch in {" ", "\t", "\n", "\r"}
+
+
+def extract_types_from_generic_alias(UnknownType: Type) -> Tuple[Set[Type], Set[Type]]:
+    Origin: Type | None = get_origin(UnknownType)
+    if Origin is None:
+        return {UnknownType}, set()
+    if Origin is UnionType or Origin is Union:
+        return deconstruct_type(UnknownType), set()
+
+    type_args = get_args(UnknownType)
+    if len(type_args) != 1:
+        raise TypeError(
+            f"Only single type generics can be deconstruct with this function, "
+            f"got {type_args}."
+        )
+
+    Generic: Type = type_args[0]
+    type_set = deconstruct_type(Generic)
+    if len(type_set) == 1:
+        return {Origin}, type_set
+    if len(type_set) != 2:
+        raise TypeError(
+            f"Union type must have exactly two types in its union, "
+            f"got {get_args(Generic)}."
+        )
+    if NoneType not in get_args(Generic):
+        raise TypeError(
+            "Union type must include NoneType if it is used as a generic argument."
+        )
+    return {Origin}, type_set
+
+
+def deconstruct_type(UnknownType: Type) -> Set[Type]:
+    Origin: Type | None = get_origin(UnknownType)
+    if UnknownType is None:
+        return {NoneType}
+    if Origin is None:
+        return {UnknownType}
+    if not (Origin is UnionType or Origin is Union):
+        return {Origin}
+    type_args = get_args(UnknownType)
+    return set(type_args)
+
+
+def get_main_type(type_set: Set[Type]) -> Type:
+    type_set_copy: Set[Type] = type_set.copy()
+    if NoneType in type_set_copy and len(type_set_copy) == 2:
+        type_set_copy.remove(NoneType)
+    if len(type_set_copy) != 1:
+        raise TypeError(
+            f"Expected exactly one type, got {type_set_copy}. If you want to allow "
+            "NoneType, use Union[int, NoneType]."
+        )
+    return type_set_copy.pop()

jmux/pda.py

@@ -0,0 +1,32 @@
+from typing import List
+
+
+class PushDownAutomata[Context, State]:
+    def __init__(self, start_state: State) -> None:
+        self._stack: List[Context] = []
+        self._state: State = start_state
+
+    @property
+    def state(self) -> State:
+        return self._state
+
+    @property
+    def stack(self) -> List[Context]:
+        return self._stack
+
+    @property
+    def top(self) -> Context | None:
+        if not self._stack:
+            return None
+        return self._stack[-1]
+
+    def set_state(self, new_state: State) -> None:
+        self._state = new_state
+
+    def push(self, mode: Context) -> None:
+        self._stack.append(mode)
+
+    def pop(self) -> Context:
+        if not self._stack:
+            raise IndexError("PDA stack is empty.")
+        return self._stack.pop()

jmux/types.py

@@ -0,0 +1,57 @@
+from enum import Enum
+from typing import Set
+
+
+class State(Enum):
+    START = "start"
+    END = "end"
+    ERROR = "error"
+    # expect
+    EXPECT_KEY = "expect_key"
+    EXPECT_COLON = "expect_colon"
+    EXPECT_VALUE = "expect_value"
+    EXPECT_COMMA_OR_EOC = "expect_comma_or_eoc"
+    # parsing
+    PARSING_KEY = "parsing_key"
+    PARSING_STRING = "parsing_string"
+    PARSING_INTEGER = "parsing_integer"
+    PARSING_FLOAT = "parsing_float"
+    PARSING_BOOLEAN = "parsing_boolean"
+    PARSING_NULL = "parsing_null"
+    PARSING_OBJECT = "parsing_object"
+
+
+PRIMITIVE_STATES: Set[State] = {
+    State.PARSING_INTEGER,
+    State.PARSING_FLOAT,
+    State.PARSING_BOOLEAN,
+    State.PARSING_NULL,
+}
+
+
+class Mode(Enum):
+    ROOT = "$"
+    OBJECT = "object"
+    ARRAY = "array"
+
+
+OBJECT_OPEN = set("{")
+OBJECT_CLOSE = set("}")
+COLON = set(":")
+ARRAY_OPEN = set("[")
+ARRAY_CLOSE = set("]")
+COMMA = set(",")
+QUOTE = set('"')
+
+NUMBER_OPEN = set("0123456789-")
+BOOLEAN_OPEN = set("tf")
+NULL_OPEN = set("n")
+
+INTERGER_ALLOWED = set("0123456789")
+FLOAT_ALLOWED = set("0123456789-+eE.")
+BOOLEAN_ALLOWED = set("truefals")
+NULL_ALLOWED = set("nul")
+
+JSON_FALSE = "false"
+JSON_TRUE = "true"
+JSON_NULL = "null"

jmux-0.0.1.dist-info/METADATA

@@ -0,0 +1,364 @@
+Metadata-Version: 2.4
+Name: jmux
+Version: 0.0.1
+Summary: JMux: A Python package for demultiplexing a JSON string into multiple awaitable variables.
+Author-email: "Johannes A.I. Unruh" <johannes@unruh.ai>
+License: MIT License
+        
+        Copyright (c) 2025 Johannes A.I. Unruh
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, subject to the following conditions:
+        
+        1. The above copyright notice and this permission notice shall be included in all
+           copies or substantial portions of the Software.
+        
+        2. You may use this software, including for commercial purposes, as long as
+           you do not sell, license, or otherwise distribute the original or
+           substantially similar versions of this software for a fee.
+        
+        3. This restriction does not apply to using this software as a dependency in
+           your own commercial applications or products, provided you are not selling
+           this software itself as a standalone product.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+Project-URL: Homepage, https://github.com/jaunruh/jmux
+Project-URL: Repository, https://github.com/jaunruh/jmux
+Keywords: demultiplexer,python,package,json
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: anyio>=4.0.0
+Requires-Dist: pydantic>=2.0.0
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest-anyio; extra == "test"
+Provides-Extra: dev
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-anyio; extra == "dev"
+Requires-Dist: uv; extra == "dev"
+Requires-Dist: build; extra == "dev"
+Requires-Dist: twine; extra == "dev"
+Requires-Dist: setuptools; extra == "dev"
+Requires-Dist: setuptools_scm[toml]; extra == "dev"
+Dynamic: license-file
+
+# JMux: A Python package for demultiplexing a JSON string into multiple awaitable variables.
+
+JMux is a powerful Python package that allows you to demultiplex a JSON stream into multiple awaitable variables. It is specifically designed for asynchronous applications that interact with Large Language Models (LLMs) using libraries like `litellm`. When an LLM streams a JSON response, `jmux` enables you to parse and use parts of the JSON object _before_ the complete response has been received, significantly improving responsiveness.
+
+## Inspiration
+
+This package is inspired by `Snapshot Streaming` mentioned in the [`WWDC25: Meet the Foundation Models framework`](https://youtu.be/mJMvFyBvZEk?si=DVIvxzuJOA87lb7I&t=465) keynote by Apple.
+
+## Features
+
+- **Asynchronous by Design**: Built on top of `asyncio`, JMux is perfect for modern, high-performance Python applications.
+- **Pydantic Integration**: Validate your `JMux` classes against Pydantic models to ensure type safety and consistency.
+- **Awaitable and Streamable Sinks**: Use `AwaitableValue` for single values and `StreamableValues` for streams of values.
+- **Robust Error Handling**: JMux provides a comprehensive set of exceptions to handle parsing errors and other issues.
+- **Lightweight**: JMux has only a few external dependencies, making it easy to integrate into any project.
+
+## Installation
+
+You can install JMux from PyPI using pip:
+
+```bash
+pip install jmux
+```
+
+## Usage with LLMs (e.g., `litellm`)
+
+The primary use case for `jmux` is to process streaming JSON responses from LLMs. This allows you to react to parts of the data as it arrives, rather than waiting for the entire JSON object to be transmitted. While this should be obvious, I should mention, that **the order in which the pydantic model defines the properties, defines which stream is filled first**.
+
+Here’s a conceptual example of how you might integrate `jmux` with an LLM call, such as one made with `litellm`:
+
+```python
+import asyncio
+from pydantic import BaseModel
+from jmux import JMux, AwaitableValue, StreamableValues
+# litellm is used conceptually here
+# from litellm import acompletion
+
+# 1. Define the Pydantic model for the expected JSON response
+class LlmResponse(BaseModel):
+    thought: str # **This property is filled first**
+    tool_code: str
+
+# 2. Define the corresponding JMux class
+class LlmResponseMux(JMux):
+    thought: AwaitableValue[str]
+    tool_code: StreamableValues[str] # Stream the code as it's generated
+
+# 3. Validate that the JMux class matches the Pydantic model
+LlmResponseMux.assert_conforms_to(LlmResponse)
+
+# A mock function that simulates a streaming LLM call
+async def mock_llm_stream():
+    json_stream = '{"thought": "I need to write some code.", "tool_code": "print(\'Hello, World!\')"}'
+    for char in json_stream:
+        yield char
+        await asyncio.sleep(0.01) # Simulate network latency
+
+# Main function to orchestrate the call and processing
+async def process_llm_response():
+    jmux_instance = LlmResponseMux()
+
+    # This task will consume the LLM stream and feed it to jmux
+    async def feed_stream():
+        async for chunk in mock_llm_stream():
+            await jmux_instance.feed_chunks(chunk)
+
+    # These tasks will consume the demultiplexed data from jmux
+    async def consume_thought():
+        thought = await jmux_instance.thought
+        print(f"LLM's thought received: '{thought}'")
+        # You can act on the thought immediately
+        # without waiting for the tool_code to finish streaming.
+
+    async def consume_tool_code():
+        print("Receiving tool code...")
+        full_code = ""
+        async for code_fragment in jmux_instance.tool_code:
+            full_code += code_fragment
+            print(f"  -> Received fragment: {code_fragment}")
+        print(f"Full tool code received: {full_code}")
+
+    # Run all tasks concurrently
+    await asyncio.gather(
+        feed_stream(),
+        consume_thought(),
+        consume_tool_code()
+    )
+
+if __name__ == "__main__":
+    asyncio.run(process_llm_response())
+```
+
+## Example Implementation
+
+<details>
+<summary>Python Code</summary>
+
+```python
+def create_json_streaming_completion[T: BaseModel, J: IJsonDemuxer](
+        self,
+        messages: List[ILlmMessage],
+        ReturnType: Type[T],
+        JMux: Type[J],
+        retries: int = 3,
+    ) -> StreamResponseTuple[T, J]:
+        try:
+            JMux.assert_conforms_to(ReturnType)
+            litellm_messages = self._convert_messages(messages)
+            jmux_instance: J = JMux()
+
+            async def stream_feeding_llm_call() -> T:
+                nonlocal jmux_instance
+                buffer = ""
+                stream: CustomStreamWrapper = await self._router.acompletion( # see litellm `router`
+                    model=self._internal_model_name.value,
+                    messages=litellm_messages,
+                    stream=True,
+                    num_retries=retries,
+                    response_format=ReturnType,
+                    **self._maybe_google_credentials_param,
+                    **self._model_params.model_dump(exclude_none=True),
+                    **self._additional_params,
+                )
+
+                async for chunk in stream:
+                    content_fragment: str | None = None
+
+                    tool_calls = chunk.choices[0].delta.tool_calls
+                    if tool_calls:
+                        content_fragment = tool_calls[0].function.arguments
+                    elif chunk.choices[0].delta.content:
+                        content_fragment = chunk.choices[0].delta.content
+
+                    if content_fragment:
+                        try:
+                            buffer += content_fragment
+                            await jmux_instance.feed_chunks(content_fragment)
+                        except Exception as e:
+                            logger.warning(f"error in JMux feed_chunks: {e}")
+                            raise e
+
+                return ReturnType.model_validate_json(buffer)
+
+            awaitable_llm_result = create_task(stream_feeding_llm_call())
+            return (awaitable_llm_result, jmux_instance)
+        except Exception as e:
+            logger.warning(f"error in create_json_streaming_completion: {e}")
+            raise e
+```
+
+The code above shows an example implementation that uses a `litellm` router for `acompletion`.
+
+You can either `await awaitable_llm_result` if you need the full result, or use `await jmux_instance.your_awaitable_value` or `async for ele in jmux_instance.your_streamable_values` to access partial results.
+
+</details>
+
+## Basic Usage
+
+Here is a simple example of how to use JMux to parse a JSON stream:
+
+```python
+import asyncio
+from enum import Enum
+from types import NoneType
+from pydantic import BaseModel
+
+from jmux import JMux, AwaitableValue, StreamableValues
+
+# 1. Define your JMux class
+class SObject(JMux):
+    class SNested(JMux):
+        key_str: AwaitableValue[str]
+
+    class SEnum(Enum):
+        VALUE1 = "value1"
+        VALUE2 = "value2"
+
+    key_str: AwaitableValue[str]
+    key_int: AwaitableValue[int]
+    key_float: AwaitableValue[float]
+    key_bool: AwaitableValue[bool]
+    key_none: AwaitableValue[NoneType]
+    key_stream: StreamableValues[str]
+    key_enum: AwaitableValue[SEnum]
+    key_nested: AwaitableValue[SNested]
+
+# 2. (Optional) Define a Pydantic model for validation
+class PObject(BaseModel):
+    class PNested(BaseModel):
+        key_str: str
+
+    class PEnum(Enum):
+        VALUE1 = "value1"
+        VALUE2 = "value2"
+
+    key_str: str
+    key_int: int
+    key_float: float
+    key_bool: bool
+    key_none: NoneType
+    key_stream: str
+    key_enum: PEnum
+    key_nested: PNested
+
+# 3. Validate the JMux class against the Pydantic model
+SObject.assert_conforms_to(PObject)
+
+# 4. Create an instance of your JMux class
+s_object = SObject()
+
+# 5. Feed the JSON stream to the JMux instance
+async def main():
+    json_stream = '{"key_str": "hello", "key_int": 42, "key_float": 3.14, "key_bool": true, "key_none": null, "key_stream": "world", "key_enum": "value1", "key_nested": {"key_str": "nested"}}'
+
+    async def produce():
+        for char in json_stream:
+            await s_object.feed_char(char)
+
+    async def consume():
+        key_str = await s_object.key_str
+        print(f"key_str: {key_str}")
+
+        key_int = await s_object.key_int
+        print(f"key_int: {key_int}")
+
+        key_float = await s_object.key_float
+        print(f"key_float: {key_float}")
+
+        key_bool = await s_object.key_bool
+        print(f"key_bool: {key_bool}")
+
+        key_none = await s_object.key_none
+        print(f"key_none: {key_none}")
+
+        key_stream = ""
+        async for char in s_object.key_stream:
+            key_stream += char
+        print(f"key_stream: {key_stream}")
+
+        key_enum = await s_object.key_enum
+        print(f"key_enum: {key_enum}")
+
+        key_nested = await s_object.key_nested
+        nested_key_str = await key_nested.key_str
+        print(f"nested_key_str: {nested_key_str}")
+
+    await asyncio.gather(produce(), consume())
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```
+
+## API Reference
+
+### Abstract Calss `jmux.JMux`
+
+The abstract base class for creating JSON demultiplexers.
+
+> `JMux.assert_conforms_to(pydantic_model: Type[BaseModel]) -> None`
+
+Asserts that the JMux class conforms to a given Pydantic model.
+
+> `async JMux.feed_char(ch: str) -> None`
+
+Feeds a character to the JMux parser.
+
+> `async JMux.feed_chunks(chunks: str) -> None`
+
+Feeds a string of characters to the JMux parser.
+
+### Class `jmux.AwaitableValue[T]`
+
+A class that represents a value that will be available in the future. You are awaiting the full value and do not get partial results.
+
+Allowed types here are (they can all be combined with `Optional`):
+
+- `int`, `float`, `str`, `bool`, `NoneType`
+- `JMux`
+- `Enum`
+
+In all cases, the corresponding `pydantic.BaseModel` should **not** be `list`
+
+### Class `jmux.StreamableValues[T]`
+
+A class that represents a stream of values that can be asynchronously iterated over.
+
+Allowed types are listed below and should all be wrapped in a `list` on the pydantic model:
+
+- `int`, `float`, `str`, `bool`, `NoneType`
+- `JMux`
+- `Enum`
+
+Additionally the following type is supported without being wrapped into `list`:
+
+- `str`
+
+This allows you to fully stream strings directly to a sink.
+
+## License
+
+This project is licensed under the terms of the MIT license. See the [LICENSE](LICENSE) file for details.
+
+## Planned Improvements
+
+- Add support for older Python versions
+
+## Contributions
+
+As you might see, this repo has only been created recently and so far I am the only developer working on it. If you want to contribute, reach out via `johannes@unruh.ai` or `johannes.a.unruh@gmail.com`.
+
+If you have suggestions or find any errors in my implementation, feel free to create an issue or also reach out via email.

jmux-0.0.1.dist-info/RECORD

@@ -0,0 +1,13 @@
+jmux/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+jmux/awaitable.py,sha256=gceBygIf3fAIWLsN1lWxsz9ExWNasDuk1WaGz8d9FAc,8427
+jmux/decoder.py,sha256=Y6KVryRDLvGV5nBsneXpTvC0WUGhR5Z89Dvqz4HMAgg,1562
+jmux/demux.py,sha256=ouhu2PIJkaEIsToNfDhJWGQYjFYULjqFSEX5ElmAjPs,34861
+jmux/error.py,sha256=VZJYivt8RPfjcF2bs-T7_UkH3dVA3xH-xGbZggQV14k,4665
+jmux/helpers.py,sha256=6y33RohUGVvGGaAREyRQTmEWfgV4w295SyqDwIMnUNs,1982
+jmux/pda.py,sha256=81gnh0eWGsgd_SrHkqjRQy_KkOSlBf5nor7pqKGgYjw,791
+jmux/types.py,sha256=V63wMx1I7l_P83JAqzOQ7H7B-xKrNUuGshJ3NjKsdeQ,1192
+jmux-0.0.1.dist-info/licenses/LICENSE,sha256=y0qnwaAe4bEqzNPyq4M_VZA2I2mQly8MawajyZhqw0k,1169
+jmux-0.0.1.dist-info/METADATA,sha256=Y4E3GeCPbNYvUv6B877m_c7DZUXlZO5voWCbxXVtwSg,13330
+jmux-0.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+jmux-0.0.1.dist-info/top_level.txt,sha256=TF2N6kHqLghfOkCiNlCueMDX4l5rPn_5MSPNtYrS1-o,5
+jmux-0.0.1.dist-info/RECORD,,

jmux-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+