dg-dataagent 0.1.0__py3-none-any.whl

dataagent/actions/environment/__init__.py

@@ -0,0 +1,21 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+__all__ = [
+    "Env",
+    "CompoundEnv",
+    "from_config",
+]
+
+from dataagent.actions.environment.compound_env import CompoundEnv
+from dataagent.actions.environment.env import Env
+from dataagent.actions.environment.env_config import from_config

dataagent/actions/environment/compound_env.py

@@ -0,0 +1,89 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+from dataagent.actions.environment.env import Env
+
+
+class CompoundEnv(Env):
+    """
+    Combines multiple Env instances into a single unified environment.
+
+    The CompoundEnv merges all tools from the component environments,
+    and delegates init/close operations to all sub-environments.
+
+    Example usage:
+        >>> class MathEnv(Env):
+        ...     @Env.tool
+        ...     def add(self, a, b):
+        ...         return a + b
+        ...
+        >>> class StringEnv(Env):
+        ...     @Env.tool
+        ...     def concat(self, a, b):
+        ...         return a + b
+        ...
+        >>> math_env = MathEnv()
+        >>> string_env = StringEnv()
+        >>> compound = CompoundEnv([math_env, string_env])
+        >>> compound.tools['add'](5, 3)
+        8
+        >>> compound.tools['concat']('Hello', ' World')
+        'Hello World'
+
+    Notes:
+        - Tools from later environments override tools with the same name
+          from earlier environments
+        - Calling init() on CompoundEnv initializes all sub-environments
+        - Calling close() on CompoundEnv closes all sub-environments
+    """
+
+    def __init__(self, envs: list[Env]):
+        """
+        Initialize CompoundEnv with a list of Env instances.
+
+        Args:
+            envs: List of Env instances to combine
+        """
+        # Set envs BEFORE calling super().__init__() because _register_tools() needs it
+        self.envs = envs
+        super().__init__()
+
+    def __repr__(self) -> str:
+        """String representation showing number of environments and tools"""
+        return f"CompoundEnv(envs={len(self.envs)}, tools={len(self.tools)})"
+
+    def init(self):
+        """Initialize all component environments"""
+        for env in self.envs:
+            env.init()
+
+    def close(self):
+        """Close all component environments"""
+        for env in self.envs:
+            env.close()
+
+    def get_description(self) -> str:
+        """Combine all descriptions"""
+        return "\n".join([env.get_description() for env in self.envs])
+
+    def _register_tools(self):
+        """
+        Override parent's _register_tools to merge tools from all sub-environments.
+
+        Tools from later environments will override tools with the same name
+        from earlier environments.
+        """
+        super()._register_tools()
+
+        for env in self.envs:
+            env_tools = env.get_tools()
+            self.tools.update(env_tools)

dataagent/actions/environment/env.py

@@ -0,0 +1,52 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+from collections.abc import Callable
+
+
+class Env:
+    def __init__(self):
+        self.tools: dict[str, Callable] = {}
+        self.init()
+        self._register_tools()
+
+    def __del__(self):
+        self.close()
+
+    @staticmethod
+    def tool(func: Callable) -> Callable:
+        """Decorator to mark a method as a tool"""
+        func._is_tool = True
+        return func
+
+    def init(self):
+        """Initialize the environment, override in subclass"""
+        pass
+
+    def close(self):
+        """Close the environment, override in subclass"""
+        pass
+
+    def get_tools(self) -> dict[str, Callable]:
+        """Return the tools bound the the environment"""
+        return self.tools
+
+    def get_description(self) -> str:
+        """Return the description of the environment"""
+        return ""
+
+    def _register_tools(self):
+        """Register all methods marked with the @tool decorator"""
+        for name in dir(self):
+            attr = getattr(self, name)
+            if callable(attr) and hasattr(attr, "_is_tool"):
+                self.tools[name] = attr

dataagent/actions/environment/env_config.py

@@ -0,0 +1,101 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+from typing import Any
+
+from dataagent.actions.environment.compound_env import CompoundEnv
+from dataagent.actions.environment.env import Env
+from dataagent.utils.import_utils import import_class
+
+
+def from_config(config: list[dict[str, Any]] | dict[str, Any]) -> Env:
+    """
+    Create an Env instance from configuration.
+
+    Args:
+        config: Either a single environment configuration dict or a list of dicts.
+                Each dict must contain a "module" field specifying the Env subclass
+                to instantiate (e.g., "mypackage.envs.MyEnv").
+                Other fields in the dict are passed as keyword arguments to __init__.
+
+    Returns:
+        An Env instance. If config is a list, returns a CompoundEnv.
+
+    Raises:
+        ValueError: If configuration is invalid
+        ImportError: If the specified module cannot be imported
+        AttributeError: If the specified class doesn't exist in the module
+
+    Examples:
+        Single environment:
+        >>> config = {
+        ...     "module": "mypackage.envs.MathEnv",
+        ...     "precision": 10
+        ... }
+        >>> env = from_config(config)
+
+        Multiple environments (creates CompoundEnv):
+        >>> config = [
+        ...     {"module": "mypackage.envs.MathEnv"},
+        ...     {"module": "mypackage.envs.StringEnv", "encoding": "utf-8"}
+        ... ]
+        >>> compound_env = from_config(config)
+    """
+    if isinstance(config, list):
+        # Create multiple environments and combine them
+        if not config:
+            raise ValueError("Configuration list cannot be empty")
+
+        envs = [_create_single_env(env_config) for env_config in config]
+        return CompoundEnv(envs)
+    if isinstance(config, dict):
+        # Create a single environment
+        return _create_single_env(config)
+    raise ValueError(f"Configuration must be a dict or list of dicts, got {type(config)}")
+
+
+def _create_single_env(config: dict[str, Any]) -> Env:
+    """
+    Create a single Env instance from a configuration dict.
+
+    Args:
+        config: Configuration dict with "module" field and optional init parameters
+
+    Returns:
+        An Env instance
+
+    Raises:
+        ValueError: If "module" field is missing
+        ImportError: If module cannot be imported
+        AttributeError: If class doesn't exist
+        TypeError: If the class is not a subclass of Env
+    """
+    if "module" not in config:
+        raise ValueError("Configuration dict must contain a 'module' field")
+
+    class_path = config["module"]
+
+    # Load the class using the utility function
+    env_class = import_class(class_path)
+
+    # Verify it's an Env subclass
+    if not issubclass(env_class, Env):
+        raise TypeError(f"Class {class_path} must be a subclass of Env, got {env_class}")
+
+    # Extract init parameters (everything except "module")
+    init_params = {k: v for k, v in config.items() if k != "module"}
+
+    # Instantiate the environment
+    try:
+        return env_class(**init_params)
+    except TypeError as e:
+        raise TypeError(f"Failed to instantiate {class_path} with parameters {init_params}: {e}") from e

dataagent/actions/gym/__init__.py

@@ -0,0 +1,32 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+__all__ = ["ArithmeticEnv", "ICBCEnv", "ManufacureEnv", "OntologyEnv", "SQLiteEnv"]
+
+import importlib
+
+_ENV_REGISTRY = {
+    "ArithmeticEnv": "dataagent.actions.gym.arithmetic.arithmetic",
+    "ICBCEnv": "dataagent.actions.gym.icbc.icbc_env",
+    "ManufacureEnv": "dataagent.actions.gym.manufacturing.manufacturing_env",
+    "OntologyEnv": "dataagent.actions.gym.ontology_env",
+    "SQLiteEnv": "dataagent.actions.gym.nl2sql.sqlite_env",
+}
+
+
+def __getattr__(name: str):
+    if name not in _ENV_REGISTRY:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    module = importlib.import_module(_ENV_REGISTRY[name])
+    return getattr(module, name)

dataagent/actions/gym/arithmetic/arithmetic.py

@@ -0,0 +1,124 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+from dataagent.actions.environment.env import Env
+
+
+class ArithmeticEnv(Env):
+    """
+    Environment providing basic arithmetic operations.
+
+    This environment provides four fundamental mathematical operations:
+    addition, subtraction, multiplication, and division.
+
+    Example:
+        >>> env = ArithmeticEnv()
+        >>> env.tools['add'](5, 3)
+        8
+        >>> env.tools['multiply'](4, 7)
+        28
+        >>> env.tools['divide'](10, 2)
+        5.0
+    """
+
+    def __init__(self):
+        """
+        Initialize the arithmetic environment.
+
+        Args:
+            precision: Number of decimal places to round results to (default: 2)
+        """
+        super().__init__()
+
+    @Env.tool
+    def add(self, a: float, b: float) -> float:
+        """
+        Add two numbers.
+
+        Args:
+            a: First number
+            b: Second number
+
+        Returns:
+            The sum of a and b
+
+        Example:
+            >>> add(5, 3)
+            8.0
+            >>> add(1.5, 2.7)
+            4.2
+        """
+        return a + b
+
+    @Env.tool
+    def subtract(self, a: float, b: float) -> float:
+        """
+        Subtract b from a.
+
+        Args:
+            a: Number to subtract from
+            b: Number to subtract
+
+        Returns:
+            The difference (a - b)
+
+        Example:
+            >>> subtract(10, 3)
+            7.0
+            >>> subtract(5.5, 2.3)
+            3.2
+        """
+        return a - b
+
+    @Env.tool
+    def multiply(self, a: float, b: float) -> float:
+        """
+        Multiply two numbers.
+
+        Args:
+            a: First number
+            b: Second number
+
+        Returns:
+            The product of a and b
+
+        Example:
+            >>> multiply(4, 7)
+            28.0
+            >>> multiply(2.5, 3.2)
+            8.0
+        """
+        return a * b
+
+    @Env.tool
+    def divide(self, a: float, b: float) -> float:
+        """
+        Divide a by b.
+
+        Args:
+            a: Dividend (number to be divided)
+            b: Divisor (number to divide by)
+
+        Returns:
+            The quotient (a / b), or NaN if b is zero
+
+        Example:
+            >>> divide(10, 2)
+            5.0
+            >>> divide(7, 3)
+            2.33
+            >>> divide(5, 0)
+            nan
+        """
+        if b == 0:
+            return float("nan")
+        return a / b