Merge pull request #3 from DeanNeaht/fix/replace-asserts-with-proper-errors

Replace assert statements with proper exception handling

src/investments_review/exceptions.py

@@ -0,0 +1,22 @@
+class LlamaCloudAPIError(Exception):
+    """Base exception for all LlamaCloud API-related errors."""
+
+    pass
+
+
+class ClassificationError(LlamaCloudAPIError):
+    """Exception raised when document classification fails or returns invalid data."""
+
+    pass
+
+
+class ExtractionError(LlamaCloudAPIError):
+    """Exception raised when data extraction fails or returns invalid data."""
+
+    pass
+
+
+class SheetParsingError(LlamaCloudAPIError):
+    """Exception raised when sheet parsing fails or returns invalid data."""
+
+    pass

src/investments_review/presentations/workflow.py

@@ -9,6 +9,7 @@ from workflows import Context, Workflow, step
 from workflows.events import Event, StopEvent
 from workflows.resource import Resource
 
+from ..exceptions import ClassificationError, ExtractionError
 from ..shared import FileEvent, FileUploadedEvent, get_llama_cloud_client
 from .models import BoardUpdateDeck, ManagementPresentation, rules
 
@@ -81,9 +82,8 @@ class PresentationWorkflow(Workflow):
         logging.info("Finished classification")
         result_item = result.items[0]  # there is only one classified file
         if result_item.result is not None:
-            assert result_item.result.type is not None, (
-                "Classification type should not be None"
-            )
+            if result_item.result.type is None:
+                raise ClassificationError("Classification type should not be None")
             logging.info(f"Classified document as: {result_item.result.type}")
             event = ClassificationEvent(
                 category=result_item.result.type,
@@ -115,7 +115,8 @@ class PresentationWorkflow(Workflow):
         )
         logging.info("Finished extracting details from presentation file")
         if result.data is not None:
-            assert isinstance(result.data, dict), "Data should be a dictionary"
+            if not isinstance(result.data, dict):
+                raise ExtractionError("Data should be a dictionary")
             details = schema.model_validate(result.data)
             return ExtractionEvent(final_result=details.to_string())
         return ExtractionEvent(error="Could not extract details from document")

src/investments_review/sheets/workflow.py

@@ -11,6 +11,7 @@ from workflows import Context, Workflow, step
 from workflows.events import Event, StopEvent
 from workflows.resource import Resource
 
+from ..exceptions import SheetParsingError
 from ..shared import FileEvent, FileUploadedEvent, get_llama_cloud_client
 from .llm import OpenAILLM, get_llm
 from .models import InvestmentSheetAnalysis
@@ -87,11 +88,13 @@ class SheetWorkflow(Workflow):
         file_paths = []
         if result.success:
             logging.info("Starting to download Parquet files...")
-            assert result.regions is not None, (
-                "Regions should have been extracted if the job was successfull"
-            )
+            if result.regions is None:
+                raise SheetParsingError(
+                    "Regions should have been extracted if the job was successful"
+                )
             for region in result.regions:
-                assert region.region_id is not None, "Region should have an ID"
+                if region.region_id is None:
+                    raise SheetParsingError("Region should have an ID")
                 parquet_region_resp = (
                     await llama_cloud_client.beta.sheets.get_result_table(
                         region_type=region.region_type,  # type: ignore