msgflux 0.1.0__py3-none-any.whl

msgflux/core/message.py

@@ -0,0 +1,45 @@
+from typing import Any, Dict, Optional, Union
+
+from msgflux.core.dotdict import dotdict
+
+
+class Message(dotdict):
+    def __init__(
+        self,
+        *,
+        content: Optional[Union[str, Dict[str, Any]]] = None,
+        context: Optional[Dict[str, Any]] = None,
+        texts: Optional[Dict[str, Any]] = None,
+        audios: Optional[Dict[str, Any]] = None,
+        images: Optional[Dict[str, Any]] = None,
+        videos: Optional[Dict[str, Any]] = None,
+        extra: Optional[Dict[str, Any]] = None,
+    ):
+        if extra is None:
+            extra = {}
+        if videos is None:
+            videos = {}
+        if images is None:
+            images = {}
+        if audios is None:
+            audios = {}
+        if texts is None:
+            texts = {}
+        if context is None:
+            context = {}
+        super().__init__()
+        self.content = content
+        self.texts = texts
+        self.context = context
+        self.audios = audios
+        self.images = images
+        self.videos = videos
+        self.extra = extra
+        self.outputs = {}
+        self.response = {}
+
+    def get_response(self):
+        if self.get("response"):
+            return next(iter(self.get("response").values()))
+        else:
+            return self.get("response")

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/dbs/__init__.py

@@ -0,0 +1,16 @@
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from msgflux.data.dbs.db import DB
+
+__all__ = ["DB"]
+
+
+def __getattr__(name: str):
+    if name != "DB":
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    value = getattr(import_module("msgflux.data.dbs.db"), name)
+    globals()[name] = value
+    return value

msgflux/data/dbs/base.py

@@ -0,0 +1,124 @@
+from abc import abstractmethod
+from typing import TYPE_CHECKING, List, Mapping, Optional, Union
+
+import msgspec
+
+if TYPE_CHECKING:
+    import numpy as np
+
+import msgflux.nn.functional as F
+from msgflux._private.client import BaseClient
+from msgflux.data.dbs.response import DBResponse
+from msgflux.utils.convert import convert_str_to_hash
+
+
+class BaseDB(BaseClient):
+    msgflux_type = "db"
+    to_ignore = ["client", "documents"]
+
+    def instance_type(self) -> Mapping[str, str]:
+        return {"db_type": self.db_type}
+
+    @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
+
+
+class BaseKV:
+    def _single_search(self, query: str) -> str:
+        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)
+        return v
+
+    def _search(self, queries: Union[str, List[str]]):
+        if not isinstance(queries, list):
+            queries = [queries]
+        args_list = [(query,) for query in queries]
+        results = F.map_gather(
+            self._single_search,
+            args_list=args_list,
+        )
+        return results
+
+    def __call__(self, queries: Union[str, List[str]]) -> DBResponse:
+        """Executes a search in the key-value database for the given query
+        or queries.
+
+        This method searches the underlying key-value store using the
+        provided query string or list of query strings. If a single query
+        is provided, the search result is returned directly. If multiple
+        queries are provided, a list of results is returned.
+
+        Args:
+            queries:
+                A single query string or a list of query strings to search
+                for in the db.
+
+        Returns:
+            An object containing the search results.
+
+        Raises:
+            ValueError:
+                If `queries` is empty or contains unsupported types.
+        """
+        results = self._search(queries)
+        if len(results) == 1:
+            results = results[0]
+        response = DBResponse()
+        response.set_response_type("key_search")
+        response.add(results)
+        return response
+
+
+class BaseVector:
+    def __call__(
+        self,
+        queries: Union[List[List[float]], "np.ndarray"],
+        *,
+        top_k: Optional[int] = 4,
+        threshold: Optional[float] = None,
+        return_score: Optional[bool] = False,
+    ) -> DBResponse:
+        """Executes a vector similarity search against the vector database.
+
+        This method searches for the closest vectors to the given query vector(s)
+        in the underlying vector store. If a single vector is provided, it is
+        automatically wrapped into a list for processing.
+
+        Args:
+            queries:
+                A single query vector or a list/array of query vectors, where each
+                vector is represented as a list of floats or a NumPy array.
+            top_k:
+                The maximum number of closest matches to return for each query.
+            threshold:
+                A similarity score threshold; results with scores below this value
+                will be filtered out.
+            return_score:
+                If True, include similarity scores in the returned results.
+
+        Returns:
+            An object containing the search results.
+
+        Raises:
+            ValueError: If the query vectors are not in a supported format or are empty.
+        """
+        results = self._search(queries, top_k, threshold, return_score=return_score)
+        response = DBResponse()
+        response.set_response_type("vector_search")
+        response.add(results)
+        return response

msgflux/data/dbs/db.py

@@ -0,0 +1,65 @@
+from typing import TYPE_CHECKING, Any, Mapping
+
+from msgflux.data.dbs.registry import db_registry
+from msgflux.data.dbs.types import KVDB, VectorDB
+
+if TYPE_CHECKING:
+    from msgflux.data.dbs.base import BaseDB
+
+
+class DB:
+    @classmethod
+    def providers(cls):
+        return {k: list(v.keys()) for k, v in db_registry.items()}
+
+    @classmethod
+    def db_types(cls):
+        return list(db_registry.keys())
+
+    @classmethod
+    def _get_db_class(cls, db_type: str, provider: str) -> type["BaseDB"]:
+        if db_type not in db_registry:
+            raise ValueError(f"DB type `{db_type}` is not supported")
+        if provider not in db_registry[db_type]:
+            raise ValueError(
+                f"Provider `{provider}` not registered for type `{db_type}`"
+            )
+        db_cls = db_registry[db_type][provider]
+        return db_cls
+
+    @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: Mapping[str, Any]
+    ) -> type["BaseDB"]:
+        """Creates a db instance from serialized parameters.
+
+        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/dbs/providers/cachetools.py

@@ -0,0 +1,57 @@
+from typing import Any, Dict, List, Union
+
+import msgspec
+
+try:
+    from cachetools import TTLCache
+except ImportError:
+    TTLCache = None
+
+from msgflux.data.dbs.base import BaseDB, BaseKV
+from msgflux.data.dbs.registry import register_db
+from msgflux.data.dbs.types import KVDB
+from msgflux.utils.convert import convert_str_to_hash
+
+
+@register_db
+class CacheToolsKVDB(BaseKV, BaseDB, KVDB):
+    """CacheTools Key-Value DB."""
+
+    provider = "cachetools"
+
+    def __init__(
+        self,
+        *,
+        ttl: int = 3600,
+        maxsize: int = 10000,
+        hash_key: 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.
+        """
+        if TTLCache is None:
+            raise ImportError(
+                "`cachetools` client is not available. Install with "
+                "`pip install cachetools`"
+            )
+        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 key, value in document.items():
+                encoded_key = convert_str_to_hash(key) if self.hash_key else key
+                encoded_value = msgspec.msgpack.encode(value)
+                self.client[encoded_key] = encoded_value

msgflux/data/dbs/providers/diskcache.py

@@ -0,0 +1,48 @@
+from typing import Any, Dict, List, Union
+
+import msgspec
+
+try:
+    from diskcache import Cache
+except ImportError:
+    Cache = None
+
+from msgflux.data.dbs.base import BaseDB, BaseKV
+from msgflux.data.dbs.registry import register_db
+from msgflux.data.dbs.types import KVDB
+from msgflux.utils.convert import convert_str_to_hash
+
+
+@register_db
+class DiskCacheKVDB(BaseKV, BaseDB, KVDB):
+    """DiskCache Key-Value DB."""
+
+    provider = "diskcache"
+
+    def __init__(self, *, ttl: int = 3600, hash_key: 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.
+        """
+        if Cache is None:
+            raise ImportError(
+                "`diskcache` client is not available. Install with "
+                "`pip install diskcache`"
+            )
+        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():
+                encoded_k = convert_str_to_hash(k) if self.hash_key else k
+                encoded_v = msgspec.msgpack.encode(v)
+                self.client.set(encoded_k, encoded_v, expire=self.ttl)

msgflux/data/dbs/providers/faiss.py

@@ -0,0 +1,220 @@
+import os
+import pickle
+import platform
+from typing import Dict, List, Optional, Union
+
+try:
+    import faiss
+    import numpy as np
+except ImportError:
+    faiss = None
+    np = None
+
+from msgflux.data.dbs.base import BaseDB, BaseVector
+from msgflux.data.dbs.registry import register_db
+from msgflux.data.dbs.types import VectorDB
+
+if platform.system() == "Windows":
+    # Solve OpenMP duplicate versions in Windows
+    os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+
+@register_db
+class FAISSVectorDB(BaseVector, BaseDB, VectorDB):
+    """FAISS Vector DB."""
+
+    provider = "faiss"
+
+    def __init__(
+        self,
+        dimension: int,
+        metric_type: str = "cosine",
+        index_type: str = "flat",
+    ):
+        """Args:
+        dimension:
+            Dimensionality of the embeddings.
+        metric_type:
+            Distance metric for similarity search ('cosine', 'l2', 'ip').
+        index_type:
+            Type of FAISS index ('flat', 'ivf', etc.).
+        """
+        if faiss is None:
+            raise ImportError(
+                "`faiss` client is not available. "
+                "Install with `pip install faiss-cpu` "
+                "or `pip install faiss-gpu`"
+            )
+
+        valid_metrics = {
+            "cosine": faiss.METRIC_INNER_PRODUCT,
+            "l2": faiss.METRIC_L2,
+            "ip": faiss.METRIC_INNER_PRODUCT,
+        }
+        if metric_type not in valid_metrics:
+            raise ValueError(
+                f"Invalid metric type. Choose from {list(valid_metrics.keys())}"
+            )
+
+        self.dimension = dimension
+        self.index_type = index_type
+        self.metric_type = metric_type
+        self._initialize()
+
+    def _initialize(self):
+        if self.index_type == "flat":
+            if self.metric_type == "cosine":
+                # For cosine similarity, use IndexFlatIP with normalized vectors
+                self.index = faiss.IndexFlatIP(self.dimension)
+            elif self.metric_type == "l2":
+                self.index = faiss.IndexFlatL2(self.dimension)
+            else:  # inner product
+                self.index = faiss.IndexFlatIP(self.dimension)
+        else:
+            raise ValueError("Currently only 'flat' index type is supported")
+        self.documents = []
+        # TODO: support to restore data
+
+    def add(
+        self,
+        embeddings: Union[List[List[float]], "np.ndarray"],
+        documents: List[Dict],
+    ):
+        """Add embeddings and their corresponding metadata to the vector db.
+
+        Args:
+            embeddings:
+                List or NumPy array of embeddings (N x dimension).
+            documents:
+                List of metadata.
+        """
+        # Convert to NumPy array if needed
+        if not isinstance(embeddings, np.ndarray):
+            embeddings = np.array(embeddings, dtype=np.float32)
+
+        # Ensure correct dtype and dimensionality
+        if embeddings.dtype != np.float32:
+            embeddings = embeddings.astype(np.float32)
+
+        if embeddings.shape[1] != self.dimension:
+            raise ValueError(f"Embedding dimension must be {self.dimension}")
+
+        # Normalize vectors if using cosine similarity
+        if self.metric_type == "cosine":
+            # Create a copy to avoid modifying the original array
+            embeddings = embeddings.copy()
+
+            # Use numpy for normalization to avoid FAISS normalization issues
+            norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
+            # Avoid division by zero
+            norms[norms == 0] = 1
+            embeddings /= norms
+
+        if len(documents) != embeddings.shape[0]:
+            raise ValueError("Number of documents must match number of embeddings")
+
+        self.index.add(embeddings)
+        self.documents.extend(documents)
+
+    def _search(  # noqa: C901
+        self,
+        queries: Union[List[List[float]], "np.ndarray"],
+        top_k: int = 4,
+        threshold: Optional[float] = None,
+        *,
+        return_score: bool = False,
+    ) -> List[List[Dict]]:
+        if len(np.array(queries).shape) == 1:
+            queries = [queries]
+
+        if not isinstance(queries, np.ndarray):
+            queries = np.array(queries, dtype=np.float32)
+
+        if queries.dtype != np.float32:
+            queries = queries.astype(np.float32)
+
+        if queries.shape[1] != self.dimension:
+            raise ValueError(f"Query embedding dimension must be {self.dimension}")
+
+        if self.metric_type == "cosine":
+            queries = queries.copy()
+            norms = np.linalg.norm(queries, axis=1, keepdims=True)
+            norms[norms == 0] = 1
+            queries /= norms
+
+        all_results = []
+
+        for query_embedding in queries:
+            query_vec = query_embedding.reshape(1, -1)
+            distances, indices = self.index.search(query_vec, top_k)
+
+            query_results = []
+            for dist, idx in zip(distances[0], indices[0]):
+                if idx == -1:
+                    continue
+
+                if threshold is not None and dist < threshold:
+                    continue
+
+                result = {"data": self.documents[idx]}
+                if return_score:
+                    result["score"] = float(dist)
+                query_results.append(result)
+
+            all_results.append(query_results)
+
+        return all_results
+
+    def save(self, path: str):
+        """Save the vector store to disk.
+
+        Args:
+            path:
+                Path where the index and metadata will be saved.
+        """
+        # Create directory if it doesn't exist
+        os.makedirs(path, exist_ok=True)
+
+        # Save FAISS index
+        index_path = os.path.join(path, "faiss_index.bin")
+        faiss.write_index(self.index, index_path)
+
+        # Save metadata and configuration
+        metadata = {
+            "dimension": self.dimension,
+            "metric_type": self.metric_type,
+            "documents": self.documents,
+        }
+        metadata_path = os.path.join(path, "metadata.pkl")
+        with open(metadata_path, "wb") as f:
+            pickle.dump(metadata, f)
+
+    @classmethod
+    def load(cls, path: str) -> "FAISSVectorDB":
+        """Load a vector store from disk.
+
+        Args:
+            path:
+                Path where the index and metadata are saved
+
+        Returns:
+            FAISSVectorDB instance
+        """
+        # Load metadata
+        metadata_path = os.path.join(path, "metadata.pkl")
+        with open(metadata_path, "rb") as f:
+            metadata = pickle.load(f)  # noqa: S301
+
+        # Create instance
+        instance = cls(
+            dimension=metadata["dimension"], metric_type=metadata["metric_type"]
+        )
+
+        # Load FAISS index
+        index_path = os.path.join(path, "faiss_index.bin")
+        instance.index = faiss.read_index(index_path)
+
+        # Restore metadata
+        instance.documents = metadata["documents"]
+
+        return instance

msgflux/data/dbs/registry.py

@@ -0,0 +1,19 @@
+from typing import TYPE_CHECKING
+
+from msgflux.utils.imports import AutoloadRegistry
+
+if TYPE_CHECKING:
+    from msgflux.data.dbs.base import BaseDB
+
+db_registry = AutoloadRegistry("msgflux.data.dbs.providers")
+
+
+def register_db(cls: type["BaseDB"]):
+    db_type = getattr(cls, "db_type", None)
+    provider = getattr(cls, "provider", None)
+
+    if not db_type or not provider:
+        raise ValueError(f"{cls.__name__} must define `db_type` and `provider`.")
+
+    db_registry.setdefault(db_type, {})[provider] = cls
+    return cls

msgflux/data/dbs/response.py

@@ -0,0 +1,10 @@
+from typing import Literal
+
+from msgflux._private.response import BaseResponse
+
+
+class DBResponse(BaseResponse):
+    response_type: Literal[
+        "key_search",
+        "vector_search",
+    ]

msgflux/data/dbs/types.py

@@ -0,0 +1,26 @@
+class GeoDB:
+    db_type = "geo"
+
+
+class GraphDB:
+    db_type = "graph"
+
+
+class KVDB:
+    db_type = "kv"
+
+
+class NoSQLDB:
+    db_type = "nosql"
+
+
+class RelationalDB:
+    db_type = "relational"
+
+
+class TimeSeriesDB:
+    db_type = "time_series"
+
+
+class VectorDB:
+    db_type = "vector"

msgflux/data/parsers/__init__.py

@@ -0,0 +1,16 @@
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from msgflux.data.parsers.parser import Parser
+
+__all__ = ["Parser"]
+
+
+def __getattr__(name: str):
+    if name != "Parser":
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    value = getattr(import_module("msgflux.data.parsers.parser"), name)
+    globals()[name] = value
+    return value