From 9804f078d070c3b6c9fc8e59141e6f6b3b55a180 Mon Sep 17 00:00:00 2001
From: Jaimin Godhani <112328542+Jai0401@users.noreply.github.com>
Date: Mon, 2 Sep 2024 20:23:47 +0530
Subject: [PATCH] feat: automates package installation (#1780)

---
 mem0/configs/vector_stores/chroma.py | 16 +++++++++++++---
 mem0/llms/aws_bedrock.py             | 15 ++++++++++++---
 mem0/llms/groq.py                    | 15 ++++++++++++---
 mem0/llms/litellm.py                 | 15 ++++++++++++---
 mem0/llms/ollama.py                  | 16 +++++++++++++---
 mem0/llms/together.py                | 16 +++++++++++++---
 6 files changed, 75 insertions(+), 18 deletions(-)

diff --git a/mem0/configs/vector_stores/chroma.py b/mem0/configs/vector_stores/chroma.py
index 8054f3b5..23e182f1 100644
--- a/mem0/configs/vector_stores/chroma.py
+++ b/mem0/configs/vector_stores/chroma.py
@@ -1,3 +1,5 @@
+import subprocess
+import sys
 from typing import Optional, ClassVar, Dict, Any
 
 from pydantic import BaseModel, Field, model_validator
@@ -7,9 +9,17 @@ class ChromaDbConfig(BaseModel):
     try:
         from chromadb.api.client import Client
     except ImportError:
-        raise ImportError(
-            "Chromadb requires extra dependencies. Install with `pip install chromadb`"
-        ) from None
+        user_input: Any = input("The 'chromadb' library is required. Install it now? [y/N]: ")
+        if user_input.lower() == 'y':
+            try:
+                subprocess.check_call([sys.executable, "-m", "pip", "install", "chromadb"])
+                from chromadb.api.client import Client
+            except subprocess.CalledProcessError:
+                print("Failed to install 'chromadb'. Please install it manually using 'pip install chromadb'.")
+                sys.exit(1)
+        else:
+            print("The required 'chromadb' library is not installed.")
+            sys.exit(1)
     Client: ClassVar[type] = Client
 
     collection_name: str = Field("mem0", description="Default name for the collection")
diff --git a/mem0/llms/aws_bedrock.py b/mem0/llms/aws_bedrock.py
index 801be578..e2c1ea0f 100644
--- a/mem0/llms/aws_bedrock.py
+++ b/mem0/llms/aws_bedrock.py
@@ -1,3 +1,5 @@
+import subprocess
+import sys
 import os
 import json
 from typing import Dict, List, Optional, Any
@@ -5,9 +7,16 @@ from typing import Dict, List, Optional, Any
 try:
     import boto3
 except ImportError:
-    raise ImportError(
-        "AWS Bedrock requires extra dependencies. Install with `pip install boto3`"
-    ) from None
+    user_input = input("The 'boto3' library is required. Install it now? [y/N]: ")
+    if user_input.lower() == 'y':
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", "boto3"])
+            import boto3
+        except subprocess.CalledProcessError:
+            print("Failed to install 'boto3'. Please install it manually using 'pip install boto3'")
+            sys.exit(1)
+    else:
+        raise ImportError("The required 'boto3' library is not installed.")
 
 from mem0.llms.base import LLMBase
 from mem0.configs.llms.base import BaseLlmConfig
diff --git a/mem0/llms/groq.py b/mem0/llms/groq.py
index e5b9015e..c08bb511 100644
--- a/mem0/llms/groq.py
+++ b/mem0/llms/groq.py
@@ -1,3 +1,5 @@
+import subprocess
+import sys
 import os
 import json
 from typing import Dict, List, Optional
@@ -5,9 +7,16 @@ from typing import Dict, List, Optional
 try:
     from groq import Groq
 except ImportError:
-    raise ImportError(
-        "Groq requires extra dependencies. Install with `pip install groq`"
-    ) from None
+    user_input = input("The 'groq' library is required. Install it now? [y/N]: ")
+    if user_input.lower() == 'y':
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", "groq"])
+            from groq import Groq
+        except subprocess.CalledProcessError:
+            print("Failed to install 'groq'. Please install it manually using 'pip install groq'.")
+            sys.exit(1)
+    else:
+        raise ImportError("The required 'groq' library is not installed.")
 
 from mem0.llms.base import LLMBase
 from mem0.configs.llms.base import BaseLlmConfig
diff --git a/mem0/llms/litellm.py b/mem0/llms/litellm.py
index f3b32169..9d910370 100644
--- a/mem0/llms/litellm.py
+++ b/mem0/llms/litellm.py
@@ -1,12 +1,21 @@
+import subprocess
+import sys
 import json
 from typing import Dict, List, Optional
 
 try:
     import litellm
 except ImportError:
-    raise ImportError(
-        "litellm requires extra dependencies. Install with `pip install litellm`"
-    ) from None
+    user_input = input("The 'litellm' library is required. Install it now? [y/N]: ")
+    if user_input.lower() == 'y':
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", "litellm"])
+            import litellm
+        except subprocess.CalledProcessError:
+            print("Failed to install 'litellm'. Please install it manually using 'pip install litellm'.")
+            sys.exit(1)
+    else:
+        raise ImportError("The required 'litellm' library is not installed.")
 
 from mem0.llms.base import LLMBase
 from mem0.configs.llms.base import BaseLlmConfig
diff --git a/mem0/llms/ollama.py b/mem0/llms/ollama.py
index 144c0718..e86cac5e 100644
--- a/mem0/llms/ollama.py
+++ b/mem0/llms/ollama.py
@@ -1,11 +1,21 @@
+import subprocess
+import sys
 from typing import Dict, List, Optional
 
 try:
     from ollama import Client
 except ImportError:
-    raise ImportError(
-        "Ollama requires extra dependencies. Install with `pip install ollama`"
-    ) from None
+    user_input = input("The 'ollama' library is required. Install it now? [y/N]: ")
+    if user_input.lower() == 'y':
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", "ollama"])
+            from ollama import Client
+        except subprocess.CalledProcessError:
+            print("Failed to install 'ollama'. Please install it manually using 'pip install ollama'.")
+            sys.exit(1)
+    else:
+        print("The required 'ollama' library is not installed.")
+        sys.exit(1)
 
 from mem0.llms.base import LLMBase
 from mem0.configs.llms.base import BaseLlmConfig
diff --git a/mem0/llms/together.py b/mem0/llms/together.py
index 1d827d3e..046c02be 100644
--- a/mem0/llms/together.py
+++ b/mem0/llms/together.py
@@ -1,3 +1,5 @@
+import subprocess
+import sys
 import os
 import json
 from typing import Dict, List, Optional
@@ -5,9 +7,17 @@ from typing import Dict, List, Optional
 try:
     from together import Together
 except ImportError:
-    raise ImportError(
-        "Together requires extra dependencies. Install with `pip install together`"
-    ) from None
+    user_input = input("The 'together' library is required. Install it now? [y/N]: ")
+    if user_input.lower() == 'y':
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", "together"])
+            from together import Together
+        except subprocess.CalledProcessError:
+            print("Failed to install 'together'. Please install it manually using 'pip install together'.")
+            sys.exit(1)
+    else:
+        print("The required 'together' library is not installed.")
+        sys.exit(1)
 
 from mem0.llms.base import LLMBase
 from mem0.configs.llms.base import BaseLlmConfig