File size: 865 Bytes
472e1d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from typing import List, Optional, Dict
import pandas as pd
from typing import TypedDict


class SQLAgentState(TypedDict):
    """
    Carries context through the text→SQL and plotting pipeline:
      - question: original NL input
      - sql_query: generated SQL
      - sql_result: raw query result DataFrame  
      - answer: final answer
      - error: any error messages
    """
    question: str
    db_info: dict
    sql_query: str
    sql_result: Optional[pd.DataFrame] = None
    answer: str = ""
    error: Optional[str] = None
    plot_path: Optional[str] = None
    response_md: str = ""
    step: Optional[str] = None
    visualization: Optional[str] = None
    visualization_reason: Optional[str] = None
    formatted_data_for_visualization: Optional[dict] = None
    visualization_output: Optional[str] = None
    off_topic: Optional[str] = None