Dynamic column retrieval for metadata

elm/wizard.py

@@ -405,7 +405,8 @@ class EnergyWizardPostgres(EnergyWizardBase):
     """Optional mappings for weird azure names to tiktoken/openai names."""
 
     def __init__(self, db_host, db_port, db_name,
-                 db_schema, db_table, cursor=None, boto_client=None,
+                 db_schema, db_table, meta_columns=None,
+                 cursor=None, boto_client=None,
                  model=None, token_budget=3500):
         """
         Parameters
@@ -421,6 +422,9 @@ def __init__(self, db_host, db_port, db_name,
         db_table : str
             Table to query in Postgres database. Necessary columns: id,
             chunks, embedding, title, and url.
+        meta_columns : list
+            List of metadata columns to retrieve from database. Default
+            query returns title and url.
         cursor : psycopg2.extensions.cursor
             PostgreSQL database cursor used to execute queries.
         boto_client: botocore.client.BedrockRuntime
@@ -437,6 +441,10 @@ def __init__(self, db_host, db_port, db_name,
 
         self.db_schema = db_schema
         self.db_table = db_table
+        if meta_columns is None:
+            self.meta_columns = ['title', 'url']
+        else:
+            self.meta_columns = meta_columns
 
         if cursor is None:
             db_user = os.getenv("EWIZ_DB_USER")
@@ -559,20 +567,29 @@ def make_ref_list(self, ids):
         """
 
         placeholders = ', '.join(['%s'] * len(ids))
+        columns_str = ', '.join([f"{self.db_table}.{c}"
+                                 for c in self.meta_columns])
 
-        sql_query = (f"SELECT {self.db_table}.title, {self.db_table}.url "
+        sql_query = (f"SELECT {columns_str} "
                      f"FROM {self.db_schema}.{self.db_table} "
                      f"WHERE {self.db_table}.id IN (" + placeholders + ")")
 
         self.cursor.execute(sql_query, ids)
 
         refs = self.cursor.fetchall()
 
-        ref_strs = (f"{{\"parentTitle\": \"{item[0]}\", "
-                    f"\"parentUrl\": \"{item[1]}\"}} " for item in refs)
+        ref_list = []
+        for item in refs:
+            ref_dict = {self.meta_columns[i]: item[i]
+                        for i in range(len(self.meta_columns))}
+            ref_str = "{"
+            ref_str += ", ".join([f"\"{key}\": \"{value}\""
+                                  for key, value in ref_dict.items()])
+            ref_str += "}"
 
-        unique_values = set(ref_strs)
+            ref_list.append(ref_str)
 
-        ref_list = list(unique_values)
+        unique_values = set(ref_list)
+        unique_list = list(unique_values)
 
-        return ref_list
+        return unique_list

tests/test_wizard_postgres.py

@@ -84,6 +84,6 @@ def test_postgres():
     ref_list = wizard.make_ref_list(ids)
 
     assert len(ref_list) > 0
-    assert 'parentTitle' in str(ref_list)
-    assert 'parentUrl' in str(ref_list)
+    assert 'title' in str(ref_list)
+    assert 'url' in str(ref_list)
     assert 'research-hub.nrel.gov' in str(ref_list)