aryn-ai · austintlee · Feb 14, 2025 · Feb 8, 2025 · Feb 11, 2025 · Feb 11, 2025
diff --git a/lib/sycamore/sycamore/connectors/aryn/ArynReader.py b/lib/sycamore/sycamore/connectors/aryn/ArynReader.py
@@ -150,10 +150,16 @@ def _to_doc(self, doc: dict[str, Any]) -> dict[str, Any]:
         client = self.Client.from_client_params(self._client_params)
         aryn_client = client._client
 
+        doc_id = doc["doc_id"]
         doc = aryn_client.get_doc(self._query_params.docset_id, doc["doc_id"])
         elements = doc.get("elements", [])
         document = Document(**doc)
-        document.data["elements"] = [create_element(**element) for element in elements]
+        document.doc_id = doc_id
+        document.data["elements"] = []
+        for json_element in elements:
+            element = create_element(**json_element)
+            element.data["doc_id"] = json_element["id"]
+            document.data["elements"].append(element)
         return {"doc": Document.serialize(document)}
 
     def execute(self, **kwargs) -> "Dataset":

diff --git a/lib/sycamore/sycamore/transforms/extract_schema.py b/lib/sycamore/sycamore/transforms/extract_schema.py
@@ -161,7 +161,7 @@ def __init__(
         llm: LLM,
         schema_name: Optional[str] = None,
         schema: Optional[Union[dict[str, str], Schema]] = None,
-        num_of_elements: int = 10,
+        num_of_elements: Optional[int] = None,
         prompt_formatter: Callable[[list[Element]], str] = element_list_formatter,
     ):
         super().__init__()
@@ -221,7 +221,8 @@ def as_llm_map(self, child: Optional[Node], **kwargs) -> Node:
 
             if self._schema_name is not None:
                 prompt = prompt.set(entity=self._schema_name)
-        prompt = prompt.set(num_elements=self._num_of_elements)
+        if self._num_of_elements is not None:
+            prompt = prompt.set(num_elements=self._num_of_elements)
         if self._prompt_formatter is not element_list_formatter:
             prompt = prompt.set(prompt_formatter=self._prompt_formatter)
 

diff --git a/lib/sycamore/sycamore/writer.py b/lib/sycamore/sycamore/writer.py
@@ -858,9 +858,9 @@ def aryn(
                 raise e
         client_params = ArynWriterClientParams(aryn_url, aryn_api_key)
         target_params = ArynWriterTargetParams(docset_id)
-        ds = ArynWriter(self.plan, client_params=client_params, target_params=target_params, **kwargs)
+        writer: Node = ArynWriter(self.plan, client_params=client_params, target_params=target_params, **kwargs)
 
-        return self._maybe_execute(ds, True)
+        return self._maybe_execute(writer, True)
 
     def _maybe_execute(
         self, node: Node, execute: bool, client: Optional[BaseDBWriter.Client] = None