parallex 0.2.1__tar.gz → 0.3.1__tar.gz

{parallex-0.2.1 → parallex-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: parallex
-Version: 0.2.1
+Version: 0.3.1
 Summary: PDF to markdown using Azure OpenAI batch processing
 Home-page: https://github.com/Summed-AI/parallex
 Author: Jeff Hostetler

{parallex-0.2.1 → parallex-0.3.1}/parallex/ai/uploader.py

@@ -9,7 +9,7 @@ from parallex.models.batch_file import BatchFile
 from parallex.models.image_file import ImageFile
 from parallex.utils.constants import CUSTOM_ID_DELINEATOR
 
-MAX_FILE_SIZE = 150 * 1024 * 1024  # 150 MB in bytes
+MAX_FILE_SIZE = 180 * 1024 * 1024  # 180 MB in bytes. Limit for Azure is 200MB.
 
 
 async def upload_images_for_processing(
@@ -23,22 +23,18 @@ async def upload_images_for_processing(
     current_index = 0
     batch_files = []
     upload_file_location = file_in_temp_dir(
-        directory=temp_directory, file_name=f"image-{trace_id}-{current_index}.jsonl"
+        directory=temp_directory, file_name=f"{trace_id}-{current_index}.jsonl"
     )
 
     for image_file in image_files:
-        if (
-            os.path.exists(upload_file_location)
-            and os.path.getsize(upload_file_location) > MAX_FILE_SIZE
-        ):
+        if await _approaching_file_size_limit(upload_file_location):
             """When approaching upload file limit, upload and start new file"""
             batch_file = await _create_batch_file(
                 client, trace_id, upload_file_location
             )
             batch_files.append(batch_file)
-            current_index += 1
-            upload_file_location = file_in_temp_dir(
-                directory=temp_directory, file_name=f"{trace_id}-{current_index}.jsonl"
+            upload_file_location = await _increment_batch_file_index(
+                current_index, temp_directory, trace_id
             )
 
         with open(image_file.path, "rb") as image:
@@ -57,21 +53,62 @@ async def upload_images_for_processing(
 
 async def upload_prompts_for_processing(
     client: OpenAIClient, prompts: list[str], temp_directory: str, trace_id: UUID
-) -> BatchFile:
+) -> list[BatchFile]:
     """Creates jsonl file and uploads for processing"""
-    upload_file_location = file_in_temp_dir(
-        directory=temp_directory, file_name=f"prompts-{trace_id}.jsonl"
+    current_index = 0
+    batch_files = []
+
+    upload_file_location = await set_file_location(
+        current_index, temp_directory, trace_id
     )
     for index, prompt in enumerate(prompts):
+        if await _approaching_file_size_limit(upload_file_location):
+            """When approaching upload file limit, upload and start new file"""
+            batch_file = await _create_batch_file(
+                client, trace_id, upload_file_location
+            )
+            batch_files.append(batch_file)
+            upload_file_location = await _increment_batch_file_index(
+                current_index, temp_directory, trace_id
+            )
+
         prompt_custom_id = f"{trace_id}{CUSTOM_ID_DELINEATOR}{index}.jsonl"
         jsonl = _simple_jsonl_format(prompt_custom_id, prompt)
         with open(upload_file_location, "a") as jsonl_file:
             jsonl_file.write(json.dumps(jsonl) + "\n")
     batch_file = await _create_batch_file(client, trace_id, upload_file_location)
-    return batch_file
+    batch_files.append(batch_file)
+    return batch_files
+
+
+async def set_file_location(
+    current_index: int, temp_directory: str, trace_id: UUID
+) -> str:
+    return file_in_temp_dir(
+        directory=temp_directory, file_name=f"{trace_id}-{current_index}.jsonl"
+    )
+
+
+async def _approaching_file_size_limit(upload_file_location: str) -> bool:
+    return (
+        os.path.exists(upload_file_location)
+        and os.path.getsize(upload_file_location) > MAX_FILE_SIZE
+    )
+
+
+async def _increment_batch_file_index(
+    current_index: int, temp_directory: str, trace_id: UUID
+) -> str:
+    current_index += 1
+    upload_file_location = await set_file_location(
+        current_index, temp_directory, trace_id
+    )
+    return upload_file_location
 
 
-async def _create_batch_file(client, trace_id, upload_file_location):
+async def _create_batch_file(
+    client: OpenAIClient, trace_id: UUID, upload_file_location: str
+) -> BatchFile:
     file_response = await client.upload(upload_file_location)
     return BatchFile(
         id=file_response.id,
@@ -82,7 +119,7 @@ async def _create_batch_file(client, trace_id, upload_file_location):
     )
 
 
-def _simple_jsonl_format(prompt_custom_id: str, prompt_text: str):
+def _simple_jsonl_format(prompt_custom_id: str, prompt_text: str) -> dict:
     return {
         "custom_id": prompt_custom_id,
         "method": "POST",

{parallex-0.2.1 → parallex-0.3.1}/parallex/parallex.py

@@ -46,6 +46,7 @@ async def parallex(
         )
     except Exception as e:
         logger.error(f"Error occurred: {e}")
+        raise e
     finally:
         await _delete_associated_files(open_ai_client, remote_file_handler)
 
@@ -55,6 +56,7 @@ async def parallex_simple_prompts(
     prompts: list[str],
     post_process_callable: Optional[Callable[..., None]] = None,
     log_level: Optional[str] = "ERROR",
+    concurrency: Optional[int] = 20,
 ) -> ParallexPromptsCallableOutput:
     setup_logger(log_level)
     remote_file_handler = RemoteFileHandler()
@@ -64,9 +66,11 @@ async def parallex_simple_prompts(
             open_ai_client=open_ai_client,
             prompts=prompts,
             post_process_callable=post_process_callable,
+            concurrency=concurrency,
         )
     except Exception as e:
         logger.error(f"Error occurred: {e}")
+        raise e
     finally:
         await _delete_associated_files(open_ai_client, remote_file_handler)
 
@@ -75,27 +79,42 @@ async def _prompts_execute(
     open_ai_client: OpenAIClient,
     prompts: list[str],
     post_process_callable: Optional[Callable[..., None]] = None,
+    concurrency: Optional[int] = 20,
 ):
     with tempfile.TemporaryDirectory() as temp_directory:
         trace_id = uuid.uuid4()
-        batch_file = await upload_prompts_for_processing(
+        batch_files = await upload_prompts_for_processing(
             client=open_ai_client,
             prompts=prompts,
             temp_directory=temp_directory,
             trace_id=trace_id,
         )
-        batch = await create_batch(
-            client=open_ai_client, file_id=batch_file.id, trace_id=trace_id
-        )
-        logger.info(f"waiting for batch to complete - {batch.id} - {batch.trace_id}")
-        output_file_id = await wait_for_batch_completion(
-            client=open_ai_client, batch=batch
-        )
-        logger.info(f"batch completed - {batch.id} - {batch.trace_id}")
-        prompts_output = await process_prompts_output(
-            client=open_ai_client, output_file_id=output_file_id
-        )
-        sorted_responses = sorted(prompts_output, key=lambda x: x.prompt_index)
+        start_batch_semaphore = asyncio.Semaphore(concurrency)
+        start_batch_tasks = []
+        for file in batch_files:
+            batch_task = asyncio.create_task(
+                _create_batch_jobs(
+                    batch_file=file,
+                    client=open_ai_client,
+                    trace_id=trace_id,
+                    semaphore=start_batch_semaphore,
+                )
+            )
+            start_batch_tasks.append(batch_task)
+        batch_jobs = await asyncio.gather(*start_batch_tasks)
+
+        process_semaphore = asyncio.Semaphore(concurrency)
+        prompt_tasks = []
+        for batch in batch_jobs:
+            prompt_task = asyncio.create_task(
+                _wait_and_create_prompt_responses(batch=batch, client=open_ai_client, semaphore=process_semaphore)
+            )
+            prompt_tasks.append(prompt_task)
+        prompt_response_groups = await asyncio.gather(*prompt_tasks)
+
+        flat_responses = [response for batch in prompt_response_groups for response in batch]
+
+        sorted_responses = sorted(flat_responses, key=lambda x: x.prompt_index)
         callable_output = ParallexPromptsCallableOutput(
             original_prompts=prompts,
             trace_id=trace_id,
@@ -155,7 +174,7 @@ async def _execute(
 
         pages = [page for batch_pages in page_groups for page in batch_pages]
         logger.info(f"pages done. total pages- {len(pages)} - {trace_id}")
-        sorted_pages = sorted(pages, key=lambda x: x.page_number)
+        sorted_pages = sorted(pages, key=lambda x: x.prompt_index)
 
         # TODO add combined version of MD to output / save to file system
         callable_output = ParallexCallableOutput(
@@ -182,6 +201,19 @@ async def _wait_and_create_pages(
         return page_responses
 
 
+async def _wait_and_create_prompt_responses(
+    batch: UploadBatch, client: OpenAIClient, semaphore: asyncio.Semaphore
+):
+    async with semaphore:
+        logger.info(f"waiting for batch to complete - {batch.id} - {batch.trace_id}")
+        output_file_id = await wait_for_batch_completion(client=client, batch=batch)
+        logger.info(f"batch completed - {batch.id} - {batch.trace_id}")
+        prompt_responses = await process_prompts_output(
+            client=client, output_file_id=output_file_id
+        )
+        return prompt_responses
+
+
 async def _create_batch_jobs(
     batch_file: BatchFile,
     client: OpenAIClient,

{parallex-0.2.1 → parallex-0.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "parallex"
-version = "0.2.1"
+version = "0.3.1"
 description = "PDF to markdown using Azure OpenAI batch processing"
 authors = ["Jeff Hostetler <jeff@summed.ai>", "Kevin Bao <kevin@summed.ai>"]
 repository = "https://github.com/Summed-AI/parallex"