langgraph-sql-agent/app.py at main · raphaelDuff/langgraph-sql-agent · GitHub

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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
from dotenv import load_dotenv

load_dotenv()

import pandas as pd
import plotly.express as px
import streamlit as st

from langchain_core.runnables import RunnableConfig

from franq_agent.graph import build_graph


st.set_page_config(
    page_title="Dr. Frankstein - Data Assistant", page_icon="🧟", layout="wide"
)
st.title("🧟 Dr. Frankstein - Data Assistant")
st.caption(
    "**F**ramework for **R**easoning and **A**nalytics using **N**atural **K**nowledge from **S**tructured **T**ables with **E**xplainable **I**nsights **N**arratively"
)


def _render_chart(
    df: pd.DataFrame, viz_type: str, title: str, viz_config: dict | None
) -> None:
    cfg = viz_config or {}

    if viz_type == "table":
        st.dataframe(df, use_container_width=True)

    elif viz_type == "bar" and cfg.get("x") and cfg.get("y"):
        fig = px.bar(df, x=cfg["x"], y=cfg["y"], color=cfg.get("color"), title=title)
        st.plotly_chart(fig, use_container_width=True)

    elif viz_type == "line" and cfg.get("x") and cfg.get("y"):
        fig = px.line(df, x=cfg["x"], y=cfg["y"], color=cfg.get("color"), title=title)
        st.plotly_chart(fig, use_container_width=True)

    elif viz_type == "pie" and cfg.get("x") and cfg.get("y"):
        fig = px.pie(df, names=cfg["x"], values=cfg["y"], title=title)
        st.plotly_chart(fig, use_container_width=True)

    else:
        st.dataframe(df, use_container_width=True)


if "graph" not in st.session_state:
    st.session_state.graph = build_graph()

if "thread_id" not in st.session_state:
    st.session_state.thread_id = "streamlit-session"

if "history" not in st.session_state:
    st.session_state.history = []


for turn in st.session_state.history:
    with st.chat_message("user"):
        st.write(turn["question"])

    with st.chat_message("assistant"):
        st.write(turn["answer"])

        if turn.get("sql"):
            with st.expander("🔍 SQL executed"):
                st.code(turn["sql"], language="sql")

        data = turn.get("data")
        viz_type = turn.get("data_viz_type")
        viz_config = turn.get("viz_config")

        if data and viz_type and viz_type != "none":
            df = pd.DataFrame(data)
            _render_chart(df, viz_type, turn["question"], viz_config)


question = st.chat_input("Ask a question about the data…")

if question:
    with st.chat_message("user"):
        st.write(question)

    with st.chat_message("assistant"):
        with st.spinner("Thinking…"):
            config = RunnableConfig(
                configurable={"thread_id": st.session_state.thread_id}
            )
            result = st.session_state.graph.invoke(
                {"question": question}, config=config
            )

        answer = result.get("final_answer", "No answer generated.")
        sql = result.get("last_sql_query")
        data_viz_type = result.get("data_viz_type")
        viz_config = result.get("viz_config")
        data = result.get("query_result") or []

        st.write(answer)

        if sql:
            with st.expander("🔍 SQL executed"):
                st.code(sql, language="sql")

        if data and data_viz_type and data_viz_type != "none":
            df = pd.DataFrame(data)
            _render_chart(df, data_viz_type, question, viz_config)

    st.session_state.history.append(
        {
            "question": question,
            "answer": answer,
            "sql": sql,
            "data_viz_type": data_viz_type,
            "viz_config": viz_config,
            "data": data,
        }
    )