fix: check for common datascience packages, added them as optional deps, warn if not available

Author: ErikBjare

gptme/cli.py

@@ -20,6 +20,7 @@
 to do so, it needs to be able to store and query past conversations in a database.
 """
 # The above may be used as a prompt for the agent.
+
 import atexit
 import importlib.metadata
 import io
@@ -50,7 +51,12 @@
 )
 from .prompts import initial_prompt_single_message
 from .tabcomplete import register_tabcomplete
-from .tools import execute_msg, execute_python, execute_shell
+from .tools import (
+    execute_msg,
+    execute_python,
+    execute_shell,
+    init_tools,
+)
 from .tools.shell import get_shell
 from .tools.summarize import summarize
 from .tools.useredit import edit_text_with_editor
@@ -266,6 +272,7 @@ def main(
     load_dotenv()
     _load_readline_history()
     init_llm(llm)  # set up API_KEY and API_BASE
+    init_tools()
 
     if "PYTEST_CURRENT_TEST" in os.environ:
         interactive = False

gptme/tools/__init__.py

@@ -2,9 +2,9 @@
 from typing import Generator
 
 from ..message import Message
-from .python import execute_python
-from .shell import execute_shell
+from .python import execute_python, init_python
 from .save import execute_save
+from .shell import execute_shell
 from .summarize import summarize
 
 logger = logging.getLogger(__name__)
@@ -46,3 +46,8 @@ def execute_codeblock(codeblock: str, ask: bool) -> Generator[Message, None, Non
         logger.warning(
             f"Unknown codeblock type '{lang_or_fn}', neither supported language or filename."
         )
+
+
+def init_tools() -> None:
+    """Runs initialization logic for tools."""
+    init_python()

gptme/tools/python.py

@@ -1,29 +1,34 @@
 import ast
 import io
 from contextlib import redirect_stderr, redirect_stdout
+from logging import getLogger
 from typing import Generator
 
 from ..message import Message
 from ..util import ask_execute, print_preview
 
+logger = getLogger(__name__)
 locals_ = {}  # type: ignore
 
 
+def init_python():
+    check_available_packages()
+
+
 def execute_python(code: str, ask: bool) -> Generator[Message, None, None]:
     """Executes a python codeblock and returns the output."""
     code = code.strip()
     if ask:
         print_preview(code, "python")
         confirm = ask_execute()
         print()
+        if not confirm:
+            # early return
+            yield Message("system", "Aborted, user chose not to run command.")
+            return
     else:
         print("Skipping confirmation")
 
-    if ask and not confirm:
-        # early return
-        yield Message("system", "Aborted, user chose not to run command.")
-        return
-
     # remove blank lines
     code = "\n".join([line for line in code.split("\n") if line.strip()])
 
@@ -50,6 +55,21 @@ def execute_python(code: str, ask: bool) -> Generator[Message, None, None]:
     yield Message("system", "Executed code block.\n\n" + output)
 
 
+def check_available_packages():
+    """Checks that essentials like numpy, pandas, matplotlib are available."""
+    expected = ["numpy", "pandas", "matplotlib"]
+    missing = []
+    for package in expected:
+        try:
+            __import__(package)
+        except ImportError:
+            missing.append(package)
+    if missing:
+        logger.warning(
+            f"Missing packages: {', '.join(missing)}. Install them with `pip install gptme-python -E datascience`"
+        )
+
+
 def test_execute_python():
     assert execute_python("1 + 1", ask=False) == ">>> 1 + 1\n2\n"
     assert execute_python("a = 2\na", ask=False) == ">>> a = 2\n>>> a\n2\n"