groundx 2.2.8__tar.gz → 2.3.0__tar.gz

{groundx-2.2.8 → groundx-2.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: groundx
-Version: 2.2.8
+Version: 2.3.0
 Summary: 
 License: MIT
 Requires-Python: >=3.8,<4.0

{groundx-2.2.8 → groundx-2.3.0}/pyproject.toml

@@ -3,7 +3,7 @@ name = "groundx"
 
 [tool.poetry]
 name = "groundx"
-version = "2.2.8"
+version = "2.3.0"
 description = ""
 readme = "README.md"
 authors = []

{groundx-2.2.8 → groundx-2.3.0}/src/groundx/core/client_wrapper.py

@@ -16,7 +16,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "groundx",
-            "X-Fern-SDK-Version": "2.2.8",
+            "X-Fern-SDK-Version": "2.3.0",
         }
         headers["X-API-Key"] = self.api_key
         return headers

groundx-2.3.0/src/groundx/csv_splitter.py

@@ -0,0 +1,64 @@
+import csv, math, os, tempfile, typing
+from pathlib import Path
+
+
+class CSVSplitter:
+    def __init__(self, filepath, delimiter=','):
+        self.filepath = filepath
+        self.delimiter = delimiter
+        self.filename = os.path.basename(filepath)
+        self.file_size = os.path.getsize(filepath)
+        self.rows_count = self.get_row_count()
+
+    def get_row_count(self):
+        with open(self.filepath, "r", newline="", encoding="utf-8") as csvfile:
+            return sum(1 for _ in csvfile) - 1
+
+    def determine_splits(self):
+        row_mod = int(self.rows_count / 1000) + 1
+        file_mod = int(self.file_size / 1024 / 1024) + 1
+
+        return max(row_mod, file_mod)
+
+    def split(self):
+        splits = self.determine_splits()
+        if splits < 2:
+            return [Path(self.filepath)]
+
+        rows_per_file = math.ceil(self.rows_count / splits)
+
+        split_files: typing.List[Path] = []
+        with open(self.filepath, "r", newline="", encoding="utf-8") as csvfile:
+            reader = csv.reader(csvfile, delimiter=self.delimiter)
+            headers = next(reader)
+
+            temp_dir = tempfile.mkdtemp()
+
+            current_file_number = 1
+            current_row = 0
+            current_writer = None
+            current_output_file = None
+
+            for row in reader:
+                if current_row % rows_per_file == 0:
+                    if current_output_file:
+                        current_output_file.close()
+                    output_file_path = os.path.join(
+                        temp_dir,
+                        f"{os.path.splitext(self.filename)[0]}_{current_file_number}.csv",
+                    )
+                    split_files.append(Path(output_file_path))
+                    current_output_file = open(
+                        output_file_path, "w", newline="", encoding="utf-8"
+                    )
+                    current_writer = csv.writer(current_output_file, delimiter=self.delimiter)
+                    current_writer.writerow(headers)
+                    current_file_number += 1
+
+                current_writer.writerow(row)
+                current_row += 1
+
+            if current_output_file:
+                current_output_file.close()
+
+        return split_files

{groundx-2.2.8 → groundx-2.3.0}/src/groundx/ingest.py

@@ -5,10 +5,11 @@ from urllib.parse import urlparse, urlunparse
 
 from .client import GroundXBase, AsyncGroundXBase
 from .core.request_options import RequestOptions
+from .csv_splitter import CSVSplitter
 from .types.document import Document
-from .types.document_type import DocumentType
 from .types.ingest_remote_document import IngestRemoteDocument
 from .types.ingest_response import IngestResponse
+from .types.ingest_response_ingest import IngestResponseIngest
 
 # this is used as the default value for optional parameters
 OMIT = typing.cast(typing.Any, ...)
@@ -38,10 +39,18 @@ MIME_TO_DOCUMENT_TYPE = {v: k for k, v in DOCUMENT_TYPE_TO_MIME.items()}
 
 ALLOWED_SUFFIXES = {f".{k}": v for k, v in DOCUMENT_TYPE_TO_MIME.items()}
 
+CSV_SPLITS = {
+    ".csv": True,
+}
+TSV_SPLITS = {
+    ".tsv": True,
+}
+
 SUFFIX_ALIASES = {
-    ".jpeg": ".jpg",
-    ".heic": ".heif",
-    ".tif": ".tiff",
+    ".jpeg": "jpg",
+    ".heic": "heif",
+    ".tif": "tiff",
+    ".md": "txt",
 }
 
 MAX_BATCH_SIZE = 50
@@ -115,11 +124,25 @@ def prep_documents(
     return remote_documents, local_documents
 
 
+def split_doc(file):
+    if file.is_file() and (
+        file.suffix.lower() in ALLOWED_SUFFIXES
+        or file.suffix.lower() in SUFFIX_ALIASES
+    ):
+        if file.suffix.lower() in CSV_SPLITS:
+            return CSVSplitter(filepath=file).split()
+        elif file.suffix.lower() in TSV_SPLITS:
+            return CSVSplitter(filepath=file, delimiter='\t').split()
+        return [file]
+    return []
+
 class GroundX(GroundXBase):
     def ingest(
         self,
         *,
         documents: typing.Sequence[Document],
+        batch_size: typing.Optional[int] = 10,
+        wait_for_complete: typing.Optional[bool] = False,
         upload_api: typing.Optional[str] = "https://api.eyelevel.ai/upload/file",
         request_options: typing.Optional[RequestOptions] = None,
     ) -> IngestResponse:
@@ -130,6 +153,13 @@ class GroundX(GroundXBase):
         ----------
         documents : typing.Sequence[Document]
 
+        # defines how many files to send per batch
+        # ignored unless wait_for_complete is True
+        batch_size : typing.Optional[int]
+
+        # will turn on progress bar and wait for ingestion to complete
+        wait_for_complete : typing.Optional[bool]
+
         # an endpoint that accepts 'name' and 'type' query params
         # and returns a presigned URL in a JSON dictionary with key 'URL'
         upload_api : typing.Optional[str]
@@ -163,41 +193,87 @@ class GroundX(GroundXBase):
         """
         remote_documents, local_documents = prep_documents(documents)
 
-        if local_documents and remote_documents:
-            raise ValueError("Documents must all be either local or remote, not a mix.")
+        if len(remote_documents) + len(local_documents) == 0:
+            raise ValueError("No valid documents were provided")
 
-        if len(remote_documents) > 0:
-            if len(remote_documents) > MAX_BATCH_SIZE:
-                raise ValueError("You have sent too many documents in this request")
+        if wait_for_complete:
+            with tqdm(total=len(remote_documents) + len(local_documents), desc="Ingesting Files", unit="file") as pbar:
+                n = max(MIN_BATCH_SIZE, min(batch_size or MIN_BATCH_SIZE, MAX_BATCH_SIZE))
 
-            return self.documents.ingest_remote(
-                documents=remote_documents,
-                request_options=request_options,
-            )
+                remote_batch: typing.List[IngestRemoteDocument] = []
+                ingest = IngestResponse(ingest=IngestResponseIngest(process_id="",status="queued"))
+
+                progress = float(len(remote_documents))
+                for rd in remote_documents:
+                    if len(remote_batch) >= n:
+                        ingest = self.documents.ingest_remote(
+                            documents=remote_batch,
+                            request_options=request_options,
+                        )
+                        ingest, progress = self._monitor_batch(ingest, progress, pbar)
+
+                        remote_batch = []
+
+                    remote_batch.append(rd)
+                    pbar.update(0.25)
+                    progress -= 0.25
+
+                if remote_batch:
+                    ingest = self.documents.ingest_remote(
+                        documents=remote_batch,
+                        request_options=request_options,
+                    )
+                    ingest, progress = self._monitor_batch(ingest, progress, pbar)
+
+
+                if progress > 0:
+                    pbar.update(progress)
+
+                current_batch_size = 0
+                local_batch: typing.List[Document] = []
+
+                progress = float(len(local_documents))
+                for ld in local_documents:
+                    fp = Path(os.path.expanduser(ld.file_path))
+                    file_size = fp.stat().st_size
+
+                    if (current_batch_size + file_size > MAX_BATCH_SIZE_BYTES) or (len(local_batch) >= n):
+                        up_docs, progress = self._process_local(local_batch, upload_api, progress, pbar)
 
-        if len(local_documents) > MAX_BATCH_SIZE:
+                        ingest = self.documents.ingest_remote(
+                            documents=up_docs,
+                            request_options=request_options,
+                        )
+                        ingest, progress = self._monitor_batch(ingest, progress, pbar)
+
+                        local_batch = []
+                        current_batch_size = 0
+
+                    local_batch.append(ld)
+                    current_batch_size += file_size
+
+                if local_batch:
+                    up_docs, progress = self._process_local(local_batch, upload_api, progress, pbar)
+
+                    ingest = self.documents.ingest_remote(
+                        documents=up_docs,
+                        request_options=request_options,
+                    )
+                    ingest, progress = self._monitor_batch(ingest, progress, pbar)
+
+                if progress > 0:
+                    pbar.update(progress)
+
+                return ingest
+        elif len(remote_documents) + len(local_documents) > MAX_BATCH_SIZE:
             raise ValueError("You have sent too many documents in this request")
 
-        if len(local_documents) == 0:
-            raise ValueError("No valid documents were provided")
 
-        docs: typing.List[IngestRemoteDocument] = []
-        for d in local_documents:
-            url = self._upload_file(upload_api, Path(os.path.expanduser(d.file_path)))
-
-            docs.append(
-                IngestRemoteDocument(
-                    bucket_id=d.bucket_id,
-                    file_name=d.file_name,
-                    file_type=d.file_type,
-                    process_level=d.process_level,
-                    search_data=d.search_data,
-                    source_url=url,
-                )
-            )
+        up_docs, _ = self._process_local(local_documents, upload_api)
+        remote_documents.extend(up_docs)
 
         return self.documents.ingest_remote(
-            documents=docs,
+            documents=remote_documents,
             request_options=request_options,
         )
 
@@ -252,14 +328,10 @@ class GroundX(GroundXBase):
         def load_directory_files(directory: str) -> typing.List[Path]:
             dir_path = Path(directory)
 
-            matched_files = [
-                file
-                for file in dir_path.rglob("*")
-                if file.is_file() and (
-                    file.suffix.lower() in ALLOWED_SUFFIXES
-                    or file.suffix.lower() in SUFFIX_ALIASES
-                )
-            ]
+            matched_files: typing.List[Path] = []
+            for file in dir_path.rglob("*"):
+                for sd in split_doc(file):
+                    matched_files.append(sd)    
 
             return matched_files      
 
@@ -301,6 +373,8 @@ class GroundX(GroundXBase):
     ):
         file_name = os.path.basename(file_path)
         file_extension = os.path.splitext(file_name)[1][1:].lower()
+        if f".{file_extension}" in SUFFIX_ALIASES:
+            file_extension = SUFFIX_ALIASES[f".{file_extension}"]
 
         presigned_info = get_presigned_url(endpoint, file_name, file_extension)
 
@@ -330,6 +404,92 @@ class GroundX(GroundXBase):
 
         return strip_query_params(upload_url)
 
+    def _process_local(
+        self,
+        local_docs,
+        upload_api,
+        progress = None,
+        pbar = None,
+    ):
+        remote_docs = []
+        for d in local_docs:
+            splits = split_doc(Path(os.path.expanduser(d.file_path)))
+
+            for sd in splits:
+                url = self._upload_file(upload_api, sd)
+
+                ft = d.file_type
+                if sd.suffix.lower() in SUFFIX_ALIASES:
+                    ft = SUFFIX_ALIASES[sd.suffix.lower()]
+
+                fn = sd.name
+                if len(splits) == 1 and d.file_name:
+                    fn = d.file_name
+
+                remote_docs.append(
+                    IngestRemoteDocument(
+                        bucket_id=d.bucket_id,
+                        file_name=fn,
+                        file_type=ft,
+                        process_level=d.process_level,
+                        search_data=d.search_data,
+                        source_url=url,
+                    )
+                )
+
+                if progress is not None and pbar is not None and pbar.update is not None:
+                    pbar.update(0.25)
+                    progress -= 0.25
+
+        return remote_docs, progress
+
+    def _monitor_batch(
+        self,
+        ingest,
+        progress,
+        pbar,
+    ):
+        completed_files = set()
+
+        while (
+            ingest is not None
+            and ingest.ingest.status not in ["complete", "error", "cancelled"]
+        ):
+            time.sleep(3)
+            ingest = self.documents.get_processing_status_by_id(ingest.ingest.process_id)
+
+            if ingest.ingest.progress:
+                if ingest.ingest.progress.processing and ingest.ingest.progress.processing.documents:
+                    for doc in ingest.ingest.progress.processing.documents:
+                        if doc.status in ["complete", "error", "cancelled"] and doc.document_id not in completed_files:
+                            pbar.update(0.75)
+                            progress -= 0.75
+                            completed_files.add(doc.document_id)
+                if ingest.ingest.progress.complete and ingest.ingest.progress.complete.documents:
+                    for doc in ingest.ingest.progress.complete.documents:
+                        if doc.status in ["complete", "error", "cancelled"] and doc.document_id not in completed_files:
+                            pbar.update(0.75)
+                            progress -= 0.75
+                            completed_files.add(doc.document_id)
+                if ingest.ingest.progress.cancelled and ingest.ingest.progress.cancelled.documents:
+                    for doc in ingest.ingest.progress.cancelled.documents:
+                        if doc.status in ["complete", "error", "cancelled"] and doc.document_id not in completed_files:
+                            pbar.update(0.75)
+                            progress -= 0.75
+                            completed_files.add(doc.document_id)
+                if ingest.ingest.progress.errors and ingest.ingest.progress.errors.documents:
+                    for doc in ingest.ingest.progress.errors.documents:
+                        if doc.status in ["complete", "error", "cancelled"] and doc.document_id not in completed_files:
+                            pbar.update(0.75)
+                            progress -= 0.75
+                            completed_files.add(doc.document_id)
+
+
+        if ingest.ingest.status in ["error", "cancelled"]:
+            raise ValueError(f"Ingest failed with status: {ingest.ingest.status}")
+
+        return ingest, progress
+
     def _upload_file_batch(
         self,
         bucket_id,
@@ -340,41 +500,35 @@ class GroundX(GroundXBase):
     ):
         docs = []
 
-        progress = len(batch)
+        progress =  float(len(batch))
         for file in batch:
             url = self._upload_file(upload_api, file)
-            docs.append(
-                Document(
-                    bucket_id=bucket_id,
-                    file_path=url,
-                ),
-            )
+            if file.suffix.lower() in SUFFIX_ALIASES:
+                docs.append(
+                    Document(
+                        bucket_id=bucket_id,
+                        file_name=file.name,
+                        file_path=url,
+                        file_type=SUFFIX_ALIASES[file.suffix.lower()],
+                    ),
+                )
+            else:
+                docs.append(
+                    Document(
+                        bucket_id=bucket_id,
+                        file_name=file.name,
+                        file_path=url,
+                    ),
+                )
             pbar.update(0.25)
             progress -= 0.25
 
         if docs:
             ingest = self.ingest(documents=docs, request_options=request_options)
+            ingest, progress = self._monitor_batch(ingest, progress, pbar)
 
-            completed_files = set()
-
-            while (
-                ingest is not None
-                and ingest.ingest.status not in ["complete", "error", "cancelled"]
-            ):
-                time.sleep(3)
-                ingest = self.documents.get_processing_status_by_id(ingest.ingest.process_id)
-
-                if ingest.ingest.progress and ingest.ingest.progress.processing:
-                    for doc in ingest.ingest.progress.processing.documents:
-                        if doc.status == "complete" and doc.document_id not in completed_files:
-                            pbar.update(0.75)
-                            progress -= 0.75
-
-            if ingest.ingest.status in ["error", "cancelled"]:
-                raise ValueError(f"Ingest failed with status: {ingest.ingest.status}")
-
-            if progress > 0:
-                pbar.update(progress)
+        if progress > 0:
+            pbar.update(progress)
 
 
 
@@ -431,41 +585,39 @@ class AsyncGroundX(AsyncGroundXBase):
         """
         remote_documents, local_documents = prep_documents(documents)
 
-        if local_documents and remote_documents:
-            raise ValueError("Documents must all be either local or remote, not a mix.")
-
-        if len(remote_documents) > 0:
-            if len(remote_documents) > MAX_BATCH_SIZE:
-                raise ValueError("You have sent too many documents in this request")
-
-            return await self.documents.ingest_remote(
-                documents=remote_documents,
-                request_options=request_options,
-            )
-
-        if len(local_documents) > MAX_BATCH_SIZE:
+        if len(remote_documents) + len(local_documents) > MAX_BATCH_SIZE:
             raise ValueError("You have sent too many documents in this request")
 
-        if len(local_documents) == 0:
+        if len(remote_documents) + len(local_documents) == 0:
             raise ValueError("No valid documents were provided")
 
-        docs: typing.List[IngestRemoteDocument] = []
         for d in local_documents:
-            url = self._upload_file(upload_api, Path(os.path.expanduser(d.file_path)))
-
-            docs.append(
-                IngestRemoteDocument(
-                    bucket_id=d.bucket_id,
-                    file_name=d.file_name,
-                    file_type=d.file_type,
-                    process_level=d.process_level,
-                    search_data=d.search_data,
-                    source_url=url,
+            splits = split_doc(Path(os.path.expanduser(d.file_path)))
+
+            for sd in splits:
+                url = self._upload_file(upload_api, sd)
+
+                ft = d.file_type
+                if sd.suffix.lower() in SUFFIX_ALIASES:
+                    ft = SUFFIX_ALIASES[sd.suffix.lower()]
+
+                fn = sd.name
+                if len(splits) == 1 and d.file_name:
+                    fn = d.file_name
+
+                remote_documents.append(
+                    IngestRemoteDocument(
+                        bucket_id=d.bucket_id,
+                        file_name=fn,
+                        file_type=ft,
+                        process_level=d.process_level,
+                        search_data=d.search_data,
+                        source_url=url,
+                    )
                 )
-            )
 
         return await self.documents.ingest_remote(
-            documents=docs,
+            documents=remote_documents,
             request_options=request_options,
         )
 
@@ -476,6 +628,8 @@ class AsyncGroundX(AsyncGroundXBase):
     ):
         file_name = os.path.basename(file_path)
         file_extension = os.path.splitext(file_name)[1][1:].lower()
+        if f".{file_extension}" in SUFFIX_ALIASES:
+            file_extension = SUFFIX_ALIASES[f".{file_extension}"]
 
         presigned_info = get_presigned_url(endpoint, file_name, file_extension)