fix: import not on top since community package

libs/community/langchain_community/document_loaders/needle.py

@@ -2,8 +2,6 @@
 
 from langchain_core.document_loaders.base import BaseLoader
 from langchain_core.documents import Document
-from needle.v1 import NeedleClient
-from needle.v1.models import FileToAdd
 
 
 class NeedleLoader(BaseLoader):
@@ -57,6 +55,7 @@ def __init__(
         Raises:
             ValueError: If the collection ID is not provided.
         """
+        from needle.v1 import NeedleClient
         super().__init__()
         self.needle_api_key = needle_api_key
         self.collection_id = collection_id
@@ -92,6 +91,8 @@ def add_files(self, files: dict) -> None:
         Raises:
             ValueError: If the collection is not properly initialized.
         """
+        from needle.v1.models import FileToAdd
+
         self._get_collection()
         assert self.client is not None, "NeedleClient must be initialized."
 

libs/community/langchain_community/retrievers/needle.py

@@ -3,7 +3,6 @@
 from langchain_core.callbacks import CallbackManagerForRetrieverRun
 from langchain_core.documents import Document
 from langchain_core.retrievers import BaseRetriever
-from needle.v1 import NeedleClient
 from pydantic import BaseModel, Field
 
 
@@ -39,6 +38,7 @@ class NeedleRetriever(BaseRetriever, BaseModel):
             for doc in results:
                 print(doc.page_content)
     """
+    from needle.v1 import NeedleClient
 
     needle_api_key: Optional[str] = Field(None, description="Needle API Key")
     collection_id: Optional[str] = Field(
@@ -52,6 +52,8 @@ def _initialize_client(self) -> None:
 
         If a client instance is already provided, this method does nothing.
         """
+        from needle.v1 import NeedleClient
+
         if not self.client:
             self.client = NeedleClient(api_key=self.needle_api_key)
 

libs/community/tests/unit_tests/document_loaders/test_needle.py

@@ -38,7 +38,7 @@ def list(self, collection_id: str):
 
 @pytest.mark.requires("needle-python")
 def test_add_and_fetch_files(mocker: MockerFixture):
-    """Test adding and fetching files using the NeedleLoader with a mock NeedleClient."""
+    """Test adding and fetching files using the NeedleLoader with a mock."""
     # Mock the NeedleClient to use the mock implementation
     mocker.patch("needle.v1.NeedleClient", new=MockNeedleClient)