fixes

Author: SwetnaT-ITS

policyengine_api/ai_prompts/simulation_analysis_prompt.py

@@ -90,9 +90,9 @@ def generate_simulation_analysis_prompt(params: InboundParameters) -> str:
 
     audience_description: str = audience_descriptions[parameters.audience]
 
-    country_id_uppercase: Annotated[
-        str, "Uppercase two-letter country ID"
-    ] = parameters.country_id.upper()
+    country_id_uppercase: Annotated[str, "Uppercase two-letter country ID"] = (
+        parameters.country_id.upper()
+    )
 
     impact_budget: str = json.dumps(parameters.impact["budget"])
     impact_intra_decile: dict[str, Any] = json.dumps(

policyengine_api/jobs/calculate_economy_simulation_job.py

@@ -254,22 +254,22 @@ def run(
                             **comparison_data,
                         }
                     )
-                    sim_config: dict[
-                        str, Any
-                    ] = self.api_v2._setup_sim_options(
-                        country_id=country_id,
-                        scope="macro",
-                        reform_policy=reform_policy,
-                        baseline_policy=baseline_policy,
-                        time_period=time_period,
-                        region=region,
-                        dataset=dataset,
-                        model_version=COUNTRY_PACKAGE_VERSIONS[country_id],
-                        data_version=(
-                            uk_dataset_version
-                            if country_id == "uk"
-                            else us_dataset_version
-                        ),
+                    sim_config: dict[str, Any] = (
+                        self.api_v2._setup_sim_options(
+                            country_id=country_id,
+                            scope="macro",
+                            reform_policy=reform_policy,
+                            baseline_policy=baseline_policy,
+                            time_period=time_period,
+                            region=region,
+                            dataset=dataset,
+                            model_version=COUNTRY_PACKAGE_VERSIONS[country_id],
+                            data_version=(
+                                uk_dataset_version
+                                if country_id == "uk"
+                                else us_dataset_version
+                            ),
+                        )
                     )
 
                     api_v2_execution = self.api_v2.run(sim_config)
@@ -369,15 +369,17 @@ def run(
 
                     v2_country_package_version = api_v2_output["model_version"]
 
-                    completion_log: V2V1Comparison = V2V1Comparison.model_validate(
-                        {
-                            **comparison_data,
-                            "v1_impact": impact,
-                            "v2_impact": api_v2_output,
-                            "v1_v2_diff": None,
-                            "v2_country_package_version": v2_country_package_version,
-                            "message": "APIv2 job completed",
-                        }
+                    completion_log: V2V1Comparison = (
+                        V2V1Comparison.model_validate(
+                            {
+                                **comparison_data,
+                                "v1_impact": impact,
+                                "v2_impact": api_v2_output,
+                                "v1_v2_diff": None,
+                                "v2_country_package_version": v2_country_package_version,
+                                "message": "APIv2 job completed",
+                            }
+                        )
                     )
 
                     logger.log_struct(
@@ -390,15 +392,17 @@ def run(
                         y=api_v2_output,
                     )
                     # Push relevant info into logging schema
-                    comparison_log: V2V1Comparison = V2V1Comparison.model_validate(
-                        {
-                            **comparison_data,
-                            "v1_impact": impact,
-                            "v2_impact": api_v2_output,
-                            "v1_v2_diff": v1_v2_diff,
-                            "v2_country_package_version": v2_country_package_version,
-                            "message": "APIv2 job comparison with APIv1 completed",
-                        }
+                    comparison_log: V2V1Comparison = (
+                        V2V1Comparison.model_validate(
+                            {
+                                **comparison_data,
+                                "v1_impact": impact,
+                                "v2_impact": api_v2_output,
+                                "v1_v2_diff": v1_v2_diff,
+                                "v2_country_package_version": v2_country_package_version,
+                                "message": "APIv2 job comparison with APIv1 completed",
+                            }
+                        )
                     )
                     logger.log_struct(
                         comparison_log.model_dump(mode="json"),
@@ -552,9 +556,9 @@ def _compute_economy(
     def _create_simulation_uk(
         self, country, reform, region, time_period
     ) -> Microsimulation:
-        CountryMicrosimulation: Type[
-            Microsimulation
-        ] = country.country_package.Microsimulation
+        CountryMicrosimulation: Type[Microsimulation] = (
+            country.country_package.Microsimulation
+        )
 
         simulation = CountryMicrosimulation(
             reform=reform,