msgflux 0.1.0a1__py3-none-any.whl

msgflux/__init__.py

@@ -0,0 +1,37 @@
+from .cache import response_cache
+from .data.databases.database import DataBase
+from .data.retrievers.retriever import Retriever
+from .dotdict import dotdict
+from .dsl.inline import inline
+from .dsl.signature import InputField, OutputField, Signature
+from .envs import set_envs
+from .message import Message
+from .models.gateway import ModelGateway
+from .models.model import Model
+from .utils.chat import ChatML
+from .utils.inspect import get_fn_name
+from .utils.msgspec import load, save
+from .utils.tool import tool_config
+from .telemetry.span import instrument
+
+
+__all__ = [
+    "ChatML",
+    "DataBase",  
+    "InputField",
+    "Message",
+    "Model",
+    "ModelGateway",
+    "OutputField",
+    "Retriever",    
+    "Signature",
+    "dotdict",
+    "get_fn_name",
+    "inline",
+    "instrument",  
+    "load",
+    "response_cache",
+    "save",
+    "set_envs",
+    "tool_config"
+]

msgflux/_private/client.py

@@ -0,0 +1,91 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+from msgflux._private.core import Core
+
+
+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__ """
+        return self.__call__(*args, **kwargs)
+
+    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,77 @@
+import asyncio
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from typing import Callable
+import uvloop
+from msgflux.envs import envs
+
+
+class AsyncWorker:
+    def __init__(self):
+        """Initializes a worker with its own event loop in a separate thread."""
+        self.loop = uvloop.new_event_loop()
+        self.thread = threading.Thread(target=self.loop.run_forever)
+        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_pool = ThreadPoolExecutor(max_workers=self.num_threads)
+        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.
+        """
+        if hasattr(f, "acall"):
+            coro = f.acall(*args, **kwargs)
+            return self._submit_to_async_worker(coro)
+        elif asyncio.iscoroutinefunction(f):
+            coro = f(*args, **kwargs)
+            return self._submit_to_async_worker(coro)
+        else:
+            return self.thread_pool.submit(f, *args, **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 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/cache.py

@@ -0,0 +1,50 @@
+from functools import wraps
+from typing import Optional
+from msgflux.data.databases.base import BaseDB
+from msgflux.message import Message
+from msgflux.models.types import TextEmbedderModel
+
+
+def response_cache(
+    db: BaseDB, 
+    task_inputs: Optional[str] = None, 
+    response_mode: Optional[str] = None,
+    model: Optional[TextEmbedderModel] = None,
+):
+    def decorator(func):        
+        @wraps(func)
+        def wrapper(msg):
+            if isinstance(msg, Message):
+                if task_inputs is None:
+                    raise ValueError(
+                        "`task_inputs` is required when input "
+                        "is a Message instance"
+                    )
+                key = msg.get(task_inputs)
+            else:
+                key = msg
+
+            if model:
+                key = model(key)            
+
+            cached_response = db(key)
+            
+            if cached_response is not None:
+                if isinstance(msg, Message):
+                    if response_mode is None:
+                        raise ValueError(
+                            "`response_mode` is required when "
+                            "input is a Message instance"
+                        )
+                    msg.set(response_mode, cached_response)
+                    return msg
+                return cached_response
+            
+            response = func(msg)
+            
+            if response is not None:
+                db.add(key, response)
+
+            return response
+        return wrapper
+    return decorator

msgflux/data/audio.py

@@ -0,0 +1 @@
+#https://github.com/openai/openai-python/blob/384e6b23ce0366d6b2f31cc98d35525da5b22c10/src/openai/helpers/local_audio_player.py#L19

msgflux/data/databases/base.py

@@ -0,0 +1,32 @@
+from abc import abstractmethod
+from typing import Dict
+from msgflux._private.client import BaseClient
+
+
+class BaseDB(BaseClient):
+
+    msgflux_type = "db"
+    to_ignore = ["client", "documents"]
+
+    def instance_type(self) -> Dict[str, str]:
+         return {"db_type": self.db_type}  
+    
+    @abstractmethod
+    def add(self):
+        """
+        Insert data on database
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _initialize(self):
+        """
+        Initialize the class. This method must be implemented by subclasses.
+
+        This method is called during the deserialization process to ensure that the client
+        is properly initialized after its state has been restored.
+
+        Raises:
+            NotImplementedError: If the method is not implemented by the subclass.
+        """
+        raise NotImplementedError

msgflux/data/databases/database.py

@@ -0,0 +1,125 @@
+from typing import Any, Dict, Type
+from msgflux.data.databases.base import BaseDB
+from msgflux.data.databases.types import (
+    #GeoDB,
+    #GraphDB,
+    KVDB,
+    #NoSQLDB,
+    #RelationalDB,
+    #TimeSeriesDB,
+    VectorDB,
+)
+from msgflux.utils.imports import import_module_from_lib
+
+
+_SUPPORTED_DB_TYPES = [
+    "kv",    
+    "vector",
+    #"graph",
+    #"relational",
+    #"nosql",
+    #"time_series",
+    #"geo",
+]
+_DB_NAMESPACE_TRANSLATOR = {
+    "cachetools": "CacheTools",
+    "diskmanager": "DiskManager",
+    "faiss": "FAISS",
+    #"qdrant": "Qdrant",
+    #"postgres": "Postgres",
+    #"weavite": "Weavite",
+    #"milvus": "Milvus",
+    #"lance_db": "LanceDB",
+    #"vespa": "Vespa",
+    #"pinecone": "Pinecone",
+    #"cassandra": "Cassandra",
+    #"redis": "Redis",
+    #"deep_lake": "DeepLake",
+    #"marqo": "Marqo",
+    #"azure": "AzureAISearch",
+    #"supabase": "SupaBase",
+    #"elastic_search": "ElasticSearch",
+    #"chroma": "Chroma",
+    #"sqlite": "SQLite",
+    #"neo4j": "Neo4J",
+}
+_VECTOR_DB_PROVIDERS = [
+    #"qdrant",
+    #"postgres",
+    #"weavite",
+    #"milvus",
+    #"lance_db",
+    #"vespa",
+    #"redis",
+    #"cassandra",
+    #"deep_lake",
+    #"marqo",
+    #"elastic_search",
+    #"sqlite",
+    "faiss",
+    #"chroma",
+]
+#_GRAPH_DB_PROVIDERS = ["neo4j"]
+_KV_DB_PROVIDERS = ["cachetools", "diskmanager"]
+
+_PROVIDERS_BY_DB_TYPE = {
+    "kv": _KV_DB_PROVIDERS,
+    "vector": _VECTOR_DB_PROVIDERS,
+}
+
+
+class DataBase:
+    supported_db_types = _SUPPORTED_DB_TYPES
+    providers_by_db_type = _PROVIDERS_BY_DB_TYPE
+
+    @classmethod
+    def _get_db_class(cls, db_type: str, provider: str) -> Type[BaseDB]:
+        if db_type not in cls.supported_db_types:
+            raise ValueError(f"DB type `{db_type}` is not supported")
+            
+        providers = cls.providers_by_db_type[db_type]
+        if provider not in providers:
+            raise ValueError(f"Provider `{provider}` is not supported for {db_type}")
+
+        if len(db_type) <= 3:
+            db_type = db_type.upper()
+        else:
+            db_type = db_type.title().replace("_", "")
+
+        provider_class_name = f"{_DB_NAMESPACE_TRANSLATOR[provider]}{db_type}DB"                
+        module_name = f"msgflux.data.databases.providers.{provider}"                
+        return import_module_from_lib(provider_class_name, module_name)
+
+    @classmethod
+    def _create_db(cls, db_type: str, provider: str, **kwargs) -> Type[BaseDB]:
+        db_cls = cls._get_db_class(db_type, provider)
+        return db_cls(**kwargs)
+
+    @classmethod
+    def from_serialized(cls, provider: str, db_type: str, params: Dict[str, Any]) -> Type[BaseDB]:
+        """
+        Creates a db instance from serialized parameters without calling __init__.
+        
+        Args:
+            provider: The db provider (e.g., "faiss", "cachetools")
+            db_type: The type of db (e.g., "vector", "kv")
+            params: Dictionary containing the serialized db parameters
+            
+        Returns:
+            An instance of the appropriate db class with restored state
+        """
+        db_cls = cls._get_db_class(db_type, provider)
+        # Create instance without calling __init__
+        instance = object.__new__(db_cls)
+        # Restore the instance state
+        instance.from_serialized(params)
+        return instance
+
+    @classmethod
+    def kv(cls, provider: str, **kwargs) -> KVDB:
+        return cls._create_db("kv", provider, **kwargs)
+
+    @classmethod
+    def vector(cls, provider: str, **kwargs) -> VectorDB:
+        return cls._create_db("vector", provider, **kwargs)
+ 

msgflux/data/databases/providers/cachetools.py

@@ -0,0 +1,64 @@
+from cachetools import TTLCache
+from typing import Any, Dict, List, Optional, Union
+import msgspec
+from msgflux.data.databases.base import BaseDB
+from msgflux.data.databases.types import KVDB
+from msgflux.utils.convert import convert_str_to_hash
+
+
+class CacheToolsKVDB(BaseDB, KVDB):
+    """CacheTools Key-Value DB."""
+
+    provider = "cachetools"
+
+    def __init__(
+        self, 
+        ttl: Optional[int] = 3600, 
+        maxsize: Optional[int] = 10000, 
+        hash_key: Optional[bool] = True
+    ):
+        """
+        Args:
+            ttl: 
+                The time-to-live (TTL) for each cache entry in seconds.
+            maxsize: 
+                The maximum number of items the cache can store.
+            hash_key: 
+                Whether to hash the keys before storing them in the cache.
+        """
+        self.hash_key = hash_key
+        self.maxsize = maxsize
+        self.ttl = ttl
+        self._initialize()
+
+    def _initialize(self):
+        self.client = TTLCache(maxsize=self.maxsize, ttl=self.ttl)
+        
+    def add(self, documents: Union[List[Dict[str, Any]], Dict[str, Any]]):
+        if not isinstance(documents, list):
+            documents = [documents]
+        for document in documents:            
+            for k, v in document.items():
+                if self.hash_key:
+                    k = convert_str_to_hash(k)
+                v = msgspec.msgpack.encode(v)
+                self.client[k] = v
+            
+    def _search(self, queries: Union[str, List[str]]):
+        if not isinstance(queries, list):
+            queries = [queries]
+        query_results = []
+        for query in queries:
+            if self.hash_key:
+                query = convert_str_to_hash(query)
+            v = self.client.get(query)
+            if v is not None:
+                v = msgspec.msgpack.decode(v)
+            query_results.append(v)
+        return query_results
+
+    def __call__(self, queries: Union[str, List[str]]) -> Union[Any, List[Any]]:
+        query_results = self._search(queries)
+        if len(query_results) == 1:
+            return query_results[0]
+        return query_results

msgflux/data/databases/providers/diskcache.py

@@ -0,0 +1,60 @@
+from typing import Any, Dict, List, Optional, Union
+import msgspec
+try:
+    from diskcache import Cache
+except:
+    raise ImportError("`diskcache` is not detected, please install"
+                      "using `pip install diskcache`")
+from msgflux.data.databases.base import BaseDB
+from msgflux.data.databases.types import KVDB
+from msgflux.utils.convert import convert_str_to_hash
+
+
+class DiskCacheKVDB(BaseDB, KVDB):
+    """DiskCache Key-Value DB."""
+
+    provider = "diskcache"
+
+    def __init__(self, ttl: Optional[int] = 3600, hash_key: Optional[bool] = True):
+        """
+        Args:
+            ttl:
+                The time-to-live (TTL) for each cache entry in seconds.
+            hash_key:
+                Whether to hash the keys before storing them in the cache.
+        """
+        self.hash_key = hash_key
+        self.ttl = ttl
+        self._initialize()
+
+    def _initialize(self):
+        self.client = Cache(timeout=1)
+        
+    def add(self, documents: Union[List[Dict[str, Any]], Dict[str, Any]]):
+        if not isinstance(documents, list):
+            documents = [documents]
+        for document in documents:            
+            for k, v in document.items():
+                if self.hash_key:
+                    k = convert_str_to_hash(k)
+                v = msgspec.msgpack.encode(v)
+                self.client.set(k, v, expire=self.ttl)
+            
+    def _search(self, queries: Union[str, List[str]]):
+        if not isinstance(queries, list):
+            queries = [queries]
+        query_results = []
+        for query in queries:
+            if self.hash_key:
+                query = convert_str_to_hash(query)
+            v = self.client.get(query)
+            if v is not None:
+                v = msgspec.msgpack.decode(v)
+            query_results.append(v)
+        return query_results
+
+    def __call__(self, queries: Union[str, List[str]]) -> Union[Any, List[Any]]:
+        query_results = self._search(queries)
+        if len(query_results) == 1:
+            return query_results[0]
+        return query_results