From 1d7c33fb4d36d73a6328b6a949ebfa6631ffb696 Mon Sep 17 00:00:00 2001
From: Felix Altenberger <felix@zenml.io>
Date: Wed, 22 Nov 2023 09:18:50 +0100
Subject: [PATCH] Fix PyYAML and linter errors

---
 pyproject.toml                               | 1 +
 scripts/install-zenml-dev.sh                 | 9 +++++++--
 src/zenml/integrations/langchain/__init__.py | 2 +-
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index e710fba3473..5451ab4c993 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -349,6 +349,7 @@ module = [
     "datasets.*",
     "pyngrok.*",
     "cloudpickle.*",
+    "matplotlib.*",
 ]
 ignore_missing_imports = true
 
diff --git a/scripts/install-zenml-dev.sh b/scripts/install-zenml-dev.sh
index 842cb57d6fc..7a442b94171 100755
--- a/scripts/install-zenml-dev.sh
+++ b/scripts/install-zenml-dev.sh
@@ -48,11 +48,16 @@ install_integrations() {
     zenml integration export-requirements \
         --output-file integration-requirements.txt \
         $ignore_integrations_args
+
+    # pin pyyaml>=6.0.1
+    echo "" >> integration-requirements.txt
+    echo "pyyaml>=6.0.1" >> integration-requirements.txt
+
     pip install -r integration-requirements.txt
     rm integration-requirements.txt
 
-    # install langchain and llama_index integrations separately
-    zenml integration install -y langchain llama_index
+    # install langchain separately
+    zenml integration install -y langchain
 }
 
 
diff --git a/src/zenml/integrations/langchain/__init__.py b/src/zenml/integrations/langchain/__init__.py
index 7e3d6217b7e..6c01c2968b3 100644
--- a/src/zenml/integrations/langchain/__init__.py
+++ b/src/zenml/integrations/langchain/__init__.py
@@ -25,7 +25,7 @@ class LangchainIntegration(Integration):
     """Definition of langchain integration for ZenML."""
 
     NAME = LANGCHAIN
-    REQUIREMENTS = ["langchain>=0.0.116"]
+    REQUIREMENTS = ["langchain>=0.0.116", "pyyaml>=6.0.1"]
 
     @classmethod
     def activate(cls) -> None: