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

symai/backend/base.py

@@ -1,115 +1,125 @@
 import logging
-import os
 import time
+from abc import ABC, abstractmethod
+from typing import Any
 
-from abc import ABC
-from typing import Any, List, Tuple
+from ..collect import CollectionRepository, rec_serialize
+from ..utils import UserMessage
 from .settings import HOME_PATH
 
-from ..collect import CollectionRepository, rec_serialize
+ENGINE_UNREGISTERED = "<UNREGISTERED/>"
 
+COLLECTION_LOGGING_ENGINES = {
+    "GPTXChatEngine",
+    "GPTXCompletionEngine",
+    "SerpApiEngine",
+    "WolframAlphaEngine",
+    "SeleniumEngine",
+    "OCREngine",
+}
 
-ENGINE_UNREGISTERED = '<UNREGISTERED/>'
 
 class Engine(ABC):
     def __init__(self) -> None:
         super().__init__()
-        self.verbose    = False
-        self.logging    = False
-        self.log_level  = logging.DEBUG
+        self.verbose = False
+        self.logging = False
+        self.log_level = logging.DEBUG
         self.time_clock = False
         self.collection = CollectionRepository()
         self.collection.connect()
         # create formatter
-        __root_dir__  = HOME_PATH
-        os.makedirs(__root_dir__, exist_ok=True)
+        __root_dir__ = HOME_PATH
+        __root_dir__.mkdir(parents=True, exist_ok=True)
         __file_path__ = __root_dir__ / "engine.log"
-        logging.basicConfig(filename=__file_path__, filemode="a", format='%(asctime)s %(name)s %(levelname)s %(message)s')
-        self.logger     = logging.getLogger()
+        logging.basicConfig(
+            filename=__file_path__,
+            filemode="a",
+            format="%(asctime)s %(name)s %(levelname)s %(message)s",
+        )
+        self.logger = logging.getLogger()
         self.logger.setLevel(logging.DEBUG)
         # logging to console
-        stream          = logging.StreamHandler()
-        streamformat    = logging.Formatter("%(asctime)s %(message)s")
+        stream = logging.StreamHandler()
+        streamformat = logging.Formatter("%(asctime)s %(message)s")
         stream.setLevel(logging.INFO)
         stream.setFormatter(streamformat)
         self.logger.addHandler(stream)
 
-    def __call__(self, argument: Any) -> Tuple[List[str], dict]:
-        log = {
-            'Input': {
-                'self': self,
-                'args': argument.args,
-                **argument.kwargs
-            }
-        }
+    def __call__(self, argument: Any) -> tuple[list[str], dict]:
+        log = {"Input": {"self": self, "args": argument.args, **argument.kwargs}}
         start_time = time.time()
 
-        # check for global object based input handler
-        if hasattr(argument.prop.instance, '_metadata') and hasattr(argument.prop.instance._metadata, 'input_handler'):
-            input_handler = argument.prop.instance._metadata.input_handler if hasattr(argument.prop.instance._metadata, 'input_handler') else None
-            if input_handler is not None:
-                input_handler((argument.prop.processed_input, argument))
-        # check for kwargs based input handler
-        if argument.prop.input_handler is not None:
-            argument.prop.input_handler((argument.prop.processed_input, argument))
+        self._trigger_input_handlers(argument)
 
-        # execute the engine
         res, metadata = self.forward(argument)
 
-        # compute time
         req_time = time.time() - start_time
-        metadata['time'] = req_time
+        metadata["time"] = req_time
         if self.time_clock:
-            print(f"{argument.prop.func}: {req_time} sec")
-        log['Output'] = res
+            UserMessage(f"{argument.prop.func}: {req_time} sec")
+        log["Output"] = res
         if self.verbose:
-            view   = {k: v for k, v in list(log['Input'].items()) if k != 'self'}
-            input_ = f"{str(log['Input']['self'])[:50]}, {str(argument.prop.func)}, {str(view)}"
-            print(input_[:150], str(log['Output'])[:100])
+            view = {k: v for k, v in list(log["Input"].items()) if k != "self"}
+            input_ = f"{str(log['Input']['self'])[:50]}, {argument.prop.func!s}, {view!s}"
+            UserMessage(f"{input_[:150]} {str(log['Output'])[:100]}")
         if self.logging:
             self.logger.log(self.log_level, log)
 
-        # share data statistics with the collection repository
-        if     str(self) == 'GPTXChatEngine' \
-            or str(self) == 'GPTXCompletionEngine' \
-            or str(self) == 'SerpApiEngine' \
-            or str(self) == 'WolframAlphaEngine' \
-            or str(self) == 'SeleniumEngine' \
-            or str(self) == 'OCREngine':
-            self.collection.add(
-                forward={'args': rec_serialize(argument.args), 'kwds': rec_serialize(argument.kwargs)},
-                engine=str(self),
-                metadata={
-                    'time': req_time,
-                    'data': rec_serialize(metadata),
-                    'argument': rec_serialize(argument)
-                }
-            )
-
-        # check for global object based output handler
-        if hasattr(argument.prop.instance, '_metadata') and hasattr(argument.prop.instance._metadata, 'output_handler'):
-            output_handler = argument.prop.instance._metadata.output_handler if hasattr(argument.prop.instance._metadata, 'output_handler') else None
-            if output_handler:
-                output_handler(res)
-        # check for kwargs based output handler
-        if argument.prop.output_handler:
-            argument.prop.output_handler((res, metadata))
+        if str(self) in COLLECTION_LOGGING_ENGINES:
+            self._record_collection_entry(argument, metadata, req_time)
+
+        self._trigger_output_handlers(argument, res, metadata)
         return res, metadata
 
+    def _trigger_input_handlers(self, argument: Any) -> None:
+        instance_metadata = getattr(argument.prop.instance, "_metadata", None)
+        if instance_metadata is not None:
+            input_handler = getattr(instance_metadata, "input_handler", None)
+            if input_handler is not None:
+                input_handler((argument.prop.processed_input, argument))
+        argument_handler = argument.prop.input_handler
+        if argument_handler is not None:
+            argument_handler((argument.prop.processed_input, argument))
+
+    def _trigger_output_handlers(self, argument: Any, result: Any, metadata: dict | None) -> None:
+        instance_metadata = getattr(argument.prop.instance, "_metadata", None)
+        if instance_metadata is not None:
+            output_handler = getattr(instance_metadata, "output_handler", None)
+            if output_handler:
+                output_handler(result)
+        argument_handler = argument.prop.output_handler
+        if argument_handler:
+            argument_handler((result, metadata))
+
+    def _record_collection_entry(self, argument: Any, metadata: dict, req_time: float) -> None:
+        self.collection.add(
+            forward={"args": rec_serialize(argument.args), "kwds": rec_serialize(argument.kwargs)},
+            engine=str(self),
+            metadata={
+                "time": req_time,
+                "data": rec_serialize(metadata),
+                "argument": rec_serialize(argument),
+            },
+        )
+
     def id(self) -> str:
         return ENGINE_UNREGISTERED
 
     def preview(self, argument):
-        # used here to avoid circular import
-        from ..symbol import Symbol
+        # Used here to avoid backend.base <-> symbol circular import.
+        from ..symbol import (  # noqa
+            Symbol,
+        )
+
         class Preview(Symbol):
             def __repr__(self) -> str:
-                '''
+                """
                 Get the representation of the Symbol object as a string.
 
                 Returns:
                     str: The representation of the Symbol object.
-                '''
+                """
                 return str(self.value.prop.prepared_input)
 
             def prepared_input(self):
@@ -117,36 +127,38 @@ class Engine(ABC):
 
         return Preview(argument), {}
 
-    def forward(self, *args: Any, **kwds: Any) -> List[str]:
-        raise NotADirectoryError()
+    @abstractmethod
+    def forward(self, *args: Any, **kwds: Any) -> list[str]:
+        raise NotADirectoryError
 
+    @abstractmethod
     def prepare(self, argument):
-        raise NotImplementedError()
-
-    def command(self, *args, **kwargs):
-        if kwargs.get('verbose', None):
-            self.verbose = kwargs['verbose']
-        if kwargs.get('logging', None):
-            self.logging = kwargs['logging']
-        if kwargs.get('log_level', None):
-            self.log_level = kwargs['log_level']
-        if kwargs.get('time_clock', None):
-            self.time_clock = kwargs['time_clock']
+        raise NotImplementedError
+
+    def command(self, *_args, **kwargs):
+        if kwargs.get("verbose"):
+            self.verbose = kwargs["verbose"]
+        if kwargs.get("logging"):
+            self.logging = kwargs["logging"]
+        if kwargs.get("log_level"):
+            self.log_level = kwargs["log_level"]
+        if kwargs.get("time_clock"):
+            self.time_clock = kwargs["time_clock"]
 
     def __str__(self) -> str:
         return self.__class__.__name__
 
     def __repr__(self) -> str:
-        '''
+        """
         Get the representation of the Symbol object as a string.
 
         Returns:
             str: The representation of the Symbol object.
-        '''
+        """
         # class with full path
-        class_ = self.__class__.__module__ + '.' + self.__class__.__name__
-        hex_   = hex(id(self))
-        return f'<class {class_} at {hex_}>'
+        class_ = self.__class__.__module__ + "." + self.__class__.__name__
+        hex_ = hex(id(self))
+        return f"<class {class_} at {hex_}>"
 
 
 class BatchEngine(Engine):
@@ -155,41 +167,42 @@ class BatchEngine(Engine):
         self.time_clock = True
         self.allows_batching = True
 
-    def __call__(self, arguments: List[Any]) -> List[Tuple[Any, dict]]:
+    def __call__(self, arguments: list[Any]) -> list[tuple[Any, dict]]:
         start_time = time.time()
         for arg in arguments:
-            if hasattr(arg.prop.instance, '_metadata') and hasattr(arg.prop.instance._metadata, 'input_handler'):
-                input_handler = getattr(arg.prop.instance._metadata, 'input_handler', None)
-                if input_handler is not None:
-                    input_handler((arg.prop.processed_input, arg))
-            if arg.prop.input_handler is not None:
-                arg.prop.input_handler((arg.prop.processed_input, arg))
+            self._trigger_input_handlers(arg)
 
-        try:
-            results, metadata_list = self.forward(arguments)
-        except Exception as e:
-            results = [e] * len(arguments)
-            metadata_list = [None] * len(arguments)
+        results, metadata_list = self._execute_batch(arguments)
 
         total_time = time.time() - start_time
         if self.time_clock:
-            print(f"Total execution time: {total_time} sec")
+            UserMessage(f"Total execution time: {total_time} sec")
 
-        return_list = []
+        return self._prepare_batch_results(arguments, results, metadata_list, total_time)
 
-        for arg, result, metadata in zip(arguments, results, metadata_list):
+    def _execute_batch(self, arguments: list[Any]) -> tuple[list[Any], list[dict | None]]:
+        try:
+            return self.forward(arguments)
+        except Exception as error:
+            return [error] * len(arguments), [None] * len(arguments)
+
+    def _prepare_batch_results(
+        self,
+        arguments: list[Any],
+        results: list[Any],
+        metadata_list: list[dict | None],
+        total_time: float,
+    ) -> list[tuple[Any, dict | None]]:
+        return_list = []
+        for arg, result, metadata in zip(arguments, results, metadata_list, strict=False):
             if metadata is not None:
-                metadata['time'] = total_time / len(arguments)
-
-            if hasattr(arg.prop.instance, '_metadata') and hasattr(arg.prop.instance._metadata, 'output_handler'):
-                output_handler = getattr(arg.prop.instance._metadata, 'output_handler', None)
-                if output_handler:
-                    output_handler(result)
-            if arg.prop.output_handler:
-                arg.prop.output_handler((result, metadata))
+                metadata["time"] = total_time / len(arguments)
 
+            self._trigger_output_handlers(arg, result, metadata)
             return_list.append((result, metadata))
         return return_list
 
-    def forward(self, arguments: List[Any]) -> Tuple[List[Any], List[dict]]:
-        raise NotImplementedError("Subclasses must implement forward method")
+    def forward(self, _arguments: list[Any]) -> tuple[list[Any], list[dict]]:
+        msg = "Subclasses must implement forward method"
+        UserMessage(msg)
+        raise NotImplementedError(msg)

symai/backend/engines/drawing/engine_bfl.py

@@ -1,11 +1,12 @@
 import logging
 import tempfile
 import time
-from typing import Optional
+from pathlib import Path
 
 import requests
 
 from ....symbol import Result
+from ....utils import UserMessage
 from ...base import Engine
 from ...settings import SYMAI_CONFIG
 
@@ -19,89 +20,89 @@ class FluxResult(Result):
     def __init__(self, value, **kwargs):
         super().__init__(value, **kwargs)
         # unpack the result
-        path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
-        url = value.get('result').get('sample')
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
+            path = tmp_file.name
+        url = value.get("result").get("sample")
         request = requests.get(url, allow_redirects=True)
         request.raise_for_status()
-        with open(path, "wb") as f:
+        with Path(path).open("wb") as f:
             f.write(request.content)
         self._value = [path]
 
 
 class DrawingEngine(Engine):
-    def __init__(self, api_key: Optional[str] = None, model: Optional[str] = None):
+    def __init__(self, api_key: str | None = None, model: str | None = None):
         super().__init__()
         self.config = SYMAI_CONFIG
-        self.api_key = self.config['DRAWING_ENGINE_API_KEY'] if api_key is None else api_key
-        self.model = self.config['DRAWING_ENGINE_MODEL'] if model is None else model
+        self.api_key = self.config["DRAWING_ENGINE_API_KEY"] if api_key is None else api_key
+        self.model = self.config["DRAWING_ENGINE_MODEL"] if model is None else model
         self.name = self.__class__.__name__
 
     def id(self) -> str:
-        if  self.config['DRAWING_ENGINE_API_KEY'] and self.config['DRAWING_ENGINE_MODEL'].startswith("flux"):
-            return 'drawing'
-        return super().id() # default to unregistered
+        if self.config["DRAWING_ENGINE_API_KEY"] and self.config["DRAWING_ENGINE_MODEL"].startswith(
+            "flux"
+        ):
+            return "drawing"
+        return super().id()  # default to unregistered
 
     def command(self, *args, **kwargs):
         super().command(*args, **kwargs)
-        if 'DRAWING_ENGINE_API_KEY' in kwargs:
-            self.api_key = kwargs['DRAWING_ENGINE_API_KEY']
-        if 'DRAWING_ENGINE_MODEL' in kwargs:
-            self.model = kwargs['DRAWING_ENGINE_MODEL']
+        if "DRAWING_ENGINE_API_KEY" in kwargs:
+            self.api_key = kwargs["DRAWING_ENGINE_API_KEY"]
+        if "DRAWING_ENGINE_MODEL" in kwargs:
+            self.model = kwargs["DRAWING_ENGINE_MODEL"]
 
     def forward(self, argument):
         prompt = argument.prop.prepared_input
         kwargs = argument.kwargs
-        model = kwargs.get('model', self.model)
-        width = kwargs.get('width', 1024)
-        height = kwargs.get('height', 768)
-        steps = kwargs.get('steps', 40)
-        prompt_upsampling = kwargs.get('prompt_upsampling', False)
-        seed = kwargs.get('seed', None)
-        guidance = kwargs.get('guidance', None)
-        safety_tolerance = kwargs.get('safety_tolerance', 2)
-        interval = kwargs.get('interval', None)
-        output_format = kwargs.get('output_format', 'png')
-        except_remedy = kwargs.get('except_remedy', None)
+        width = kwargs.get("width", 1024)
+        height = kwargs.get("height", 768)
+        steps = kwargs.get("steps", 40)
+        seed = kwargs.get("seed", None)
+        guidance = kwargs.get("guidance", None)
+        safety_tolerance = kwargs.get("safety_tolerance", 2)
+        except_remedy = kwargs.get("except_remedy", None)
 
         headers = {
-            'accept': 'application/json',
-            'x-key': self.api_key,
-            'Content-Type': 'application/json',
+            "accept": "application/json",
+            "x-key": self.api_key,
+            "Content-Type": "application/json",
         }
 
         payload = {
-            'prompt': prompt,
-            'width': width,
-            'height': height,
-            'num_inference_steps': steps,
-            'guidance_scale': guidance,
-            'seed': seed,
-            'safety_tolerance': safety_tolerance,
+            "prompt": prompt,
+            "width": width,
+            "height": height,
+            "num_inference_steps": steps,
+            "guidance_scale": guidance,
+            "seed": seed,
+            "safety_tolerance": safety_tolerance,
         }
         # drop any None values so Flux API won't return 500
         payload = {k: v for k, v in payload.items() if v is not None}
 
-        if kwargs.get('operation') == 'create':
+        if kwargs.get("operation") == "create":
             try:
                 response = requests.post(
-                    f'https://api.us1.bfl.ai/v1/{self.model}',
-                    headers=headers,
-                    json=payload
+                    f"https://api.us1.bfl.ai/v1/{self.model}", headers=headers, json=payload
                 )
                 # fail early on HTTP errors
                 response.raise_for_status()
                 data = response.json()
                 request_id = data.get("id")
                 if not request_id:
-                    raise Exception(f"Failed to get request ID! Response payload: {data}")
+                    UserMessage(
+                        f"Failed to get request ID! Response payload: {data}",
+                        raise_with=Exception,
+                    )
 
                 while True:
                     time.sleep(5)
 
                     result = requests.get(
-                        'https://api.us1.bfl.ai/v1/get_result',
+                        "https://api.us1.bfl.ai/v1/get_result",
                         headers=headers,
-                        params={'id': request_id}
+                        params={"id": request_id},
                     )
 
                     result.raise_for_status()
@@ -118,8 +119,8 @@ class DrawingEngine(Engine):
 
             metadata = {}
             return [rsp], metadata
-        else:
-            raise Exception(f"Unknown operation: {kwargs['operation']}")
+        UserMessage(f"Unknown operation: {kwargs['operation']}", raise_with=Exception)
+        return [], {}
 
     def prepare(self, argument):
         argument.prop.prepared_input = str(argument.prop.processed_input)