synth-ai 0.2.4.dev6__py3-none-any.whl → 0.2.4.dev7__py3-none-any.whl

synth_ai/environments/environment/registry.py

@@ -27,19 +27,18 @@ Example:
     >>> env_instance = env_cls(task_config)
 """
 
-from typing import Type, Dict, List
-import logging
 import importlib.metadata
+import logging
 
 from synth_ai.environments.stateful.core import StatefulEnvironment
 
 logger = logging.getLogger(__name__)
 
 # Global registry for environment types
-ENV_REGISTRY: Dict[str, Type[StatefulEnvironment]] = {}
+ENV_REGISTRY: dict[str, type[StatefulEnvironment]] = {}
 
 
-def register_environment(name: str, cls: Type[StatefulEnvironment]) -> None:
+def register_environment(name: str, cls: type[StatefulEnvironment]) -> None:
     """
     Register an environment class under a unique name.
 
@@ -78,7 +77,7 @@ def register_environment(name: str, cls: Type[StatefulEnvironment]) -> None:
     ENV_REGISTRY[name] = cls
 
 
-def get_environment_cls(env_type: str) -> Type[StatefulEnvironment]:
+def get_environment_cls(env_type: str) -> type[StatefulEnvironment]:
     """
     Retrieve a registered environment class by name.
 
@@ -122,10 +121,10 @@ def get_environment_cls(env_type: str) -> Type[StatefulEnvironment]:
         available_types = list(ENV_REGISTRY.keys())
         raise ValueError(
             f"Unsupported environment type: '{env_type}'. Available types: {available_types}"
-        )
+        ) from None
 
 
-def list_supported_env_types() -> List[str]:
+def list_supported_env_types() -> list[str]:
     """
     List all registered environment type names.
 
@@ -165,36 +164,36 @@ def list_supported_env_types() -> List[str]:
         This function returns a copy of the environment names, so modifying
         the returned list will not affect the registry.
     """
-    return sorted(list(ENV_REGISTRY.keys()))
+    return sorted(ENV_REGISTRY.keys())
 
 
 def discover_entry_point_environments() -> None:
     """
     Discover and register environments from entry points.
-    
+
     This function scans for environments registered via setuptools entry points
     under the group 'synth_ai.environments'. This allows third-party packages
     to register environments by declaring them in their pyproject.toml:
-    
+
     [project.entry-points."synth_ai.environments"]
     my_env = "my_package.my_env:MyEnvironment"
     another_env = "my_package.other:AnotherEnv"
-    
+
     The function will automatically import and register all discovered environments.
     """
     try:
         entry_points = importlib.metadata.entry_points()
-        if hasattr(entry_points, 'select'):
+        if hasattr(entry_points, "select"):
             # Python 3.10+
-            env_entry_points = entry_points.select(group='synth_ai.environments')
+            env_entry_points = entry_points.select(group="synth_ai.environments")
         else:
             # Python 3.9 and below
-            env_entry_points = entry_points.get('synth_ai.environments', [])
-        
+            env_entry_points = entry_points.get("synth_ai.environments", [])
+
         for entry_point in env_entry_points:
             try:
                 env_cls = entry_point.load()
-                
+
                 # Validate that it's a StatefulEnvironment subclass
                 if not issubclass(env_cls, StatefulEnvironment):
                     logger.warning(
@@ -202,10 +201,10 @@ def discover_entry_point_environments() -> None:
                         f"is not a StatefulEnvironment subclass. Skipping."
                     )
                     continue
-                
+
                 register_environment(entry_point.name, env_cls)
                 logger.info(f"Registered environment '{entry_point.name}' from entry point")
-                
+
             except Exception as e:
                 logger.error(
                     f"Failed to load environment entry point '{entry_point.name}' "
@@ -218,11 +217,11 @@ def discover_entry_point_environments() -> None:
 def auto_discover_environments() -> None:
     """
     Automatically discover and register environments from multiple sources.
-    
+
     This function combines multiple discovery mechanisms:
     1. Entry points (setuptools plugins)
     2. Could be extended with more discovery methods in the future
-    
+
     This should be called once at application startup to populate the registry
     with all available environments.
     """

synth_ai/environments/environment/resources/sqlite.py

@@ -1,8 +1,7 @@
+import contextlib
 import sqlite3
 import threading
-import contextlib
 from pathlib import Path
-from typing import Optional
 
 
 class SQLiteManager:
@@ -36,7 +35,7 @@ class SQLiteManager:
                 self._conn.close()
                 self._conn = None
 
-    def reset(self, new_db_path: Optional[Path] = None):
+    def reset(self, new_db_path: Path | None = None):
         """Resets the connection, optionally switching to a new database path for a fresh session."""
         with self._lock:
             if self._conn:

synth_ai/environments/environment/rewards/core.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from typing import List, Any
+from typing import Any
 
 
 class RewardComponent(ABC):
@@ -17,7 +18,7 @@ class RewardComponent(ABC):
 class RewardStack:
     """Aggregates multiple RewardComponent instances to compute a total reward."""
 
-    def __init__(self, components: List[RewardComponent]):
+    def __init__(self, components: list[RewardComponent]):
         self.components = components
 
     async def step_reward(self, state: Any, action: Any) -> float:

synth_ai/environments/environment/tools/__init__.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
+
 from abc import ABC, abstractmethod
+from typing import Any
+
 from pydantic import BaseModel, Field
-from typing import Any, Dict, Type
 
 
 class EnvToolCall(BaseModel):
@@ -38,7 +40,7 @@ class EnvToolCall(BaseModel):
     """
 
     tool: str
-    args: Dict[str, Any] = Field(default_factory=dict)
+    args: dict[str, Any] = Field(default_factory=dict)
 
 
 class ToolResult(BaseModel):
@@ -123,8 +125,8 @@ class AbstractTool(ABC):
     """
 
     name: str
-    call_schema: Type[BaseModel]
-    result_schema: Type[BaseModel] = ToolResult
+    call_schema: type[BaseModel]
+    result_schema: type[BaseModel] = ToolResult
 
     @abstractmethod
     async def __call__(self, call: EnvToolCall) -> ToolResult:

synth_ai/environments/examples/crafter_classic/__init__.py

@@ -1,6 +1,6 @@
 from .config_logging import configure_logging
-from .environment import CrafterClassicEnvironment
 from .engine import CrafterEngine
+from .environment import CrafterClassicEnvironment
 
 # Configure logging when crafter_classic module is imported
 configure_logging()

synth_ai/environments/examples/crafter_classic/engine.py

@@ -4,32 +4,32 @@ This file follows the same structure as the SokobanEngine shown earlier.
 
 from __future__ import annotations
 
-# Import logging configuration first to suppress JAX debug messages
-from .config_logging import safe_compare
-
-# Import patches
-from . import engine_deterministic_patch  # Ensures deterministic behavior
-from . import engine_serialization_patch_v3 as engine_serialization_patch  # Adds save/load methods
-from . import world_config_patch_simple as world_config_patch  # Adds configurable world generation
-
+import copy
+import dataclasses
 import logging
 import time
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, Tuple, Union
 
-import numpy as np
 import crafter  # type: ignore
-import copy
-import dataclasses
+import numpy as np
 
+from synth_ai.environments.environment.rewards.core import RewardComponent, RewardStack  # Added
 from synth_ai.environments.environment.shared_engine import (
     GetObservationCallable,
     InternalObservation,
 )
+from synth_ai.environments.reproducibility.core import IReproducibleEngine
 from synth_ai.environments.stateful.engine import StatefulEngine, StatefulEngineSnapshot
 from synth_ai.environments.tasks.core import TaskInstance
-from synth_ai.environments.reproducibility.core import IReproducibleEngine
-from synth_ai.environments.environment.rewards.core import RewardStack, RewardComponent  # Added
+
+# Import patches
+from . import engine_deterministic_patch  # Ensures deterministic behavior
+from . import engine_serialization_patch_v3 as engine_serialization_patch  # Adds save/load methods
+from . import world_config_patch_simple as world_config_patch  # Adds configurable world generation
+
+# Import logging configuration first to suppress JAX debug messages
+from .config_logging import safe_compare
 
 # Local helper imports (must exist relative to this file)
 from .engine_helpers.action_map import CRAFTER_ACTION_MAP  # action‑name → int

synth_ai/environments/examples/crafter_classic/engine_deterministic_patch.py

@@ -6,6 +6,7 @@ the hash-based set-iteration nondeterminism that caused the drift.
 """
 
 import collections
+
 import crafter
 
 print("[PATCH] Attempting to apply Crafter deterministic patch...")

synth_ai/environments/examples/crafter_classic/engine_helpers/action_map.py

@@ -1,5 +1,6 @@
-import crafter.constants as C
 from typing import Dict
 
+import crafter.constants as C
+
 # Map each action name to its corresponding index in the crafter package
 CRAFTER_ACTION_MAP: Dict[str, int] = {action_name: idx for idx, action_name in enumerate(C.actions)}

synth_ai/environments/examples/crafter_classic/engine_helpers/serialization.py

@@ -1,7 +1,8 @@
 import importlib
-import numpy as np
 from typing import Any, Dict
 
+import numpy as np
+
 # Minimal attributes to serialize per object type
 BASIC_ATTRS: Dict[str, list] = {
     "Player": [

synth_ai/environments/examples/crafter_classic/engine_serialization_patch_v3.py

@@ -5,10 +5,11 @@ This version handles player references for Zombie and Skeleton objects.
 
 import collections
 import pickle
-import numpy as np
+from typing import Any, Dict, Optional, Set
+
 import crafter
+import numpy as np
 from crafter import objects
-from typing import Dict, Any, Optional, Set
 
 print("[PATCH] Attempting to apply Crafter serialization patch v3...")
 

synth_ai/environments/examples/crafter_classic/environment.py

@@ -2,13 +2,10 @@
 
 from __future__ import annotations
 
-from typing import List, Optional, Any, Dict, Union
 import dataclasses
 import logging
 import time
-
-# Import logging configuration to suppress JAX debug messages
-from .config_logging import safe_compare
+from typing import Any, Dict, List, Optional, Union
 
 # Import tracing abstractions
 from synth_ai.tracing_v3.abstractions import (
@@ -17,29 +14,33 @@ from synth_ai.tracing_v3.abstractions import (
     TimeRecord,
 )
 
+# Import logging configuration to suppress JAX debug messages
+from .config_logging import safe_compare
+
 logger = logging.getLogger(__name__)
 
-from synth_ai.environments.examples.crafter_classic.engine import (
-    CrafterEngine,
-    CrafterPrivateState,
-    CrafterPublicState,
-    CrafterEngineSnapshot,
-)
-from synth_ai.environments.examples.crafter_classic.taskset import CrafterTaskInstance
+from pydantic import BaseModel, Field
+
 from synth_ai.environments.environment.shared_engine import (
     GetObservationCallable,
     InternalObservation,
 )
-from synth_ai.environments.reproducibility.core import ReproducibleEnvironment
-from synth_ai.environments.stateful.core import StatefulEnvironment
 from synth_ai.environments.environment.tools import (
+    TOOL_REGISTRY,
     AbstractTool,
     EnvToolCall,
     ToolResult,
-    TOOL_REGISTRY,
     register_tool,
 )
-from pydantic import BaseModel, Field
+from synth_ai.environments.examples.crafter_classic.engine import (
+    CrafterEngine,
+    CrafterEngineSnapshot,
+    CrafterPrivateState,
+    CrafterPublicState,
+)
+from synth_ai.environments.examples.crafter_classic.taskset import CrafterTaskInstance
+from synth_ai.environments.reproducibility.core import ReproducibleEnvironment
+from synth_ai.environments.stateful.core import StatefulEnvironment
 
 
 # --- Tool Definition ---

synth_ai/environments/examples/crafter_classic/taskset.py

@@ -6,12 +6,12 @@ from __future__ import annotations
 
 import asyncio
 import random
-from dataclasses import dataclass, asdict, fields
+from dataclasses import asdict, dataclass, fields
 from typing import Dict, List
 from uuid import UUID, uuid4
 
-import numpy as np
 import crafter
+import numpy as np
 from crafter import objects
 
 from synth_ai.environments.tasks.core import (

synth_ai/environments/examples/crafter_classic/trace_hooks_v3.py

@@ -3,12 +3,11 @@ Trace hooks for Crafter Classic environment - v3 version.
 Updated to use the v3 tracing architecture with async support.
 """
 
-from typing import Any, Optional, Dict, Set
 from datetime import datetime
+from typing import Any, Dict, Optional, Set
 
-from synth_ai.tracing_v3.hooks import HookManager
 from synth_ai.tracing_v3.abstractions import BaseEvent, EnvironmentEvent
-
+from synth_ai.tracing_v3.hooks import HookManager
 
 # Achievement categorization by difficulty
 EASY_ACHIEVEMENTS = {

synth_ai/environments/examples/crafter_classic/world_config_patch_simple.py

@@ -5,7 +5,8 @@ This version modifies generation parameters rather than rewriting functions.
 
 import json
 import os
-from typing import Dict, Any, Optional
+from typing import Any, Dict, Optional
+
 import crafter
 
 print("[PATCH] Attempting to apply simplified Crafter world configuration patch...")

synth_ai/environments/examples/crafter_custom/crafter/__init__.py

@@ -1,7 +1,7 @@
+from . import constants
+from .config import PRESETS, WorldGenConfig
 from .env import Env
 from .recorder import Recorder
-from .config import WorldGenConfig, PRESETS
-from . import constants
 
 # Note: We don't register with gym since this is a custom version
 # Users should import directly from this module

synth_ai/environments/examples/crafter_custom/crafter/config.py

@@ -4,8 +4,8 @@ Configuration system for customizable Crafter environment.
 
 import json
 import os
-from dataclasses import dataclass, field, asdict
-from typing import Dict, Any, Optional
+from dataclasses import asdict, dataclass, field
+from typing import Any, Dict, Optional
 
 
 @dataclass

synth_ai/environments/examples/crafter_custom/crafter/env.py

@@ -2,13 +2,9 @@ import collections
 
 import numpy as np
 
-from . import constants
-from . import engine
-from . import objects
-from . import worldgen
+from . import constants, engine, objects, worldgen
 from .config import WorldGenConfig
 
-
 # Gym is an optional dependency.
 try:
     import gym

synth_ai/environments/examples/crafter_custom/crafter/objects.py

@@ -1,7 +1,6 @@
 import numpy as np
 
-from . import constants
-from . import engine
+from . import constants, engine
 
 
 class Object:

synth_ai/environments/examples/crafter_custom/crafter/worldgen.py

@@ -3,8 +3,7 @@ import functools
 import numpy as np
 import opensimplex
 
-from . import constants
-from . import objects
+from . import constants, objects
 from .config import WorldGenConfig
 
 

synth_ai/environments/examples/crafter_custom/dataset_builder.py

@@ -3,14 +3,14 @@
 Dataset builder for Crafter Custom Environments
 """
 
+import gzip
 import json
+import pickle
+import random
 import uuid
-from dataclasses import dataclass, asdict
+from dataclasses import asdict, dataclass
 from pathlib import Path
-from typing import Any, Dict, Optional, List, Set
-import random
-import gzip
-import pickle
+from typing import Any, Dict, List, Optional, Set
 
 
 @dataclass

synth_ai/environments/examples/crafter_custom/environment.py

@@ -2,10 +2,10 @@
 
 from __future__ import annotations
 
-from typing import List, Optional, Any, Dict, Union
 import dataclasses
 import logging
 import time
+from typing import Any, Dict, List, Optional, Union
 
 # Import logging configuration to suppress JAX debug messages
 from synth_ai.environments.examples.crafter_classic.config_logging import safe_compare
@@ -20,28 +20,25 @@ from synth_ai.tracing_v3.abstractions import (
 logger = logging.getLogger(__name__)
 
 # Import the base Crafter components
-from synth_ai.environments.examples.crafter_classic.engine import (
-    CrafterEngine,
-    CrafterPrivateState,
-    CrafterPublicState,
-    CrafterEngineSnapshot,
-)
-from synth_ai.environments.examples.crafter_classic.taskset import CrafterTaskInstance
+from pydantic import BaseModel, Field
+
 from synth_ai.environments.environment.shared_engine import (
     GetObservationCallable,
     InternalObservation,
 )
-from synth_ai.environments.reproducibility.core import ReproducibleEnvironment
-from synth_ai.environments.stateful.core import StatefulEnvironment
 from synth_ai.environments.environment.tools import (
+    TOOL_REGISTRY,
     AbstractTool,
     EnvToolCall,
     ToolResult,
-    TOOL_REGISTRY,
     register_tool,
 )
-from pydantic import BaseModel, Field
-
+from synth_ai.environments.examples.crafter_classic.engine import (
+    CrafterEngine,
+    CrafterEngineSnapshot,
+    CrafterPrivateState,
+    CrafterPublicState,
+)
 
 # Use the same tool and observation classes as CrafterClassic
 from synth_ai.environments.examples.crafter_classic.environment import (
@@ -49,6 +46,9 @@ from synth_ai.environments.examples.crafter_classic.environment import (
     CrafterInteractTool,
     SynthCrafterObservationCallable,
 )
+from synth_ai.environments.examples.crafter_classic.taskset import CrafterTaskInstance
+from synth_ai.environments.reproducibility.core import ReproducibleEnvironment
+from synth_ai.environments.stateful.core import StatefulEnvironment
 
 
 class CrafterCustomEnvironment(StatefulEnvironment, ReproducibleEnvironment[CrafterEngine]):

synth_ai/environments/examples/crafter_custom/run_dataset.py

@@ -3,14 +3,14 @@
 Run script for Crafter dataset instances
 """
 
-import json
 import argparse
-import random
-from pathlib import Path
-from typing import List, Optional, Dict, Any
-import uuid
+import json
 import os
+import random
 import sys
+import uuid
+from pathlib import Path
+from typing import Any, Dict, List, Optional
 
 sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))))
 

synth_ai/environments/examples/enron/art_helpers/email_search_tools.py

@@ -1,8 +1,8 @@
-import sqlite3
 import logging
+import sqlite3
 import textwrap
-from typing import List, Optional
 from dataclasses import dataclass
+from typing import List, Optional
 
 from synth_ai.environments.environment.db.sqlite import SQLiteManager
 from synth_ai.environments.examples.enron.art_helpers.types_enron import Email

synth_ai/environments/examples/enron/art_helpers/local_email_db.py

@@ -1,10 +1,11 @@
-import sqlite3
-import os
 import logging
-from datasets import load_dataset, Dataset, Features, Value, Sequence
-from tqdm import tqdm
+import os
+import sqlite3
 from datetime import datetime
 
+from datasets import Dataset, Features, Sequence, Value, load_dataset
+from tqdm import tqdm
+
 # Resolve paths relative to this file so it works regardless of the current working directory
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 # Database will live in "../data/enron_emails.db" relative to project root

synth_ai/environments/examples/enron/art_helpers/types_enron.py

@@ -1,6 +1,7 @@
-from pydantic import BaseModel
 from typing import List, Optional
 
+from pydantic import BaseModel
+
 
 class SyntheticQuery(BaseModel):
     id: int

synth_ai/environments/examples/enron/engine.py

@@ -1,28 +1,32 @@
 # engine.py
 from __future__ import annotations
-from dataclasses import dataclass, asdict
-from typing import Any, Dict, Tuple, Optional, List
-from pydantic import BaseModel
+
+from dataclasses import asdict, dataclass
 from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
 
-from synth_ai.environments.examples.enron.art_helpers.types_enron import Email
+from pydantic import BaseModel
+from synth_ai.environments.environment.db.sqlite import SQLiteManager
+from synth_ai.environments.environment.rewards.core import RewardComponent, RewardStack
 from synth_ai.environments.examples.enron.art_helpers.email_search_tools import (
-    search_emails as helper_search_emails,
-    read_email as helper_read_email,
     SearchResult,
 )
-
-# SQLite-backed helpers
-from synth_ai.environments.stateful.engine import StatefulEngine, StatefulEngineSnapshot
-from synth_ai.environments.examples.enron.taskset import EnronTaskInstance
-from synth_ai.zyk import LM  # Import LM class
-
-from synth_ai.environments.environment.db.sqlite import SQLiteManager
-from synth_ai.environments.environment.rewards.core import RewardStack, RewardComponent
+from synth_ai.environments.examples.enron.art_helpers.email_search_tools import (
+    read_email as helper_read_email,
+)
+from synth_ai.environments.examples.enron.art_helpers.email_search_tools import (
+    search_emails as helper_search_emails,
+)
 from synth_ai.environments.examples.enron.art_helpers.local_email_db import (
     DEFAULT_DB_PATH,
     generate_database,
 )
+from synth_ai.environments.examples.enron.art_helpers.types_enron import Email
+from synth_ai.environments.examples.enron.taskset import EnronTaskInstance
+
+# SQLite-backed helpers
+from synth_ai.environments.stateful.engine import StatefulEngine, StatefulEngineSnapshot
+from synth_ai.zyk import LM  # Import LM class
 
 # --------------------------------------------------------------------------- actions
 ACTION_SEARCH = "search"