File size: 1,147 Bytes
38830c1 | 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 | from __future__ import annotations
from typing import Any, Literal, TypedDict
class BIState(TypedDict, total=False):
question: str
next_step: Literal["sql", "analysis", "reflection", "visualization", "end"] | None
plan: dict[str, Any] | None
sql: dict[str, Any] | None
analysis: dict[str, Any] | None
reflection: dict[str, Any] | None
visualization: dict[str, Any] | None
memories: list[dict[str, Any]]
def build_initial_state(question: str, memories: list[dict[str, Any]] | None = None) -> BIState:
return {
"question": question,
"next_step": None,
"plan": None,
"sql": None,
"analysis": None,
"reflection": None,
"visualization": None,
"memories": memories or [],
}
def final_answer(state: BIState) -> dict[str, Any]:
return {
"question": state["question"],
"plan": state.get("plan"),
"sql": state.get("sql"),
"analysis": state.get("analysis"),
"reflection": state.get("reflection"),
"visualization": state.get("visualization"),
"memories": state.get("memories", []),
}
|