@@ -90,7 +90,7 @@ def selector(messages):
9090 "cached_questions_and_schemas"
9191 ) is not None and cache_result .get ("contains_pre_run_results" ):
9292 decision = "sql_query_correction_agent"
93- if (
93+ elif (
9494 cache_result .get ("cached_questions_and_schemas" ) is not None
9595 and cache_result .get ("contains_pre_run_results" ) is False
9696 ):
@@ -99,12 +99,20 @@ def selector(messages):
9999 decision = "question_decomposition_agent"
100100
101101 elif messages [- 1 ].source == "question_decomposition_agent" :
102- decomposition_result = json .loads (messages [- 1 ].content )
103-
104- if len (decomposition_result ["entities" ]) == 1 :
105- decision = "sql_schema_selection_agent"
106- else :
107- decision = "parallel_sql_flow_agent"
102+ # Fix: Ensure decomposition_result is properly handled as a list of dictionaries
103+ try :
104+ decomposition_result = json .loads (messages [- 1 ].content )
105+
106+ if isinstance (decomposition_result , list ) and len (decomposition_result ) == 1 :
107+ decision = "sql_schema_selection_agent"
108+ elif isinstance (decomposition_result , list ) and len (decomposition_result ) > 1 :
109+ decision = "parallel_sql_flow_agent"
110+ else :
111+ logging .error ("Unexpected format in decomposition result: %s" , decomposition_result )
112+ raise ValueError ("Invalid decomposition result format" )
113+ except (json .JSONDecodeError , ValueError ) as e :
114+ logging .error ("Error processing decomposition result: %s" , e )
115+ decision = "question_decomposition_agent" # Retry or handle gracefully
108116
109117 elif messages [- 1 ].source == "sql_schema_selection_agent" :
110118 decision = "sql_query_generation_agent"
0 commit comments