File size: 1,366 Bytes
e8e5773
 
da99b31
 
e8e5773
 
 
da99b31
 
 
 
 
e8e5773
 
da99b31
 
 
 
e8e5773
 
 
da99b31
 
 
 
 
 
 
 
 
e8e5773
da99b31
 
 
 
 
e8e5773
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
import pandas as pd
import numpy as np
import math, json, re, ast
from schemas import ScenarioPlan, TaskSpec

class ScenarioEngine:
    @staticmethod
    def render_plan(plan: ScenarioPlan, results: dict) -> str:
        sections = ["# Scenario Output\n"]
        for t in plan.tasks:
            sections.append(ScenarioEngine._render_task(t, results))
        return "\n".join(sections)

    @staticmethod
    def _df(v):
        if isinstance(v, pd.DataFrame): return v
        if isinstance(v, list) and v and isinstance(v[0], dict): return pd.DataFrame(v)
        if isinstance(v, dict): return pd.DataFrame([v])
        return None

    @staticmethod
    def _render_task(t: TaskSpec, results: dict) -> str:
        out=[f"## {t.title}\n"]
        df=ScenarioEngine._df(results.get(t.data_key)) if t.data_key else None
        if df is None: return "\n".join(out+["_No data_"])
        if t.filter: df=df.query(t.filter)
        if t.group_by or t.agg: df=df.groupby(t.group_by).agg("first").reset_index()
        if t.sort_by: df=df.sort_values(by=t.sort_by, ascending=(t.sort_dir=="asc"))
        if t.top: df=df.head(t.top)
        if t.fields: df=df[t.fields]

        if t.format=="list":
            out += [f"- {row.to_dict()}" for _,row in df.iterrows()]
        else:
            out.append(df.to_markdown(index=False))
        return "\n".join(out)