unstructured-ingest 1.0.7__tar.gz → 1.0.11__tar.gz

{unstructured_ingest-1.0.7 → unstructured_ingest-1.0.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unstructured_ingest
-Version: 1.0.7
+Version: 1.0.11
 Summary: Local ETL data pipeline to get data RAG ready
 Author-email: Unstructured Technologies <devops@unstructuredai.io>
 License-Expression: Apache-2.0

{unstructured_ingest-1.0.7 → unstructured_ingest-1.0.11}/pyproject.toml

@@ -143,6 +143,7 @@ test = [
     "vertexai",
     "pyiceberg",
     "pyarrow",
+    "networkx"
 ]
 # Add constraints needed for CI
 ci = [

unstructured_ingest-1.0.11/unstructured_ingest/__version__.py

@@ -0,0 +1 @@
+__version__ = "1.0.11"  # pragma: no cover

unstructured_ingest-1.0.11/unstructured_ingest/data_types/entities.py

@@ -0,0 +1,17 @@
+from pydantic import BaseModel, Field
+
+
+class Entity(BaseModel):
+    type: str
+    entity: str
+
+
+class EntityRelationship(BaseModel):
+    to: str
+    from_: str = Field(..., alias="from")
+    relationship: str
+
+
+class EntitiesData(BaseModel):
+    items: list[Entity] = Field(default_factory=list)
+    relationships: list[EntityRelationship] = Field(default_factory=list)

{unstructured_ingest-1.0.7 → unstructured_ingest-1.0.11}/unstructured_ingest/processes/connectors/google_drive.py

@@ -153,7 +153,13 @@ class GoogleDriveIndexer(Indexer):
         """
         try:
             # A very minimal call: list 1 file from the drive.
-            client.list(spaces="drive", pageSize=1, fields="files(id)").execute()
+            client.list(
+                supportsAllDrives=True,
+                includeItemsFromAllDrives=True,
+                spaces="drive",
+                pageSize=1,
+                fields="files(id)",
+            ).execute()
         except HttpError as e:
             error_content = e.content.decode() if hasattr(e, "content") else ""
             lower_error = error_content.lower()
@@ -183,6 +189,8 @@ class GoogleDriveIndexer(Indexer):
             page_token = None
             while True:
                 response = files_client.list(
+                    supportsAllDrives=True,
+                    includeItemsFromAllDrives=True,
                     spaces="drive",
                     q=query,
                     fields="nextPageToken, files(id, mimeType, fileExtension)",
@@ -251,6 +259,8 @@ class GoogleDriveIndexer(Indexer):
                 else:
                     # Non-recursive: check for at least one immediate non-folder child.
                     response = client.list(
+                        supportsAllDrives=True,
+                        includeItemsFromAllDrives=True,
                         spaces="drive",
                         fields="files(id)",
                         pageSize=1,
@@ -348,6 +358,8 @@ class GoogleDriveIndexer(Indexer):
         files_response = []
         while not done:
             response: dict = files_client.list(
+                supportsAllDrives=True,
+                includeItemsFromAllDrives=True,
                 spaces="drive",
                 fields=fields_input,
                 corpora="user",
@@ -381,7 +393,9 @@ class GoogleDriveIndexer(Indexer):
         return files_response
 
     def get_root_info(self, files_client, object_id: str) -> dict:
-        return files_client.get(fileId=object_id, fields=",".join(self.fields)).execute()
+        return files_client.get(
+            supportsAllDrives=True, fileId=object_id, fields=",".join(self.fields)
+        ).execute()
 
     def get_files(
         self,

{unstructured_ingest-1.0.7 → unstructured_ingest-1.0.11}/unstructured_ingest/processes/connectors/ibm_watsonx/ibm_watsonx_s3.py

@@ -172,7 +172,7 @@ class IbmWatsonxUploaderConfig(UploaderConfig):
     namespace: str = Field(description="Namespace name")
     table: str = Field(description="Table name")
     max_retries: int = Field(
-        default=5, description="Maximum number of retries to upload data", ge=2, le=10
+        default=5, description="Maximum number of retries to upload data", ge=2, le=500
     )
     record_id_key: str = Field(
         default=RECORD_ID_LABEL,

{unstructured_ingest-1.0.7 → unstructured_ingest-1.0.11}/unstructured_ingest/processes/connectors/neo4j.py

@@ -9,8 +9,9 @@ from enum import Enum
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, AsyncGenerator, Literal, Optional
 
-from pydantic import BaseModel, ConfigDict, Field, Secret, field_validator
+from pydantic import BaseModel, ConfigDict, Field, Secret, ValidationError, field_validator
 
+from unstructured_ingest.data_types.entities import EntitiesData, Entity, EntityRelationship
 from unstructured_ingest.data_types.file_data import FileData
 from unstructured_ingest.error import DestinationConnectionError
 from unstructured_ingest.interfaces import (
@@ -97,7 +98,6 @@ class Neo4jUploadStager(UploadStager):
         **kwargs: Any,
     ) -> Path:
         elements = get_json_data(elements_filepath)
-
         nx_graph = self._create_lexical_graph(
             elements, self._create_document_node(file_data=file_data)
         )
@@ -109,28 +109,54 @@ class Neo4jUploadStager(UploadStager):
 
         return output_filepath
 
-    def _add_entities(self, element: dict, graph: "Graph", element_node: _Node) -> None:
-        entities = element.get("metadata", {}).get("entities", [])
-        if not entities:
-            return None
-        if not isinstance(entities, list):
-            return None
-
+    def _add_entities(self, entities: list[Entity], graph: "Graph", element_node: _Node) -> None:
         for entity in entities:
-            if not isinstance(entity, dict):
-                continue
-            if "entity" not in entity or "type" not in entity:
-                continue
             entity_node = _Node(
-                labels=[Label.ENTITY], properties={"id": entity["entity"]}, id_=entity["entity"]
+                labels=[Label.ENTITY], properties={"id": entity.entity}, id_=entity.entity
             )
             graph.add_edge(
                 entity_node,
-                _Node(labels=[Label.ENTITY], properties={"id": entity["type"]}, id_=entity["type"]),
+                _Node(labels=[Label.ENTITY], properties={"id": entity.type}, id_=entity.type),
                 relationship=Relationship.ENTITY_TYPE,
             )
             graph.add_edge(element_node, entity_node, relationship=Relationship.HAS_ENTITY)
 
+    def _add_entity_relationships(
+        self, relationships: list[EntityRelationship], graph: "Graph"
+    ) -> None:
+        for relationship in relationships:
+            from_node = _Node(
+                labels=[Label.ENTITY],
+                properties={"id": relationship.from_},
+                id_=relationship.from_,
+            )
+            to_node = _Node(
+                labels=[Label.ENTITY], properties={"id": relationship.to}, id_=relationship.to
+            )
+            graph.add_edge(from_node, to_node, relationship=relationship.relationship)
+
+    def _add_entity_data(self, element: dict, graph: "Graph", element_node: _Node) -> None:
+        entities = element.get("metadata", {}).get("entities", {})
+        if not entities:
+            return None
+        try:
+            if isinstance(entities, list):
+                self._add_entities(
+                    [Entity.model_validate(e) for e in entities if isinstance(e, dict)],
+                    graph,
+                    element_node,
+                )
+            elif isinstance(entities, dict):
+                entity_data = EntitiesData.model_validate(entities)
+                self._add_entities(entity_data.items, graph, element_node)
+                self._add_entity_relationships(entity_data.relationships, graph)
+        except ValidationError:
+            logger.warning(
+                "Failed to add entities to the graph. "
+                "Please check the format of the entities in the input data."
+            )
+        return None
+
     def _create_lexical_graph(self, elements: list[dict], document_node: _Node) -> "Graph":
         import networkx as nx
 
@@ -149,7 +175,7 @@ class Neo4jUploadStager(UploadStager):
             previous_node = element_node
             graph.add_edge(element_node, document_node, relationship=Relationship.PART_OF_DOCUMENT)
 
-            self._add_entities(element, graph, element_node)
+            self._add_entity_data(element, graph, element_node)
 
             if self._is_chunk(element):
                 for origin_element in format_and_truncate_orig_elements(element, include_text=True):
@@ -165,7 +191,7 @@ class Neo4jUploadStager(UploadStager):
                         document_node,
                         relationship=Relationship.PART_OF_DOCUMENT,
                     )
-                    self._add_entities(origin_element, graph, origin_element_node)
+                    self._add_entity_data(origin_element, graph, origin_element_node)
 
         return graph
 
@@ -208,7 +234,9 @@ class _GraphData(BaseModel):
             _Edge(
                 source=u,
                 destination=v,
-                relationship=Relationship(data_dict["relationship"]),
+                relationship=Relationship(data_dict["relationship"])
+                if data_dict["relationship"] in Relationship
+                else data_dict["relationship"],
             )
             for u, v, data_dict in nx_graph.edges(data=True)
         ]
@@ -242,7 +270,7 @@ class _Edge(BaseModel):
 
     source: _Node
     destination: _Node
-    relationship: Relationship
+    relationship: Relationship | str
 
 
 class Label(Enum):
@@ -380,7 +408,7 @@ class Neo4jUploader(Uploader):
         )
         logger.info(f"Finished merging {len(graph_data.nodes)} graph nodes.")
 
-        edges_by_relationship: defaultdict[tuple[Relationship, Label, Label], list[_Edge]] = (
+        edges_by_relationship: defaultdict[tuple[Relationship | str, Label, Label], list[_Edge]] = (
             defaultdict(list)
         )
         for edge in graph_data.edges:
@@ -463,16 +491,19 @@ class Neo4jUploader(Uploader):
     @staticmethod
     def _create_edges_query(
         edges: list[_Edge],
-        relationship: Relationship,
+        relationship: Relationship | str,
         source_label: Label,
         destination_label: Label,
     ) -> tuple[str, dict]:
         logger.info(f"Preparing MERGE query for {len(edges)} {relationship} relationships.")
+        relationship = (
+            relationship.value if isinstance(relationship, Relationship) else relationship
+        )
         query_string = f"""
             UNWIND $edges AS edge
             MATCH (u: `{source_label.value}` {{id: edge.source}})
             MATCH (v: `{destination_label.value}` {{id: edge.destination}})
-            MERGE (u)-[:`{relationship.value}`]->(v)
+            MERGE (u)-[:`{relationship}`]->(v)
             """
         parameters = {
             "edges": [

{unstructured_ingest-1.0.7 → unstructured_ingest-1.0.11}/unstructured_ingest/processes/connectors/pinecone.py

@@ -1,6 +1,7 @@
 import json
 import re
 from dataclasses import dataclass, field
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, Optional
 
 from pydantic import Field, Secret
@@ -18,11 +19,14 @@ from unstructured_ingest.interfaces import (
 )
 from unstructured_ingest.logger import logger
 from unstructured_ingest.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.utils import ndjson
 from unstructured_ingest.utils.constants import RECORD_ID_LABEL
 from unstructured_ingest.utils.data_prep import (
     flatten_dict,
     generator_batching_wbytes,
     get_enhanced_element_id,
+    get_json_data,
+    write_data,
 )
 from unstructured_ingest.utils.dep_check import requires_dependencies
 
@@ -162,6 +166,28 @@ class PineconeUploadStager(UploadStager):
             "metadata": metadata,
         }
 
+    def stream_update(self, input_file: Path, output_file: Path, file_data: FileData) -> None:
+        with input_file.open() as in_f:
+            reader = ndjson.reader(in_f)
+            with output_file.open("w") as out_f:
+                writer = ndjson.writer(out_f)
+                for element in reader:
+                    if "embeddings" not in element:
+                        continue
+                    conformed_element = self.conform_dict(element_dict=element, file_data=file_data)
+                    writer.write(row=conformed_element)
+                    writer.f.flush()
+
+    def process_whole(self, input_file: Path, output_file: Path, file_data: FileData) -> None:
+        elements_contents = get_json_data(path=input_file)
+
+        conformed_elements = [
+            self.conform_dict(element_dict=element, file_data=file_data)
+            for element in elements_contents
+            if "embeddings" in element
+        ]
+        write_data(path=output_file, data=conformed_elements)
+
 
 @dataclass
 class PineconeUploader(VectorDBUploader):

unstructured_ingest-1.0.7/unstructured_ingest/__version__.py

@@ -1 +0,0 @@
-__version__ = "1.0.7"  # pragma: no cover