Skip to content

API Documentation

GeoDataFrame AI

GeoDataFrameAI

Bases: GeoDataFrame

A class to represent a GeoDataFrame with AI capabilities. It is a proxy for the GeoPandas GeoDataFrame class, allowing for additional functionality related to AI and machine learning tasks.

Source code in geopandasai/geodataframe_ai.py
 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
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
class GeoDataFrameAI(GeoDataFrame):
    """
    A class to represent a GeoDataFrame with AI capabilities. It is a proxy for
    the GeoPandas GeoDataFrame class, allowing for additional functionality
    related to AI and machine learning tasks.
    """

    def __init__(self, *args, description: str = None, **kwargs):
        """
        Initialize the GeoDataFrameAI class.
        """
        super().__init__(*args, **kwargs)
        # User provided description through the describe method
        self.description = description or ""

        # The state of the conversation, calling chat initializes this
        # while calling improve will update the state.
        self.state: Union[MagicReturn, Any] = None

        # A helper storing previous conversation to ensure the reset
        # method delete entire conversation history, even if multiple
        # '.chat' were called.
        self._memories = set()

    def set_description(self, description: str):
        """
        Describe the GeoDataFrameAI. This is a user-provided description that
        can be used to provide context for the AI.
        """
        assert (
            self.state is None
        ), "You cannot set a description after running a chat or improve method. Please reset the state first."
        self.description = description
        return self

    def chat(
        self,
        prompt: str,
        *other_dfs: Sequence[GeoOrDataFrame | "GeoDataFrameAI"],
        return_type: Type = None,
        provided_libraries: List[str] = None,
    ) -> Union[Any, MagicReturn]:
        """
        This method is used to start a conversation with the AI.
        It takes a prompt and any number of other GeoDataFrames as input.
        The prompt is a string that describes the task or question to be answered.
        The other_dfs are additional GeoDataFrames that can be used to provide
        context for the AI.

        :param prompt: The prompt to start the conversation.
        :param other_dfs: Additional GeoDataFrames to provide context.
        :param return_type: The type of the return value. If None, it will be inferred.
        :param provided_libraries: A list of libraries that are provided by the user.
        :return: The result of the conversation, which can be any type or a MagicReturn object.
        """

        # Reset the state if it was previously set, and initialize a new memory
        self.state = magic_prompt_with_dataframes(
            prompt,
            *([self] + list(other_dfs)),
            return_type=return_type,
            provided_libraries=provided_libraries,
        )
        # Only for caching purposes, we store the memory of the conversation to be able to clean it up later.
        self._memories.add(self.state.memory)

        return self.return_value()

    def return_value(self):
        if (
            self.state.memory.return_type == GeoDataFrame
            or self.state.memory.return_type == DataFrame
        ):
            return self.state.internal
        return self.state.execute()

    def improve(
        self,
        prompt: str,
        *other_dfs: List[GeoOrDataFrame | "GeoDataFrameAI"],
        return_type: Type = None,
        provided_libraries: List[str] = None,
    ) -> Union[Any, MagicReturn]:
        if self.state is None:
            raise ValueError("No code has been generated yet. Please run a chat first.")

        self.state = self.state.improve(
            prompt,
            *(self.state.memory.dfs if not other_dfs else ([self] + list(other_dfs))),
            return_type=return_type,
            provided_libraries=provided_libraries,
        )

        return self.return_value()

    @property
    def code(self) -> str:
        if self.state is None:
            raise ValueError("No code has been generated yet. Please run a chat first.")
        return self.state.code

    def inspect(self):
        """
        Print the history of the last output.
        """
        if self.state is None:
            raise ValueError("No code has been generated yet. Please run a chat first.")
        self.state.inspect()

    def reset(self):
        """
        Reset the state of the GeoDataFrameAI.
        """
        self.state = None
        for memory in self._memories:
            memory.reset()

    def inject(self, function_name: str, ai_module="ai", ai_module_path="ai"):
        """
        Inject the state of the GeoDataFrameAI into the current context.
        """
        if self.state is None:
            raise ValueError("No code has been generated yet. Please run a chat first.")
        return self.state.inject(function_name, ai_module, ai_module_path)
__init__(*args, description=None, **kwargs)

Initialize the GeoDataFrameAI class.

Source code in geopandasai/geodataframe_ai.py
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def __init__(self, *args, description: str = None, **kwargs):
    """
    Initialize the GeoDataFrameAI class.
    """
    super().__init__(*args, **kwargs)
    # User provided description through the describe method
    self.description = description or ""

    # The state of the conversation, calling chat initializes this
    # while calling improve will update the state.
    self.state: Union[MagicReturn, Any] = None

    # A helper storing previous conversation to ensure the reset
    # method delete entire conversation history, even if multiple
    # '.chat' were called.
    self._memories = set()
chat(prompt, *other_dfs, return_type=None, provided_libraries=None)

This method is used to start a conversation with the AI. It takes a prompt and any number of other GeoDataFrames as input. The prompt is a string that describes the task or question to be answered. The other_dfs are additional GeoDataFrames that can be used to provide context for the AI.

:param prompt: The prompt to start the conversation. :param other_dfs: Additional GeoDataFrames to provide context. :param return_type: The type of the return value. If None, it will be inferred. :param provided_libraries: A list of libraries that are provided by the user. :return: The result of the conversation, which can be any type or a MagicReturn object.

Source code in geopandasai/geodataframe_ai.py
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def chat(
    self,
    prompt: str,
    *other_dfs: Sequence[GeoOrDataFrame | "GeoDataFrameAI"],
    return_type: Type = None,
    provided_libraries: List[str] = None,
) -> Union[Any, MagicReturn]:
    """
    This method is used to start a conversation with the AI.
    It takes a prompt and any number of other GeoDataFrames as input.
    The prompt is a string that describes the task or question to be answered.
    The other_dfs are additional GeoDataFrames that can be used to provide
    context for the AI.

    :param prompt: The prompt to start the conversation.
    :param other_dfs: Additional GeoDataFrames to provide context.
    :param return_type: The type of the return value. If None, it will be inferred.
    :param provided_libraries: A list of libraries that are provided by the user.
    :return: The result of the conversation, which can be any type or a MagicReturn object.
    """

    # Reset the state if it was previously set, and initialize a new memory
    self.state = magic_prompt_with_dataframes(
        prompt,
        *([self] + list(other_dfs)),
        return_type=return_type,
        provided_libraries=provided_libraries,
    )
    # Only for caching purposes, we store the memory of the conversation to be able to clean it up later.
    self._memories.add(self.state.memory)

    return self.return_value()
inject(function_name, ai_module='ai', ai_module_path='ai')

Inject the state of the GeoDataFrameAI into the current context.

Source code in geopandasai/geodataframe_ai.py
130
131
132
133
134
135
136
def inject(self, function_name: str, ai_module="ai", ai_module_path="ai"):
    """
    Inject the state of the GeoDataFrameAI into the current context.
    """
    if self.state is None:
        raise ValueError("No code has been generated yet. Please run a chat first.")
    return self.state.inject(function_name, ai_module, ai_module_path)
inspect()

Print the history of the last output.

Source code in geopandasai/geodataframe_ai.py
114
115
116
117
118
119
120
def inspect(self):
    """
    Print the history of the last output.
    """
    if self.state is None:
        raise ValueError("No code has been generated yet. Please run a chat first.")
    self.state.inspect()
reset()

Reset the state of the GeoDataFrameAI.

Source code in geopandasai/geodataframe_ai.py
122
123
124
125
126
127
128
def reset(self):
    """
    Reset the state of the GeoDataFrameAI.
    """
    self.state = None
    for memory in self._memories:
        memory.reset()
set_description(description)

Describe the GeoDataFrameAI. This is a user-provided description that can be used to provide context for the AI.

Source code in geopandasai/geodataframe_ai.py
37
38
39
40
41
42
43
44
45
46
def set_description(self, description: str):
    """
    Describe the GeoDataFrameAI. This is a user-provided description that
    can be used to provide context for the AI.
    """
    assert (
        self.state is None
    ), "You cannot set a description after running a chat or improve method. Please reset the state first."
    self.description = description
    return self

Configuration

GeoPandasAIConfig dataclass

Configuration class for GeoPandasAI.

This class holds the configuration settings for GeoPandasAI, including the lite LLM configuration, libraries to be used, cache backend, data descriptor, code injector, code executor, and return types. It is designed to be immutable after creation, ensuring that the configuration remains consistent throughout the application.

Source code in geopandasai/config.py
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
@dataclass(frozen=True, eq=True)
class GeoPandasAIConfig:
    """
    Configuration class for GeoPandasAI.

    This class holds the configuration settings for GeoPandasAI, including
    the lite LLM configuration, libraries to be used, cache backend,
    data descriptor, code injector, code executor, and return types.
    It is designed to be immutable after creation, ensuring that the configuration
    remains consistent throughout the application.
    """

    lite_llm_config: Optional[dict] = field(
        default_factory=lambda: _load_default_lite_llm_config()
    )
    libraries: List[str] = field(
        default_factory=lambda: [
            "pandas",
            "matplotlib.pyplot",
            "folium",
            "geopandas",
            "contextily",
        ]
    )

    cache_backend: ACacheBackend = FileSystemCacheBackend()

    descriptor: ADescriptor = PublicDataDescriptor()

    injector: ACodeInjector = PrintCodeInjector()

    executor: ACodeExecutor = TrustedCodeExecutor()

    return_types: Set[Type] = field(
        default_factory=lambda: {
            int,
            float,
            str,
            bool,
            list,
            dict,
            GeoDataFrame,
            DataFrame,
            folium.Map,
            Figure,
        }
    )

set_active_lite_llm_config(lite_llm_config)

Set the active lite LLM configuration.

Source code in geopandasai/config.py
89
90
91
92
93
94
95
96
@deprecated(
    "This function is deprecated and will be removed in future versions. Use `update_geopandas_ai_config` instead.",
)
def set_active_lite_llm_config(lite_llm_config: dict) -> None:
    """Set the active lite LLM configuration."""
    update_geopandasai_config(
        lite_llm_config=lite_llm_config,
    )

update_geopandasai_config(lite_llm_config=None, libraries=None, cache_backend=None, descriptor=None, return_types=None, injector=None, executor=None)

Update the GeoPandasAI configuration.

This function allows you to update the configuration of GeoPandasAI, including the lite LLM configuration, libraries, cache backend, descriptor, return types, and code injector. If a parameter is not provided, the current value will be retained. :param lite_llm_config: The configuration for the lite LLM, if any. :param libraries: A list of libraries to be used in the GeoPandasAI environment. :param cache_backend: The cache backend to be used. :param descriptor: The data descriptor to be used. :param return_types: A set of types that can be returned by the AI. :param injector: The code injector to be used. :param executor: The code executor to be used.

Source code in geopandasai/config.py
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
def update_geopandasai_config(
    lite_llm_config: Optional[dict] = None,
    libraries: Optional[List[str]] = None,
    cache_backend: Optional[ACacheBackend] = None,
    descriptor: Optional[ADescriptor] = None,
    return_types: Optional[Set[Type]] = None,
    injector: Optional[ACodeInjector] = None,
    executor: Optional[ACodeExecutor] = None,
) -> None:
    """Update the GeoPandasAI configuration.

    This function allows you to update the configuration of GeoPandasAI, including
    the lite LLM configuration, libraries, cache backend, descriptor, return types,
    and code injector. If a parameter is not provided, the current value will be retained.
    :param lite_llm_config: The configuration for the lite LLM, if any.
    :param libraries: A list of libraries to be used in the GeoPandasAI environment.
    :param cache_backend: The cache backend to be used.
    :param descriptor: The data descriptor to be used.
    :param return_types: A set of types that can be returned by the AI.
    :param injector: The code injector to be used.
    :param executor: The code executor to be used.
    """

    current_config = container.config()
    updated_config = replace(
        current_config,
        lite_llm_config=(
            lite_llm_config
            if lite_llm_config is not None
            else current_config.lite_llm_config
        ),
        libraries=libraries if libraries is not None else current_config.libraries,
        cache_backend=(
            cache_backend if cache_backend is not None else current_config.cache_backend
        ),
        descriptor=(
            descriptor if descriptor is not None else current_config.descriptor
        ),
        return_types=(
            return_types if return_types is not None else current_config.return_types
        ),
        injector=(injector if injector is not None else current_config.injector),
        executor=(executor if executor is not None else current_config.executor),
    )
    container.config.override(updated_config)

GeoPandas Wrapper

read_arrow(source, *args, **kwargs)

Read Arrow (e.g. pyarrow.Table) data and return a GeoDataFrameAI.

Source code in geopandasai/wrapper.py
62
63
64
65
66
67
68
69
70
def read_arrow(
    source,
    *args,
    **kwargs,
) -> GeoDataFrameAI:
    """
    Read Arrow (e.g. pyarrow.Table) data and return a GeoDataFrameAI.
    """
    return GeoDataFrameAI(gpd.read_arrow(source, *args, **kwargs))

read_feather(path, *args, **kwargs)

Read a Feather file and return a GeoDataFrameAI.

Source code in geopandasai/wrapper.py
28
29
30
31
32
33
34
35
36
def read_feather(
    path: str,
    *args,
    **kwargs,
) -> GeoDataFrameAI:
    """
    Read a Feather file and return a GeoDataFrameAI.
    """
    return GeoDataFrameAI(gpd.read_feather(path, *args, **kwargs))

read_file(filename, *args, **kwargs)

Read a vector file (shapefile, GeoJSON, etc.) and return a GeoDataFrameAI.

Source code in geopandasai/wrapper.py
 6
 7
 8
 9
10
11
12
13
14
def read_file(
    filename: str,
    *args,
    **kwargs,
) -> GeoDataFrameAI:
    """
    Read a vector file (shapefile, GeoJSON, etc.) and return a GeoDataFrameAI.
    """
    return GeoDataFrameAI(gpd.read_file(filename, *args, **kwargs))

read_fileobj(fileobj, *args, **kwargs)

Read a file-like object and return a GeoDataFrameAI.

Source code in geopandasai/wrapper.py
51
52
53
54
55
56
57
58
59
def read_fileobj(
    fileobj,
    *args,
    **kwargs,
) -> GeoDataFrameAI:
    """
    Read a file-like object and return a GeoDataFrameAI.
    """
    return GeoDataFrameAI(gpd.read_file(fileobj, *args, **kwargs))

read_parquet(path, *args, **kwargs)

Read a Parquet file and return a GeoDataFrameAI.

Source code in geopandasai/wrapper.py
17
18
19
20
21
22
23
24
25
def read_parquet(
    path: str,
    *args,
    **kwargs,
) -> GeoDataFrameAI:
    """
    Read a Parquet file and return a GeoDataFrameAI.
    """
    return GeoDataFrameAI(gpd.read_parquet(path, *args, **kwargs))

read_postgis(sql, con, *args, **kwargs)

Read data from a PostGIS-enabled database and return a GeoDataFrameAI.

Source code in geopandasai/wrapper.py
39
40
41
42
43
44
45
46
47
48
def read_postgis(
    sql: str,
    con,
    *args,
    **kwargs,
) -> GeoDataFrameAI:
    """
    Read data from a PostGIS-enabled database and return a GeoDataFrameAI.
    """
    return GeoDataFrameAI(gpd.read_postgis(sql, con, *args, **kwargs))

to_feather(gdf, path, *args, **kwargs)

Write a GeoDataFrameAI to a Feather file.

Source code in geopandasai/wrapper.py
86
87
88
89
90
91
92
93
94
95
def to_feather(
    gdf: GeoDataFrameAI,
    path: str,
    *args,
    **kwargs,
) -> None:
    """
    Write a GeoDataFrameAI to a Feather file.
    """
    gdf.to_feather(path, *args, **kwargs)

to_file(gdf, filename, *args, **kwargs)

Write a GeoDataFrameAI to a file (shapefile, GeoPackage, GeoJSON, etc.).

Source code in geopandasai/wrapper.py
 98
 99
100
101
102
103
104
105
106
107
def to_file(
    gdf: GeoDataFrameAI,
    filename: str,
    *args,
    **kwargs,
) -> None:
    """
    Write a GeoDataFrameAI to a file (shapefile, GeoPackage, GeoJSON, etc.).
    """
    gdf.to_file(filename, *args, **kwargs)

to_parquet(gdf, path, *args, **kwargs)

Write a GeoDataFrameAI to a Parquet file.

Source code in geopandasai/wrapper.py
74
75
76
77
78
79
80
81
82
83
def to_parquet(
    gdf: GeoDataFrameAI,
    path: str,
    *args,
    **kwargs,
) -> None:
    """
    Write a GeoDataFrameAI to a Parquet file.
    """
    gdf.to_parquet(path, *args, **kwargs)

to_postgis(gdf, name, con, *args, **kwargs)

Write a GeoDataFrameAI to a PostGIS-enabled database.

Source code in geopandasai/wrapper.py
110
111
112
113
114
115
116
117
118
119
120
def to_postgis(
    gdf: GeoDataFrameAI,
    name: str,
    con,
    *args,
    **kwargs,
) -> None:
    """
    Write a GeoDataFrameAI to a PostGIS-enabled database.
    """
    gdf.to_postgis(name, con, *args, **kwargs)

Services

Cache Backend

ACacheBackend

Bases: ABC

Base class for all cache backends. This class defines the interface for caching mechanisms used in GeoPandasAI. It provides methods to get, set, clear, and reset cache entries. Subclasses should implement these methods to provide specific caching functionality.

Source code in geopandasai/external/cache/backend/base.py
 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
class ACacheBackend(abc.ABC):
    """
    Base class for all cache backends.
    This class defines the interface for caching mechanisms used in GeoPandasAI.
    It provides methods to get, set, clear, and reset cache entries.
    Subclasses should implement these methods to provide specific caching functionality.
    """

    def get_cache(self, key: str) -> Optional[bytes]:
        """
        Get the cached result for the given key.
        """
        pass

    def set_cache(self, key: str, value: bytes) -> None:
        """
        Set the cached result for the given key.
        """
        pass

    def clear_cache(self, key: str) -> None:
        """
        Clear the cached result for the given key.
        """
        pass

    def reset_cache(self) -> None:
        """
        Reset the cache.
        """
        pass
clear_cache(key)

Clear the cached result for the given key.

Source code in geopandasai/external/cache/backend/base.py
25
26
27
28
29
def clear_cache(self, key: str) -> None:
    """
    Clear the cached result for the given key.
    """
    pass
get_cache(key)

Get the cached result for the given key.

Source code in geopandasai/external/cache/backend/base.py
13
14
15
16
17
def get_cache(self, key: str) -> Optional[bytes]:
    """
    Get the cached result for the given key.
    """
    pass
reset_cache()

Reset the cache.

Source code in geopandasai/external/cache/backend/base.py
31
32
33
34
35
def reset_cache(self) -> None:
    """
    Reset the cache.
    """
    pass
set_cache(key, value)

Set the cached result for the given key.

Source code in geopandasai/external/cache/backend/base.py
19
20
21
22
23
def set_cache(self, key: str, value: bytes) -> None:
    """
    Set the cached result for the given key.
    """
    pass

FileSystemCacheBackend

Bases: ACacheBackend

FileSystemCacheBackend is a cache backend that stores cached results in the file system.

Source code in geopandasai/external/cache/backend/file_system.py
 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
38
39
40
41
42
class FileSystemCacheBackend(ACacheBackend):
    """
    FileSystemCacheBackend is a cache backend that stores cached results in the file system.
    """

    def __init__(self, cache_dir: str = "./.geopandasai_cache"):
        self.cache_dir = cache_dir
        os.makedirs(cache_dir, exist_ok=True)

    def get_cache(self, key: str) -> Optional[bytes]:
        try:
            with open(os.path.join(self.cache_dir, key), "rb") as f:
                return f.read()
        except FileNotFoundError:
            return None

    def set_cache(self, key: str, value: bytes) -> None:
        with open(os.path.join(self.cache_dir, key), "wb") as f:
            f.write(value)

    def clear_cache(self, key: str) -> None:
        try:
            os.remove(os.path.join(self.cache_dir, key))
        except FileNotFoundError:
            pass

    def reset_cache(self) -> None:
        if not os.path.exists(self.cache_dir):
            return
        for filename in os.listdir(self.cache_dir):
            file_path = os.path.join(self.cache_dir, filename)
            try:
                if os.path.isfile(file_path):
                    os.remove(file_path)
            except Exception as e:
                print(f"Error deleting file {file_path}: {e}")

InMemoryCacheBackend

Bases: ACacheBackend

InMemoryCacheBackend is a cache backend that stores cached results in memory.

Source code in geopandasai/external/cache/backend/in_memory.py
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class InMemoryCacheBackend(ACacheBackend):
    """
    InMemoryCacheBackend is a cache backend that stores cached results in memory.
    """

    def __init__(self):
        self.cache = {}

    def get_cache(self, key: str) -> Optional[bytes]:
        return self.cache.get(key)

    def set_cache(self, key: str, value: bytes) -> None:
        self.cache[key] = value

    def clear_cache(self, key: str) -> None:
        if key in self.cache:
            del self.cache[key]

    def reset_cache(self) -> None:
        self.cache.clear()

Code Execution Backend

ACodeExecutor

Bases: ABC

Abstract base class for executing Python code.

Source code in geopandasai/services/code/executor/base.py
 7
 8
 9
10
11
12
13
14
class ACodeExecutor(abc.ABC):
    """
    Abstract base class for executing Python code.
    """

    @abc.abstractmethod
    def execute(self, code: str, return_type: Type, *dfs: Iterable[GeoOrDataFrame]):
        pass

TrustedCodeExecutor

Bases: ACodeExecutor

A class for executing trusted Python code in a temporary file.

This class inherits from ACodeExecutor and provides functionality to safely execute Python code passed as a string. It ensures the result matches the expected return type and wraps GeoDataFrame results with GeoDataFrameAI for additional functionality.

Source code in geopandasai/services/code/executor/trusted.py
 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
38
39
40
41
42
43
44
45
46
47
class TrustedCodeExecutor(ACodeExecutor):
    """
    A class for executing trusted Python code in a temporary file.

    This class inherits from `ACodeExecutor` and provides functionality
    to safely execute Python code passed as a string. It ensures the
    result matches the expected return type and wraps GeoDataFrame results
    with `GeoDataFrameAI` for additional functionality.
    """

    def execute(self, code: str, return_type: Type, *dfs: Iterable[GeoOrDataFrame]):
        """
        Executes the provided Python code in a temporary file and returns the result.
        :param code: The Python code to execute as a string.
        :param return_type: The expected return type of the code execution.
        :param dfs: An iterable of GeoDataFrame or DataFrame objects to be passed as arguments to the executed code.
        :return:            The result of the executed code, wrapped in GeoDataFrameAI if it is a GeoDataFrame.
        """
        with tempfile.NamedTemporaryFile(delete=True, suffix=".py", mode="w") as f:
            f.write(code)
            f.flush()
            spec = importlib.util.spec_from_file_location("output", f.name)
            output_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(output_module)
            result = output_module.execute(
                **{f"df_{i+1}": df for i, df in enumerate(dfs)}
            )

            if not isinstance(result, return_type):
                raise TypeError(
                    f"Expected return type {return_type}, but got {type(result)}\n\nCode:\n\n{code}\n\n"
                )

            if isinstance(result, GeoOrDataFrame):
                from ....geodataframe_ai import GeoDataFrameAI

                result = GeoDataFrameAI(result)

            return result
execute(code, return_type, *dfs)

Executes the provided Python code in a temporary file and returns the result. :param code: The Python code to execute as a string. :param return_type: The expected return type of the code execution. :param dfs: An iterable of GeoDataFrame or DataFrame objects to be passed as arguments to the executed code. :return: The result of the executed code, wrapped in GeoDataFrameAI if it is a GeoDataFrame.

Source code in geopandasai/services/code/executor/trusted.py
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def execute(self, code: str, return_type: Type, *dfs: Iterable[GeoOrDataFrame]):
    """
    Executes the provided Python code in a temporary file and returns the result.
    :param code: The Python code to execute as a string.
    :param return_type: The expected return type of the code execution.
    :param dfs: An iterable of GeoDataFrame or DataFrame objects to be passed as arguments to the executed code.
    :return:            The result of the executed code, wrapped in GeoDataFrameAI if it is a GeoDataFrame.
    """
    with tempfile.NamedTemporaryFile(delete=True, suffix=".py", mode="w") as f:
        f.write(code)
        f.flush()
        spec = importlib.util.spec_from_file_location("output", f.name)
        output_module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(output_module)
        result = output_module.execute(
            **{f"df_{i+1}": df for i, df in enumerate(dfs)}
        )

        if not isinstance(result, return_type):
            raise TypeError(
                f"Expected return type {return_type}, but got {type(result)}\n\nCode:\n\n{code}\n\n"
            )

        if isinstance(result, GeoOrDataFrame):
            from ....geodataframe_ai import GeoDataFrameAI

            result = GeoDataFrameAI(result)

        return result

UntrustedCodeExecutor

Bases: TrustedCodeExecutor

A class for executing untrusted Python code with user confirmation.

This class extends TrustedCodeExecutor and adds an interactive confirmation step before executing the provided code. It highlights the code for better readability and ensures the user explicitly approves execution, making it suitable for scenarios where code safety is a concern.

Source code in geopandasai/services/code/executor/untrusted.py
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
class UntrustedCodeExecutor(TrustedCodeExecutor):
    """
    A class for executing untrusted Python code with user confirmation.

    This class extends `TrustedCodeExecutor` and adds an interactive confirmation
    step before executing the provided code. It highlights the code for better
    readability and ensures the user explicitly approves execution, making it
    suitable for scenarios where code safety is a concern.
    """

    def execute(self, code: str, return_type: Type, *dfs: Iterable[GeoOrDataFrame]):
        print("\n" + "-" * 30 + " Code Preview " + "-" * 30)
        print(_highlight_code(code))
        print("-" * 75 + "\n")

        confirmation = (
            input("⚠️  Do you want to execute this code? [y/N]: ").strip().lower()
        )
        if confirmation != "y":
            print("🚫 Execution aborted.")
            raise ValueError(
                "Code execution aborted by user because unsafe or untrusted code was signaled."
            )

        return super().execute(code, return_type, *dfs)