msgflux 0.1.0__py3-none-any.whl

msgflux/__init__.py

@@ -0,0 +1,104 @@
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from msgspec_ext import load_dotenv
+
+    from msgflux.cache import response_cache
+    from msgflux.core.dotdict import dotdict
+    from msgflux.core.examples import Example
+    from msgflux.core.message import Message
+    from msgflux.data.dbs import DB
+    from msgflux.data.parsers import Parser
+    from msgflux.data.retrievers import Retriever
+    from msgflux.data.types import Audio, File, Image, Video
+    from msgflux.dsl.inline import Inline
+    from msgflux.dsl.signature import InputField, OutputField, Signature
+    from msgflux.envs import set_envs
+    from msgflux.exceptions import TaskError
+    from msgflux.models import Model
+    from msgflux.models.gateway import ModelGateway
+    from msgflux.telemetry import Spans
+    from msgflux.tools.config import tool_config
+    from msgflux.utils.chat import ChatBlock, ChatML
+    from msgflux.utils.console import cprint
+    from msgflux.utils.inspect import get_fn_name
+    from msgflux.utils.msgspec import load, msgspec_dumps, save
+
+__all__ = [
+    "DB",
+    "Audio",
+    "ChatBlock",
+    "ChatML",
+    "Example",
+    "File",
+    "Image",
+    "Inline",
+    "InputField",
+    "Message",
+    "Model",
+    "ModelGateway",
+    "OutputField",
+    "Parser",
+    "Retriever",
+    "Signature",
+    "Spans",
+    "TaskError",
+    "Video",
+    "cprint",
+    "dotdict",
+    "get_fn_name",
+    "load",
+    "load_dotenv",
+    "msgspec_dumps",
+    "response_cache",
+    "save",
+    "set_envs",
+    "tool_config",
+]
+
+_LAZY_IMPORTS = {
+    "Audio": ("msgflux.data.types", "Audio"),
+    "ChatBlock": ("msgflux.utils.chat", "ChatBlock"),
+    "ChatML": ("msgflux.utils.chat", "ChatML"),
+    "DB": ("msgflux.data.dbs", "DB"),
+    "Example": ("msgflux.core.examples", "Example"),
+    "File": ("msgflux.data.types", "File"),
+    "Image": ("msgflux.data.types", "Image"),
+    "Inline": ("msgflux.dsl.inline", "Inline"),
+    "InputField": ("msgflux.dsl.signature", "InputField"),
+    "Message": ("msgflux.core.message", "Message"),
+    "Model": ("msgflux.models", "Model"),
+    "ModelGateway": ("msgflux.models.gateway", "ModelGateway"),
+    "OutputField": ("msgflux.dsl.signature", "OutputField"),
+    "Parser": ("msgflux.data.parsers", "Parser"),
+    "Retriever": ("msgflux.data.retrievers", "Retriever"),
+    "Signature": ("msgflux.dsl.signature", "Signature"),
+    "Spans": ("msgflux.telemetry", "Spans"),
+    "TaskError": ("msgflux.exceptions", "TaskError"),
+    "Video": ("msgflux.data.types", "Video"),
+    "cprint": ("msgflux.utils.console", "cprint"),
+    "dotdict": ("msgflux.core.dotdict", "dotdict"),
+    "get_fn_name": ("msgflux.utils.inspect", "get_fn_name"),
+    "load": ("msgflux.utils.msgspec", "load"),
+    "load_dotenv": ("msgspec_ext", "load_dotenv"),
+    "msgspec_dumps": ("msgflux.utils.msgspec", "msgspec_dumps"),
+    "response_cache": ("msgflux.cache", "response_cache"),
+    "save": ("msgflux.utils.msgspec", "save"),
+    "set_envs": ("msgflux.envs", "set_envs"),
+    "tool_config": ("msgflux.tools.config", "tool_config"),
+}
+
+
+def __getattr__(name: str):
+    if name not in _LAZY_IMPORTS:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    module_name, attr_name = _LAZY_IMPORTS[name]
+    value = getattr(import_module(module_name), attr_name)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(__all__))

msgflux/_private/client.py

@@ -0,0 +1,105 @@
+import asyncio
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+from msgflux._private.core import Core
+from msgflux._private.executor import Executor
+
+
+class BaseClient(ABC, Core):
+    """BaseClient is an abstract base class that defines the interface and
+    common behavior for client classes in the msgflux system.
+
+    This class inherits from `Core` to leverage its state management
+    capabilities and extends it with additional functionality specific to
+    clients. It requires subclasses to implement certain methods, ensuring
+    a consistent interface across different client implementations.
+
+    Subclasses must implement the `_initialize` and `__call__` methods,
+    which are essential for client initialization and operation,
+    respectively. Additionally, this class provides methods for
+    serialization and deserialization, allowing client instances to be
+    easily saved and restored.
+
+    Attributes:
+        to_ignore ([optional[list[str]]): A list of attribute keys to be
+            ignored during serialization.
+        msgflux_type (str): A string identifying the type of the client
+            in the msgflux system.
+        provider (optional[str]): The provider associated with the client,
+            if applicable.
+        instance_type (optional[callable]): A method or callable that
+            returns additional instance-specific data to be included during
+            serialization.
+    """
+
+    @abstractmethod
+    def __call__(self):
+        """Execute the client's main functionality. This method must be
+        implemented by subclasses.
+
+        This method defines the primary behavior of the client and is
+        called when the instance is invoked as a function.
+
+        Raises:
+            NotImplementedError: If the method is not implemented by the
+                subclass.
+        """
+        raise NotImplementedError
+
+    async def acall(self, *args, **kwargs):
+        """Async interface to __call__."""
+        executor = Executor.get_instance()
+        future = executor.submit(self.__call__, *args, **kwargs)
+        done, _ = await asyncio.wait([future])
+        response = await done[0]
+        return response
+
+    def serialize(self) -> Dict[str, Any]:
+        """Serialize the client instance into a dictionary.
+
+        This method prepares the client's state for serialization by
+        removing unnecessary attributes and including additional metadata
+        such as the client type and provider.
+
+        Returns:
+            Dict[str, Any]: A dictionary containing the serialized state
+                of the client, including metadata and the client's
+                internal state.
+        """
+        state = self.__getstate__()
+        if hasattr(self, "to_ignore"):
+            for key in self.to_ignore:
+                state.pop(key, None)
+        data = {"msgflux_type": self.msgflux_type}
+        if hasattr(self, "provider"):
+            data["provider"] = self.provider
+        if hasattr(self, "instance_type"):
+            instance_type = self.instance_type()
+            data.update(instance_type)
+        data["state"] = state
+        return data
+
+    @classmethod
+    def from_serialized(cls, state: Dict[str, Any]):
+        """Deserialize a client instance from a dictionary.
+
+        This method creates a new instance of the client and restores its
+        state from the provided dictionary. It also ensures that the
+        client is properly initialized after deserialization.
+
+        Args:
+            state (Dict[str, Any]): A dictionary containing the serialized
+                state of the client.
+
+        Returns:
+            BaseClient: A new instance of the client with its state
+                restored.
+        """
+        if "msgflux_type" in state:
+            state.pop("msgflux_type")
+        instance = cls()
+        instance.__setstate__(state)
+        if hasattr(instance, "_initialize"):
+            instance._initialize()
+        return instance

msgflux/_private/core.py

@@ -0,0 +1,14 @@
+class Core:
+    """Core is a set of methods that are common among some classes in msgflux.
+
+    This class provides basic functionality for state management, allowing instances
+    to be serialized and deserialized. It is intended to be inherited by other classes
+    that require these common features.
+    """
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)

msgflux/_private/executor.py

@@ -0,0 +1,134 @@
+import asyncio
+import threading
+from concurrent.futures import Future, ThreadPoolExecutor
+from functools import partial
+from typing import Any, Callable, Tuple
+
+try:
+    import platform
+
+    import uvloop
+
+    if platform.python_implementation() != "PyPy":
+        asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+except ImportError:
+    pass
+
+from msgflux.envs import envs
+
+
+class AsyncWorker:
+    def __init__(self):
+        """Initializes a worker with its own event loop in a separate thread."""
+        self.loop = asyncio.new_event_loop()
+        self.thread = threading.Thread(target=self.loop.run_forever, daemon=True)
+        self.thread.start()
+
+    def submit(self, coro):
+        """Submits a coroutine to the worker's event loop."""
+        return asyncio.run_coroutine_threadsafe(coro, self.loop)
+
+    def shutdown(self):
+        """Terminates the event loop and worker thread."""
+        self.loop.call_soon_threadsafe(self.loop.stop)
+        self.thread.join()
+
+
+class Executor:
+    """Async pool to manage the execution of synchronous and asynchronous code.
+    Executor distributes tasks among a ThreadPool or among Async Workers.
+    """
+
+    _instance = None
+    _lock = threading.Lock()
+
+    def __init__(self):
+        self.num_threads = envs.executor_num_threads
+        self.num_async_workers = envs.executor_num_async_workers
+        self.async_worker_index = 0
+        self._thread_local = threading.local()
+        self.thread_pool = ThreadPoolExecutor(
+            max_workers=self.num_threads,
+            initializer=self._mark_thread_pool_worker,
+        )
+        self.async_workers = [AsyncWorker() for _ in range(self.num_async_workers)]
+
+    @classmethod
+    def get_instance(cls):
+        """Returns the singleton instance in a thread-safe manner."""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = cls()
+        return cls._instance
+
+    def submit(self, f: Callable, *args, **kwargs):
+        """Submits a task to the appropriate pool based on the function type.
+        Returns a Future to track the result.
+        """
+        mode, target = self._resolve_submission_target(f, args, kwargs)
+        if mode == "async":
+            return self._submit_to_async_worker(target)
+
+        sync_callable, sync_args, sync_kwargs = target
+        if self._is_in_thread_pool_worker():
+            return self._execute_inline(sync_callable, *sync_args, **sync_kwargs)
+
+        return self.thread_pool.submit(sync_callable, *sync_args, **sync_kwargs)
+
+    def _submit_to_async_worker(self, coro):
+        """Distribute a coroutine to an asynchronous worker using round-robin."""
+        worker = self.async_workers[self.async_worker_index]
+        self.async_worker_index = (self.async_worker_index + 1) % self.num_async_workers
+        return worker.submit(coro)
+
+    def _resolve_submission_target(
+        self,
+        f: Callable,
+        args: Tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Tuple[str, Any]:
+        """Resolve wrappers like functools.partial before deciding execution mode.
+
+        This preserves `.acall()` routing for partially-bound Modules/Tools and
+        avoids sending nested sync work back into the same saturated pool.
+        """
+        if isinstance(f, partial):
+            merged_kwargs = dict(f.keywords or {})
+            merged_kwargs.update(kwargs)
+            merged_args = (*f.args, *args)
+            return self._resolve_submission_target(f.func, merged_args, merged_kwargs)
+
+        if hasattr(f, "acall"):
+            return ("async", f.acall(*args, **kwargs))
+
+        if asyncio.iscoroutinefunction(f):
+            return ("async", f(*args, **kwargs))
+
+        return ("sync", (f, args, kwargs))
+
+    def _mark_thread_pool_worker(self) -> None:
+        self._thread_local.in_thread_pool_worker = True
+
+    def _is_in_thread_pool_worker(self) -> bool:
+        return bool(getattr(self._thread_local, "in_thread_pool_worker", False))
+
+    def _execute_inline(self, f: Callable, *args, **kwargs) -> Future:
+        """Execute nested sync work inline and wrap it in a completed Future."""
+        future: Future = Future()
+        try:
+            result = f(*args, **kwargs)
+        except Exception as exc:
+            future.set_exception(exc)
+        else:
+            future.set_result(result)
+        return future
+
+    def shutdown(self):
+        """Shutdown the executor, closing the pools."""
+        self.thread_pool.shutdown()
+        for worker in self.async_workers:
+            worker.shutdown()
+
+    def __del__(self):
+        self.shutdown()

msgflux/_private/response.py

@@ -0,0 +1,230 @@
+import asyncio
+import threading
+from collections import deque
+from typing import Any, AsyncGenerator, Literal, Optional, Union
+
+
+class CoreResponse:
+    def set_metadata(self, metadata: Any):
+        self.metadata = metadata
+
+    def set_response_type(self, response_type: str):
+        if isinstance(response_type, str):
+            self.response_type = response_type
+        else:
+            raise TypeError(
+                f"`response_type` requires strgiven `{type(response_type)}`"
+            )
+
+
+class BaseResponse(CoreResponse):
+    def __init__(self):
+        self.data = None
+        self.reasoning = None
+        self.metadata = None
+        self.response_type = None
+
+    @property
+    def has_reasoning(self) -> bool:
+        return self.reasoning is not None
+
+    def add(self, data: Any):
+        self.data = data
+
+    def consume(self) -> Any:
+        return self.data
+
+    def consume_reasoning(self) -> Optional[str]:
+        return self.reasoning
+
+
+class BaseStreamResponse(CoreResponse):
+    def __init__(self, mode: Literal["sync", "async"] = "sync"):
+        if mode not in {"sync", "async"}:
+            raise ValueError("`mode` must be `sync` or `async`")
+        self.mode = mode
+        if mode == "async":
+            self.first_chunk_event = asyncio.Event()
+            self._response_type_event = asyncio.Event()
+        else:
+            self.first_chunk_event = threading.Event()
+            self._response_type_event = threading.Event()
+        self.data = None
+        self.reasoning = None
+        self.has_reasoning = False
+
+        # Content queue
+        self._queue = None
+        self._queue_loop = None
+        self._pending_chunks = deque()
+        self._queue_lock = threading.Lock()
+
+        # Reasoning queue
+        self._reasoning_queue = None
+        self._reasoning_queue_loop = None
+        self._reasoning_pending_chunks = deque()
+        self._reasoning_queue_lock = threading.Lock()
+
+        self.metadata = None
+        self.response_type = None
+        self.error = None
+
+    def _finish_queue_with_none(
+        self,
+        *,
+        queue_attr: str,
+        loop_attr: str,
+        pending_attr: str,
+        lock_attr: str,
+    ) -> None:
+        with getattr(self, lock_attr):
+            queue = getattr(self, queue_attr)
+            loop = getattr(self, loop_attr)
+            pending = getattr(self, pending_attr)
+            if queue is None or loop is None or loop.is_closed():
+                pending.append(None)
+                return
+
+        loop.call_soon_threadsafe(queue.put_nowait, None)
+
+    def _fail_stream(self, error: Exception) -> None:
+        self.set_error(error)
+        self._finish_queue_with_none(
+            queue_attr="_queue",
+            loop_attr="_queue_loop",
+            pending_attr="_pending_chunks",
+            lock_attr="_queue_lock",
+        )
+        self._finish_queue_with_none(
+            queue_attr="_reasoning_queue",
+            loop_attr="_reasoning_queue_loop",
+            pending_attr="_reasoning_pending_chunks",
+            lock_attr="_reasoning_queue_lock",
+        )
+
+    def _accumulate_data(self, data: Any) -> None:
+        if data is None:
+            return
+
+        if not isinstance(data, (str, bytes)):
+            raise TypeError(
+                "ModelStreamResponse only supports `str` or `bytes` chunks, "
+                f"got `{type(data).__name__}`."
+            )
+
+        if self.data is None:
+            self.data = data
+            return
+
+        if isinstance(self.data, str) and isinstance(data, str):
+            self.data += data
+            return
+
+        if isinstance(self.data, bytes) and isinstance(data, bytes):
+            self.data += data
+            return
+
+        raise TypeError(
+            "ModelStreamResponse received mixed chunk types: "
+            f"`{type(self.data).__name__}` then `{type(data).__name__}`."
+        )
+
+    def set_response_type(self, response_type: str):
+        super().set_response_type(response_type)
+        if not self._response_type_event.is_set():
+            self._response_type_event.set()
+
+    def set_error(self, error: Exception):
+        self.error = error
+        if not self.first_chunk_event.is_set():
+            self.first_chunk_event.set()
+        if not self._response_type_event.is_set():
+            self._response_type_event.set()
+
+    def add(self, data: Any):
+        """Add data to the content stream queue in a thread-safe way."""
+        if not self.first_chunk_event.is_set():
+            self.first_chunk_event.set()
+
+        try:
+            self._accumulate_data(data)
+        except Exception as e:
+            self._fail_stream(e)
+            raise
+
+        with self._queue_lock:
+            queue = self._queue
+            loop = self._queue_loop
+            if queue is None or loop is None or loop.is_closed():
+                self._pending_chunks.append(data)
+                return
+
+        loop.call_soon_threadsafe(queue.put_nowait, data)
+
+    def add_reasoning(self, data: Any):
+        """Add data to the reasoning stream queue in a thread-safe way."""
+        if data is not None and not self.has_reasoning:
+            self.has_reasoning = True
+        if not self.first_chunk_event.is_set():
+            self.first_chunk_event.set()
+        with self._reasoning_queue_lock:
+            queue = self._reasoning_queue
+            loop = self._reasoning_queue_loop
+            if queue is None or loop is None or loop.is_closed():
+                self._reasoning_pending_chunks.append(data)
+                return
+
+        loop.call_soon_threadsafe(queue.put_nowait, data)
+
+    def _bind_consumer_queue(self) -> asyncio.Queue:
+        loop = asyncio.get_running_loop()
+        with self._queue_lock:
+            if self._queue is None:
+                self._queue = asyncio.Queue()
+                self._queue_loop = loop
+                while self._pending_chunks:
+                    self._queue.put_nowait(self._pending_chunks.popleft())
+            elif self._queue_loop is not loop:
+                raise RuntimeError(
+                    "BaseStreamResponse.consume() must run on the same event loop."
+                )
+            return self._queue
+
+    def _bind_reasoning_queue(self) -> asyncio.Queue:
+        loop = asyncio.get_running_loop()
+        with self._reasoning_queue_lock:
+            if self._reasoning_queue is None:
+                self._reasoning_queue = asyncio.Queue()
+                self._reasoning_queue_loop = loop
+                while self._reasoning_pending_chunks:
+                    self._reasoning_queue.put_nowait(
+                        self._reasoning_pending_chunks.popleft()
+                    )
+            elif self._reasoning_queue_loop is not loop:
+                raise RuntimeError(
+                    "BaseStreamResponse.consume_reasoning() "
+                    "must run on the same event loop."
+                )
+            return self._reasoning_queue
+
+    async def consume(self) -> AsyncGenerator[Union[bytes, str], None]:
+        """Async generator that yields content chunks until None is received."""
+        queue = self._bind_consumer_queue()
+        while True:
+            chunk = await queue.get()
+            if chunk is None:
+                if self.error is not None:
+                    raise self.error
+                break
+            yield chunk
+
+    async def consume_reasoning(self) -> AsyncGenerator[str, None]:
+        """Async generator that yields reasoning chunks until None is received."""
+        queue = self._bind_reasoning_queue()
+        while True:
+            chunk = await queue.get()
+            if chunk is None:
+                if self.error is not None:
+                    raise self.error
+                break
+            yield chunk

msgflux/auto/__init__.py

@@ -0,0 +1,5 @@
+"""Auto module for automatic parameter and configuration management."""
+
+from msgflux.auto.params import AutoParams
+
+__all__ = ["AutoParams"]