aryn-ai · mdwelsh · Oct 18, 2024 · Oct 18, 2024 · Oct 18, 2024 · Oct 18, 2024
diff --git a/lib/sycamore/sycamore/materialize.py b/lib/sycamore/sycamore/materialize.py
@@ -97,7 +97,7 @@ def __init__(
 
         self._maybe_anonymous()
 
-    def _maybe_anonymous(self):
+    def _maybe_anonymous(self) -> None:
         if self._root is None:
             return
         from pyarrow.fs import S3FileSystem

diff --git a/lib/sycamore/sycamore/tests/unit/transforms/test_graph_entity_extractor.py b/lib/sycamore/sycamore/tests/unit/transforms/test_graph_entity_extractor.py
@@ -1,7 +1,7 @@
 from typing import Optional
 from pydantic import BaseModel
 import sycamore
-from sycamore.data.document import Document
+from sycamore.data.document import Document, HierarchicalDocument
 from sycamore.data.element import Element
 from sycamore.llms.llms import LLM
 from sycamore.reader import DocSetReader
@@ -74,7 +74,7 @@ async def generate_async(self, *, prompt_kwargs: dict, llm_kwargs: Optional[dict
             }
             """
 
-    def test_entity_extractor(self):
+    def test_entity_extractor(self) -> None:
         context = sycamore.init()
         reader = DocSetReader(context)
         ds = reader.document(self.docs)
@@ -83,10 +83,10 @@ class Company(BaseModel):
             name: str
 
         llm = self.MockLLM()
-        ds = ds.extract_document_structure(structure=StructureBySection).extract_graph_entities(
+        ds = ds.extract_document_structure(structure=StructureBySection()).extract_graph_entities(
             [EntityExtractor(llm=llm, entities=[Company])]
         )
-        docs = ds.take_all()
+        docs = [HierarchicalDocument(doc.data) for doc in ds.take_all()]
 
         for doc in docs:
             for section in doc.children:

diff --git a/lib/sycamore/sycamore/tests/unit/transforms/test_graph_relationship_extractor.py b/lib/sycamore/sycamore/tests/unit/transforms/test_graph_relationship_extractor.py
@@ -1,7 +1,7 @@
 from typing import Optional
 from pydantic import BaseModel
 import sycamore
-from sycamore.data.document import Document
+from sycamore.data.document import Document, HierarchicalDocument
 from sycamore.data.element import Element
 from sycamore.llms.llms import LLM
 from sycamore.reader import DocSetReader
@@ -80,8 +80,8 @@ class MockRelationshipLLM(LLM):
         def __init__(self):
             super().__init__(model_name="mock_model")
 
-        def generate(self, *, prompt_kwargs: dict, llm_kwargs: Optional[dict] = None):
-            pass
+        def generate(self, *, prompt_kwargs: dict, llm_kwargs: Optional[dict] = None) -> str:
+            return ""
 
         def is_chat_mode(self):
             return True
@@ -94,7 +94,7 @@ async def generate_async(self, *, prompt_kwargs: dict, llm_kwargs: Optional[dict
             }
             """
 
-    def test_relationship_extractor(self):
+    def test_relationship_extractor(self) -> None:
         context = sycamore.init()
         reader = DocSetReader(context)
         ds = reader.document(self.docs)
@@ -107,11 +107,11 @@ class Competes(BaseModel):
             end: Company
 
         ds = (
-            ds.extract_document_structure(structure=StructureBySection)
+            ds.extract_document_structure(structure=StructureBySection())
             .extract_graph_entities([EntityExtractor(self.MockEntityLLM(), [Company])])
             .extract_graph_relationships([RelationshipExtractor(self.MockRelationshipLLM(), [Competes])])
         )
-        docs = ds.take_all()
+        docs = [HierarchicalDocument(doc.data) for doc in ds.take_all()]
 
         for doc in docs:
             for section in doc.children:

diff --git a/lib/sycamore/sycamore/tests/unit/transforms/test_resolve_graph_entities.py b/lib/sycamore/sycamore/tests/unit/transforms/test_resolve_graph_entities.py
@@ -1,7 +1,7 @@
 from typing import Optional
 from pydantic import BaseModel
 import sycamore
-from sycamore.data.document import Document
+from sycamore.data.document import Document, HierarchicalDocument
 from sycamore.data.element import Element
 from sycamore.llms.llms import LLM
 from sycamore.reader import DocSetReader
@@ -93,7 +93,7 @@ async def generate_async(self, *, prompt_kwargs: dict, llm_kwargs: Optional[dict
             }
             """
 
-    def test_resolve_entities(self):
+    def test_resolve_entities(self) -> None:
         context = sycamore.init()
         reader = DocSetReader(context)
         ds = reader.document(self.docs)
@@ -106,12 +106,12 @@ class Competes(BaseModel):
             end: Company
 
         ds = (
-            ds.extract_document_structure(structure=StructureBySection)
+            ds.extract_document_structure(structure=StructureBySection())
             .extract_graph_entities([EntityExtractor(self.MockEntityLLM(), [Company])])
             .extract_graph_relationships([RelationshipExtractor(self.MockRelationshipLLM(), [Competes])])
             .resolve_graph_entities(resolvers=[], resolve_duplicates=True)
         )
-        docs = ds.take_all()
+        docs = [HierarchicalDocument(doc.data) for doc in ds.take_all()]
 
         for doc in docs:
             if doc.data.get("EXTRACTED_NODES", False) is True:

diff --git a/lib/sycamore/sycamore/tests/unit/utils/test_import_utils.py b/lib/sycamore/sycamore/tests/unit/utils/test_import_utils.py
@@ -1,5 +1,4 @@
-from typing import TYPE_CHECKING
-from mypy import api
+from typing_extensions import assert_type
 from sycamore.utils.import_utils import requires_modules
 
 
@@ -9,13 +8,8 @@ def require_fn() -> int:
     return 42
 
 
-res = require_fn()
-if TYPE_CHECKING:
-    reveal_type(res)  # noqa: F821
-
-
 # This test fails prior to adding generic (ParamSpec and TypeVar) type annotations to the
 # requires_modules decorator, as the revealed type is "Any".
-def test_mypy_type():
-    mypy_res = api.run([__file__])
-    assert 'Revealed type is "builtins.int"' in mypy_res[0]
+def test_mypy_type() -> None:
+    res = require_fn()
+    assert_type(res, int)
diff --git a/lib/sycamore/sycamore/tests/unit/utils/test_pydantic_pickling.py b/lib/sycamore/sycamore/tests/unit/utils/test_pydantic_pickling.py
@@ -4,7 +4,7 @@
 from pydantic import BaseModel, Field
 
 
-def test_pydantic_picklng():
+def test_pydantic_picklng() -> None:
     class BoardMember(BaseModel):
         name: str
         votes_for: Optional[int]

diff --git a/lib/sycamore/sycamore/transforms/extract_graph_entities.py b/lib/sycamore/sycamore/transforms/extract_graph_entities.py
@@ -3,7 +3,7 @@
 from collections import defaultdict
 import hashlib
 import io
-from typing import Dict, Any, List
+from typing import Dict, Any, List, Type
 from sycamore.plan_nodes import Node
 from sycamore.transforms.map import Map
 from sycamore.data import HierarchicalDocument
@@ -43,7 +43,7 @@ class EntityExtractor(GraphEntityExtractor):
     def __init__(
         self,
         llm: LLM,
-        entities: list[BaseModel],
+        entities: List[Type[BaseModel]],
         prompt: str = GraphEntityExtractorPrompt.user,
         split_calls: bool = False,
     ):

diff --git a/lib/sycamore/sycamore/transforms/extract_graph_relationships.py b/lib/sycamore/sycamore/transforms/extract_graph_relationships.py
@@ -1,22 +1,23 @@
 from abc import ABC, abstractmethod
+import asyncio
 import base64
 from collections import defaultdict
 from enum import Enum
 import hashlib
+import json
 import io
+import logging
+from typing import Dict, Any, List, Type
+import uuid
+
 from PIL import Image
-from typing import Dict, Any, List
+from pydantic import BaseModel, create_model
+
 from sycamore.plan_nodes import Node
 from sycamore.transforms.map import Map
 from sycamore.data import HierarchicalDocument
 from sycamore.llms import LLM
 from sycamore.llms.prompts import GraphRelationshipExtractorPrompt
-from pydantic import BaseModel, create_model
-import asyncio
-
-import json
-import uuid
-import logging
 
 logger = logging.getLogger(__name__)
 
@@ -45,7 +46,7 @@ class RelationshipExtractor(GraphRelationshipExtractor):
     def __init__(
         self,
         llm: LLM,
-        relationships: list[BaseModel],
+        relationships: List[Type[BaseModel]],
         prompt: str = GraphRelationshipExtractorPrompt.user,
         split_calls: bool = False,
     ):