Skip to content

API Reference#

Creates a state graph for processing datasets.

This function orchestrates the creation of a workflow for handling table data. It sets up nodes for reading files and analyzing data based on provided parameters. The graph dynamically routes based on the presence of file attachments in the input state.

Parameters:

Name Type Description Default
llm Runnable

The primary language model for processing user input.

 required
pybox_manager BasePyBoxManager

A python code sandbox delegator, used to execute the data analysis code generated by llm.

 required
session_id str | None

An optional session identifier used to associate with pybox. Defaults to None.

 None
workdir Path | None

The working directory for pybox operations. Defaults to None.

 None
error_trace_cleanup bool

Flag to clean up error traces. Defaults to False.

 False
nlines int | None

Number of lines to read for preview. Defaults to None.

 None
vlm BaseLanguageModel | None

Optional vision language model for processing images. Defaults to None.

 None
safety_llm Runnable | None

Model used for safety classification of inputs. Defaults to None.

 None
dataset_retriever BaseRetriever | None

Component to retrieve datasets. Defaults to None.

 None
normalize_llm BaseLanguageModel | None

Model for data normalization tasks. Defaults to None.

 None
locate str | None

The locale of the user. Defaults to None.

 required
checkpointer BaseCheckpointSaver | None

Component for saving checkpoints. Defaults to None.

 None
verbose bool

Flag to enable verbose logging. Defaults to False.

 False

Returns:

Name Type Description
CompiledStateGraph CompiledStateGraph

A compiled state graph representing the table processing workflow.
Source code in src/tablegpt/agent/__init__.py 
 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
def create_tablegpt_graph(
    llm: BaseLanguageModel,
    pybox_manager: BasePyBoxManager,
    *,
    session_id: str | None = None,
    workdir: Path | None = None,
    error_trace_cleanup: bool = False,
    nlines: int | None = None,
    vlm: BaseLanguageModel | None = None,
    safety_llm: Runnable | None = None,
    dataset_retriever: BaseRetriever | None = None,
    normalize_llm: BaseLanguageModel | None = None,
    locale: str | None = None,
    checkpointer: BaseCheckpointSaver | None = None,
    verbose: bool = False,
) -> CompiledStateGraph:
    """Creates a state graph for processing datasets.

    This function orchestrates the creation of a workflow for handling table data.
    It sets up nodes for reading files and analyzing data based on provided parameters.
    The graph dynamically routes based on the presence of file attachments in the input state.

    Args:
        llm (Runnable): The primary language model for processing user input.
        pybox_manager (BasePyBoxManager):  A python code sandbox delegator, used to execute the data analysis code generated by llm.
        session_id (str | None, optional): An optional session identifier used to associate with `pybox`. Defaults to None.
        workdir (Path | None, optional): The working directory for `pybox` operations. Defaults to None.
        error_trace_cleanup (bool, optional): Flag to clean up error traces. Defaults to False.
        nlines (int | None, optional): Number of lines to read for preview. Defaults to None.
        vlm (BaseLanguageModel | None, optional): Optional vision language model for processing images. Defaults to None.
        safety_llm (Runnable | None, optional): Model used for safety classification of inputs. Defaults to None.
        dataset_retriever (BaseRetriever | None, optional): Component to retrieve datasets. Defaults to None.
        normalize_llm (BaseLanguageModel | None, optional): Model for data normalization tasks. Defaults to None.
        locate (str | None, optional): The locale of the user. Defaults to None.
        checkpointer (BaseCheckpointSaver | None, optional): Component for saving checkpoints. Defaults to None.
        verbose (bool, optional): Flag to enable verbose logging. Defaults to False.

    Returns:
        CompiledStateGraph: A compiled state graph representing the table processing workflow.
    """
    workflow = StateGraph(AgentState)
    file_reading_graph = create_file_reading_workflow(
        nlines=nlines,
        llm=llm,
        pybox_manager=pybox_manager,
        workdir=workdir,
        session_id=session_id,
        normalize_llm=normalize_llm,
        locale=locale,
        verbose=verbose,
    )
    data_analyze_graph = create_data_analyze_workflow(
        llm=llm,
        pybox_manager=pybox_manager,
        workdir=workdir,
        session_id=session_id,
        error_trace_cleanup=error_trace_cleanup,
        vlm=vlm,
        safety_llm=safety_llm,
        dataset_retriever=dataset_retriever,
        verbose=verbose,
    )

    def router(state: AgentState) -> str:
        # Must have at least one message when entering this router
        last_message = state["messages"][-1]
        if last_message.additional_kwargs.get("attachments"):
            return "file_reading_graph"
        return "data_analyze_graph"

    workflow.add_node("file_reading_graph", file_reading_graph)
    workflow.add_node("data_analyze_graph", data_analyze_graph)

    workflow.add_conditional_edges(START, router)
    workflow.add_edge("file_reading_graph", END)
    workflow.add_edge("data_analyze_graph", END)

    return workflow.compile(checkpointer=checkpointer, debug=verbose)