tools/state.py

from __future__ import annotations

from typing import Literal, TypedDict


QueryIntent = Literal["data_retrieval", "parameter_info", "both", "out_of_scope"]


class GraphState(TypedDict):
    user_query: str
    query_intent: QueryIntent
    is_valid: bool
    filters: dict
    extracted_data: list[dict]
    leading_country_result: list[dict]
    parameter_data: list[dict]
    error_message: str
    final_response: str
    conversation_history: list[dict]


def create_initial_state(
    user_query: str,
    conversation_history: list[dict] | None = None,
    prior_filters: dict | None = None,
) -> GraphState:
    return GraphState(
        user_query=user_query,
        query_intent="out_of_scope",
        is_valid=True,
        filters=prior_filters or {},
        extracted_data=[],
        leading_country_result=[],
        parameter_data=[],
        error_message="",
        final_response="",
        conversation_history=conversation_history or [],
    )