symbolicai 0.20.2__py3-none-any.whl → 1.0.0__py3-none-any.whl

symai/collect/pipeline.py

@@ -1,14 +1,21 @@
+from __future__ import annotations
+
 import json
 import logging
 from datetime import datetime
-from typing import Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any
 
 from bson.objectid import ObjectId
-from pymongo.collection import Collection
-from pymongo.database import Database
 from pymongo.mongo_client import MongoClient
 
 from ..backend.settings import SYMAI_CONFIG
+from ..utils import UserMessage
+
+if TYPE_CHECKING:
+    from pymongo.collection import Collection
+    from pymongo.database import Database
+else:
+    Collection = Database = Any
 
 logger = logging.getLogger(__name__)
 
@@ -24,18 +31,17 @@ def rec_serialize(obj):
     if isinstance(obj, (int, float, bool)):
         # For simple types, return the string representation directly.
         return obj
-    elif isinstance(obj, dict):
+    if isinstance(obj, dict):
         # For dictionaries, serialize each value. Keep keys as strings.
         return {str(key): rec_serialize(value) for key, value in obj.items()}
-    elif isinstance(obj, (list, tuple, set)):
+    if isinstance(obj, (list, tuple, set)):
         # For lists, tuples, and sets, serialize each element.
         return [rec_serialize(elem) for elem in obj]
-    else:
-        # Attempt JSON serialization first, then fall back to str(...)
-        try:
-            return json.dumps(obj)
-        except TypeError:
-            return str(obj)
+    # Attempt JSON serialization first, then fall back to str(...)
+    try:
+        return json.dumps(obj)
+    except TypeError:
+        return str(obj)
 
 
 class CollectionRepository:
@@ -44,15 +50,15 @@ class CollectionRepository:
         self.uri: str                           = SYMAI_CONFIG["COLLECTION_URI"]
         self.db_name: str                       = SYMAI_CONFIG["COLLECTION_DB"]
         self.collection_name: str               = SYMAI_CONFIG["COLLECTION_STORAGE"]
-        self.client: Optional[MongoClient]      = None
-        self.db: Optional[Database]             = None
-        self.collection: Optional[Collection]   = None
+        self.client: MongoClient | None      = None
+        self.db: Database | None             = None
+        self.collection: Collection | None   = None
 
-    def __enter__(self) -> 'CollectionRepository':
+    def __enter__(self) -> CollectionRepository:
         self.connect()
         return self
 
-    def __exit__(self, exc_type: Optional[type], exc_val: Optional[Exception], exc_tb: Optional[Any]) -> None:
+    def __exit__(self, exc_type: type | None, exc_val: Exception | None, exc_tb: Any | None) -> None:
         self.close()
 
     def ping(self) -> bool:
@@ -63,10 +69,12 @@ class CollectionRepository:
             self.client.admin.command('ping')
             return True
         except Exception as e:
-            print("Connection failed: " + str(e))
+            UserMessage(f"Connection failed: {e}")
             return False
 
-    def add(self, forward: Any, engine: Any, metadata: Dict[str, Any] = {}) -> Any:
+    def add(self, forward: Any, engine: Any, metadata: dict[str, Any] | None = None) -> Any:
+        if metadata is None:
+            metadata = {}
         if not self.support_community:
             return None
         record = {
@@ -78,22 +86,22 @@ class CollectionRepository:
         }
         try: # assure that adding a record does never cause a system error
             return self.collection.insert_one(record).inserted_id if self.collection else None
-        except Exception as e:
+        except Exception:
             return None
 
-    def get(self, record_id: str) -> Optional[Dict[str, Any]]:
+    def get(self, record_id: str) -> dict[str, Any] | None:
         if not self.support_community:
             return None
         return self.collection.find_one({'_id': ObjectId(record_id)}) if self.collection else None
 
     def update(self,
                record_id: str,
-               forward: Optional[Any]             = None,
-               engine: Optional[str]              = None,
-               metadata: Optional[Dict[str, Any]] = None) -> Any:
+               forward: Any | None             = None,
+               engine: str | None              = None,
+               metadata: dict[str, Any] | None = None) -> Any:
         if not self.support_community:
             return None
-        updates: Dict[str, Any] = {'updated_at': datetime.now()}
+        updates: dict[str, Any] = {'updated_at': datetime.now()}
         if forward is not None:
             updates['forward']  = forward
         if engine is not None:
@@ -108,14 +116,14 @@ class CollectionRepository:
             return None
         return self.collection.delete_one({'_id': ObjectId(record_id)}) if self.collection else None
 
-    def list(self, filters: Optional[Dict[str, Any]] = None, limit: int = 0) -> List[Dict[str, Any]]:
+    def list(self, filters: dict[str, Any] | None = None, limit: int = 0) -> list[dict[str, Any]]:
         if not self.support_community:
             return None
         if filters is None:
             filters = {}
         return list(self.collection.find(filters).limit(limit)) if self.collection else []
 
-    def count(self, filters: Optional[Dict[str, Any]] = None) -> int:
+    def count(self, filters: dict[str, Any] | None = None) -> int:
         if not self.support_community:
             return None
         if filters is None:
@@ -133,7 +141,7 @@ class CollectionRepository:
             self.client     = False
             self.db         = None
             self.collection = None
-            print("[WARN] MongoClient: Connection failed: " + str(e))
+            UserMessage(f"[WARN] MongoClient: Connection failed: {e}")
 
     def close(self) -> None:
         if self.client is not None:

symai/collect/stats.py

@@ -1,13 +1,16 @@
 import json
 import re
+from collections.abc import Callable
 from json import JSONEncoder
-from typing import Any, Callable, List, Optional, Tuple, Type, Union
+from pathlib import Path
+from typing import Any, Union
 
 import numpy as np
 import torch
 
 from ..ops.primitives import OperatorPrimitives
 from ..symbol import Symbol
+from ..utils import UserMessage
 
 SPECIAL_CONSTANT = '__aggregate_'
 EXCLUDE_LIST     = ['_ipython_canary_method_should_not_exist_', '__custom_documentations__']
@@ -24,7 +27,7 @@ class AggregatorJSONEncoder(JSONEncoder):
         if isinstance(obj, np.ndarray):
             return obj.tolist()
         # drop active from state
-        elif isinstance(obj, Aggregator):
+        if isinstance(obj, Aggregator):
             state = obj.__dict__.copy()
             state.pop('_raise_error', None)
             state.pop('_active', None)
@@ -32,8 +35,8 @@ class AggregatorJSONEncoder(JSONEncoder):
             state.pop('_map', None)
             # drop everything that starts with SPECIAL_CONSTANT
             for key in list(state.keys()):
-                if  not key.startswith(SPECIAL_CONSTANT) and key != '_value' or \
-                    key == '_value' and obj._value == [] or \
+                if  (not key.startswith(SPECIAL_CONSTANT) and key != '_value') or \
+                    (key == '_value' and obj._value == []) or \
                     key.replace(SPECIAL_CONSTANT, '') in EXCLUDE_LIST:
                     state.pop(key, None)
             return state
@@ -42,8 +45,8 @@ class AggregatorJSONEncoder(JSONEncoder):
 
 class Aggregator(Symbol):
     def __init__(self,
-                 value: Optional[Union["Aggregator", Symbol]] = None,
-                 path: Optional[str] = None,
+                 value: Union["Aggregator", Symbol] | None = None,
+                 path: str | None = None,
                  active: bool = True,
                  raise_error: bool = False,
                  *args, **kwargs):
@@ -61,7 +64,7 @@ class Aggregator(Symbol):
             elif not isinstance(self._value, (list, tuple)):
                 self._value = [self._value]
         elif value is not None:
-            raise Exception(f'Aggregator object must be of type Aggregator or Symbol! Got: {type(value)}')
+            UserMessage(f'Aggregator object must be of type Aggregator or Symbol! Got: {type(value)}', raise_with=Exception)
         else:
             self._value = []
         self._raise_error   = raise_error
@@ -71,11 +74,13 @@ class Aggregator(Symbol):
         self._path      = path
 
     def __new__(cls, *args,
-            mixin: Optional[bool] = None,
-            primitives: Optional[List[Type]] = [OperatorPrimitives], # only inherit arithmetic primitives
-            callables: Optional[List[Tuple[str, Callable]]] = None,
+            mixin: bool | None = None,
+            primitives: list[type] | None = None, # only inherit arithmetic primitives
+            callables: list[tuple[str, Callable]] | None = None,
             semantic: bool = False,
             **kwargs) -> "Symbol":
+        if primitives is None:
+            primitives = [OperatorPrimitives]
         return super().__new__(cls, *args,
                          mixin=mixin,
                          primitives=primitives,
@@ -97,8 +102,8 @@ class Aggregator(Symbol):
             # add also a property with the same name but without the SPECIAL_CONSTANT prefix as a shortcut
             self.__dict__[name] = self.__dict__[f'{SPECIAL_CONSTANT}{name}']
             return self.__dict__.get(name)
-        elif not self._active and name not in self.__dict__:
-            raise Exception(f'Aggregator object is frozen! No attribute {name} found!')
+        if not self._active and name not in self.__dict__:
+            UserMessage(f'Aggregator object is frozen! No attribute {name} found!', raise_with=Exception)
         return self.__dict__.get(name)
 
     def __setattr__(self, name, value):
@@ -139,19 +144,23 @@ class Aggregator(Symbol):
     def _set_values(obj, dictionary, parent, strict: bool = True):
         # recursively reconstruct the object
         for key, value in dictionary.items():
-            if isinstance(value, dict):
+            attr_key = key
+            attr_value = value
+            if isinstance(attr_value, dict):
                 if parent is not None:
-                    obj._path = key
-                value = Aggregator._reconstruct(value, parent=parent, strict=strict)
-            if key.startswith(SPECIAL_CONSTANT):
-                key = key.replace(SPECIAL_CONSTANT, '')
-            if key == '_value':
+                    obj._path = attr_key
+                attr_value = Aggregator._reconstruct(attr_value, parent=parent, strict=strict)
+            if attr_key.startswith(SPECIAL_CONSTANT):
+                attr_key = attr_key.replace(SPECIAL_CONSTANT, '')
+            if attr_key == '_value':
                 try:
-                    value = np.asarray(value, dtype=np.float32)
+                    attr_value = np.asarray(attr_value, dtype=np.float32)
                 except Exception as e:
                     if strict:
-                        raise Exception(f'Could not set value of Aggregator object: {obj.path}! ERROR: {e}') from e
-            obj.__setattr__(key, value)
+                        msg = f'Could not set value of Aggregator object: {obj.path}! ERROR: {e}'
+                        UserMessage(msg)
+                        raise Exception(msg) from e
+            obj.__setattr__(attr_key, attr_value)
 
     @staticmethod
     def _reconstruct(dictionary, parent = None, strict: bool = True):
@@ -186,8 +195,7 @@ class Aggregator(Symbol):
             if obj._path is not None:
                 path = obj._path.replace(SPECIAL_CONSTANT, '') + '.' + path
             obj = obj._parent
-        path = path[:-1] # remove last dot
-        return path
+        return path[:-1] # remove last dot
 
     def __or__(self, other: Any) -> Any:
         self.add(other)
@@ -211,8 +219,7 @@ class Aggregator(Symbol):
     @property
     def value(self):
         if self.map is not None:
-            res = np.asarray(self.map(np.asarray(self._value, dtype=np.float32)))
-            return res
+            return np.asarray(self.map(np.asarray(self._value, dtype=np.float32)))
         return np.asarray(self._value, dtype=np.float32)
 
     @property
@@ -232,19 +239,18 @@ class Aggregator(Symbol):
     def shape(self):
         if len(self.entries) > 0:
             return np.asarray(self.entries).shape
-        else:
-            return ()
+        return ()
 
     def serialize(self):
         return json.dumps(self, cls=AggregatorJSONEncoder)
 
     def save(self, path: str):
-        with open(path, 'w') as f:
+        with Path(path).open('w') as f:
             json.dump(self, f, cls=AggregatorJSONEncoder)
 
     @staticmethod
     def load(path: str, strict: bool = True):
-        with open(path, 'r') as f:
+        with Path(path).open() as f:
             json_ = json.load(f)
         return Aggregator._reconstruct(json_, strict=strict)
 
@@ -253,44 +259,50 @@ class Aggregator(Symbol):
 
     def add(self, entries):
         # Add entries to the aggregator
-        if not self.active:
-            if self._finalized:
-                raise Exception('Aggregator object is frozen!')
+        if not self.active and self._finalized:
+            UserMessage('Aggregator object is frozen!', raise_with=Exception)
             return
         try:
-            # Append a new entry to the aggregator
-            assert type(entries) in [tuple, list, np.float32, np.float64, np.ndarray, torch.Tensor, int, float, bool, str] or isinstance(entries, Symbol), 'Entries must be a tuple, list, numpy array, torch tensor, integer, float, boolean, string, or Symbol! Got: {}'.format(type(entries))
-            if type(entries) == torch.Tensor:
-                entries = entries.detach().cpu().numpy().astype(np.float32)
-            elif type(entries) in [tuple, list]:
-                entries = np.asarray(entries, dtype=np.float32)
-            elif type(entries) in [int, float]:
-                entries = entries
-            elif type(entries) == bool:
-                entries = int(entries)
-            elif type(entries) == str:
-                entries = Symbol(entries).embedding.astype(np.float32)
-            elif isinstance(entries, Symbol):
-                # use this to avoid recursion on map setter
-                self.add(entries._value)
-                return
-            elif isinstance(entries, Aggregator):
-                self.add(entries.get())
+            processed_entries = self._prepare_entries(entries)
+            if processed_entries is None:
                 return
-
-            if isinstance(entries, np.ndarray) or isinstance(entries, np.float32):
-                entries = entries.squeeze()
-
-            self.entries.append(entries)
+            processed_entries = self._squeeze_entries(processed_entries)
+            self.entries.append(processed_entries)
         except Exception as e:
+            msg = f'Could not add entries to Aggregator object! Please verify type or original error: {e}'
             if self._raise_error:
-                raise Exception(f'Could not add entries to Aggregator object! Please verify type or original error: {e}') from e
-            else:
-                print('Could not add entries to Aggregator object! Please verify type or original error: {}'.format(e))
+                UserMessage(msg)
+                raise Exception(msg) from e
+            UserMessage(msg)
+
+    def _prepare_entries(self, entries):
+        valid_types = (tuple, list, np.float32, np.float64, np.ndarray, torch.Tensor, int, float, bool, str, Symbol)
+        assert isinstance(entries, valid_types), f'Entries must be a tuple, list, numpy array, torch tensor, integer, float, boolean, string, or Symbol! Got: {type(entries)}'
+        if isinstance(entries, torch.Tensor):
+            return entries.detach().cpu().numpy().astype(np.float32)
+        if isinstance(entries, (tuple, list)):
+            return np.asarray(entries, dtype=np.float32)
+        if isinstance(entries, bool):
+            return int(entries)
+        if isinstance(entries, str):
+            return Symbol(entries).embedding.astype(np.float32)
+        if isinstance(entries, Symbol):
+            # Use this to avoid recursion on map setter
+            self.add(entries._value)
+            return None
+        if isinstance(entries, Aggregator):
+            self.add(entries.get())
+            return None
+        return entries
+
+    def _squeeze_entries(self, entries):
+        if isinstance(entries, (np.ndarray, np.float32)):
+            return entries.squeeze()
+        return entries
 
     def keys(self):
         # Get all key names of items that have the SPECIAL_CONSTANT prefix
-        return [key.replace(SPECIAL_CONSTANT, '') for key in self.__dict__.keys() if not key.startswith('_') and \
+        return [key.replace(SPECIAL_CONSTANT, '') for key in self.__dict__ if not key.startswith('_') and \
                 key.replace(SPECIAL_CONSTANT, '') not in EXCLUDE_LIST]
 
     @property
@@ -301,7 +313,7 @@ class Aggregator(Symbol):
     @active.setter
     def active(self, value):
         # Set the active status of the aggregator
-        assert isinstance(value, bool), 'Active status must be a boolean! Got: {}'.format(type(value))
+        assert isinstance(value, bool), f'Active status must be a boolean! Got: {type(value)}'
         self._active = value
 
     @property
@@ -312,7 +324,7 @@ class Aggregator(Symbol):
     @finalized.setter
     def finalized(self, value):
         # Set the finalized status of the aggregator
-        assert isinstance(value, bool), 'Finalized status must be a boolean! Got: {}'.format(type(value))
+        assert isinstance(value, bool), f'Finalized status must be a boolean! Got: {type(value)}'
         self._finalized = value
 
     def finalize(self):
@@ -323,7 +335,7 @@ class Aggregator(Symbol):
             if name == 'map':
                 self.__setattr__(name, value)
             else:
-                raise Exception('Aggregator object is frozen!')
+                UserMessage('Aggregator object is frozen!', raise_with=Exception)
         self.__setattr__ = raise_exception
         def get_attribute(*args, **kwargs):
             return self.__dict__.get(*args, **kwargs)
@@ -342,7 +354,7 @@ class Aggregator(Symbol):
     def clear(self):
         # Clear the entries of the aggregator
         if self._finalized:
-            raise Exception('Aggregator object is frozen!')
+            UserMessage('Aggregator object is frozen!', raise_with=Exception)
         self._value = []
 
     def sum(self, axis=0):