From 82bead23929a23998f4d7504f31d83bf253927bf Mon Sep 17 00:00:00 2001 From: Marco Perini Date: Mon, 11 Nov 2024 12:25:50 +0100 Subject: [PATCH] perf: add custom logging level --- brickllm/__init__.py | 2 ++ brickllm/logger.py | 38 +++++++++++++++++++++++++++++ brickllm/nodes/generation_local.py | 3 ++- brickllm/nodes/get_elem_children.py | 5 +++- brickllm/nodes/get_elements.py | 3 ++- brickllm/nodes/get_relationships.py | 3 ++- brickllm/nodes/get_sensors.py | 3 ++- brickllm/nodes/schema_to_ttl.py | 3 ++- brickllm/nodes/validate_schema.py | 3 ++- 9 files changed, 56 insertions(+), 7 deletions(-) create mode 100644 brickllm/logger.py diff --git a/brickllm/__init__.py b/brickllm/__init__.py index 1a4e70a..748d6b9 100644 --- a/brickllm/__init__.py +++ b/brickllm/__init__.py @@ -1,4 +1,5 @@ from .configs import GraphConfig +from .logger import custom_logger from .schemas import ( ElemListSchema, RelationshipsSchema, @@ -15,4 +16,5 @@ "State", "StateLocal", "GraphConfig", + "custom_logger", ] diff --git a/brickllm/logger.py b/brickllm/logger.py new file mode 100644 index 0000000..dc93b86 --- /dev/null +++ b/brickllm/logger.py @@ -0,0 +1,38 @@ +import logging + +# Create custom log level +EURAC_LEVEL = 25 +logging.addLevelName(EURAC_LEVEL, "EURAC") + + +def eurac(self, message, *args, **kwargs): + """ + Log with custom EURAC level + """ + if self.isEnabledFor(EURAC_LEVEL): + self._log(EURAC_LEVEL, message, args, **kwargs) + + +# Add eurac method to Logger class +logging.Logger.eurac = eurac + + +# Create and configure logger +def get_logger(name="BrickLLM"): + logger = logging.getLogger(name) + + # Create handler if none exists + if not logger.handlers: + handler = logging.StreamHandler() + formatter = logging.Formatter( + "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + ) + handler.setFormatter(formatter) + logger.addHandler(handler) + + logger.setLevel(EURAC_LEVEL) + return logger + + +# Create default logger instance +custom_logger = get_logger() diff --git a/brickllm/nodes/generation_local.py b/brickllm/nodes/generation_local.py index 3c3d089..7a8a8e4 100644 --- a/brickllm/nodes/generation_local.py +++ b/brickllm/nodes/generation_local.py @@ -2,6 +2,7 @@ from .. import StateLocal from ..helpers import prompt_template_local +from ..logger import custom_logger from ..utils import extract_rdf_graph @@ -17,7 +18,7 @@ def generation_local(state: StateLocal, config: Dict[str, Any]) -> Dict[str, Any dict: A dictionary containing the output generated. """ - print("---One shot generation with local LLM Node---") + custom_logger.eurac("🤖 Starting one shot generation with local LLM") instructions = state["instructions"] user_prompt = state["user_prompt"] diff --git a/brickllm/nodes/get_elem_children.py b/brickllm/nodes/get_elem_children.py index 0c8c70b..bc383b4 100644 --- a/brickllm/nodes/get_elem_children.py +++ b/brickllm/nodes/get_elem_children.py @@ -4,6 +4,7 @@ from .. import ElemListSchema, State from ..helpers import get_elem_children_instructions +from ..logger import custom_logger from ..utils import create_hierarchical_dict, filter_elements, get_children_hierarchy @@ -18,7 +19,9 @@ def get_elem_children(state: State, config: Dict[str, Any]) -> Dict[str, Any]: Returns: dict: A dictionary containing the hierarchical structure of identified elements. """ - print("---Get Elem Children Node---") + custom_logger.eurac( + "📊 Getting children for each BrickSchema category in the element list" + ) user_prompt = state["user_prompt"] categories = state["elem_list"] diff --git a/brickllm/nodes/get_elements.py b/brickllm/nodes/get_elements.py index beb32d0..fc33b30 100644 --- a/brickllm/nodes/get_elements.py +++ b/brickllm/nodes/get_elements.py @@ -4,6 +4,7 @@ from .. import ElemListSchema, State from ..helpers import get_elem_instructions +from ..logger import custom_logger from ..utils import get_brick_definition, get_hierarchical_info @@ -19,7 +20,7 @@ def get_elements(state: State, config: Dict[str, Any]) -> Dict[str, Any]: Returns: dict: A dictionary containing the list of identified elements. """ - print("---Get Elements Node---") + custom_logger.eurac("🔍 Getting elements from user prompt") user_prompt = state["user_prompt"] diff --git a/brickllm/nodes/get_relationships.py b/brickllm/nodes/get_relationships.py index 81eb08b..70fb157 100644 --- a/brickllm/nodes/get_relationships.py +++ b/brickllm/nodes/get_relationships.py @@ -6,6 +6,7 @@ from .. import RelationshipsSchema, State from ..helpers import get_relationships_instructions +from ..logger import custom_logger from ..utils import build_hierarchy, find_sensor_paths @@ -20,7 +21,7 @@ def get_relationships(state: State, config: Dict[str, Any]) -> Dict[str, Any]: Returns: dict: A dictionary containing the grouped sensor paths. """ - print("---Get Relationships Node---") + custom_logger.eurac("🔗 Getting relationships between building components") user_prompt = state["user_prompt"] building_structure = state["elem_hierarchy"] diff --git a/brickllm/nodes/get_sensors.py b/brickllm/nodes/get_sensors.py index 0a386a4..ac6c207 100644 --- a/brickllm/nodes/get_sensors.py +++ b/brickllm/nodes/get_sensors.py @@ -1,6 +1,7 @@ from typing import Any, Dict from .. import State +from ..logger import custom_logger def get_sensors(state: State) -> Dict[str, Any]: @@ -13,7 +14,7 @@ def get_sensors(state: State) -> Dict[str, Any]: Returns: dict: A dictionary containing sensor UUIDs mapped to their locations. """ - print("---Get Sensor Node---") + custom_logger.eurac("📡 Getting sensors information") uuid_dict = { "Building#1>Floor#1>Office#1>Room#1": [ diff --git a/brickllm/nodes/schema_to_ttl.py b/brickllm/nodes/schema_to_ttl.py index 741a3e4..f9486b1 100644 --- a/brickllm/nodes/schema_to_ttl.py +++ b/brickllm/nodes/schema_to_ttl.py @@ -5,6 +5,7 @@ from .. import State, TTLSchema from ..helpers import schema_to_ttl_instructions, ttl_example +from ..logger import custom_logger def schema_to_ttl(state: State, config: Dict[str, Any]) -> Dict[str, Any]: @@ -18,7 +19,7 @@ def schema_to_ttl(state: State, config: Dict[str, Any]) -> Dict[str, Any]: Returns: dict: A dictionary containing the generated TTL output. """ - print("---Schema To TTL Node---") + custom_logger.eurac("📝 Generating TTL from schema") user_prompt = state["user_prompt"] sensors_dict = state["sensors_dict"] diff --git a/brickllm/nodes/validate_schema.py b/brickllm/nodes/validate_schema.py index 913f1d8..76d91d0 100644 --- a/brickllm/nodes/validate_schema.py +++ b/brickllm/nodes/validate_schema.py @@ -1,5 +1,6 @@ from typing import Any, Dict +from ..logger import custom_logger from ..utils import validate_ttl @@ -13,7 +14,7 @@ def validate_schema(state) -> Dict[str, Any]: Returns: dict: A dictionary containing the validation status and report. """ - print("---Validate Schema Node---") + custom_logger.eurac("✅ Validating TTL schema") ttl_output = state.get("ttl_output", None) max_iter = state.get("validation_max_iter", 2)