File size: 1,462 Bytes
5dd1bb4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
"""SQLEnv: Interactive Database Query Environment for the OpenEnv Challenge."""

# ---------------------------------------------------------------------------
# Pydantic / TypedDict compatibility shim
# ---------------------------------------------------------------------------
# The openenv library defines ``Message`` with ``typing.TypedDict``.
# On Python < 3.12, Pydantic 2.x rejects ``typing.TypedDict`` in model
# fields; it requires ``typing_extensions.TypedDict`` instead.  We patch
# ``typing.TypedDict`` early so that all downstream imports see the
# compatible version before any Pydantic model is constructed.
import sys

if sys.version_info < (3, 12):
    import typing
    import typing_extensions

    typing.TypedDict = typing_extensions.TypedDict  # type: ignore[attr-defined]

try:
    from .models import SQLAction, SQLObservation, SQLState
except ImportError:
    # When pytest imports this file standalone (not as part of the sql_env
    # package), relative imports fail. Fall back to absolute imports.
    try:
        from sql_env.models import SQLAction, SQLObservation, SQLState  # type: ignore[no-redef]
    except ImportError:
        pass  # Imports not available; this file is being collected, not used.

# Client is not imported at package level to avoid loading torch unnecessarily.
# Import it explicitly when needed: from sql_env.client import SQLEnvClient

__all__ = [
    "SQLAction",
    "SQLObservation",
    "SQLState",
]