fix self-query retriever

app.py

@@ -35,8 +35,6 @@ st.set_page_config(page_title="ChatData")
 
 st.header("ChatData")
 
-columns = ['ref_id', 'title', 'id', 'categories', 'abstract', 'authors', 'pubdate']
-
 
 def try_eval(x):
     try:
@@ -45,12 +43,12 @@ def try_eval(x):
         return x
 
 
-def display(dataframe, columns=None, index=None):
+def display(dataframe, columns_=None, index=None):
     if index:
         dataframe.set_index(index)
     if len(dataframe) > 0:
-        if columns:
-            st.dataframe(dataframe[columns])
+        if columns_:
+            st.dataframe(dataframe[columns_])
         else:
             st.dataframe(dataframe)
     else:
@@ -261,7 +259,7 @@ ENGINE = ReplacingMergeTree ORDER BY id
                     f"### Answer from LLM\n{ret['answer']}\n### References")
                 docs = ret['sources']
                 docs = pd.DataFrame([{**d.metadata, 'abstract': d.page_content} for d in docs])
-                display(docs, columns, index='ref_id')
+                display(docs, ['ref_id', 'title', 'id', 'categories', 'abstract', 'authors', 'pubdate'], index='ref_id')
             except Exception as e:
                 st.write('Oops 😵 Something bad happened...')
                 raise e
@@ -288,7 +286,7 @@ with tab_self_query:
                 docs = pd.DataFrame(
                     [{**d.metadata, 'abstract': d.page_content} for d in docs])
 
-                display(docs, columns)
+                display(docs, ['title', 'id', 'categories', 'abstract', 'authors', 'pubdate'])
             except Exception as e:
                 st.write('Oops 😵 Something bad happened...')
                 raise e
@@ -307,7 +305,7 @@ with tab_self_query:
                     f"### Answer from LLM\n{ret['answer']}\n### References")
                 docs = ret['sources']
                 docs = pd.DataFrame([{**d.metadata, 'abstract': d.page_content} for d in docs])
-                display(docs, columns, index='ref_id')
+                display(docs, ['title', 'id', 'categories', 'abstract', 'authors', 'pubdate'], index='ref_id')
             except Exception as e:
                 st.write('Oops 😵 Something bad happened...')
                 raise e