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

symai/misc/console.py

@@ -1,24 +1,24 @@
-import re
-import pygments
 import logging
+import re
 
 #@TODO: refactor to use rich instead of prompt_toolkit
 from html import escape as escape_html
-from pygments.lexers.python import PythonLexer
-from pygments.lexers.javascript import JavascriptLexer
+from typing import ClassVar
+
+import pygments
+from prompt_toolkit import HTML, print_formatted_text
+from prompt_toolkit.formatted_text import PygmentsTokens
 from pygments.lexers.c_cpp import CppLexer
+from pygments.lexers.javascript import JavascriptLexer
+from pygments.lexers.python import PythonLexer
 from pygments.lexers.shell import BashLexer
-from prompt_toolkit import print_formatted_text
-from prompt_toolkit import HTML
-from prompt_toolkit.formatted_text import PygmentsTokens
-
 
 logger = logging.getLogger(__name__)
-print = print_formatted_text
+print = print_formatted_text # noqa
 
 
-class ConsoleStyle(object):
-    style_types = {
+class ConsoleStyle:
+    style_types: ClassVar[dict[str, str]] = {
         'alert':     'orange',
         'error':     'ansired',
         'warn':      'ansiyellow',
@@ -55,32 +55,42 @@ class ConsoleStyle(object):
         style = self.style_types.get(self.style_type, self.style_types['default'])
 
         if style == self.style_types['code']:
-            # Split the message using ``` as delimiter
-            segments = re.split(r'(```)', message)
-            is_code = False
-            for segment in segments:
-                if segment == '```':
-                    is_code = not is_code
-                    continue
-                if is_code:
-                    # Determine lexer
-                    if 'python' in segment.lower():
-                        lexer = PythonLexer()
-                    elif 'javascript' in segment.lower() or 'typescript' in segment.lower():
-                        lexer = JavascriptLexer()
-                    elif 'c++' in segment.lower():
-                        lexer = CppLexer()
-                    else:
-                        lexer = BashLexer()
-                    # Print highlighted code
-                    tokens = list(pygments.lex("```" + segment + "```", lexer))
-                    print(PygmentsTokens(tokens), end='')
-                else:
-                    # Print the segment normally
-                    print(segment, end='\n')
-        elif style == self.style_types['default']:
+            self._print_code_message(message)
+            return
+        if style == self.style_types['default']:
             print(message)
-        elif style == self.style_types['custom']:
-            print(HTML(f'<style fg="{self.color}">{content_for_html}</style>'))
-        else:
-            print(HTML(f'<style fg="{style}">{content_for_html}</style>'))
+            return
+        if style == self.style_types['custom']:
+            self._print_html(self.color, content_for_html)
+            return
+        self._print_html(style, content_for_html)
+
+    def _print_code_message(self, message: str) -> None:
+        segments = re.split(r'(```)', message)
+        is_code_segment = False
+        for segment in segments:
+            if segment == '```':
+                is_code_segment = not is_code_segment
+                continue
+            if is_code_segment:
+                self._print_code_segment(segment)
+                continue
+            print(segment, end='\n')
+
+    def _print_code_segment(self, segment: str) -> None:
+        lexer = self._select_lexer(segment)
+        tokens = list(pygments.lex("```" + segment + "```", lexer))
+        print(PygmentsTokens(tokens), end='')
+
+    def _select_lexer(self, segment: str):
+        lowered_segment = segment.lower()
+        if 'python' in lowered_segment:
+            return PythonLexer()
+        if 'javascript' in lowered_segment or 'typescript' in lowered_segment:
+            return JavascriptLexer()
+        if 'c++' in lowered_segment:
+            return CppLexer()
+        return BashLexer()
+
+    def _print_html(self, color: str, content: str) -> None:
+        print(HTML(f'<style fg="{color}">{content}</style>'))

symai/misc/loader.py

@@ -1,12 +1,14 @@
 import sys
 from itertools import cycle
 from shutil import get_terminal_size
-from threading import Thread, Event
+from threading import Event, Thread
 from time import sleep
+
 from prompt_toolkit import print_formatted_text
+
 from .console import ConsoleStyle
 
-print = print_formatted_text
+print = print_formatted_text # noqa
 
 class Loader:
     def __init__(self, desc="Loading...", end="\n", timeout=0.1):
@@ -56,4 +58,4 @@ class Loader:
         self.stop()
 
     def print(self, message):
-        print(message, style='ansigray')
+        print(message, style="ansigray")

symai/models/__init__.py

@@ -1,2 +1,18 @@
-from .base import Const, LengthConstraint, LLMDataModel, CustomConstraint, build_dynamic_llm_datamodel
+from .base import (
+    Const,
+    CustomConstraint,
+    LengthConstraint,
+    LLMDataModel,
+    build_dynamic_llm_datamodel,
+)
 from .errors import ExceptionWithUsage, TypeValidationError
+
+__all__ = [
+    "Const",
+    "CustomConstraint",
+    "ExceptionWithUsage",
+    "LLMDataModel",
+    "LengthConstraint",
+    "TypeValidationError",
+    "build_dynamic_llm_datamodel",
+]