refactor: enhance equipment simulation with normalized tag keys, expanded distribution parameters, and improved error logging

src/aeros_equipment/service.py

@@ -167,14 +167,15 @@ async def update_node(
             json=updateNodeReq,
             headers={"Content-Type": "application/json"},
         )
-
+        response.raise_for_status()
         result = response.json()
-
+        log.info("Aeros UpdateEquipmentDistributions response: %s", result)
         return result
 
     except Exception as e:
+        log.error("Failed to update equipment nodes in Aeros: %s", str(e))
         raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Aeros update failed: {str(e)}"
         )
 
 
@@ -259,21 +260,26 @@ async def get_asset_batch(location_tags: List[str], nr_location_tags: List[str],
         
         for row in repairable_result:
             location_tag = row[0]
+            if not location_tag:
+                continue
+            
+            # Normalize key
+            tag_key = location_tag.strip().upper()
             raw_parameters = row[1]
             parameters = raw_parameters if isinstance(raw_parameters, dict) else json.loads(raw_parameters) if raw_parameters else {}
             
             try:
                 distribution, reldisp1, reldisp2 = get_distribution_from_db(parameters, is_repairable=True)
                 
-                results[location_tag]["cmDisType"] = "Normal"
+                results[tag_key]["cmDisType"] = "Normal"
-                results[location_tag]["cmDisP1"] = 6
+                results[tag_key]["cmDisP1"] = 6
-                results[location_tag]["cmDisP2"] = 3
+                results[tag_key]["cmDisP2"] = 3
-                results[location_tag]["relDisType"] = distribution
+                results[tag_key]["relDisType"] = distribution
-                results[location_tag]["relDisP1"] = reldisp1
+                results[tag_key]["relDisP1"] = reldisp1
-                results[location_tag]["relDisP2"] = reldisp2
+                results[tag_key]["relDisP2"] = reldisp2
-                results[location_tag]["parameters"] = parameters
+                results[tag_key]["parameters"] = parameters
             except Exception as e:
-                raise Exception(f"Error processing repairable item {location_tag}: {e}")
+                log.error(f"Error processing repairable item {location_tag}: {e}")
 
     # Query non-repairable equipment data
     if nr_location_tags:
@@ -287,6 +293,11 @@ async def get_asset_batch(location_tags: List[str], nr_location_tags: List[str],
         
         for row in non_repairable_result:
             location_tag = row[0]
+            if not location_tag:
+                continue
+                
+            # Normalize key
+            tag_key = location_tag.strip().upper()
             raw_parameters = row[1]
             parameters = raw_parameters if isinstance(raw_parameters, dict) else json.loads(raw_parameters) if raw_parameters else {}
             best_distribution = row[2]
@@ -294,15 +305,15 @@ async def get_asset_batch(location_tags: List[str], nr_location_tags: List[str],
             try:
                 distribution, reldisp1, reldisp2 = get_distribution_from_db(parameters, is_repairable=False, best_distribution=best_distribution)
                 
-                results[location_tag]["cmDisType"] = "Normal"
+                results[tag_key]["cmDisType"] = "Normal"
-                results[location_tag]["cmDisP1"] = 6
+                results[tag_key]["cmDisP1"] = 6
-                results[location_tag]["cmDisP2"] = 3
+                results[tag_key]["cmDisP2"] = 3
-                results[location_tag]["relDisType"] = distribution
+                results[tag_key]["relDisType"] = distribution
-                results[location_tag]["relDisP1"] = reldisp1
+                results[tag_key]["relDisP1"] = reldisp1
-                results[location_tag]["relDisP2"] = reldisp2
+                results[tag_key]["relDisP2"] = reldisp2
-                results[location_tag]["parameters"] = parameters
+                results[tag_key]["parameters"] = parameters
             except Exception as e:
-                raise Exception(f"Error processing non-repairable item {location_tag}: {e}")
+                log.error(f"Error processing non-repairable item {location_tag}: {e}")
 
     return results
 
@@ -366,7 +377,7 @@ def get_distribution_from_db(parameters: dict, is_repairable: bool = True, best_
             return "Normal", mu, sigma
     
 async def update_oh_interval_offset(*, aeros_db_session, overhaul_offset, overhaul_interval, project_name):
-    query = text("""
+    stmt = text("""
         UPDATE public."RegularNodes" rn
         SET "OHInterval" = :new_oh_interval,
             "OHOffset" = :new_oh_offset
@@ -377,13 +388,14 @@ async def update_oh_interval_offset(*, aeros_db_session, overhaul_offset, overha
         AND p."ProjectName" = :project_name
     """)
 
-    await aeros_db_session.execute(query, {
+    result = await aeros_db_session.execute(stmt, {
         "new_oh_interval": overhaul_interval,
         "new_oh_offset": overhaul_offset,
         "project_name": project_name
     })
     
     await aeros_db_session.commit()
+    log.info("Updated OH Interval/Offset for project %s. Rows affected: %s", project_name, result.rowcount)
 
 
 async def update_equipment_for_simulation(
@@ -528,7 +540,15 @@ async def update_equipment_for_simulation(
                         eq["relDisP1"] = 1.0
                         eq["relDisP2"] = MTBF
 
+                eq["relDisP3"] = 0
+                eq["relDisUnitCode"] = "UHour"
+                eq["cmDisP2"] = 0
+                eq["cmDisP3"] = 0
+                eq["cmDisUnitCode"] = "UHour"
+                eq["ohDisType"] = "Fixed"
                 eq["ohDisP1"] = overhaul_duration
+                eq["ohDisP2"] = 0
+                eq["ohDisP3"] = 0
                 eq["ohDisUnitCode"] = "UHour"
 
                 reqNodeInputs.append(eq)
@@ -547,24 +567,48 @@ async def update_equipment_for_simulation(
             if eq_name_raw in results:
                 continue
 
-            reliability = reliability_data.get(eq_name_raw, {})
+            # Lookup reliability with normalized key
+            reliability = reliability_data.get(eq_name, {})
 
+            if reliability:
+                log.info("Found reliability data for %s", eq_name_raw)
                 eq["cmDisType"] = reliability.get("cmDisType", "Fixed")
                 eq["cmDisP1"] = reliability.get("cmDisP1", 0)
                 eq["cmDisP2"] = reliability.get("cmDisP2", 0)
+                eq["cmDisP3"] = 0
+                eq["cmDisUnitCode"] = "UHour"
                 eq["relDisType"] = reliability.get("relDisType", "Fixed")
                 eq["relDisP1"] = reliability.get("relDisP1", 0)
                 eq["relDisP2"] = reliability.get("relDisP2", 0)
+                eq["relDisP3"] = 0
+                eq["relDisUnitCode"] = "UHour"
+            else:
+                log.debug("No reliability data for %s, keeping existing or using defaults", eq_name_raw)
+                # Ensure existing unit codes and types are at least present if not already
+                eq.setdefault("cmDisType", "Fixed")
+                eq.setdefault("cmDisUnitCode", "UHour")
+                eq.setdefault("relDisType", "Fixed")
+                eq.setdefault("relDisUnitCode", "UHour")
+                eq.setdefault("cmDisP1", 0)
+                eq.setdefault("cmDisP2", 0)
+                eq.setdefault("cmDisP3", 0)
+                eq.setdefault("relDisP1", 0)
+                eq.setdefault("relDisP2", 0)
+                eq.setdefault("relDisP3", 0)
+
+            eq["ohDisType"] = "Fixed"
             eq["ohDisP1"] = overhaul_duration
+            eq["ohDisP2"] = 0
+            eq["ohDisP3"] = 0
             eq["ohDisUnitCode"] = "UHour"
 
             reqNodeInputs.append(eq)
             results[eq_name_raw] = {
-                "mttr": eq["cmDisP1"],
+                "mttr": eq.get("cmDisP1", 0),
-                "distribution": eq["relDisType"],
+                "distribution": eq.get("relDisType", "Fixed"),
-                "beta": eq["relDisP1"],
+                "beta": eq.get("relDisP1", 0),
-                "eta": eq["relDisP2"],
+                "eta": eq.get("relDisP2", 0),
-                "parameters": eq.get("parameters", {}),
+                "parameters": reliability.get("parameters", {}),
                 "oh_duration": overhaul_duration
             }
 

src/aeros_simulation/router.py

@@ -648,6 +648,7 @@ airflow_router = APIRouter()
 @airflow_router.post("/calculate_eaf_contribution", response_model=StandardResponse[dict])
 async def calculate_contribution(
     db_session: DbSession,
+    aeros_db_session: CollectorDbSession,
     simulation_in: SimulationInput,
     batch_num: int = Query(0, ge=0)
 ):
@@ -700,7 +701,15 @@ async def calculate_contribution(
             }
 
             results = await update_equipment_for_simulation(
-                db_session=db_session, project_name=project.project_name, schematic_name=simulation_in.SchematicName, custom_input=custom_input
+                db_session=db_session,
+                aeros_db_session=aeros_db_session,
+                project_name=project.project_name,
+                simulation_duration=simulation_in.SimDuration,
+                overhaul_duration=simulation_in.OverhaulDuration,
+                overhaul_interval=simulation_in.OverhaulInterval,
+                offset=simulation_in.OffSet,
+                schematic_name=simulation_in.SchematicName,
+                custom_input=custom_input
             )
 
             # await update_simulation(