File size: 549 Bytes
34a5262
 
 
 
 
 
 
 
 
2080a31
34a5262
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
from pydantic import BaseModel, Field
from langgraph.graph import MessagesState

class SQLAgentState(MessagesState):
    """State for the agent"""
    next_tool : str = ""
    tables_list: str = ""
    schema_of_table: str = ""
    query_gen : str= ""
    # check_query: str = ""
    execute_query : str = ""
    task_complete: bool = False
    response_to_user: str= ""
    current_task: str = ""
    query: str = "" ## query of the human stored in it 

class DBQuery(BaseModel):
    query: str = Field(..., description="The SQL query to execute")