parallex 0.1.3__py3-none-any.whl → 0.2.0__py3-none-any.whl

parallex/ai/batch_processor.py

@@ -35,8 +35,6 @@ async def wait_for_batch_completion(client: OpenAIClient, batch: UploadBatch) ->
         await asyncio.sleep(delay)
         batch_response = await client.retrieve_batch(batch.id)
         status = batch_response.status
-        batch.output_file_id = batch_response.output_file_id
-        batch.error_file_id = batch_response.error_file_id
         delay = 30
         if status == "completed":
             return batch_response.output_file_id

parallex/ai/open_ai_client.py

@@ -4,34 +4,46 @@ from openai import AsyncAzureOpenAI
 from openai._legacy_response import HttpxBinaryResponseContent
 from openai.types import FileObject, Batch, FileDeleted
 
+from parallex.file_management.remote_file_handler import RemoteFileHandler
 from parallex.utils.logger import logger
 
 
 # Exceptions for missing keys, etc
 class OpenAIClient:
-    def __init__(self, model: str):
+    def __init__(self, model: str, remote_file_handler: RemoteFileHandler):
         self.model = model
+        self.file_handler = remote_file_handler
 
         self._client = AsyncAzureOpenAI(
-            azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
-            api_key=os.getenv("AZURE_OPENAI_API_KEY"),
-            api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
+            azure_endpoint=os.getenv("AZURE_API_BASE"),
+            api_key=os.getenv("AZURE_API_KEY"),
+            api_version=os.getenv("AZURE_API_VERSION"),
         )
 
     async def upload(self, file_path: str) -> FileObject:
-        return await self._client.files.create(
+        file = await self._client.files.create(
             file=open(file_path, "rb"), purpose="batch"
         )
+        self.file_handler.add_file(file.id)
+        return file
 
     async def create_batch(self, upload_file_id: str) -> Batch:
-        return await self._client.batches.create(
+        batch = await self._client.batches.create(
             input_file_id=upload_file_id,
             endpoint="/chat/completions",
             completion_window="24h",
         )
+        self.file_handler.add_file(batch.input_file_id)
+        self.file_handler.add_file(batch.output_file_id)
+        self.file_handler.add_file(batch.error_file_id)
+        return batch
 
     async def retrieve_batch(self, batch_id: str) -> Batch:
-        return await self._client.batches.retrieve(batch_id)
+        batch = await self._client.batches.retrieve(batch_id)
+        self.file_handler.add_file(batch.input_file_id)
+        self.file_handler.add_file(batch.output_file_id)
+        self.file_handler.add_file(batch.error_file_id)
+        return batch
 
     async def retrieve_file(self, file_id: str) -> HttpxBinaryResponseContent:
         return await self._client.files.content(file_id)

parallex/ai/uploader.py

@@ -16,7 +16,7 @@ async def upload_images_for_processing(
     image_files: list[ImageFile],
     temp_directory: str,
     prompt_text: str,
-):
+) -> list[BatchFile]:
     """Base64 encodes image, converts to expected jsonl format and uploads"""
     trace_id = image_files[0].trace_id
     current_index = 0
@@ -71,7 +71,7 @@ def _jsonl_format(prompt_custom_id: str, encoded_image: str, prompt_text: str):
         "method": "POST",
         "url": "/chat/completions",
         "body": {
-            "model": os.getenv("AZURE_OPENAI_API_DEPLOYMENT"),
+            "model": os.getenv("AZURE_API_DEPLOYMENT"),
             "messages": [
                 {
                     "role": "user",

parallex/file_management/remote_file_handler.py

@@ -0,0 +1,7 @@
+class RemoteFileHandler:
+    def __init__(self):
+        self.created_files = set()
+
+    def add_file(self, file_name: str) -> None:
+        if file_name is not None:
+            self.created_files.add(file_name)

parallex/parallex.py

@@ -9,7 +9,8 @@ from parallex.ai.output_processor import process_output
 from parallex.ai.uploader import upload_images_for_processing
 from parallex.file_management.converter import convert_pdf_to_images
 from parallex.file_management.file_finder import add_file_to_temp_directory
-from parallex.models.image_file import ImageFile
+from parallex.file_management.remote_file_handler import RemoteFileHandler
+from parallex.models.batch_file import BatchFile
 from parallex.models.parallex_callable_output import ParallexCallableOutput
 from parallex.models.upload_batch import UploadBatch
 from parallex.utils.constants import DEFAULT_PROMPT
@@ -26,9 +27,32 @@ async def parallex(
     log_level: Optional[str] = "ERROR",
 ) -> ParallexCallableOutput:
     setup_logger(log_level)
-    with tempfile.TemporaryDirectory() as temp_directory:
-        open_ai_client = OpenAIClient(model=model)
+    remote_file_handler = RemoteFileHandler()
+    open_ai_client = OpenAIClient(model=model, remote_file_handler=remote_file_handler)
+    try:
+        return await _execute(
+            open_ai_client=open_ai_client,
+            pdf_source_url=pdf_source_url,
+            post_process_callable=post_process_callable,
+            concurrency=concurrency,
+            prompt_text=prompt_text,
+        )
+    except Exception as e:
+        logger.error(f"Error occurred: {e}")
+    finally:
+        for file in remote_file_handler.created_files:
+            logger.info(f"deleting - {file}")
+            await open_ai_client.delete_file(file)
 
+
+async def _execute(
+    open_ai_client: OpenAIClient,
+    pdf_source_url: str,
+    post_process_callable: Optional[Callable[..., None]] = None,
+    concurrency: Optional[int] = 20,
+    prompt_text: Optional[str] = DEFAULT_PROMPT,
+) -> ParallexCallableOutput:
+    with tempfile.TemporaryDirectory() as temp_directory:
         raw_file = await add_file_to_temp_directory(
             pdf_source_url=pdf_source_url, temp_directory=temp_directory
         )
@@ -47,7 +71,7 @@ async def parallex(
         start_batch_tasks = []
         for file in batch_files:
             batch_task = asyncio.create_task(
-                _create_images_and_batch_jobs(
+                _create_batch_jobs(
                     batch_file=file,
                     client=open_ai_client,
                     trace_id=trace_id,
@@ -55,11 +79,11 @@ async def parallex(
                 )
             )
             start_batch_tasks.append(batch_task)
-        batches = await asyncio.gather(*start_batch_tasks)
+        batch_jobs = await asyncio.gather(*start_batch_tasks)
 
         pages_tasks = []
         process_semaphore = asyncio.Semaphore(concurrency)
-        for batch in batches:
+        for batch in batch_jobs:
             page_task = asyncio.create_task(
                 _wait_and_create_pages(
                     batch=batch, client=open_ai_client, semaphore=process_semaphore
@@ -69,7 +93,7 @@ async def parallex(
         page_groups = await asyncio.gather(*pages_tasks)
 
         pages = [page for batch_pages in page_groups for page in batch_pages]
-        logger.debug(f"pages done. total pages- {len(pages)} - {trace_id}")
+        logger.info(f"pages done. total pages- {len(pages)} - {trace_id}")
         sorted_pages = sorted(pages, key=lambda x: x.page_number)
 
         # TODO add combined version of MD to output / save to file system
@@ -88,30 +112,23 @@ async def _wait_and_create_pages(
     batch: UploadBatch, client: OpenAIClient, semaphore: asyncio.Semaphore
 ):
     async with semaphore:
-        logger.debug(f"waiting for batch to complete - {batch.id} - {batch.trace_id}")
+        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.debug(f"batch completed - {batch.id} - {batch.trace_id}")
+        logger.info(f"batch completed - {batch.id} - {batch.trace_id}")
         page_responses = await process_output(
             client=client, output_file_id=output_file_id
         )
-        await _remove_global_batch_files(client=client, batch=batch)
         return page_responses
 
 
-async def _remove_global_batch_files(client: OpenAIClient, batch: UploadBatch):
-    file_ids = [batch.input_file_id, batch.output_file_id, batch.error_file_id]
-    for file_id in file_ids:
-        await client.delete_file(file_id)
-
-
-async def _create_images_and_batch_jobs(
-    batch_file: ImageFile,
+async def _create_batch_jobs(
+    batch_file: BatchFile,
     client: OpenAIClient,
     trace_id: UUID,
     semaphore: asyncio.Semaphore,
 ):
     async with semaphore:
-        batch = await create_batch(
+        upload_batch = await create_batch(
             client=client, file_id=batch_file.id, trace_id=trace_id
         )
-        return batch
+        return upload_batch

parallex/utils/logger.py

@@ -1,8 +1,8 @@
 import logging
 
-from aiologger import Logger
+from aiologger.loggers.json import JsonLogger
 
-logger = Logger.with_default_handlers(name="parallex")
+logger = JsonLogger.with_default_handlers(name="parallex")
 
 
 def setup_logger(level: str = "ERROR"):
@@ -15,6 +15,4 @@ def setup_logger(level: str = "ERROR"):
         "NOTSET": logging.NOTSET,
     }.get(level, logging.INFO)
 
-    logging.basicConfig(
-        level=level, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-    )
+    logger.setLevel = level

{parallex-0.1.3.dist-info → parallex-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: parallex
-Version: 0.1.3
+Version: 0.2.0
 Summary: PDF to markdown using Azure OpenAI batch processing
 Home-page: https://github.com/Summed-AI/parallex
 Author: Jeff Hostetler
@@ -45,10 +45,10 @@ import os
 from parallex.models.parallex_callable_output import ParallexCallableOutput
 from parallex.parallex import parallex
 
-os.environ["AZURE_OPENAI_API_KEY"] = "key"
-os.environ["AZURE_OPENAI_ENDPOINT"] = "your-endpoint.com"
-os.environ["AZURE_OPENAI_API_VERSION"] = "deployment_version"
-os.environ["AZURE_OPENAI_API_DEPLOYMENT"] = "deployment_name"
+os.environ["AZURE_API_KEY"] = "key"
+os.environ["AZURE_API_BASE"] = "your-endpoint.com"
+os.environ["AZURE_API_VERSION"] = "deployment_version"
+os.environ["AZURE_API_DEPLOYMENT"] = "deployment_name"
 
 model = "gpt-4o"
 

{parallex-0.1.3.dist-info → parallex-0.2.0.dist-info}/RECORD

@@ -1,10 +1,11 @@
 parallex/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-parallex/ai/batch_processor.py,sha256=wftq7-EKWbWO_tzz2PqVZa8XRVf7MoJlr7EcllX4-5I,1481
-parallex/ai/open_ai_client.py,sha256=Yvnvg5MGEyQrmN3HF5k8fEWse9Slthy3J-oumO6ZKkQ,1459
+parallex/ai/batch_processor.py,sha256=O5q_jaIU0VI93p7Riq4aZ_qUiN9Omxp5GOfn0IqEYgo,1361
+parallex/ai/open_ai_client.py,sha256=TRH78oYod_EWpp3hjEh097OT7hwsQmtv44_j3X9Frxo,2047
 parallex/ai/output_processor.py,sha256=P6ak7cblRHnsR1W7oEtbOGM7zd7tzZbRKigixQaXWyw,966
-parallex/ai/uploader.py,sha256=M8g8dC_bwiGNDI_S5qxcRqJljDu6KSan_eIcQWA-ERA,3162
+parallex/ai/uploader.py,sha256=92P0LxLuRgtjtD4kLtM0n8WUww_8-GImLxb3pbl-kkg,3174
 parallex/file_management/converter.py,sha256=Rj-93LXNl2gCY-XUOCZv7DdCNI2-GyRpS5FobnTqwzo,1111
 parallex/file_management/file_finder.py,sha256=BPvrkxZlwOYmRXzzS138wGTsVzuhDIKfQZn0CISUj3o,1598
+parallex/file_management/remote_file_handler.py,sha256=jsI9NhOrKQR8K3yo536lGplVBGis9XY0G4dRpumgWFM,213
 parallex/file_management/utils.py,sha256=WMdXd9UOFbJDHnL2IWfDXyyD2jhwnGtpCVI_npiSlIk,98
 parallex/models/batch_file.py,sha256=JwARFB48sMOTN-wf7J5YbsWIac2rxXnZ4fBABFESA0M,405
 parallex/models/image_file.py,sha256=LjQne2b6rIDWpQpdYT41KXNDWpg5kv9bkM1SCx6jnAI,402
@@ -12,10 +13,10 @@ parallex/models/page_response.py,sha256=KADCAV3XnkqWm-q_FBCfbt5nqDbiHg9MroZvFXaB
 parallex/models/parallex_callable_output.py,sha256=CkJKA8mwsc5olNnG1K6nrWUu4xTkJvp8bp3SSPQEX5c,465
 parallex/models/raw_file.py,sha256=Nlv6u_jlDCXDgU2_Ff7DRbDCx27pB1NZugNhEoaBMQU,483
 parallex/models/upload_batch.py,sha256=jrnds9ryXg9drL4TF8TGimMVTCDfKaWsBzFv_ed0i88,2068
-parallex/parallex.py,sha256=RvSDNyjrE7Fd3LuatEt18_1sR_btPJ-i6uxjNr_dGh0,4502
+parallex/parallex.py,sha256=EkD_kZevDu0UBpRet3nsvIr826f7uBHiT0JA5hR3E8c,5117
 parallex/utils/constants.py,sha256=c6i_-OSfCXAzW9ILzddSSHfldqHnsPEID3G3VYGYXUg,362
-parallex/utils/logger.py,sha256=5dpTogztRq4NCgYWnbbkFNx3V2sFCN-Mtoagwj8i18Q,505
-parallex-0.1.3.dist-info/LICENSE,sha256=wPwCqGrisXnEcpaUxSO79C2mdOUTbtjhLjyy8mVW6p8,1046
-parallex-0.1.3.dist-info/METADATA,sha256=oOHuVOKgITHhEjzSQQW4Odo06tZjFfY1B2AVR8q9SDg,3444
-parallex-0.1.3.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-parallex-0.1.3.dist-info/RECORD,,
+parallex/utils/logger.py,sha256=i3ZZ7YTUmhUStbvVME67F9ffnkLOv5ijm7wVUyJT8Ys,440
+parallex-0.2.0.dist-info/LICENSE,sha256=wPwCqGrisXnEcpaUxSO79C2mdOUTbtjhLjyy8mVW6p8,1046
+parallex-0.2.0.dist-info/METADATA,sha256=Aq2RRlLwkXcLZ_wNXLZAsydFYJqTSe47eVhAq78oja8,3416
+parallex-0.2.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+parallex-0.2.0.dist-info/RECORD,,