configurable config/split

app.py

@@ -72,14 +72,12 @@ def get_table_name(config: str | None, split: str | None, config_choices: list[s
     return table_name.lower()
 
 
-def get_prompt_messages(card_data: dict[str, Any], natural_language_query: str):
+def get_prompt_messages(card_data: dict[str, Any], natural_language_query: str, config: str | None, split: str | None):
     config_choices = get_config_choices(card_data)
     split_choices = get_split_choices(card_data)
 
-    chosen_config = config_choices[0] if len(config_choices) > 0 else None
-    chosen_split = split_choices[0] if len(split_choices) > 0 else None
-    table_name = get_table_name(chosen_config, chosen_split, config_choices, split_choices)
-    features = card_data[chosen_config]["features"]
+    table_name = get_table_name(config, split, config_choices, split_choices)
+    features = card_data[config]["features"]
     messages = [
         {
             "role": "system",
@@ -111,9 +109,9 @@ def get_split_choices(card_data: dict[str, Any]) -> list[str]:
     return list(splits)
 
 
-def query_dataset(hub_repo_id, card_data, query):
+def query_dataset(hub_repo_id, card_data, query, config, split):
     card_data = json.loads(card_data)
-    messages = get_prompt_messages(card_data, query)
+    messages = get_prompt_messages(card_data, query, config, split)
     api_key = os.environ["API_KEY_TOGETHER_AI"].strip()
     response = requests.post(
         "https://api.together.xyz/v1/chat/completions",
@@ -180,9 +178,27 @@ with gr.Blocks() as demo:
             lines=1,
             visible=False,
         )
+    with gr.Row():
+        card_data = gr.Code(label="Card data", language="json", visible=False)
+
+    @gr.render(inputs=[card_data])
+    def show_config_split_choices(data):
+        try:
+            data = json.loads(data.strip())
+            config_choices = get_config_choices(data)
+            split_choices = get_split_choices(data)
+        except Exception:
+            config_choices = []
+            split_choices = []
+        
+        initial_config = config_choices[0] if len(config_choices) > 0 else None
+        initial_split = split_choices[0] if len(split_choices) > 0 else None
+        with gr.Row():
+            with gr.Column():
+                config_selection = gr.Dropdown(label="Config Name", choices=config_choices, value=initial_config)
+            with gr.Column():
+                split_selection = gr.Dropdown(label="Split Name", choices=split_choices, value=initial_split)
 
-    @gr.render(triggers=[search_in.submit])
-    def show_config_split_choices():
         with gr.Row():
             with gr.Column():
                 btn = gr.Button("Show Dataset")
@@ -190,8 +206,6 @@ with gr.Blocks() as demo:
                 btn2 = gr.Button("Query Dataset")
         with gr.Row():
             search_out = gr.HTML(label="Search Results")
-        with gr.Row():
-            card_data = gr.Code(label="Card data", language="json", visible=False)
         gr.on(
             [btn.click, search_in.submit],
             fn=get_iframe,
@@ -205,7 +219,7 @@ with gr.Blocks() as demo:
         gr.on(
             [btn2.click, query.submit],
             fn=query_dataset,
-            inputs=[search_in, card_data, query],
+            inputs=[search_in, card_data, query, config_selection, split_selection],
             outputs=[sql_out, search_out],
         )