symbolicai 0.21.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

symai/core_ext.py

@@ -3,22 +3,20 @@ import atexit
 import functools
 import logging
 import multiprocessing as mp
-import os
 import pickle
 import random
-import sys
-import time
 import threading
+import time
 import traceback
+from collections.abc import Callable
 from pathlib import Path
-from typing import Callable, List
 
 import dill
 from loguru import logger
-from pathos.multiprocessing import ProcessingPool as PPool
 
 from . import __root_dir__
 from .functional import EngineRepository
+from .utils import UserMessage
 
 logging.getLogger("multiprocessing").setLevel(logging.ERROR)
 
@@ -28,6 +26,7 @@ logging.getLogger("multiprocessing").setLevel(logging.ERROR)
 _pool_registry: dict[int, mp.pool.Pool] = {}
 _pool_lock = threading.Lock()
 
+
 def _get_pool(workers: int) -> mp.pool.Pool:
     pool = _pool_registry.get(workers)
     if pool is not None:
@@ -39,66 +38,67 @@ def _get_pool(workers: int) -> mp.pool.Pool:
             _pool_registry[workers] = pool
     return pool
 
+
 @atexit.register
 def _shutdown_pools() -> None:
     for pool in _pool_registry.values():
         pool.close()
         pool.join()
 
+
 def _run_in_process(expr, func, args, kwargs):
     expr = dill.loads(expr)
     func = dill.loads(func)
     return func(expr, *args, **kwargs)
 
-def _parallel(func: Callable, expressions: List[Callable], worker: int = mp.cpu_count() // 2):
+
+def _parallel(func: Callable, expressions: list[Callable], worker: int = mp.cpu_count() // 2):
     pickled_exprs = [dill.dumps(expr) for expr in expressions]
-    pickled_func  = dill.dumps(func)
+    pickled_func = dill.dumps(func)
     pool = _get_pool(worker)
+
     def proxy_function(*args, **kwargs):
         return pool.starmap(
-                _run_in_process,
-                [(expr, pickled_func, args, kwargs) for expr in pickled_exprs]
-            )
+            _run_in_process, [(expr, pickled_func, args, kwargs) for expr in pickled_exprs]
+        )
+
     return proxy_function
 
+
 # Decorator
-def parallel(expressions: List[Callable], worker: int = mp.cpu_count() // 2):
+def parallel(expressions: list[Callable], worker: int = mp.cpu_count() // 2):
     def decorator_parallel(func):
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
             # Run expressions in parallel
             parallel_func = _parallel(func, expressions, worker=worker)
             # Call the proxy function to execute in parallel and capture results
-            results = parallel_func(*args, **kwargs)
-            return results
+            return parallel_func(*args, **kwargs)
+
         return wrapper
+
     return decorator_parallel
 
+
 def bind(engine: str, property: str):
-    '''
+    """
     Bind to an engine and retrieve one of its properties.
-    '''
+    """
+
     def decorator(func):
         @functools.wraps(func)
-        def wrapper(*args, **kwargs):
-            return EngineRepository.bind_property(
-                    engine=engine,
-                    property=property
-                )
+        def wrapper(*_args, **_kwargs):
+            return EngineRepository.bind_property(engine=engine, property=property)
+
         return wrapper
+
     return decorator
 
 
 def retry(
-    exceptions=Exception,
-    tries=-1,
-    delay=0,
-    max_delay=-1,
-    backoff=1,
-    jitter=0,
-    graceful=False
+    exceptions=Exception, tries=-1, delay=0, max_delay=-1, backoff=1, jitter=0, graceful=False
 ):
-    '''
+    """
     Returns a retry decorator for both async and sync functions.
 
     :param exceptions: an exception or a tuple of exceptions to catch. default: Exception.
@@ -111,9 +111,11 @@ def retry(
     :param graceful:   whether to raise an exception if all attempts fail or return with None. default: False.
 
     Credits to invlpg (https://pypi.org/project/retry)
-    '''
+    """
+
     def decorator(func):
         if asyncio.iscoroutinefunction(func):
+
             @functools.wraps(func)
             async def async_wrapper(*args, **kwargs):
                 return await _aretry_func(
@@ -124,23 +126,26 @@ def retry(
                     max_delay=max_delay,
                     backoff=backoff,
                     jitter=jitter,
-                    graceful=graceful
+                    graceful=graceful,
                 )
+
             return async_wrapper
-        else:
-            @functools.wraps(func)
-            def sync_wrapper(*args, **kwargs):
-                return _retry_func(
-                    functools.partial(func, *args, **kwargs),
-                    exceptions=exceptions,
-                    tries=tries,
-                    delay=delay,
-                    max_delay=max_delay,
-                    backoff=backoff,
-                    jitter=jitter,
-                    graceful=graceful
-                )
-            return sync_wrapper
+
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            return _retry_func(
+                functools.partial(func, *args, **kwargs),
+                exceptions=exceptions,
+                tries=tries,
+                delay=delay,
+                max_delay=max_delay,
+                backoff=backoff,
+                jitter=jitter,
+                graceful=graceful,
+            )
+
+        return sync_wrapper
+
     return decorator
 
 
@@ -152,13 +157,13 @@ def _retry_func(
     max_delay: int,
     backoff: int,
     jitter: int,
-    graceful: bool
+    graceful: bool,
 ):
     _tries, _delay = tries, delay
     while _tries:
         try:
             return func()
-        except exceptions as e:
+        except exceptions:
             _tries -= 1
             if not _tries:
                 if graceful:
@@ -175,6 +180,7 @@ def _retry_func(
 
             if max_delay >= 0:
                 _delay = min(_delay, max_delay)
+    return None
 
 
 async def _aretry_func(
@@ -185,13 +191,13 @@ async def _aretry_func(
     max_delay: int,
     backoff: int,
     jitter: int,
-    graceful: bool
+    graceful: bool,
 ):
     _tries, _delay = tries, delay
     while _tries:
         try:
             return await func()
-        except exceptions as e:
+        except exceptions:
             _tries -= 1
             if not _tries:
                 if graceful:
@@ -208,54 +214,45 @@ async def _aretry_func(
 
             if max_delay >= 0:
                 _delay = min(_delay, max_delay)
+    return None
 
 
-def cache(
-    in_memory:  bool,
-    cache_path: str = __root_dir__ / 'cache'
-):
-    '''
+def cache(in_memory: bool, cache_path: str = __root_dir__ / "cache"):
+    """
     Cache the result of a *any* function call. This is very useful in cost optimization (e.g. computing embeddings).
-    '''
+    """
+
     def decorator(func):
         @functools.wraps(func)
         def wrapper(instance):
-            return _cache_registry_func(
-                    in_memory,
-                    cache_path,
-                    func,
-                    instance
-                )
-        return wrapper
-    return decorator
+            return _cache_registry_func(in_memory, cache_path, func, instance)
 
+        return wrapper
 
-def _cache_registry_func(
-        in_memory: bool,
-        cache_path: str,
-        func: Callable,
-        *args, **kwargs
-    ):
+    return decorator
 
-    if not os.path.exists(cache_path): os.makedirs(cache_path)
 
-    if in_memory and os.path.exists(Path(cache_path) / func.__qualname__):
-        with open(Path(cache_path) / func.__qualname__, 'rb') as f:
-            call = pickle.load(f)
+def _cache_registry_func(in_memory: bool, cache_path: str, func: Callable, *args, **kwargs):
+    cache_dir = Path(cache_path)
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    cache_file = cache_dir / func.__qualname__
 
-        return call
+    if in_memory and cache_file.exists():
+        with cache_file.open("rb") as f:
+            return pickle.load(f)
 
     call = func(*args, **kwargs)
-    with open(Path(cache_path) / func.__qualname__, 'wb') as f:
-        pickle.dump(call , f)
+    with cache_file.open("wb") as f:
+        pickle.dump(call, f)
 
     return call
 
 
 def error_logging(debug: bool = False):
-    '''
+    """
     Log the error of a function call.
-    '''
+    """
+
     def dec(func):
         @functools.wraps(func)
         def _dec(*args, **kwargs):
@@ -265,11 +262,13 @@ def error_logging(debug: bool = False):
                 logger.error(e)
                 if debug:
                     # Simple message:
-                    print('Function: {} call failed. Error: {}'.format(func.__name__, e))
+                    UserMessage(f"Function: {func.__name__} call failed. Error: {e}")
                     # print traceback
                     traceback.print_exc()
                 raise e
+
         return _dec
+
     return dec
 
 
@@ -291,18 +290,21 @@ def deprecated(reason: str = ""):
             @deprecated("Use new_method() instead")
             def old_method(self): pass
     """
+
     def decorator(obj):
         if isinstance(obj, type):
             # If obj is a class
             original_init = obj.__init__
+
             @functools.wraps(original_init)
             def new_init(self, *args, **kwargs):
                 logger.warning(
                     f"{obj.__name__} is deprecated and will be removed in future versions. {reason}",
                     category=DeprecationWarning,
-                    stacklevel=2
+                    stacklevel=2,
                 )
                 original_init(self, *args, **kwargs)
+
             obj.__init__ = new_init
             return obj
 
@@ -312,9 +314,10 @@ def deprecated(reason: str = ""):
             logger.warning(
                 f"{obj.__name__} is deprecated and will be removed in future versions. {reason}",
                 category=DeprecationWarning,
-                stacklevel=2
+                stacklevel=2,
             )
             return obj(*args, **kwargs)
+
         return wrapper
 
     return decorator