make langchain/chroma optional (when using rag)

Author: andreashappe

pyproject.toml

@@ -46,12 +46,6 @@ dependencies = [
     'websockets == 13.1',
     'faker',
     'fpdf',
-    'langchain_core',
-    'langchain_community',
-    'langchain_chroma',
-    'langchain_openai',
-    'markdown',
-    'chromadb',
 ]
 
 [project.urls]
@@ -70,17 +64,26 @@ where = ["src"]
 [tool.pytest.ini_options]
 pythonpath = "src"
 addopts = ["--import-mode=importlib"]
+
 [project.optional-dependencies]
-testing = ['pytest', 'pytest-mock', 'pandas', 'faker',     'langchain_core']
+testing = [
+	'pytest',
+	'pytest-mock',
+	'faker',
+	'langchain_core'
+]
+
 dev = [
 	'ruff',
 ]
-rag-usecase = [
-    'langchain-community',
+
+rag = [
+	'langchain_core',
+	'langchain-community',
+	'langchain-chroma',
 	'langchain-openai',
 	'markdown',
 	'chromadb',
-	'langchain-chroma',
 ]
 
 [project.scripts]

src/hackingBuddyGPT/usecases/linux_privesc.py

@@ -9,10 +9,13 @@
 from hackingBuddyGPT.usecases.usecase import use_case
 from hackingBuddyGPT.utils import llm_util
 from hackingBuddyGPT.utils.logging import log_conversation
-from hackingBuddyGPT.utils.rag import RagBackground
+from hackingBuddyGPT.utils.rag import has_langchain
 from hackingBuddyGPT.utils.connectors.ssh_connection import SSHConnection
 from hackingBuddyGPT.utils.shell_root_detection import got_root
 
+if has_langchain():
+    from hackingBuddyGPT.utils.rag import RagBackground
+
 template_analyze = Template("""Your task is to analyze the result of an executed command to determina 
 a way to escalate your privileges into a root shell. Describe your findings including all needed
 information while being as concise as possible.
@@ -132,6 +135,10 @@ def init(self):
         guidance = []
 
         if self.rag_path != '':
+            if not has_langchain():
+                self.log.console.print("[red]RAG path provided but langchain is not installed. Please install langchain to use RAG functionality, e.g., through `pip install -e .\[rag]`.[/red]")
+                raise ImportError("langchain is not installed")
+
             self._enable_rag = True
             self._rag_data = RagBackground(self.rag_path, self.llm)
 
@@ -254,4 +261,4 @@ def check_success(self, cmd:str, result:str) -> bool:
         ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
         last_line = result.split("\n")[-1] if result else ""
         last_line = ansi_escape.sub("", last_line)
-        return got_root(self.conn.hostname, last_line)
+        return got_root(self.conn.hostname, last_line)

src/hackingBuddyGPT/utils/rag.py

@@ -1,7 +1,15 @@
-from langchain_community.document_loaders import DirectoryLoader, TextLoader
-from langchain_chroma import Chroma
-from langchain_openai import OpenAIEmbeddings
-from langchain_text_splitters import MarkdownTextSplitter
+try:
+    from langchain_community.document_loaders import DirectoryLoader, TextLoader
+    from langchain_chroma import Chroma
+    from langchain_openai import OpenAIEmbeddings
+    from langchain_text_splitters import MarkdownTextSplitter
+except ImportError:
+    _has_langchain = False
+else:
+    _has_langchain = True
+
+def has_langchain():
+    return _has_langchain
 
 class RagBackground: