Refactor/Adapt ClassRegistry from kajson, remove sandbox manager (Pipelex#103)

- Adapt to new ClassRegistry from kajson, used with dependency injection
- Remove unused sandbox component, adapt logging, pretty printing, error reporting

Author: lchoquel

docs/pages/configuration/config-practical/logging-config.md

@@ -79,7 +79,7 @@ silenced_problem_ids = ["azure_openai_no_stream_options"]
 ### Poor Loggers
 
 ```toml
-generic_poor_logger = "#sandbox"
+generic_poor_logger = "#poor-log"
 poor_loggers = [
     "kajson.decoder.sandbox",
     "kajson.encoder.sandbox",
@@ -150,7 +150,7 @@ is_caller_info_enabled = true
 caller_info_template = "file_line_func"
 
 silenced_problem_ids = []
-generic_poor_logger = "#sandbox"
+generic_poor_logger = "#poor-log"
 poor_loggers = []
 
 [pipelex.log_config.package_log_levels]

pipelex/cogt/content_generation/assignment_models.py

@@ -1,6 +1,5 @@
 from typing import Any, Dict, Optional, Type
 
-from kajson.class_registry import class_registry
 from pydantic import BaseModel
 from typing_extensions import override
 
@@ -16,6 +15,7 @@
 from pipelex.cogt.ocr.ocr_handle import OcrHandle
 from pipelex.cogt.ocr.ocr_input import OcrInput
 from pipelex.cogt.ocr.ocr_job_components import OcrJobConfig, OcrJobParams
+from pipelex.hub import get_class_registry
 from pipelex.pipeline.job_metadata import JobMetadata
 from pipelex.tools.templating.jinja2_template_category import Jinja2TemplateCategory
 from pipelex.tools.templating.templating_models import PromptingStyle
@@ -80,7 +80,7 @@ class ObjectAssignment(BaseModel):
 
     def __init__(self, **kwargs: Any):
         super().__init__(**kwargs)
-        if not class_registry.has_class(name=self.object_class_name):
+        if not get_class_registry().has_class(name=self.object_class_name):
             error_msg = f"Could not create ObjectAssignment for class '{self.object_class_name}' because it is not in the class registry."
             raise LLMAssignmentError(error_msg)
 
@@ -90,7 +90,7 @@ def make_for_class(
         llm_assignment: LLMAssignment,
     ) -> "ObjectAssignment":
         object_class_name = object_class.__name__
-        class_registry.register_class(
+        get_class_registry().register_class(
             class_type=object_class,
             name=object_class_name,
             should_warn_if_already_registered=False,

pipelex/cogt/content_generation/llm_generate.py

@@ -1,12 +1,11 @@
 from typing import List
 
-from kajson.class_registry import class_registry
 from pydantic import BaseModel
 
 from pipelex import log
 from pipelex.cogt.content_generation.assignment_models import LLMAssignment, ObjectAssignment
 from pipelex.cogt.llm.llm_job_factory import LLMJobFactory
-from pipelex.hub import get_llm_worker
+from pipelex.hub import get_class_registry, get_llm_worker
 
 
 async def llm_gen_text(llm_assignment: LLMAssignment) -> str:
@@ -31,7 +30,7 @@ async def llm_gen_object(object_assignment: ObjectAssignment) -> BaseModel:
         llm_job_params=llm_assignment.llm_job_params,
     )
     content_class_name = object_assignment.object_class_name
-    content_class = class_registry.get_required_base_model(name=content_class_name)
+    content_class = get_class_registry().get_required_base_model(name=content_class_name)
     generated_object: BaseModel = await llm_worker.gen_object(
         llm_job=llm_job,
         schema=content_class,
@@ -49,7 +48,7 @@ async def llm_gen_object_list(object_assignment: ObjectAssignment) -> List[BaseM
         llm_job_params=llm_assignment.llm_job_params,
     )
     item_class_name = object_assignment.object_class_name
-    item_class = class_registry.get_required_class(name=item_class_name)
+    item_class = get_class_registry().get_required_class(name=item_class_name)
 
     class ListSchema(BaseModel):
         items: List[item_class]  # type: ignore

pipelex/core/concept.py

@@ -1,7 +1,7 @@
 import re
 from typing import List, Tuple
 
-from kajson.class_registry import class_registry
+from kajson.kajson_manager import KajsonManager
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 from typing_extensions import Self
 
@@ -91,10 +91,12 @@ def validate_structure_class_name(cls, value: str) -> str:
 
     @classmethod
     def is_valid_structure_class(cls, structure_class_name: str) -> bool:
-        if class_registry.has_subclass(name=structure_class_name, base_class=StuffContent):
+        # we get_class_registry directly from KajsonManager instead of pipelex hub to avoid circular import
+        if KajsonManager.get_class_registry().has_subclass(name=structure_class_name, base_class=StuffContent):
             return True
         else:
-            if class_registry.has_class(name=structure_class_name):
+            # we get_class_registry directly from KajsonManager instead of pipelex hub to avoid circular import
+            if KajsonManager.get_class_registry().has_class(name=structure_class_name):
                 log.warning(f"Concept class '{structure_class_name}' is registered but it's not a subclass of StuffContent")
             return False
 

pipelex/core/concept_factory.py

@@ -1,7 +1,6 @@
 from inspect import getsource
 from typing import Any, Dict, List, Optional, Type, Union
 
-from kajson.class_registry import class_registry
 from pydantic import BaseModel, ConfigDict, Field, ValidationError
 
 from pipelex.core.concept import Concept
@@ -10,6 +9,7 @@
 from pipelex.core.domain import SpecialDomain
 from pipelex.core.stuff_content import TextContent
 from pipelex.exceptions import ConceptFactoryError, StructureClassError
+from pipelex.hub import get_class_registry
 
 
 class ConceptBlueprint(BaseModel):
@@ -180,8 +180,8 @@ def list_native_concepts(cls) -> List[Concept]:
 
     @classmethod
     def get_concept_class_source_code(cls, concept_name: str, base_class: Type[Any]) -> str:
-        if not class_registry.has_class(concept_name):
+        if not get_class_registry().has_class(concept_name):
             raise RuntimeError(f"Class '{concept_name}' not found in registry")
 
-        cls = class_registry.get_required_subclass(name=concept_name, base_class=base_class)
+        cls = get_class_registry().get_required_subclass(name=concept_name, base_class=base_class)
         return getsource(cls)

pipelex/core/stuff_factory.py

@@ -1,7 +1,6 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import shortuuid
-from kajson.class_registry import class_registry
 from pydantic import BaseModel, Field
 
 from pipelex.config import get_config
@@ -11,7 +10,7 @@
 from pipelex.core.stuff import Stuff, StuffCreationRecord
 from pipelex.core.stuff_content import StuffContent, StuffContentInitableFromStr
 from pipelex.exceptions import ConceptError, PipelexError
-from pipelex.hub import get_required_concept
+from pipelex.hub import get_class_registry, get_required_concept
 
 
 class StuffFactoryError(PipelexError):
@@ -111,7 +110,7 @@ def make_from_str(
             raise StuffFactoryError(f"Concept '{concept_str}' does not contain a domain, could not make stuff '{stuff_ref}'")
         the_concept = get_required_concept(concept_code=concept_code)
         the_subclass_name = the_concept.structure_class_name
-        the_subclass = class_registry.get_class(name=the_subclass_name) or eval(the_subclass_name)
+        the_subclass = get_class_registry().get_class(name=the_subclass_name) or eval(the_subclass_name)
         if not issubclass(the_subclass, StuffContentInitableFromStr):
             raise StuffFactoryError(f"Concept '{concept_code}', subclass '{the_subclass}' is not InitableFromStr")
         stuff_content: StuffContent = the_subclass.make_from_str(str_value)
@@ -155,7 +154,7 @@ def combine_stuffs(cls, concept_code: str, stuff_contents: Dict[str, StuffConten
         """
         the_concept = get_required_concept(concept_code=concept_code)
         the_subclass_name = the_concept.structure_class_name
-        the_subclass = class_registry.get_required_subclass(name=the_subclass_name, base_class=StuffContent)
+        the_subclass = get_class_registry().get_required_subclass(name=the_subclass_name, base_class=StuffContent)
         the_stuff_content = the_subclass.model_validate(obj=stuff_contents)
         return cls.make_stuff(
             concept_str=concept_code,

pipelex/hub.py

@@ -1,6 +1,8 @@
 from collections import defaultdict
 from typing import ClassVar, Dict, List, Optional, Type
 
+from kajson.class_registry_abstract import ClassRegistryAbstract
+
 from pipelex import log
 from pipelex.cogt.content_generation.content_generator_protocol import ContentGeneratorProtocol
 from pipelex.cogt.imgg.imgg_worker_abstract import ImggWorkerAbstract
@@ -43,6 +45,7 @@ def __init__(self):
         self._config: Optional[ConfigRoot] = None
         self._secrets_provider: Optional[SecretsProviderAbstract] = None
         self._template_provider: Optional[TemplateProviderAbstract] = None
+        self._class_registry: Optional[ClassRegistryAbstract] = None
         # cogt
         self._llm_models_provider: Optional[LLMModelProviderAbstract] = None
         self._llm_deck_provider: Optional[LLMDeckAbstract] = None
@@ -112,6 +115,9 @@ def set_secrets_provider(self, secrets_provider: SecretsProviderAbstract):
     def set_template_provider(self, template_provider: TemplateProviderAbstract):
         self._template_provider = template_provider
 
+    def set_class_registry(self, class_registry: ClassRegistryAbstract):
+        self._class_registry = class_registry
+
     # cogt
 
     def set_llm_models_provider(self, llm_models_provider: LLMModelProviderAbstract):
@@ -187,6 +193,11 @@ def get_required_template_provider(self) -> TemplateProviderAbstract:
             raise RuntimeError("Template provider is not set. You must initialize Pipelex first.")
         return self._template_provider
 
+    def get_required_class_registry(self) -> ClassRegistryAbstract:
+        if self._class_registry is None:
+            raise RuntimeError("ClassRegistry is not initialized")
+        return self._class_registry
+
     # cogt
 
     def get_required_llm_models_provider(self) -> LLMModelProviderAbstract:
@@ -294,6 +305,10 @@ def get_template(template_name: str) -> str:
     return get_template_provider().get_template(template_name=template_name)
 
 
+def get_class_registry() -> ClassRegistryAbstract:
+    return get_pipelex_hub().get_required_class_registry()
+
+
 # cogt
 
 

pipelex/libraries/library_manager.py

@@ -2,8 +2,8 @@
 from pathlib import Path
 from typing import Any, ClassVar, Dict, List, Optional, Type
 
-from kajson.class_registry import class_registry
 from kajson.exceptions import ClassRegistryInheritanceError, ClassRegistryNotFoundError
+from kajson.kajson_manager import KajsonManager
 from pydantic import ValidationError
 
 from pipelex import log
@@ -74,13 +74,13 @@ def teardown(self) -> None:
     def load_libraries(self):
         log.debug("LibraryManager loading separate libraries")
 
-        class_registry.register_classes_in_folder(
+        KajsonManager.get_class_registry().register_classes_in_folder(
             folder_path=LibraryConfig.loaded_pipelines_path,
         )
         library_paths = [LibraryConfig.loaded_pipelines_path]
         if runtime_manager.is_unit_testing:
             log.debug("Registering test pipeline structures for unit testing")
-            class_registry.register_classes_in_folder(
+            KajsonManager.get_class_registry().register_classes_in_folder(
                 folder_path=LibraryConfig.test_pipelines_path,
             )
             library_paths += [LibraryConfig.test_pipelines_path]
@@ -283,7 +283,7 @@ def make_pipe_from_details_dict(
         # the factory class name for that specific type of Pipe is the pipe class name with "Factory" suffix
         factory_class_name = f"{pipe_class_name}Factory"
         try:
-            pipe_factory: Type[PipeSpecificFactoryProtocol[Any, Any]] = class_registry.get_required_subclass(
+            pipe_factory: Type[PipeSpecificFactoryProtocol[Any, Any]] = KajsonManager.get_class_registry().get_required_subclass(
                 name=factory_class_name,
                 base_class=PipeSpecificFactoryProtocol,
             )

pipelex/pipe_operators/pipe_llm.py

@@ -1,6 +1,5 @@
 from typing import List, Optional, Set, Type, cast
 
-from kajson.class_registry import class_registry
 from pydantic import model_validator
 from typing_extensions import Self, override
 
@@ -34,6 +33,7 @@
     StaticValidationErrorType,
 )
 from pipelex.hub import (
+    get_class_registry,
     get_concept_provider,
     get_content_generator,
     get_llm_deck,
@@ -118,7 +118,7 @@ def _validate_inputs(self):
                         # let's check at least that the input is a structured concept
                         input_concept = concept_provider.get_required_concept(concept_code=concept_code_of_declared_input)
                         input_concept_class_name = input_concept.structure_class_name
-                        input_concept_class = class_registry.get_required_subclass(name=input_concept_class_name, base_class=StuffContent)
+                        input_concept_class = get_class_registry().get_required_subclass(name=input_concept_class_name, base_class=StuffContent)
                         if issubclass(input_concept_class, StructuredContent):
                             continue
                     explanation = "The input provided for LLM Vision must be an image or a concept that refines image"
@@ -393,7 +393,7 @@ async def _llm_gen_object_stuff_content(
         llm_prompt_2_factory: Optional[LLMPromptFactoryAbstract],
         content_generator: ContentGeneratorProtocol,
     ) -> StuffContent:
-        content_class: Type[StuffContent] = class_registry.get_required_subclass(name=output_class_name, base_class=StuffContent)
+        content_class: Type[StuffContent] = get_class_registry().get_required_subclass(name=output_class_name, base_class=StuffContent)
         task_desc: str
         the_content: StuffContent
 

pipelex/pipe_operators/pipe_llm_prompt.py

@@ -1,6 +1,5 @@
 from typing import ClassVar, List, Optional, Set, cast
 
-from kajson.class_registry import class_registry
 from pydantic import model_validator
 from typing_extensions import Self, override
 
@@ -22,7 +21,7 @@
     PipeRunParamsError,
     WorkingMemoryVariableError,
 )
-from pipelex.hub import get_template
+from pipelex.hub import get_class_registry, get_template
 from pipelex.pipe_operators.pipe_jinja2 import PipeJinja2, PipeJinja2Output
 from pipelex.pipe_operators.pipe_operator import PipeOperator
 from pipelex.pipeline.job_metadata import JobCategory, JobMetadata
@@ -219,7 +218,7 @@ async def _run_operator_pipe(
     @staticmethod
     def get_output_structure_prompt(output_concept: str) -> str:
         class_name = Concept.extract_concept_name_from_str(concept_str=output_concept)
-        output_class = class_registry.get_class(class_name)
+        output_class = get_class_registry().get_class(class_name)
         if not output_class:
             return ""