diff --git a/src/extraction_review/process_file.py b/src/extraction_review/process_file.py index 6f434df..d598548 100644 --- a/src/extraction_review/process_file.py +++ b/src/extraction_review/process_file.py @@ -25,14 +25,6 @@ class FileEvent(StartEvent): file_id: str -class DownloadFileEvent(Event): - pass - - -class FileDownloadedEvent(Event): - pass - - class Status(Event): level: Literal["info", "warning", "error"] message: str @@ -42,10 +34,6 @@ class ExtractJobStartedEvent(Event): pass -class ExtractJobCompletedEvent(Event): - pass - - class ExtractedEvent(Event): data: ExtractedData @@ -68,64 +56,9 @@ class ProcessFileWorkflow(Workflow): """Extract structured data from a document and save it for review.""" @step() - async def run_file(self, event: FileEvent, ctx: Context) -> DownloadFileEvent: - """Start extraction for the uploaded document.""" - logger.info(f"Running file {event.file_id}") - async with ctx.store.edit_state() as state: - state.file_id = event.file_id - return DownloadFileEvent() - - @step() - async def download_file( + async def start_extraction( self, - event: DownloadFileEvent, - ctx: Context[ExtractionState], - llama_cloud_client: Annotated[ - AsyncLlamaCloud, Resource(get_llama_cloud_client) - ], - ) -> FileDownloadedEvent: - """Retrieve the document from cloud storage for processing.""" - state = await ctx.store.get_state() - if state.file_id is None: - raise ValueError("File ID is not set") - try: - files = await llama_cloud_client.files.query( - filter=Filter(file_ids=[state.file_id]) - ) - file_metadata = files.items[0] - file_url = await llama_cloud_client.files.get(file_id=state.file_id) - - temp_dir = tempfile.gettempdir() - filename = file_metadata.name - file_path = os.path.join(temp_dir, filename) - client = httpx.AsyncClient() - # Report progress to the UI - logger.info(f"Downloading file {file_url.url} to {file_path}") - - async with client.stream("GET", file_url.url) as response: - with open(file_path, "wb") as f: - async for chunk in response.aiter_bytes(): - f.write(chunk) - logger.info(f"Downloaded file {file_url.url} to {file_path}") - async with ctx.store.edit_state() as state: - state.file_path = file_path - state.filename = filename - return FileDownloadedEvent() - - except Exception as e: - logger.error(f"Error downloading file {state.file_id}: {e}", exc_info=True) - ctx.write_event_to_stream( - Status( - level="error", - message=f"Error downloading file {state.file_id}: {e}", - ) - ) - raise e - - @step() - async def process_file( - self, - event: FileDownloadedEvent, + event: FileEvent, ctx: Context[ExtractionState], llama_cloud_client: Annotated[ AsyncLlamaCloud, Resource(get_llama_cloud_client) @@ -140,54 +73,74 @@ class ProcessFileWorkflow(Workflow): ), ], ) -> ExtractJobStartedEvent: - """Extract structured data fields from the document.""" - state = await ctx.store.get_state() - if state.file_path is None or state.filename is None: - raise ValueError("File path or filename is not set") + """Download document and start extraction job.""" + file_id = event.file_id + logger.info(f"Running file {file_id}") + + # Download file from cloud storage + try: + files = await llama_cloud_client.files.query( + filter=Filter(file_ids=[file_id]) + ) + file_metadata = files.items[0] + file_url = await llama_cloud_client.files.get(file_id=file_id) + + temp_dir = tempfile.gettempdir() + filename = file_metadata.name + file_path = os.path.join(temp_dir, filename) + client = httpx.AsyncClient() + # Report progress to the UI + logger.info(f"Downloading file {file_url.url} to {file_path}") + + async with client.stream("GET", file_url.url) as response: + with open(file_path, "wb") as f: + async for chunk in response.aiter_bytes(): + f.write(chunk) + logger.info(f"Downloaded file {file_url.url} to {file_path}") + + except Exception as e: + logger.error(f"Error downloading file {file_id}: {e}", exc_info=True) + ctx.write_event_to_stream( + Status( + level="error", + message=f"Error downloading file {file_id}: {e}", + ) + ) + raise e + + # Start extraction job # track the content of the file, so as to be able to de-duplicate - file_content = Path(state.file_path).read_bytes() - logger.info(f"Extracting data from file {state.filename}") + file_content = Path(file_path).read_bytes() + logger.info(f"Extracting data from file {filename}") ctx.write_event_to_stream( - Status(level="info", message=f"Extracting data from file {state.filename}") + Status(level="info", message=f"Extracting data from file {filename}") ) extract_job = await llama_cloud_client.extraction.run( config=extract_config.settings.model_dump(), data_schema=extract_config.json_schema, - file_id=state.file_id, + file_id=file_id, project_id=project_id, ) - async with ctx.store.edit_state() as st: - st.file_hash = hashlib.sha256(file_content).hexdigest() - st.extract_job_id = extract_job.id + + # Save state (mutation at end of step) + async with ctx.store.edit_state() as state: + state.file_id = file_id + state.file_path = file_path + state.filename = filename + state.file_hash = hashlib.sha256(file_content).hexdigest() + state.extract_job_id = extract_job.id return ExtractJobStartedEvent() @step() - async def wait_for_extract_job( + async def complete_extraction( self, event: ExtractJobStartedEvent, ctx: Context[ExtractionState], llama_cloud_client: Annotated[ AsyncLlamaCloud, Resource(get_llama_cloud_client) ], - ) -> ExtractJobCompletedEvent: - state = await ctx.store.get_state() - if state.extract_job_id is None: - raise ValueError("Job ID cannot be null when waiting for its completion") - await llama_cloud_client.extraction.jobs.wait_for_completion( - state.extract_job_id - ) - return ExtractJobCompletedEvent() - - @step() - async def get_extraction_job_result( - self, - event: ExtractJobCompletedEvent, - ctx: Context[ExtractionState], - llama_cloud_client: Annotated[ - AsyncLlamaCloud, Resource(get_llama_cloud_client) - ], extract_config: Annotated[ ExtractConfig, ResourceConfig( @@ -197,17 +150,27 @@ class ProcessFileWorkflow(Workflow): description="Configuration for document extraction quality and features", ), ], - ) -> ExtractedEvent | ExtractedInvalidEvent: - """Process the completed extraction job and validate results.""" + ) -> StopEvent: + """Wait for extraction to complete, validate results, and save for review.""" state = await ctx.store.get_state() if state.extract_job_id is None: - raise ValueError("Job ID cannot be null when getting its result") + raise ValueError("Job ID cannot be null when waiting for its completion") + + # Wait for extraction job to complete + await llama_cloud_client.extraction.jobs.wait_for_completion( + state.extract_job_id + ) + + # Get extraction result extracted_result = await llama_cloud_client.extraction.jobs.get_result( state.extract_job_id ) extract_run = await llama_cloud_client.extraction.runs.get( run_id=extracted_result.run_id ) + + # Validate and parse extraction result + extracted_event: ExtractedEvent | ExtractedInvalidEvent try: logger.info(f"Extracted data: {extracted_result}") # Create dynamic Pydantic model from JSON schema @@ -220,10 +183,10 @@ class ProcessFileWorkflow(Workflow): file_id=state.file_id, file_hash=state.file_hash, ) - return ExtractedEvent(data=data) + extracted_event = ExtractedEvent(data=data) except InvalidExtractionData as e: logger.error(f"Error validating extracted data: {e}", exc_info=True) - return ExtractedInvalidEvent(data=e.invalid_item) + extracted_event = ExtractedInvalidEvent(data=e.invalid_item) except Exception as e: logger.error( f"Error extracting data from file {state.filename}: {e}", @@ -237,42 +200,39 @@ class ProcessFileWorkflow(Workflow): ) raise e - @step() - async def record_extracted_data( - self, - event: ExtractedEvent | ExtractedInvalidEvent, - ctx: Context, - llama_cloud_client: Annotated[ - AsyncLlamaCloud, Resource(get_llama_cloud_client) - ], - ) -> StopEvent: - """Save extracted data for human review and correction.""" - data = event.data.model_dump() + # Stream the extracted event to client + ctx.write_event_to_stream(extracted_event) + + # Save extracted data for review + extracted_data = extracted_event.data + data_dict = extracted_data.model_dump() # remove past data when reprocessing the same file - if event.data.file_hash is not None: + if extracted_data.file_hash is not None: await llama_cloud_client.beta.agent_data.delete_by_query( deployment_name=agent_name or "_public", collection=EXTRACTED_DATA_COLLECTION, filter={ "file_hash": { - "eq": event.data.file_hash, + "eq": extracted_data.file_hash, }, }, ) logger.info( - f"Removing past data for file {event.data.file_name} with hash {event.data.file_hash}" + f"Removing past data for file {extracted_data.file_name} with hash {extracted_data.file_hash}" ) # finally, save the new data item = await llama_cloud_client.beta.agent_data.agent_data( - data=data, + data=data_dict, deployment_name=agent_name or "_public", collection=EXTRACTED_DATA_COLLECTION, ) - logger.info(f"Recorded extracted data for file {event.data.file_name or ''}") + logger.info( + f"Recorded extracted data for file {extracted_data.file_name or ''}" + ) ctx.write_event_to_stream( Status( level="info", - message=f"Recorded extracted data for file {event.data.file_name or ''}", + message=f"Recorded extracted data for file {extracted_data.file_name or ''}", ) ) return StopEvent(result=item.id)