Spaces:
Sleeping
Sleeping
| import toml | |
| import yaml | |
| import asyncio | |
| import concurrent.futures | |
| import json | |
| import os | |
| from concurrent.futures import ThreadPoolExecutor | |
| from datetime import datetime | |
| from typing import Any, Dict, List, Optional, Callable | |
| import psutil | |
| try: | |
| import gzip | |
| except ImportError as error: | |
| print(f"Error importing gzip: {error}") | |
| # from pydantic import BaseModel | |
| class BaseStructure: | |
| """Base structure. | |
| Attributes: | |
| name (Optional[str]): _description_ | |
| description (Optional[str]): _description_ | |
| save_metadata (bool): _description_ | |
| save_artifact_path (Optional[str]): _description_ | |
| save_metadata_path (Optional[str]): _description_ | |
| save_error_path (Optional[str]): _description_ | |
| Methods: | |
| run: _description_ | |
| save_to_file: _description_ | |
| load_from_file: _description_ | |
| save_metadata: _description_ | |
| load_metadata: _description_ | |
| log_error: _description_ | |
| save_artifact: _description_ | |
| load_artifact: _description_ | |
| log_event: _description_ | |
| run_async: _description_ | |
| save_metadata_async: _description_ | |
| load_metadata_async: _description_ | |
| log_error_async: _description_ | |
| save_artifact_async: _description_ | |
| load_artifact_async: _description_ | |
| log_event_async: _description_ | |
| asave_to_file: _description_ | |
| aload_from_file: _description_ | |
| run_in_thread: _description_ | |
| save_metadata_in_thread: _description_ | |
| run_concurrent: _description_ | |
| compress_data: _description_ | |
| decompres_data: _description_ | |
| run_batched: _description_ | |
| load_config: _description_ | |
| backup_data: _description_ | |
| monitor_resources: _description_ | |
| run_with_resources: _description_ | |
| run_with_resources_batched: _description_ | |
| Examples: | |
| >>> base_structure = BaseStructure() | |
| >>> base_structure | |
| BaseStructure(name=None, description=None, save_metadata=True, save_artifact_path='./artifacts', save_metadata_path='./metadata', save_error_path='./errors') | |
| """ | |
| def __init__( | |
| self, | |
| name: Optional[str] = None, | |
| description: Optional[str] = None, | |
| save_metadata_on: bool = True, | |
| save_artifact_path: Optional[str] = "./artifacts", | |
| save_metadata_path: Optional[str] = "./metadata", | |
| save_error_path: Optional[str] = "./errors", | |
| workspace_dir: Optional[str] = "./workspace", | |
| ): | |
| super().__init__() | |
| self.name = name | |
| self.description = description | |
| self.save_metadata_on = save_metadata_on | |
| self.save_artifact_path = save_artifact_path | |
| self.save_metadata_path = save_metadata_path | |
| self.save_error_path = save_error_path | |
| self.workspace_dir = workspace_dir | |
| def run(self, *args, **kwargs): | |
| """Run the structure.""" | |
| def save_to_file(self, data: Any, file_path: str): | |
| """Save data to file. | |
| Args: | |
| data (Any): _description_ | |
| file_path (str): _description_ | |
| """ | |
| with open(file_path, "w") as file: | |
| json.dump(data, file) | |
| def load_from_file(self, file_path: str) -> Any: | |
| """Load data from file. | |
| Args: | |
| file_path (str): _description_ | |
| Returns: | |
| Any: _description_ | |
| """ | |
| with open(file_path) as file: | |
| return json.load(file) | |
| def save_metadata(self, metadata: Dict[str, Any]): | |
| """Save metadata to file. | |
| Args: | |
| metadata (Dict[str, Any]): _description_ | |
| """ | |
| if self.save_metadata: | |
| file_path = os.path.join( | |
| self.save_metadata_path, f"{self.name}_metadata.json" | |
| ) | |
| self.save_to_file(metadata, file_path) | |
| def load_metadata(self) -> Dict[str, Any]: | |
| """Load metadata from file. | |
| Returns: | |
| Dict[str, Any]: _description_ | |
| """ | |
| file_path = os.path.join( | |
| self.save_metadata_path, f"{self.name}_metadata.json" | |
| ) | |
| return self.load_from_file(file_path) | |
| def log_error(self, error_message: str): | |
| """Log error to file. | |
| Args: | |
| error_message (str): _description_ | |
| """ | |
| file_path = os.path.join( | |
| self.save_error_path, f"{self.name}_errors.log" | |
| ) | |
| with open(file_path, "a") as file: | |
| file.write(f"{error_message}\n") | |
| def save_artifact(self, artifact: Any, artifact_name: str): | |
| """Save artifact to file. | |
| Args: | |
| artifact (Any): _description_ | |
| artifact_name (str): _description_ | |
| """ | |
| file_path = os.path.join( | |
| self.save_artifact_path, f"{artifact_name}.json" | |
| ) | |
| self.save_to_file(artifact, file_path) | |
| def load_artifact(self, artifact_name: str) -> Any: | |
| """Load artifact from file. | |
| Args: | |
| artifact_name (str): _description_ | |
| Returns: | |
| Any: _description_ | |
| """ | |
| file_path = os.path.join( | |
| self.save_artifact_path, f"{artifact_name}.json" | |
| ) | |
| return self.load_from_file(file_path) | |
| def _current_timestamp(self): | |
| """Current timestamp. | |
| Returns: | |
| _type_: _description_ | |
| """ | |
| return datetime.now().strftime("%Y-%m-%d %H:%M:%S") | |
| def log_event( | |
| self, | |
| event: str, | |
| event_type: str = "INFO", | |
| ): | |
| """Log event to file. | |
| Args: | |
| event (str): _description_ | |
| event_type (str, optional): _description_. Defaults to "INFO". | |
| """ | |
| timestamp = self._current_timestamp() | |
| log_message = f"[{timestamp}] [{event_type}] {event}\n" | |
| file = os.path.join( | |
| self.save_metadata_path, f"{self.name}_events.log" | |
| ) | |
| with open(file, "a") as file: | |
| file.write(log_message) | |
| async def run_async(self, *args, **kwargs): | |
| """Run the structure asynchronously.""" | |
| loop = asyncio.get_event_loop() | |
| return await loop.run_in_executor( | |
| None, self.run, *args, **kwargs | |
| ) | |
| async def save_metadata_async(self, metadata: Dict[str, Any]): | |
| """Save metadata to file asynchronously. | |
| Args: | |
| metadata (Dict[str, Any]): _description_ | |
| """ | |
| loop = asyncio.get_event_loop() | |
| return await loop.run_in_executor( | |
| None, self.save_metadata, metadata | |
| ) | |
| async def load_metadata_async(self) -> Dict[str, Any]: | |
| """Load metadata from file asynchronously. | |
| Returns: | |
| Dict[str, Any]: _description_ | |
| """ | |
| loop = asyncio.get_event_loop() | |
| return await loop.run_in_executor(None, self.load_metadata) | |
| async def log_error_async(self, error_message: str): | |
| """Log error to file asynchronously. | |
| Args: | |
| error_message (str): _description_ | |
| """ | |
| loop = asyncio.get_event_loop() | |
| return await loop.run_in_executor( | |
| None, self.log_error, error_message | |
| ) | |
| async def save_artifact_async( | |
| self, artifact: Any, artifact_name: str | |
| ): | |
| """Save artifact to file asynchronously. | |
| Args: | |
| artifact (Any): _description_ | |
| artifact_name (str): _description_ | |
| """ | |
| loop = asyncio.get_event_loop() | |
| return await loop.run_in_executor( | |
| None, self.save_artifact, artifact, artifact_name | |
| ) | |
| async def load_artifact_async(self, artifact_name: str) -> Any: | |
| """Load artifact from file asynchronously. | |
| Args: | |
| artifact_name (str): _description_ | |
| Returns: | |
| Any: _description_ | |
| """ | |
| loop = asyncio.get_event_loop() | |
| return await loop.run_in_executor( | |
| None, self.load_artifact, artifact_name | |
| ) | |
| async def log_event_async( | |
| self, | |
| event: str, | |
| event_type: str = "INFO", | |
| ): | |
| """Log event to file asynchronously. | |
| Args: | |
| event (str): _description_ | |
| event_type (str, optional): _description_. Defaults to "INFO". | |
| """ | |
| loop = asyncio.get_event_loop() | |
| return await loop.run_in_executor( | |
| None, self.log_event, event, event_type | |
| ) | |
| async def asave_to_file( | |
| self, data: Any, file: str, *args, **kwargs | |
| ): | |
| """Save data to file asynchronously. | |
| Args: | |
| data (Any): _description_ | |
| file (str): _description_ | |
| """ | |
| await asyncio.to_thread( | |
| self.save_to_file, | |
| data, | |
| file, | |
| *args, | |
| ) | |
| async def aload_from_file( | |
| self, | |
| file: str, | |
| ) -> Any: | |
| """Async load data from file. | |
| Args: | |
| file (str): _description_ | |
| Returns: | |
| Any: _description_ | |
| """ | |
| return await asyncio.to_thread(self.load_from_file, file) | |
| def run_in_thread(self, *args, **kwargs): | |
| """Run the structure in a thread.""" | |
| with concurrent.futures.ThreadPoolExecutor() as executor: | |
| return executor.submit(self.run, *args, **kwargs) | |
| def save_metadata_in_thread(self, metadata: Dict[str, Any]): | |
| """Save metadata to file in a thread. | |
| Args: | |
| metadata (Dict[str, Any]): _description_ | |
| """ | |
| with concurrent.futures.ThreadPoolExecutor() as executor: | |
| return executor.submit(self.save_metadata, metadata) | |
| def run_concurrent(self, *args, **kwargs): | |
| """Run the structure concurrently.""" | |
| return asyncio.run(self.run_async(*args, **kwargs)) | |
| def compress_data( | |
| self, | |
| data: Any, | |
| ) -> bytes: | |
| """Compress data. | |
| Args: | |
| data (Any): _description_ | |
| Returns: | |
| bytes: _description_ | |
| """ | |
| return gzip.compress(json.dumps(data).encode()) | |
| def decompres_data(self, data: bytes) -> Any: | |
| """Decompress data. | |
| Args: | |
| data (bytes): _description_ | |
| Returns: | |
| Any: _description_ | |
| """ | |
| return json.loads(gzip.decompress(data).decode()) | |
| def run_batched( | |
| self, | |
| batched_data: List[Any], | |
| batch_size: int = 10, | |
| *args, | |
| **kwargs, | |
| ): | |
| """Run batched data. | |
| Args: | |
| batched_data (List[Any]): _description_ | |
| batch_size (int, optional): _description_. Defaults to 10. | |
| Returns: | |
| _type_: _description_ | |
| """ | |
| with ThreadPoolExecutor(max_workers=batch_size) as executor: | |
| futures = [ | |
| executor.submit(self.run, data) | |
| for data in batched_data | |
| ] | |
| return [future.result() for future in futures] | |
| def load_config( | |
| self, config: str = None, *args, **kwargs | |
| ) -> Dict[str, Any]: | |
| """Load config from file. | |
| Args: | |
| config (str, optional): _description_. Defaults to None. | |
| Returns: | |
| Dict[str, Any]: _description_ | |
| """ | |
| return self.load_from_file(config) | |
| def backup_data( | |
| self, data: Any, backup_path: str = None, *args, **kwargs | |
| ): | |
| """Backup data to file. | |
| Args: | |
| data (Any): _description_ | |
| backup_path (str, optional): _description_. Defaults to None. | |
| """ | |
| timestamp = self._current_timestamp() | |
| backup_file_path = f"{backup_path}/{timestamp}.json" | |
| self.save_to_file(data, backup_file_path) | |
| def monitor_resources(self): | |
| """Monitor resource usage.""" | |
| memory = psutil.virtual_memory().percent | |
| cpu_usage = psutil.cpu_percent(interval=1) | |
| self.log_event( | |
| f"Resource usage - Memory: {memory}%, CPU: {cpu_usage}%" | |
| ) | |
| def run_with_resources(self, *args, **kwargs): | |
| """Run the structure with resource monitoring.""" | |
| self.monitor_resources() | |
| return self.run(*args, **kwargs) | |
| def run_with_resources_batched( | |
| self, | |
| batched_data: List[Any], | |
| batch_size: int = 10, | |
| *args, | |
| **kwargs, | |
| ): | |
| """Run batched data with resource monitoring. | |
| Args: | |
| batched_data (List[Any]): _description_ | |
| batch_size (int, optional): _description_. Defaults to 10. | |
| Returns: | |
| _type_: _description_ | |
| """ | |
| self.monitor_resources() | |
| return self.run_batched( | |
| batched_data, batch_size, *args, **kwargs | |
| ) | |
| def _serialize_callable( | |
| self, attr_value: Callable | |
| ) -> Dict[str, Any]: | |
| """ | |
| Serializes callable attributes by extracting their name and docstring. | |
| Args: | |
| attr_value (Callable): The callable to serialize. | |
| Returns: | |
| Dict[str, Any]: Dictionary with name and docstring of the callable. | |
| """ | |
| return { | |
| "name": getattr( | |
| attr_value, "__name__", type(attr_value).__name__ | |
| ), | |
| "doc": getattr(attr_value, "__doc__", None), | |
| } | |
| def _serialize_attr(self, attr_name: str, attr_value: Any) -> Any: | |
| """ | |
| Serializes an individual attribute, handling non-serializable objects. | |
| Args: | |
| attr_name (str): The name of the attribute. | |
| attr_value (Any): The value of the attribute. | |
| Returns: | |
| Any: The serialized value of the attribute. | |
| """ | |
| try: | |
| if callable(attr_value): | |
| return self._serialize_callable(attr_value) | |
| elif hasattr(attr_value, "to_dict"): | |
| return ( | |
| attr_value.to_dict() | |
| ) # Recursive serialization for nested objects | |
| else: | |
| json.dumps( | |
| attr_value | |
| ) # Attempt to serialize to catch non-serializable objects | |
| return attr_value | |
| except (TypeError, ValueError): | |
| return f"<Non-serializable: {type(attr_value).__name__}>" | |
| def to_dict(self) -> Dict[str, Any]: | |
| """ | |
| Converts all attributes of the class, including callables, into a dictionary. | |
| Handles non-serializable attributes by converting them or skipping them. | |
| Returns: | |
| Dict[str, Any]: A dictionary representation of the class attributes. | |
| """ | |
| return { | |
| attr_name: self._serialize_attr(attr_name, attr_value) | |
| for attr_name, attr_value in self.__dict__.items() | |
| } | |
| def to_json(self, indent: int = 4, *args, **kwargs): | |
| return json.dumps( | |
| self.to_dict(), indent=indent, *args, **kwargs | |
| ) | |
| def to_yaml(self, indent: int = 4, *args, **kwargs): | |
| return yaml.dump( | |
| self.to_dict(), indent=indent, *args, **kwargs | |
| ) | |
| def to_toml(self, *args, **kwargs): | |
| return toml.dumps(self.to_dict(), *args, **kwargs) | |
| # def model_dump_json(self): | |
| # logger.info( | |
| # f"Saving {self.agent_name} model to JSON in the {self.workspace_dir} directory" | |
| # ) | |
| # create_file_in_folder( | |
| # self.workspace_dir, | |
| # f"{self.agent_name}.json", | |
| # str(self.to_json()), | |
| # ) | |
| # return ( | |
| # f"Model saved to {self.workspace_dir}/{self.agent_name}.json" | |
| # ) | |
| # def model_dump_yaml(self): | |
| # logger.info( | |
| # f"Saving {self.agent_name} model to YAML in the {self.workspace_dir} directory" | |
| # ) | |
| # create_file_in_folder( | |
| # self.workspace_dir, | |
| # f"{self.agent_name}.yaml", | |
| # self.to_yaml(), | |
| # ) | |
| # return ( | |
| # f"Model saved to {self.workspace_dir}/{self.agent_name}.yaml" | |
| # ) | |