Slim the steps down (#193)

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)