From 9ccb0a3bf4ec04df06f8be1e947750857e381da6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= <erik@bjareho.lt>
Date: Mon, 30 Oct 2023 18:06:31 +0100
Subject: [PATCH] feat: include context about included programs and packages

---
 gptme/prompts.py | 48 +++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 47 insertions(+), 1 deletion(-)

diff --git a/gptme/prompts.py b/gptme/prompts.py
index 30dc852b..541f6765 100644
--- a/gptme/prompts.py
+++ b/gptme/prompts.py
@@ -1,3 +1,4 @@
+import functools
 import logging
 import os
 import shutil
@@ -92,10 +93,16 @@ def initial_prompt(short: bool = False) -> Generator[Message, None, None]:
             hide=True,
         )
 
+    python_libraries = get_installed_python_libraries()
+    python_libraries_str = "\n".join(f"- {lib}" for lib in python_libraries)
+
+    shell_programs = get_installed_programs()
+    shell_programs_str = "\n".join(f"- {prog}" for prog in shell_programs)
+
     if include_tools:
         yield Message(
             "system",
-            """
+            f"""
 You are gptme, an AI assistant CLI tool powered powered by large language models that helps the user.
 You can run code and execute terminal commands on their local machine.
 The assistant shows the user to write code, interact with the system, and access the internet. The user will then choose to execute the suggested commands.
@@ -111,10 +118,16 @@ def initial_prompt(short: bool = False) -> Generator[Message, None, None]:
 
 When you send a message containing Python code (and is not a file block), it will be executed in a stateful environment. Python will respond with the output of the execution.
 
+The following libraries are available:
+{python_libraries_str}
+
 ## bash
 
 When you send a message containing bash code, it will be executed in a stateful bash shell. The shell will respond with the output of the execution.
 
+These programs are available, among others:
+{shell_programs_str}
+
 ## saving files
 
 When you send a message containing a code block, if the first line contains a filename, like "```hello.py" (a "file block"), the code block will be saved to that file.
@@ -208,3 +221,36 @@ def initial_prompt(short: bool = False) -> Generator[Message, None, None]:
             hide=True,
             pinned=True,
         )
+
+
+@functools.lru_cache
+def get_installed_python_libraries() -> set[str]:
+    """Check if a select list of Python libraries are installed."""
+    candidates = [
+        "numpy",
+        "pandas",
+        "matplotlib",
+        "seaborn",
+        "scipy",
+        "scikit-learn",
+        "statsmodels",
+        "pillow",
+    ]
+    installed = set()
+    for candidate in candidates:
+        try:
+            __import__(candidate)
+            installed.add(candidate)
+        except ImportError:
+            pass
+    return installed
+
+
+@functools.lru_cache
+def get_installed_programs() -> set[str]:
+    candidates = ["ffmpeg", "convert", "pandoc"]
+    installed = set()
+    for candidate in candidates:
+        if shutil.which(candidate) is not None:
+            installed.add(candidate)
+    return installed