Fix for python 3.10

climateqa/engine/talk_to_data/workflow.py

@@ -1,6 +1,6 @@
 import os
 
-from typing import Any, Callable, NotRequired, TypedDict
+from typing import Any, Callable, TypedDict, Optional
 import pandas as pd
 
 from plotly.graph_objects import Figure
@@ -35,9 +35,9 @@ class TableState(TypedDict):
     """
     table_name: str
     params: dict[str, Any]
-    sql_query: NotRequired[str]
-    dataframe: NotRequired[pd.DataFrame | None]
-    figure: NotRequired[Callable[..., Figure]]
+    sql_query: Optional[str]
+    dataframe: Optional[pd.DataFrame | None]
+    figure: Optional[Callable[..., Figure]]
     status: str
 
 class PlotState(TypedDict):
@@ -59,7 +59,7 @@ class State(TypedDict):
     user_input: str
     plots: list[str]
     plot_states: dict[str, PlotState]
-    error: NotRequired[str]
+    error: Optional[str]
     
 async def drias_workflow(user_input: str) -> State:
     """Performs the complete workflow of Talk To Drias : from user input to sql queries, dataframes and figures generated
@@ -187,8 +187,8 @@ async def find_param(state: State, param_name:str, table: str) -> dict[str, Any]
 
 class Location(TypedDict):
     location: str
-    latitude: NotRequired[str]
-    longitude: NotRequired[str]
+    latitude: Optional[str]
+    longitude: Optional[str]
 
 async def find_location(user_input: str, table: str) -> Location:
     print(f"---- Find location in table {table} ----")