Allow PlotGenerator init without suggestions; make Groq dependency optional

plotsense/explanations/explanations.py

@@ -3,7 +3,10 @@
 import matplotlib.pyplot as plt
 from typing import Union, Optional, Dict, List
 from dotenv import load_dotenv
-from groq import Groq
+try:
+    from groq import Groq
+except ImportError:
+    Groq = None
 import warnings
 import builtins
 
@@ -59,6 +62,8 @@ def _validate_keys(self):
         }
 
         for service in ['groq']:
+            if service == 'groq' and Groq is None:
+                continue
             if not self.api_keys.get(service):
                 if self.interactive:
                     try:
@@ -84,6 +89,8 @@ def _initialize_clients(self):
         self.clients = {}
         if self.api_keys.get('groq'):
             try:
+                if Groq is None:
+                    raise ImportError("groq")
                 self.clients['groq'] = Groq(api_key=self.api_keys['groq'])
             except Exception as e:
                 warnings.warn(f"Could not initialize Groq client: {e}", ImportWarning)

plotsense/plot_generator/generator.py

@@ -21,12 +21,13 @@ def __init__(self, data: pd.DataFrame, suggestions: Optional[pd.DataFrame] = Non
             raise TypeError("Data must be a pandas DataFrame")
         if data.empty:
             raise ValueError("DataFrame is empty")
-        if not isinstance(suggestions, pd.DataFrame):
-            raise TypeError("Suggestions must be a pandas DataFrame")
-        if suggestions.empty:
-            raise ValueError("Suggestions DataFrame is empty")
-        if 'plot_type' not in suggestions.columns or 'variables' not in suggestions.columns:
-            raise ValueError("Suggestions DataFrame must contain 'plot_type' and 'variables' columns")
+        if suggestions is not None:
+            if not isinstance(suggestions, pd.DataFrame):
+                raise TypeError("Suggestions must be a pandas DataFrame")
+            if suggestions.empty:
+                raise ValueError("Suggestions DataFrame is empty")
+            if 'plot_type' not in suggestions.columns or 'variables' not in suggestions.columns:
+                raise ValueError("Suggestions DataFrame must contain 'plot_type' and 'variables' columns")
 
         self.data = data.copy()
         self.suggestions = suggestions
@@ -49,6 +50,8 @@ def generate_plot(self, suggestion_index: int, **kwargs) -> plt.Figure:
             raise TypeError("Suggestion index must be an integer")
         if not isinstance(kwargs, dict):
             raise TypeError("Additional arguments must be provided as a dictionary")
+        if self.suggestions is None:
+            raise ValueError("No suggestions available to generate a plot")
         if self.suggestions.empty:
             raise ValueError("No suggestions available to generate a plot")
         if self.data.empty:

plotsense/visual_suggestion/suggestions.py

@@ -10,7 +10,10 @@
 import textwrap
 import builtins
 from pprint import pprint
-from groq import Groq
+try:
+    from groq import Groq
+except ImportError:
+    Groq = None
 
 
 load_dotenv()
@@ -73,6 +76,8 @@ def _validate_keys(self):
         }
 
         for service in ['groq']:
+            if service == 'groq' and Groq is None:
+                continue
             if not self.api_keys.get(service):
                 if self.interactive:
                     try:
@@ -98,6 +103,8 @@ def _initialize_clients(self):
         self.clients = {}
         if self.api_keys.get('groq'):
             try:
+                if Groq is None:
+                    raise ImportError("groq")
                 self.clients['groq'] = Groq(api_key=self.api_keys['groq'])
             except ImportError:
                 warnings.warn("Groq Python client not installed. pip install groq")

test/test_explanations.py

@@ -13,6 +13,7 @@
 
 # Import the class to test
 from plotsense import PlotExplainer, explainer
+import plotsense.explanations.explanations as explanations_module
 
 # Test data setup
 @pytest.fixture
@@ -46,7 +47,7 @@ def mock_groq_completion():
 @pytest.fixture
 def mock_groq_client():
     """Fixture that mocks the Groq client"""
-    with patch('groq.Groq') as mock:
+    with patch('plotsense.explanations.explanations.Groq') as mock:
         mock_instance = MagicMock()
         mock.return_value = mock_instance
         yield mock_instance
@@ -89,12 +90,18 @@ def test_init_without_api_keys_interactive(self):
             del os.environ['GROQ_API_KEY']
 
     def test_init_without_api_keys_non_interactive(self):
-        with pytest.raises(ValueError,  match="API key is required"):
+        if explanations_module.Groq is None:
             PlotExplainer(api_keys={}, interactive=False)
+        else:
+            with pytest.raises(ValueError, match="API key is required"):
+                PlotExplainer(api_keys={}, interactive=False)
 
     def test_validate_keys_missing(self):
-        with pytest.raises(ValueError,  match="API key is required"):
+        if explanations_module.Groq is None:
             PlotExplainer(api_keys={}, interactive=False)
+        else:
+            with pytest.raises(ValueError, match="API key is required"):
+                PlotExplainer(api_keys={}, interactive=False)
 
     def test_initialize_clients(self, mock_groq_client):
         explainer = PlotExplainer(api_keys={'groq': 'test_key'}, interactive=False)

test/test_plotgen.py

@@ -72,6 +72,13 @@ def reset_plot_generator_instance():
 
 # Unit Tests
 class TestPlotGeneratorUnit:
+    def test_init_plot_generator_without_suggestions(self, sample_dataframe):
+        pg = PlotGenerator(sample_dataframe, suggestions=None)
+        assert pg.data.equals(sample_dataframe)
+        assert pg.suggestions is None
+        with pytest.raises(ValueError, match="No suggestions available"):
+            pg.generate_plot(0)
+
     def test_init_plot_generator(self, sample_dataframe, sample_suggestions):
         pg = PlotGenerator(sample_dataframe, sample_suggestions)
         assert pg.data.equals(sample_dataframe)

test/test_suggestions.py

@@ -12,6 +12,7 @@
 
 # SUT
 from plotsense.visual_suggestion.suggestions import VisualizationRecommender
+import plotsense.visual_suggestion.suggestions as suggestions_module
 
 load_dotenv()  # make .env vars visible for tests
 SEED = 42
@@ -74,8 +75,11 @@ def test_init_with_keys(self):
 
     def test_missing_key_noninteractive(self):
         """Test initialization with missing API key in non-interactive mode"""
-        with pytest.raises(ValueError, match="GROQ API key is required"):
+        if suggestions_module.Groq is None:
             VisualizationRecommender(api_keys={"groq": None}, interactive=False)
+        else:
+            with pytest.raises(ValueError, match="GROQ API key is required"):
+                VisualizationRecommender(api_keys={"groq": None}, interactive=False)
 
     def test_missing_key_interactive(self, monkeypatch):
         """Test interactive key input"""