Import simplification & dev dependencies

.github/workflows/codeflash.yml

@@ -38,7 +38,9 @@ jobs:
         if: steps.bot_check.outputs.skip_remaining_steps == 'no'
         run: curl -LsSf https://astral.sh/uv/install.sh | sh
       - if: steps.bot_check.outputs.skip_remaining_steps == 'no'
-        run: uv sync
+        run: |-
+          uv sync
+          uv pip install codeflash
       - name: Run CodeFlash on fhaviary
         if: steps.bot_check.outputs.skip_remaining_steps == 'no'
         run: uv run codeflash

packages/gsm8k/pyproject.toml

@@ -30,4 +30,4 @@ where = ["src"]
 
 [tool.setuptools_scm]
 root = "../.."
-version_file = "src/aviary/gsm8k/version.py"
+version_file = "src/aviary/envs/gsm8k/version.py"

packages/gsm8k/tests/test_gsm8k_env.py

@@ -1,8 +1,7 @@
 import pytest
 
-from aviary.env import Environment, TaskDataset
-from aviary.gsm8k.env import CalculatorEnv, CalculatorEnvConfig
-from aviary.tools import ToolCall, ToolRequestMessage
+from aviary.core import Environment, TaskDataset, ToolCall, ToolRequestMessage
+from aviary.envs.gsm8k import CalculatorEnv, CalculatorEnvConfig
 
 
 @pytest.mark.asyncio

packages/hotpotqa/pyproject.toml

@@ -28,4 +28,4 @@ where = ["src"]
 
 [tool.setuptools_scm]
 root = "../.."
-version_file = "src/aviary/hotpotqa/version.py"
+version_file = "src/aviary/envs/hotpotqa/version.py"

packages/hotpotqa/tests/test_hotpotqa_env.py

@@ -1,5 +1,5 @@
-from aviary.env import Environment, TaskDataset
-from aviary.hotpotqa import HotPotQAEnv
+from aviary.core import Environment, TaskDataset
+from aviary.envs.hotpotqa import HotPotQAEnv
 
 
 def test_env_construction() -> None:

pyproject.toml

@@ -30,6 +30,27 @@ requires-python = ">=3.11"
 cloud = [
     "boto3",
 ]
+dev = [
+    "SQLAlchemy[aiosqlite]~=2.0",  # Match aviary dependencies
+    "aviary.gsm8k[typing]",  # So `uv sync` pulls this in, and for type stubs
+    "aviary.hotpotqa",  # So `uv sync` pulls this in
+    "fhaviary[image,llm,server,typing,xml]",
+    "ipython>=8",  # Pin to keep recent
+    "mypy>=1.8",  # Pin for mutable-override
+    "pre-commit~=3.4",
+    "pydantic~=2.9",  # Pydantic 2.9 changed JSON schema exports 'allOf', so ensure tests match
+    "pylint-pydantic",
+    "pylint>=3.2",
+    "pytest-asyncio",
+    "pytest-recording",
+    "pytest-subtests",
+    "pytest-sugar",
+    "pytest-timer[colorama]",
+    "pytest-xdist",
+    "pytest>=8",  # Pin to keep recent
+    "refurb>=2",  # Pin to keep recent
+    "typeguard",
+]
 gsm8k = ["aviary.gsm8k"]
 hotpotqa = ["aviary.hotpotqa"]
 image = [
@@ -394,26 +415,7 @@ trailing_comma_inline_array = true
 
 [tool.uv]
 dev-dependencies = [
-    "SQLAlchemy[aiosqlite]~=2.0",  # Match aviary dependencies
-    "aviary.gsm8k[typing]",  # So `uv sync` pulls this in, and for type stubs
-    "aviary.hotpotqa",  # So `uv sync` pulls this in
-    "codeflash",
-    "fhaviary[image,llm,server,typing,xml]",
-    "ipython>=8",  # Pin to keep recent
-    "mypy>=1.8",  # Pin for mutable-override
-    "pre-commit~=3.4",
-    "pydantic~=2.9",  # Pydantic 2.9 changed JSON schema exports 'allOf', so ensure tests match
-    "pylint-pydantic",
-    "pylint>=3.2",
-    "pytest-asyncio",
-    "pytest-recording",
-    "pytest-subtests",
-    "pytest-sugar",
-    "pytest-timer[colorama]",
-    "pytest-xdist",
-    "pytest>=8",  # Pin to keep recent
-    "refurb>=2",  # Pin to keep recent
-    "typeguard",
+    "fhaviary[dev]",
 ]
 
 [tool.uv.sources]

src/aviary/core.py

@@ -0,0 +1,69 @@
+from aviary.env import (
+    TASK_DATASET_REGISTRY,
+    DummyEnv,
+    DummyEnvState,
+    DummyTaskDataset,
+    Environment,
+    Frame,
+    TaskConfig,
+    TaskDataset,
+)
+from aviary.env_client import EnvironmentClient
+from aviary.message import MalformedMessageError, Message, join
+from aviary.render import Renderer
+from aviary.tools import (
+    INVALID_TOOL_NAME,
+    FunctionInfo,
+    Messages,
+    MessagesAdapter,
+    Parameters,
+    Tool,
+    ToolCall,
+    ToolCallFunction,
+    ToolRequestMessage,
+    ToolResponseMessage,
+    Tools,
+    ToolsAdapter,
+    ToolSelector,
+    ToolSelectorLedger,
+    argref_by_name,
+    eval_answer,
+    wraps_doc_only,
+)
+from aviary.utils import encode_image_to_base64, is_coroutine_callable, partial_format
+
+__all__ = [
+    "INVALID_TOOL_NAME",
+    "TASK_DATASET_REGISTRY",
+    "DummyEnv",
+    "DummyEnvState",
+    "DummyTaskDataset",
+    "Environment",
+    "EnvironmentClient",
+    "Frame",
+    "FunctionInfo",
+    "MalformedMessageError",
+    "Message",
+    "Messages",
+    "MessagesAdapter",
+    "Parameters",
+    "Renderer",
+    "TaskConfig",
+    "TaskDataset",
+    "Tool",
+    "ToolCall",
+    "ToolCallFunction",
+    "ToolRequestMessage",
+    "ToolResponseMessage",
+    "ToolSelector",
+    "ToolSelectorLedger",
+    "Tools",
+    "ToolsAdapter",
+    "argref_by_name",
+    "encode_image_to_base64",
+    "eval_answer",
+    "is_coroutine_callable",
+    "join",
+    "partial_format",
+    "wraps_doc_only",
+]

src/aviary/env.py

@@ -250,8 +250,8 @@ def from_name(cls, name: str, **env_kwargs) -> Self:
 # Maps baseline environment names to their module and class names
 ENV_REGISTRY: dict[str, tuple[str, str]] = {
     "dummy": ("aviary.env", "DummyEnv"),
-    "calculator": ("aviary.gsm8k.env", "CalculatorEnv"),
-    "hotpotqa": ("aviary.hotpotqa.env", "HotPotQAEnv"),
+    "calculator": ("aviary.envs.gsm8k.env", "CalculatorEnv"),
+    "hotpotqa": ("aviary.envs.hotpotqa.env", "HotPotQAEnv"),
 }
 
 TEnvironment = TypeVar("TEnvironment", bound=Environment)
@@ -319,8 +319,8 @@ def iter_batches(
 # Maps baseline task dataset names to their module and class names
 TASK_DATASET_REGISTRY: dict[str, tuple[str, str]] = {
     "dummy": ("aviary.env", "DummyTaskDataset"),
-    "gsm8k": ("aviary.gsm8k.env", "GSM8kDataset"),
-    "hotpotqa": ("aviary.hotpotqa.env", "HotPotQADataset"),
+    "gsm8k": ("aviary.envs.gsm8k.env", "GSM8kDataset"),
+    "hotpotqa": ("aviary.envs.hotpotqa.env", "HotPotQADataset"),
 }
 
 

tests/test_envs.py

@@ -8,10 +8,14 @@
 import pytest
 from pydantic import BaseModel
 
-from aviary.env import DummyEnv, DummyEnvState, Environment, Frame, TaskDataset
-from aviary.message import Message
-from aviary.render import Renderer
-from aviary.tools import (
+from aviary.core import (
+    DummyEnv,
+    DummyEnvState,
+    Environment,
+    Frame,
+    Message,
+    Renderer,
+    TaskDataset,
     Tool,
     ToolCall,
     ToolRequestMessage,

tests/test_messages.py

@@ -3,8 +3,8 @@
 import numpy as np
 import pytest
 
-from aviary.message import Message
-from aviary.tools import (
+from aviary.core import (
+    Message,
     ToolCall,
     ToolCallFunction,
     ToolRequestMessage,

tests/test_tools.py

@@ -8,9 +8,9 @@
 from pydantic import BaseModel, Field
 from pytest_subtests import SubTests
 
-from aviary.env import DummyEnv
-from aviary.tools import (
+from aviary.core import (
     INVALID_TOOL_NAME,
+    DummyEnv,
     FunctionInfo,
     Tool,
     ToolCall,

tests/test_utils.py

@@ -1,6 +1,6 @@
 import pytest
 
-from aviary.tools import eval_answer
+from aviary.core import eval_answer
 
 
 @pytest.mark.vcr