feat: Embed chat for generating complex queries (#31) (#39)

In chat_query.py:
- Calls LLM from metadata chatbot asynchronously (requires credentials
AWS Bedrock access to Claude Sonnet 3-7)
- Generates loading icon when query is loading
- Queries are kept track of in self.queries

In query.py:
- Added complex_query_builder = ComplexQueryBuilder()

Ran linters

---------

Co-authored-by: Dan Birman <danbirman@gmail.com>

Author: sreyakumar

.gitignore

@@ -137,3 +137,6 @@ dmypy.json
 
 # MacOs
 **/.DS_Store
+
+# Local UI testing
+panel_chat_query.py

pyproject.toml

@@ -25,6 +25,8 @@ dependencies = [
     'aind-data-access-api[rds]',
     'aind-metadata-validator>=0.8.3',
     'flask',
+    'langchain',
+    'langchain_aws',
 ]
 
 [project.optional-dependencies]

src/aind_metadata_viz/app.py

@@ -2,14 +2,21 @@
 import panel as pn
 from aind_data_schema import __version__ as ads_version
 from aind_metadata_viz import database
-from aind_metadata_viz.utils import AIND_COLORS, COLOR_OPTIONS, hd_style, outer_style
+from aind_metadata_viz.utils import (
+    AIND_COLORS,
+    COLOR_OPTIONS,
+    hd_style,
+    outer_style,
+)
 from aind_metadata_viz.charts import file_present_chart, modality_present_chart
 
 pn.extension("vega", design="material")
 alt.themes.enable("ggplot2")
 
 # Define CSS to set the background color and add to panel
-background_param = pn.state.location.query_params.get("background", "dark_blue")
+background_param = pn.state.location.query_params.get(
+    "background", "dark_blue"
+)
 background_color = AIND_COLORS.get(background_param, AIND_COLORS["dark_blue"])
 
 css = f"""
@@ -22,8 +29,8 @@
 pn.config.raw_css.append(css)
 
 # Get the active color list
-colors = (
-    COLOR_OPTIONS.get(pn.state.location.query_params.get("colors"), COLOR_OPTIONS["default"])
+colors = COLOR_OPTIONS.get(
+    pn.state.location.query_params.get("colors"), COLOR_OPTIONS["default"]
 )
 color_list = list(colors.values())
 
@@ -39,7 +46,9 @@
     name="Filter by core file:", options=database.ALL_FILES
 )
 
-field_selector = pn.widgets.Select(name="Filter download by field:", options=[])
+field_selector = pn.widgets.Select(
+    name="Filter download by field:", options=[]
+)
 
 missing_selector = pn.widgets.Select(
     name="Filter download by state", options=["Missing", "Valid/Present"]
@@ -118,7 +127,12 @@ def field_present_chart(selected_file, derived_filter, **args):
 
     sum_longform_df = db.get_file_field_presence()
 
-    field_selection = alt.selection_point(fields=['field'], empty='none', name='field', value=field_selector.value)
+    field_selection = alt.selection_point(
+        fields=["field"],
+        empty="none",
+        name="field",
+        value=field_selector.value,
+    )
 
     chart = (
         alt.Chart(sum_longform_df)
@@ -151,12 +165,11 @@ def field_present_chart(selected_file, derived_filter, **args):
 
     def update_selection(event):
         if len(event.new) > 0:
-            field_selector.value = event.new[0]['field']
-    pane.selection.param.watch(update_selection, 'field')
-
-    return pane
+            field_selector.value = event.new[0]["field"]
 
+    pane.selection.param.watch(update_selection, "field")
 
+    return pane
 
 
 header = (
@@ -175,7 +188,6 @@ def update_selection(event):
 """
 
 
-
 header_pane = pn.pane.Markdown(header, styles=outer_style, width=420)
 
 total_md = f"<p style=\"text-align:center\"><b>{db.get_overall_valid():1.2f}%</b> of all metadata records are fully {hd_style('valid', colors)}</p>"
@@ -209,7 +221,10 @@ def build_row(selected_modality, derived_filter):
     db.modality_filter = selected_modality
     db.derived_filter = derived_filter
 
-    return pn.Row(file_present_chart(db, colors, top_selector), modality_present_chart(db, colors, color_list, modality_selector))
+    return pn.Row(
+        file_present_chart(db, colors, top_selector),
+        modality_present_chart(db, colors, color_list, modality_selector),
+    )
 
 
 top_row = pn.bind(
@@ -228,24 +243,42 @@ def build_row(selected_modality, derived_filter):
 # Put everything in a column and buffer it
 main_col = pn.Column(top_row, mid_plot, styles=outer_style, width=515)
 
-main_row = pn.Row(pn.HSpacer(), left_col, pn.Spacer(width=20), main_col, pn.HSpacer(), margin=20)
+main_row = pn.Row(
+    pn.HSpacer(),
+    left_col,
+    pn.Spacer(width=20),
+    main_col,
+    pn.HSpacer(),
+    margin=20,
+)
 
 # Add the validator search section
-validator_name_selector = pn.widgets.TextInput(name="Enter asset name to validate:", value="", placeholder="Asset name", width=800)
+validator_name_selector = pn.widgets.TextInput(
+    name="Enter asset name to validate:",
+    value="",
+    placeholder="Asset name",
+    width=800,
+)
 pn.state.location.sync(validator_name_selector, {"value": "validator_name"})
 
 validator = database.RecordValidator(validator_name_selector.value, colors)
 
 
 def build_validator(validator_name):
     validator.update(validator_name)
-    col = pn.Column(validator_name_selector, validator.panel(), width=(515+20+420), styles=outer_style)
+    col = pn.Column(
+        validator_name_selector,
+        validator.panel(),
+        width=(515 + 20 + 420),
+        styles=outer_style,
+    )
     row = pn.Row(pn.HSpacer(), col, pn.HSpacer())
     return row
 
 
-validator_row = pn.bind(build_validator,
-                        validator_name=validator_name_selector)
+validator_row = pn.bind(
+    build_validator, validator_name=validator_name_selector
+)
 
 pn.Column(main_row, validator_row).servable(
     title="Metadata Portal",

src/aind_metadata_viz/charts.py

@@ -11,7 +11,9 @@ def file_present_chart(db, colors, selector):
     local_states = sum_longform_df["state"].unique()
     local_color_list = [colors[state] for state in local_states]
 
-    file_selection = alt.selection_point(fields=['file'], empty='none', name='file', value=selector.value)
+    file_selection = alt.selection_point(
+        fields=["file"], empty="none", name="file", value=selector.value
+    )
 
     chart = (
         alt.Chart(sum_longform_df)
@@ -40,8 +42,9 @@ def file_present_chart(db, colors, selector):
 
     def update_selection(event):
         if len(event.new) > 0:
-            selector.value = event.new[0]['file']
-    pane.selection.param.watch(update_selection, 'file')
+            selector.value = event.new[0]["file"]
+
+    pane.selection.param.watch(update_selection, "file")
 
     return pane
 
@@ -55,7 +58,12 @@ def modality_present_chart(db, colors, color_list, selector):
         df_list.append(sum_longform_df)
     df = pd.concat(df_list)
 
-    modality_selection = alt.selection_point(fields=['modality'], empty='all', name='modality', value=(selector.value if selector.value != "all" else None))
+    modality_selection = alt.selection_point(
+        fields=["modality"],
+        empty="all",
+        name="modality",
+        value=(selector.value if selector.value != "all" else None),
+    )
 
     chart = (
         alt.Chart(df)
@@ -93,9 +101,10 @@ def modality_present_chart(db, colors, color_list, selector):
 
     def update_selection(event):
         if len(event.new) > 0:
-            selector.value = event.new[0]['modality']
+            selector.value = event.new[0]["modality"]
         else:
             selector.value = "all"
-    pane.selection.param.watch(update_selection, 'modality')
 
-    return pane
+    pane.selection.param.watch(update_selection, "modality")
+
+    return pane