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

symai/components.py

@@ -1,23 +1,19 @@
 import copy
 import inspect
 import json
-import os
 import re
 import sys
-from abc import abstractmethod
 from collections import defaultdict
+from collections.abc import Callable, Iterator
 from pathlib import Path
 from random import sample
 from string import ascii_lowercase, ascii_uppercase
 from threading import Lock
-from .context import CURRENT_ENGINE_VAR
-from typing import Callable, Dict, Iterator, List, Optional, Type, Union
+from typing import TYPE_CHECKING, Union
 
 import numpy as np
-from attr import dataclass
 from box import Box
 from loguru import logger
-from pydantic import BaseModel, ValidationError
 from pyvis.network import Network
 from tqdm import tqdm
 
@@ -25,16 +21,25 @@ from . import core, core_ext
 from .backend.base import Engine
 from .backend.settings import HOME_PATH
 from .constraints import DictFormatConstraint
+from .context import CURRENT_ENGINE_VAR
 from .formatter import ParagraphFormatter
-from .post_processors import (CodeExtractPostProcessor,
-                              JsonTruncateMarkdownPostProcessor,
-                              JsonTruncatePostProcessor, PostProcessor,
-                              StripPostProcessor)
+from .post_processors import (
+    CodeExtractPostProcessor,
+    JsonTruncateMarkdownPostProcessor,
+    JsonTruncatePostProcessor,
+    PostProcessor,
+    StripPostProcessor,
+)
 from .pre_processors import JsonPreProcessor, PreProcessor
 from .processor import ProcessorPipeline
 from .prompts import JsonPromptTemplate, Prompt
 from .symbol import Expression, Metadata, Symbol
-from .utils import CustomUserWarning
+from .utils import UserMessage
+
+if TYPE_CHECKING:
+    from .backend.engines.index.engine_vectordb import VectorDBResult
+
+_DEFAULT_PARAGRAPH_FORMATTER = ParagraphFormatter()
 
 
 class GraphViz(Expression):
@@ -58,7 +63,7 @@ class GraphViz(Expression):
                             select_menu=select_menu,
                             filter_menu=filter_menu)
 
-    def forward(self, sym: Symbol, file_path: str, **kwargs):
+    def forward(self, sym: Symbol, file_path: str, **_kwargs):
         nodes = [str(n) if n.value else n.__repr__(simplified=True) for n in sym.nodes]
         edges = [(str(e[0]) if e[0].value else e[0].__repr__(simplified=True),
                   str(e[1]) if e[1].value else e[1].__repr__(simplified=True)) for e in sym.edges]
@@ -73,21 +78,21 @@ class TrackerTraceable(Expression):
 
 
 class Any(Expression):
-    def __init__(self, *expr: List[Expression], **kwargs):
+    def __init__(self, *expr: list[Expression], **kwargs):
         super().__init__(**kwargs)
-        self.expr: List[Expression] = expr
+        self.expr: list[Expression] = expr
 
     def forward(self, *args, **kwargs) -> Symbol:
-        return self.sym_return_type(any([e() for e in self.expr(*args, **kwargs)]))
+        return self.sym_return_type(any(e() for e in self.expr(*args, **kwargs)))
 
 
 class All(Expression):
-    def __init__(self, *expr: List[Expression], **kwargs):
+    def __init__(self, *expr: list[Expression], **kwargs):
         super().__init__(**kwargs)
-        self.expr: List[Expression] = expr
+        self.expr: list[Expression] = expr
 
     def forward(self, *args, **kwargs) -> Symbol:
-        return self.sym_return_type(all([e() for e in self.expr(*args, **kwargs)]))
+        return self.sym_return_type(all(e() for e in self.expr(*args, **kwargs)))
 
 
 class Try(Expression):
@@ -117,14 +122,14 @@ class Lambda(Expression):
 
 
 class Choice(Expression):
-    def __init__(self, cases: List[str], default: Optional[str] = None, **kwargs):
+    def __init__(self, cases: list[str], default: str | None = None, **kwargs):
         super().__init__(**kwargs)
-        self.cases: List[str] = cases
-        self.default: Optional[str] = default
+        self.cases: list[str] = cases
+        self.default: str | None = default
 
     def forward(self, sym: Symbol, *args, **kwargs) -> Symbol:
         sym = self._to_symbol(sym)
-        return sym.choice(cases=self.cases, default=self.default, *args, **kwargs)
+        return sym.choice(*args, cases=self.cases, default=self.default, **kwargs)
 
 
 class Output(Expression):
@@ -137,13 +142,13 @@ class Output(Expression):
     def forward(self, *args, **kwargs) -> Expression:
         kwargs['verbose'] = self.verbose
         kwargs['handler'] = self.handler
-        return self.output(expr=self.expr, *args, **kwargs)
+        return self.output(*args, expr=self.expr, **kwargs)
 
 
 class Sequence(TrackerTraceable):
-    def __init__(self, *expressions: List[Expression], **kwargs):
+    def __init__(self, *expressions: list[Expression], **kwargs):
         super().__init__(**kwargs)
-        self.expressions: List[Expression] = expressions
+        self.expressions: list[Expression] = expressions
 
     def forward(self, *args, **kwargs) -> Symbol:
         sym = self.expressions[0](*args, **kwargs)
@@ -159,11 +164,11 @@ class Sequence(TrackerTraceable):
 
 
 class Parallel(Expression):
-    def __init__(self, *expr: List[Expression | Callable], sequential: bool = False, **kwargs):
+    def __init__(self, *expr: list[Expression | Callable], sequential: bool = False, **kwargs):
         super().__init__(**kwargs)
         self.sequential: bool       = sequential
-        self.expr: List[Expression] = expr
-        self.results: List[Symbol]  = []
+        self.expr: list[Expression] = expr
+        self.results: list[Symbol]  = []
 
     def forward(self, *args, **kwargs) -> Symbol:
         # run in sequence
@@ -180,11 +185,11 @@ class Parallel(Expression):
 
 #@TODO: BinPacker(format="...") -> ensure that data packages form a "bin" that's consistent (e.g. never break a sentence in the middle)
 class Stream(Expression):
-    def __init__(self, expr: Optional[Expression] = None, retrieval: Optional[str] = None, **kwargs):
+    def __init__(self, expr: Expression | None = None, retrieval: str | None = None, **kwargs):
         super().__init__(**kwargs)
         self.char_token_ratio:    float = 0.6
-        self.expr: Optional[Expression] = expr
-        self.retrieval:   Optional[str] = retrieval
+        self.expr: Expression | None = expr
+        self.retrieval:   str | None = retrieval
         self._trace:               bool = False
         self._previous_frame            = None
 
@@ -194,19 +199,23 @@ class Stream(Expression):
         if self._trace:
             local_vars = self._previous_frame.f_locals
             vals = []
-            for key, var in local_vars.items():
+            for _key, var in local_vars.items():
                 if isinstance(var, TrackerTraceable):
                     vals.append(var)
 
             if len(vals) == 1:
                 self.expr = vals[0]
             else:
-                raise ValueError(f"This component does either not inherit from TrackerTraceable or has an invalid number of component declarations: {len(vals)}! Only one component that inherits from TrackerTraceable is allowed in the with stream clause.")
+                UserMessage(
+                    "This component does either not inherit from TrackerTraceable or has an invalid number of component "
+                    f"declarations: {len(vals)}! Only one component that inherits from TrackerTraceable is allowed in the "
+                    "with stream clause.",
+                    raise_with=ValueError,
+                )
 
         res = sym.stream(expr=self.expr,
                          char_token_ratio=self.char_token_ratio,
                          **kwargs)
-
         if self.retrieval is not None:
             res = list(res)
             if self.retrieval == 'all':
@@ -215,9 +224,8 @@ class Stream(Expression):
                 res = sorted(res, key=lambda x: len(x), reverse=True)
                 return res[0]
             if self.retrieval == 'contains':
-                res = [r for r in res if self.expr in r]
-                return res
-            raise ValueError(f"Invalid retrieval method: {self.retrieval}")
+                return [r for r in res if self.expr in r]
+            UserMessage(f"Invalid retrieval method: {self.retrieval}", raise_with=ValueError)
 
         return res
 
@@ -231,10 +239,12 @@ class Stream(Expression):
 
 
 class Trace(Expression):
-    def __init__(self, expr: Optional[Expression] = None, engines=['all'], **kwargs):
+    def __init__(self, expr: Expression | None = None, engines=None, **kwargs):
+        if engines is None:
+            engines = ['all']
         super().__init__(**kwargs)
         self.expr: Expression = expr
-        self.engines: List[str] = engines
+        self.engines: list[str] = engines
 
     def forward(self, *args, **kwargs) -> Expression:
         Expression.command(verbose=True, engines=self.engines)
@@ -252,23 +262,26 @@ class Trace(Expression):
         Expression.command(verbose=False, engines=self.engines)
         if self.expr is not None:
             return self.expr.__exit__(type, value, traceback)
+        return None
 
 
 class Analyze(Expression):
-    def __init__(self, exception: Exception, query: Optional[str] = None, **kwargs):
+    def __init__(self, exception: Exception, query: str | None = None, **kwargs):
         super().__init__(**kwargs)
         self.exception: Expression = exception
-        self.query: Optional[str] = query
+        self.query: str | None = query
 
     def forward(self, sym: Symbol, *args, **kwargs) -> Symbol:
-        return sym.analyze(exception=self.exception, query=self.query, *args, **kwargs)
+        return sym.analyze(*args, exception=self.exception, query=self.query, **kwargs)
 
 
 class Log(Expression):
-    def __init__(self, expr: Optional[Expression] = None, engines=['all'], **kwargs):
+    def __init__(self, expr: Expression | None = None, engines=None, **kwargs):
+        if engines is None:
+            engines = ['all']
         super().__init__(**kwargs)
         self.expr: Expression = expr
-        self.engines: List[str] = engines
+        self.engines: list[str] = engines
 
     def forward(self, *args, **kwargs) -> Expression:
         Expression.command(logging=True, engines=self.engines)
@@ -286,6 +299,7 @@ class Log(Expression):
         Expression.command(logging=False, engines=self.engines)
         if self.expr is not None:
             return self.expr.__exit__(type, value, traceback)
+        return None
 
 
 class Template(Expression):
@@ -331,10 +345,10 @@ class Metric(Expression):
         self.normalize  = normalize
         self.eps        = eps
 
-    def forward(self, sym: Symbol, **kwargs) -> Symbol:
+    def forward(self, sym: Symbol, **_kwargs) -> Symbol:
         sym = self._to_symbol(sym)
-        assert sym.value_type == np.ndarray or sym.value_type == list, 'Metric can only be applied to numpy arrays or lists.'
-        if sym.value_type == list:
+        assert sym.value_type is np.ndarray or sym.value_type is list, 'Metric can only be applied to numpy arrays or lists.'
+        if sym.value_type is list:
             sym._value = np.array(sym.value)
         # compute normalization between 0 and 1
         if self.normalize:
@@ -343,17 +357,19 @@ class Metric(Expression):
             elif len(sym.value.shape) == 2:
                 pass
             else:
-                raise ValueError(f'Invalid shape: {sym.value.shape}')
+                UserMessage(f'Invalid shape: {sym.value.shape}', raise_with=ValueError)
             # normalize between 0 and 1 and sum to 1
             sym._value = np.exp(sym.value) / (np.exp(sym.value).sum() + self.eps)
         return sym
 
 
 class Style(Expression):
-    def __init__(self, description: str, libraries: List[str] = [], **kwargs):
+    def __init__(self, description: str, libraries: list[str] | None = None, **kwargs):
+        if libraries is None:
+            libraries = []
         super().__init__(**kwargs)
         self.description: str = description
-        self.libraries: List[str] = libraries
+        self.libraries: list[str] = libraries
 
     def forward(self, sym: Symbol, **kwargs) -> Symbol:
         sym = self._to_symbol(sym)
@@ -365,7 +381,7 @@ class Query(TrackerTraceable):
         super().__init__(**kwargs)
         self.prompt: str = prompt
 
-    def forward(self, sym: Symbol, context: Symbol = None, *args, **kwargs) -> Symbol:
+    def forward(self, sym: Symbol, context: Symbol = None, *_args, **kwargs) -> Symbol:
         sym = self._to_symbol(sym)
         return sym.query(prompt=self.prompt, context=context, **kwargs)
 
@@ -474,11 +490,11 @@ class ExcludeFilter(Expression):
 class FileWriter(Expression):
     def __init__(self, path: str, **kwargs):
         super().__init__(**kwargs)
-        self.path = path
+        self.path = Path(path)
 
-    def forward(self, sym: Symbol, **kwargs) -> Symbol:
+    def forward(self, sym: Symbol, **_kwargs) -> Symbol:
         sym = self._to_symbol(sym)
-        with open(self.path, 'w') as f:
+        with self.path.open('w') as f:
             f.write(str(sym))
 
 
@@ -493,12 +509,10 @@ class FileReader(Expression):
         assert len(_splits) == 1 or len(_splits) == 2, 'Invalid file link format.'
         _tmp     = Path(_tmp)
         # check if file exists and is a file
-        if os.path.exists(_tmp) and os.path.isfile(_tmp):
-            return True
-        return False
+        return _tmp.is_file()
 
     @staticmethod
-    def get_files(folder_path: str, max_depth: int = 1) -> List[str]:
+    def get_files(folder_path: str, max_depth: int = 1) -> list[str]:
         accepted_formats = ['.pdf', '.md', '.txt']
 
         folder = Path(folder_path)
@@ -512,7 +526,7 @@ class FileReader(Expression):
         return files
 
     @staticmethod
-    def extract_files(cmds: str) -> Optional[List[str]]:
+    def extract_files(cmds: str) -> list[str] | None:
         # Use the updated regular expression to match quoted and non-quoted paths
         pattern = r'''(?:"((?:\\.|[^"\\])*)"|'((?:\\.|[^'\\])*)'|`((?:\\.|[^`\\])*)`|((?:\\ |[^ ])+))'''
         # Use the regular expression to split and handle quoted and non-quoted paths
@@ -551,16 +565,16 @@ class FileReader(Expression):
         return Path(path).expanduser().resolve().as_posix()
 
     @staticmethod
-    def integrity_check(files: List[str]) -> List[str]:
+    def integrity_check(files: list[str]) -> list[str]:
         not_skipped = []
         for file in tqdm(files):
             if FileReader.exists(file):
                 not_skipped.append(file)
             else:
-                CustomUserWarning(f'Skipping file: {file}')
+                UserMessage(f'Skipping file: {file}')
         return not_skipped
 
-    def forward(self, files: Union[str, List[str]], **kwargs) -> Expression:
+    def forward(self, files: str | list[str], **kwargs) -> Expression:
         if isinstance(files, str):
             # Convert to list for uniform processing; more easily downstream
             files = [files]
@@ -586,15 +600,17 @@ class FileQuery(Expression):
 
 class Function(TrackerTraceable):
     def __init__(self, prompt: str       = '',
-                 examples: Optional[str] = [],
-                 pre_processors: Optional[List[PreProcessor]]   = None,
-                 post_processors: Optional[List[PostProcessor]] = None,
-                 default: Optional[object]       = None,
-                 constraints: List[Callable]     = [],
-                 return_type: Optional[Type]     = str,
-                 sym_return_type: Optional[Type] = Symbol,
-                 origin_type: Optional[Type]     = Expression,
+                 examples: str | None = [],
+                 pre_processors: list[PreProcessor] | None   = None,
+                 post_processors: list[PostProcessor] | None = None,
+                 default: object | None       = None,
+                 constraints: list[Callable] | None     = None,
+                 return_type: type | None     = str,
+                 sym_return_type: type | None = Symbol,
+                 origin_type: type | None     = Expression,
                  *args, **kwargs):
+        if constraints is None:
+            constraints = []
         super().__init__(**kwargs)
         chars       = ascii_lowercase + ascii_uppercase
         self.name   = 'func_' + ''.join(sample(chars, 15))
@@ -629,13 +645,16 @@ class Function(TrackerTraceable):
         if 'fn' in kwargs:
             self.prompt = kwargs['fn']
             del kwargs['fn']
-        @core.few_shot(prompt=self.prompt,
-                  examples=self.examples,
-                  pre_processors=self.pre_processors,
-                  post_processors=self.post_processors,
-                  constraints=self.constraints,
-                  default=self.default,
-                  *self.args, **self.kwargs)
+        @core.few_shot(
+            *self.args,
+            prompt=self.prompt,
+            examples=self.examples,
+            pre_processors=self.pre_processors,
+            post_processors=self.post_processors,
+            constraints=self.constraints,
+            default=self.default,
+            **self.kwargs
+        )
         def _func(_, *args, **kwargs) -> self.return_type:
             pass
         _type = type(self.name, (self.origin_type, ), {
@@ -658,11 +677,11 @@ class PrepareData(Function):
             assert argument.prop.context is not None
             instruct = argument.prop.prompt
             context  = argument.prop.context
-            return """{
-    'context': '%s',
-    'instruction': '%s',
+            return f"""{{
+    'context': '{context}',
+    'instruction': '{instruct}',
     'result': 'TODO: Replace this with the expected result.'
-}""" % (context, instruct)
+}}"""
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -707,7 +726,7 @@ class ExpressionBuilder(Function):
         super().__init__('Generate the code following the instructions:', **kwargs)
         self.processors = ProcessorPipeline([StripPostProcessor(), CodeExtractPostProcessor()])
 
-    def forward(self, instruct, *args, **kwargs):
+    def forward(self, instruct, *_args, **_kwargs):
         result = super().forward(instruct)
         return self.processors(str(result), None)
 
@@ -768,21 +787,21 @@ class JsonParser(Expression):
 
 
 class SimilarityClassification(Expression):
-    def __init__(self, classes: List[str], metric: str = 'cosine', in_memory: bool = False, **kwargs):
+    def __init__(self, classes: list[str], metric: str = 'cosine', in_memory: bool = False, **kwargs):
         super().__init__(**kwargs)
         self.classes   = classes
         self.metric    = metric
         self.in_memory = in_memory
 
         if self.in_memory:
-            CustomUserWarning(f'Caching mode is enabled! It is your responsability to empty the .cache folder if you did changes to the classes. The cache is located at {HOME_PATH}/cache')
+            UserMessage(f'Caching mode is enabled! It is your responsability to empty the .cache folder if you did changes to the classes. The cache is located at {HOME_PATH}/cache')
 
     def forward(self, x: Symbol) -> Symbol:
         x            = self._to_symbol(x)
         usr_embed    = x.embed()
         embeddings   = self._dynamic_cache()
         similarities = [usr_embed.similarity(emb, metric=self.metric) for emb in embeddings]
-        similarities = sorted(zip(self.classes, similarities), key=lambda x: x[1], reverse=True)
+        similarities = sorted(zip(self.classes, similarities, strict=False), key=lambda x: x[1], reverse=True)
 
         return Symbol(similarities[0][0])
 
@@ -790,9 +809,7 @@ class SimilarityClassification(Expression):
         @core_ext.cache(in_memory=self.in_memory)
         def embed_classes(self):
             opts = map(Symbol, self.classes)
-            embeddings = [opt.embed() for opt in opts]
-
-            return embeddings
+            return [opt.embed() for opt in opts]
 
         return embed_classes(self)
 
@@ -820,19 +837,14 @@ class Indexer(Expression):
     @staticmethod
     def replace_special_chars(index: str):
         # replace special characters that are not for path
-        index = str(index)
-        index = index.replace('-', '')
-        index = index.replace('_', '')
-        index = index.replace(' ', '')
-        index = index.lower()
-        return index
+        return str(index).replace('-', '').replace('_', '').replace(' ', '').lower()
 
     def __init__(
             self,
             index_name: str = DEFAULT,
             top_k: int = 8,
             batch_size: int = 20,
-            formatter: Callable = ParagraphFormatter(),
+            formatter: Callable = _DEFAULT_PARAGRAPH_FORMATTER,
             auto_add=False,
             raw_result: bool = False,
             new_dim: int = 1536,
@@ -861,15 +873,15 @@ class Indexer(Expression):
     def register(self):
         # check if index already exists in indices.txt and append if not
         change = False
-        with open(self.path, 'r') as f:
+        with self.path.open() as f:
             indices = f.read().split('\n')
             # filter out empty strings
             indices = [i for i in indices if i]
-            if self.index_name not in indices:
+        if self.index_name not in indices:
                 indices.append(self.index_name)
                 change = True
         if change:
-            with open(self.path, 'w') as f:
+            with self.path.open('w') as f:
                 f.write('\n'.join(indices))
 
     def exists(self) -> bool:
@@ -877,15 +889,16 @@ class Indexer(Expression):
         path = HOME_PATH / 'indices.txt'
         if not path.exists():
             return False
-        with open(path, 'r') as f:
+        with path.open() as f:
             indices = f.read().split('\n')
             if self.index_name in indices:
                 return True
+        return False
 
     def forward(
             self,
-            data: Optional[Symbol] = None,
-            raw_result: bool = False,
+            data: Symbol | None = None,
+            _raw_result: bool = False,
         ) -> Symbol:
         that = self
         if data is not None:
@@ -898,15 +911,14 @@ class Indexer(Expression):
             # we save the index
             that.config(None, save=True, index_name=that.index_name, index_dims=that.new_dim)
 
-        def _func(query, *args, **kwargs) -> Union[Symbol, 'VectorDBResult']:
+        def _func(query, *_args, **kwargs) -> Union[Symbol, 'VectorDBResult']:
             raw_result = kwargs.get('raw_result') or that.raw_result
             query_emb = Symbol(query).embed(new_dim=that.new_dim).value
             res = that.get(query_emb, index_name=that.index_name, index_top_k=that.top_k, ori_query=query, index_dims=that.new_dim, **kwargs)
             that.retrieval = res
             if raw_result:
                 return res
-            rsp = Symbol(res).query(prompt='From the retrieved data, select the most relevant information.', context=query)
-            return rsp
+            return Symbol(res).query(prompt='From the retrieved data, select the most relevant information.', context=query)
         return _func
 
 
@@ -917,8 +929,8 @@ class PrimitiveDisabler(Expression):
         self._original_primitives = defaultdict(list)
 
     def __enter__(self):
-        # Avoid circular imports; import locally
-        from .symbol import Symbol
+        # Import Symbol lazily so components does not clash with symbol during load.
+        from .symbol import Symbol # noqa
 
         frame = inspect.currentframe()
         f_locals = frame.f_back.f_locals
@@ -934,7 +946,7 @@ class PrimitiveDisabler(Expression):
             for func in self._primitives:
                 if hasattr(sym, func):
                     self._original_primitives[sym_name].append((func, getattr(sym, func)))
-                    setattr(sym, func, lambda *args, **kwargs: None)
+                    setattr(sym, func, lambda *_args, **_kwargs: None)
 
     def _enable_primitives(self):
         for sym_name, sym in self._symbols.items():
@@ -968,7 +980,7 @@ class FunctionWithUsage(Function):
 
     def print_verbose(self, msg):
         if self.verbose:
-            print(msg)
+            UserMessage(msg)
 
     def _format_usage(self, prompt_tokens, completion_tokens, total_tokens):
         return Box(
@@ -1020,12 +1032,11 @@ class FunctionWithUsage(Function):
             self.completion_tokens += completion_tokens
             self.total_tokens += total_tokens
         else:
-            if self.missing_usage_exception and not "preview" in kwargs:
-                raise Exception("Missing usage in metadata of neursymbolic engine")
-            else:
-                prompt_tokens = 0
-                completion_tokens = 0
-                total_tokens = 0
+            if self.missing_usage_exception and "preview" not in kwargs:
+                UserMessage("Missing usage in metadata of neursymbolic engine", raise_with=Exception)
+            prompt_tokens = 0
+            completion_tokens = 0
+            total_tokens = 0
 
         return res, self._format_usage(prompt_tokens, completion_tokens, total_tokens)
 
@@ -1041,7 +1052,7 @@ class SelfPrompt(Expression):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-    def forward(self, existing_prompt: Dict[str, str], **kwargs) -> Dict[str, str]:
+    def forward(self, existing_prompt: dict[str, str], **kwargs) -> dict[str, str]:
         """
         Generate new system and user prompts based on the existing prompt.
 
@@ -1094,14 +1105,13 @@ class MetadataTracker(Expression):
         value = value or self.metadata
         if isinstance(value, dict):
             return '{\n\t' + ', \n\t'.join(f'"{k}": {self.__str__(v)}' for k,v in value.items()) + '\n}'
-        elif isinstance(value, list):
+        if isinstance(value, list):
             return '[' + ', '.join(self.__str__(item) for item in value) + ']'
-        elif isinstance(value, str):
+        if isinstance(value, str):
             return f'"{value}"'
-        else:
-            return f"\n\t    {value}"
+        return f"\n\t    {value}"
 
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls, *_args, **_kwargs):
         cls._lock = getattr(cls, '_lock', Lock())
         with cls._lock:
             instance = super().__new__(cls)
@@ -1122,25 +1132,26 @@ class MetadataTracker(Expression):
 
     def _trace_calls(self, frame, event, arg):
         if not self._trace:
-            return
+            return None
 
-        if event == 'return' and frame.f_code.co_name == 'forward':
-            # Check if this is an engine forward call
-            if ('self' in frame.f_locals
-                and
-                isinstance(frame.f_locals['self'], Engine)):
-                _, metadata = arg  # arg contains return value on 'return' event
-                engine_name = frame.f_locals['self'].__class__.__name__
-                model_name = frame.f_locals['self'].model
-                self._metadata[(self._metadata_id, engine_name, model_name)] = metadata
-                self._metadata_id += 1
+        if (
+            event == 'return'
+            and frame.f_code.co_name == 'forward'
+            and 'self' in frame.f_locals
+            and isinstance(frame.f_locals['self'], Engine)
+        ):
+            _, metadata = arg  # arg contains return value on 'return' event
+            engine_name = frame.f_locals['self'].__class__.__name__
+            model_name = frame.f_locals['self'].model
+            self._metadata[(self._metadata_id, engine_name, model_name)] = metadata
+            self._metadata_id += 1
 
         return self._trace_calls
 
     def _accumulate_completion_token_details(self):
         """Parses the return object and accumulates completion token details per token type"""
         if not self._metadata:
-            CustomUserWarning("No metadata available to generate usage details.")
+            UserMessage("No metadata available to generate usage details.")
             return {}
 
         token_details = defaultdict(lambda: defaultdict(lambda: defaultdict(int)))
@@ -1182,15 +1193,59 @@ class MetadataTracker(Expression):
                     logger.warning(f"Tracking {engine_name} is not supported.")
                     continue
             except Exception as e:
-                CustomUserWarning(f"Failed to parse metadata for {engine_name}: {e}", raise_with=AttributeError)
+                UserMessage(f"Failed to parse metadata for {engine_name}: {e}", raise_with=AttributeError)
 
         # Convert to normal dict
         return {**token_details}
 
+    def _can_accumulate_engine(self, engine_name: str) -> bool:
+        supported_engines = ("GPTXChatEngine", "GPTXReasoningEngine", "GPTXSearchEngine")
+        return engine_name in supported_engines
+
+    def _accumulate_time_field(self, accumulated: dict, metadata: dict) -> None:
+        if 'time' in metadata and 'time' in accumulated:
+            accumulated['time'] += metadata['time']
+
+    def _accumulate_usage_fields(self, accumulated: dict, metadata: dict) -> None:
+        if 'raw_output' not in metadata or 'raw_output' not in accumulated:
+            return
+
+        metadata_raw_output = metadata['raw_output']
+        accumulated_raw_output = accumulated['raw_output']
+        if not hasattr(metadata_raw_output, 'usage') or not hasattr(accumulated_raw_output, 'usage'):
+            return
+
+        current_usage = metadata_raw_output.usage
+        accumulated_usage = accumulated_raw_output.usage
+
+        for attr in ['completion_tokens', 'prompt_tokens', 'total_tokens']:
+            if hasattr(current_usage, attr) and hasattr(accumulated_usage, attr):
+                setattr(
+                    accumulated_usage,
+                    attr,
+                    getattr(accumulated_usage, attr) + getattr(current_usage, attr),
+                )
+
+        for detail_attr in ['completion_tokens_details', 'prompt_tokens_details']:
+            if not hasattr(current_usage, detail_attr) or not hasattr(accumulated_usage, detail_attr):
+                continue
+
+            current_details = getattr(current_usage, detail_attr)
+            accumulated_details = getattr(accumulated_usage, detail_attr)
+
+            for attr in dir(current_details):
+                if attr.startswith('_') or not hasattr(accumulated_details, attr):
+                    continue
+
+                current_val = getattr(current_details, attr)
+                accumulated_val = getattr(accumulated_details, attr)
+                if isinstance(current_val, (int, float)) and isinstance(accumulated_val, (int, float)):
+                    setattr(accumulated_details, attr, accumulated_val + current_val)
+
     def _accumulate_metadata(self):
         """Accumulates metadata across all tracked engine calls."""
         if not self._metadata:
-            CustomUserWarning("No metadata available to generate usage details.")
+            UserMessage("No metadata available to generate usage details.")
             return {}
 
         # Use first entry as base
@@ -1199,39 +1254,12 @@ class MetadataTracker(Expression):
 
         # Skipz first entry
         for (_, engine_name), metadata in list(self._metadata.items())[1:]:
-            if engine_name not in ("GPTXChatEngine", "GPTXReasoningEngine", "GPTXSearchEngine"):
+            if not self._can_accumulate_engine(engine_name):
                 logger.warning(f"Metadata accumulation for {engine_name} is not supported. Try `.usage` instead for now.")
                 continue
 
-            # Accumulate time if it exists
-            if 'time' in metadata and 'time' in accumulated:
-                accumulated['time'] += metadata['time']
-
-            # Handle usage stats accumulation
-            if 'raw_output' in metadata and 'raw_output' in accumulated:
-                if hasattr(metadata['raw_output'], 'usage') and hasattr(accumulated['raw_output'], 'usage'):
-                    current_usage = metadata['raw_output'].usage
-                    accumulated_usage = accumulated['raw_output'].usage
-
-                    # Accumulate token counts
-                    for attr in ['completion_tokens', 'prompt_tokens', 'total_tokens']:
-                        if hasattr(current_usage, attr) and hasattr(accumulated_usage, attr):
-                            setattr(accumulated_usage, attr,
-                                    getattr(accumulated_usage, attr) + getattr(current_usage, attr))
-
-                    # Handle nested token details if they exist
-                    for detail_attr in ['completion_tokens_details', 'prompt_tokens_details']:
-                        if hasattr(current_usage, detail_attr) and hasattr(accumulated_usage, detail_attr):
-                            current_details = getattr(current_usage, detail_attr)
-                            accumulated_details = getattr(accumulated_usage, detail_attr)
-
-                            # Accumulate all numeric attributes in the details
-                            for attr in dir(current_details):
-                                if not attr.startswith('_') and hasattr(accumulated_details, attr):
-                                    current_val = getattr(current_details, attr)
-                                    accumulated_val = getattr(accumulated_details, attr)
-                                    if isinstance(current_val, (int, float)) and isinstance(accumulated_val, (int, float)):
-                                        setattr(accumulated_details, attr, accumulated_val + current_val)
+            self._accumulate_time_field(accumulated, metadata)
+            self._accumulate_usage_fields(accumulated, metadata)
 
         return accumulated
 
@@ -1250,7 +1278,7 @@ class MetadataTracker(Expression):
 
 class DynamicEngine(Expression):
     """Context manager for dynamically switching neurosymbolic engine models."""
-    def __init__(self, model: str, api_key: str, debug: bool = False, **kwargs):
+    def __init__(self, model: str, api_key: str, _debug: bool = False, **_kwargs):
         super().__init__()
         self.model = model
         self.api_key = api_key
@@ -1259,7 +1287,7 @@ class DynamicEngine(Expression):
         self.engine_instance = None
         self._ctx_token = None
 
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls, *_args, **_kwargs):
         cls._lock = getattr(cls, '_lock', Lock())
         with cls._lock:
             instance = super().__new__(cls)
@@ -1293,11 +1321,12 @@ class DynamicEngine(Expression):
 
     def _create_engine_instance(self):
         """Create an engine instance based on the model name."""
-        from .backend.engines.neurosymbolic import ENGINE_MAPPING
+        # Deferred to avoid components <-> neurosymbolic engine circular imports.
+        from .backend.engines.neurosymbolic import ENGINE_MAPPING # noqa
         try:
             engine_class = ENGINE_MAPPING.get(self.model)
             if engine_class is None:
-                raise ValueError(f"Unsupported model '{self.model}'")
+                UserMessage(f"Unsupported model '{self.model}'", raise_with=ValueError)
             return engine_class(api_key=self.api_key, model=self.model)
         except Exception as e:
-            raise ValueError(f"Failed to create engine for model '{self.model}': {str(e)}")
+            UserMessage(f"Failed to create engine for model '{self.model}': {e!s}", raise_with=ValueError)