chore: make langchain optional

byaldi/RAGModel.py

@@ -1,13 +1,18 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
 
-from langchain_core.retrievers import BaseRetriever
 from PIL import Image
 
 from byaldi.colpali import ColPaliModel
-from byaldi.integrations import ByaldiLangChainRetriever
+
 from byaldi.objects import Result
 
+# Optional langchain integration
+try:
+    from byaldi.integrations import ByaldiLangChainRetriever
+except ImportError:
+    pass
+
 
 class RAGMultiModalModel:
     """
@@ -52,7 +57,10 @@ def from_pretrained(
         """
         instance = cls()
         instance.model = ColPaliModel.from_pretrained(
-            pretrained_model_name_or_path, index_root=index_root, device=device, verbose=verbose
+            pretrained_model_name_or_path,
+            index_root=index_root,
+            device=device,
+            verbose=verbose,
         )
         return instance
 
@@ -168,5 +176,5 @@ def search(
     def get_doc_ids_to_file_names(self):
         return self.model.get_doc_ids_to_file_names()
 
-    def as_langchain_retriever(self, **kwargs: Any) -> BaseRetriever:
-        return ByaldiLangChainRetriever(model=self, kwargs=kwargs)
+    def as_langchain_retriever(self, **kwargs: Any):
+        return ByaldiLangChainRetriever(model=self, kwargs=kwargs)

byaldi/integrations/__init__.py

@@ -1,3 +1,8 @@
-from byaldi.integrations._langchain import ByaldiLangChainRetriever
+_all__ = []
 
-__all__ = ["ByaldiLangChainRetriever"]
+try:
+    from byaldi.integrations._langchain import ByaldiLangChainRetriever
+
+    _all__.append("ByaldiLangChainRetriever")
+except ImportError:
+    pass

pyproject.toml

@@ -36,12 +36,12 @@ dependencies = [
     "srsly",
     "torch",
     "transformers",
-    "langchain-core",
 ]
 
 [project.optional-dependencies]
 dev = ["pytest>=7.4.0", "ruff>=0.1.9"]
 server = ["uvicorn", "fastapi"]
+langchain = ["langchain-core"]
 
 [project.urls]
 "Homepage" = "https://github.com/answerdotai/byaldi"