feat: improved agentbase platform performance and layout

interface.py

@@ -1,5 +1,5 @@
 """
-AgentBase Visualisation Demo.
+AgentBase Visualisation UI.
 
 Author: Arastun Mammadli
 Date: [Current Date]
@@ -15,51 +15,51 @@ from retrieval.models.sentence_bert import DenseRetriever
 from retrieval.utils import load_queries
 
 
-def keyword_filter(df, query, columns=["agent_name", "agent_description"], top_k=100) -> List[Tuple[str, float]]:
+@st.cache_resource()
+def load_retrievers(agentbase_path: str, index_configs: List[str]) -> Tuple[dict, dict]:
+    bm25s = {}
+    bges = {}
+    for idx_config in index_configs:
+        bm25s[idx_config] = BM25Retriever(agentbase_path, index_config=idx_config)
+        bges[idx_config] = DenseRetriever("BAAI/bge-large-en-v1.5", agentbase_path, index_config=idx_config)
+    return bm25s, bges
+
+@st.cache_resource()
+def load_agentbase_data(agentbase_path: str) -> pd.DataFrame:
+    return pd.read_csv(agentbase_path)
+
+def keyword_filter(query, top_k, df, columns=["agent_name", "agent_description"]) -> List[Tuple[str, float]]:
     """
     Simple keyword-based boolean filter across specified columns.
     """
-    if not query: return df
-
     mask = df[columns].astype(str).apply(
         lambda col: col.str.contains(query, case=False, na=False)
     ).any(axis=1)
-    filtered_df = df[mask].head(top_k)
+    filtered_df = df[mask].head(top_k).copy()
     filtered_df["scores"] = 1    
     return filtered_df
 
+
 class AgentBaseUI:
     """
     AgentBase Streamlit-based UI Components.
     """
-
     def __init__(self, agentbase_path, platforms_path):
-        self.agents_df = pd.read_csv(agentbase_path)
-        self.agents_df['scores'] = 0.0
+        self.agents_df = load_agentbase_data(agentbase_path)
         self.platforms_df = pd.read_csv(platforms_path)
+        self.bm25s, self.bges = load_retrievers(agentbase_path, index_configs=["v1", "naive"])
 
-        # initialise retrievers
+        # selection options and defaults
         self.retrieval_models = ["bge-large-en-v1.5", "bm25", "keyword"]
         self.selected_model = "bge-large-en-v1.5"
         self.indexing_configs = ["v1", "naive"]
-        self.indexing_config = "v1" # default
-
-        self.visible_configs = ["indexing1", "indexing2"]
-        self.visible_configs_to_indexing = {
-            "indexing1": "v1",
-            "indexing2": "naive"
-        }
+        self.indexing_config = "v1"
 
-        self.bm25s = {}
-        self.bges = {}
-        for idx_config in self.indexing_configs:
-            self.bm25s[idx_config] = BM25Retriever(agentbase_path, index_config=idx_config)
-            self.bges[idx_config] = DenseRetriever("BAAI/bge-large-en-v1.5", agentbase_path, index_config=idx_config)
-
-    def show_header(self):
+    def header_panel(self):
         st.set_page_config(page_title="AgentBase", layout="wide")
         st.title("AgentBase Platform Demo")
         st.write("A Large-Scale Agent Collection for Automated Agent Recommendation.")
+
         st.subheader("🔍 Retrieval")
         if "query" not in st.session_state:
             st.session_state.query = ""
@@ -71,45 +71,49 @@ class AgentBaseUI:
                 st.session_state.query = suggestion
 
         col1, col2, col3 = st.columns([4, 1, 1])
-        with col1: 
-            st.session_state.query = st.text_input("", placeholder="Type to search...", value=st.session_state.query)
+        with col1:
+            st.text_input("", placeholder="Type to search...", key="query")
         with col2:
             self.selected_model = st.selectbox("", self.retrieval_models, index=0)
         with col3:
-            self.indexing_config = st.selectbox("", self.visible_configs, index=0)
+            self.indexing_config = st.selectbox("", self.indexing_configs, index=0)
 
-    def show_agents(self):
-        st.header("Agents")
-        filtered_df = self.retrieve_agents(st.session_state.query) if st.session_state.query else self.agents_df
-        top_k = st.slider("Top K", 3, 100, 5) if st.session_state.query else len(filtered_df)
-        top_k = min(top_k, len(filtered_df))
+        _, col2 = st.columns([2, 1])
+        with col2:
+            with st.expander("See explanation"):
+                st.write('''
+                - **Retrieval Models**:
+                    - **BGE-Large**: a dense retrieval model.
+                    - **BM25**: a sparse retrieval model.
+                    - **Keyword**: simple boolean keyword matching.
+                - **Indexing Configurations**:
+                    - **v1**: using all columns with priority ordering (e.g., name, description come first).
+                    - **naive**: using agent name and description only.
+                ''')
+
+    def retrieval_panel(self):
+        top_k = st.slider("Top K", 3, 100, 5)
+        if st.session_state.query:
+            self.filtered_df = self.retrieve_agents(st.session_state.query, top_k)
+        else:
+            self.filtered_df = self.agents_df.copy()
+            self.filtered_df['scores'] = 0.0
 
-        if len(filtered_df) > 0:
+        if len(self.filtered_df) > 0:
             st.write(f"Showing {top_k} of {len(self.agents_df)} agents")
-            # cropped display
-            agent_config = {
+            agent_config = { # clean column display
                 "agent_url": st.column_config.LinkColumn("agent_url", display_text="Visit →"),
                 "agent_description": st.column_config.TextColumn("agent_description", width="large"),
                 "agent_accessibility": st.column_config.TextColumn("agent_accessibility", width="small"),
             }
-            
-            # option to show details
-            if st.toggle("Show detailed view"):
-                st.dataframe(
-                    filtered_df.head(top_k),
-                    column_config=agent_config,
-                    use_container_width=True,
-                    hide_index=True
-                )
-            else:
-                key_columns = ['agent_name', 'platform_name', 'agent_description', 'agent_url', 'scores']
-                if (filtered_df['scores'] == 0).all(): key_columns.remove("scores")
-                st.dataframe(
-                    filtered_df[key_columns].head(top_k),
-                    column_config=agent_config,
-                    use_container_width=True,
-                    hide_index=True
-                )
+            key_columns = ['agent_name', 'platform_name', 'agent_description', 'agent_url', 'scores']
+            if (self.filtered_df['scores'] == 0).all(): key_columns.remove("scores")
+            st.dataframe(
+                self.filtered_df[key_columns].head(top_k),
+                column_config=agent_config,
+                use_container_width=True,
+                hide_index=True
+            )
         else:
             st.info("No agents match your search.")
 
@@ -119,11 +123,11 @@ class AgentBaseUI:
         Default maximum top_k of 100
         """
         if self.selected_model == 'keyword':
-            return keyword_filter(self.agents_df, query, top_k)
+            return keyword_filter(query, top_k, self.agents_df)
         elif self.selected_model == 'bm25':
-            res = self.bm25s[self.visible_configs_to_indexing[self.indexing_config]].retrieve(query, top_k)
+            res = self.bm25s[self.indexing_config].retrieve(query, top_k)
         elif self.selected_model == 'bge-large-en-v1.5':
-            res = self.bges[self.visible_configs_to_indexing[self.indexing_config]].retrieve(query, top_k)
+            res = self.bges[self.indexing_config].retrieve(query, top_k)
         else:
             raise ValueError(f"Selected model must be one of {self.retrieval_models}")
 
@@ -134,29 +138,24 @@ class AgentBaseUI:
             score = dict(res).get(row['agent_id'], 0)
             filtered_df.at[index, 'scores'] = score
         return filtered_df.sort_values(by="scores", ascending=False)
-            
-    def show_platforms(self):
-        st.header("Platforms")
-        platforms_filtered = self.platforms_df
-        
-        if len(platforms_filtered) > 0:
-            st.write(f"Showing {len(platforms_filtered)} of {len(self.platforms_df)} platforms")
-            # cropped display
-            platform_config = {
-                "platform_url": st.column_config.LinkColumn("platform_url", display_text="Visit →"),
-            }
+    
+    def info_panel(self):
+        with st.expander(f"View AgentBase-v1.0"):
             st.dataframe(
-                platforms_filtered,
-                column_config=platform_config,
+                self.agents_df,
+                self.agents_df.columns,
+                use_container_width=True,
+                hide_index=True
+            )
+            st.dataframe(
+                self.platforms_df,
+                self.platforms_df.columns,
                 use_container_width=True,
                 hide_index=True
             )
-        else:
-            st.info("No platforms match your search.")
-
 
 if __name__ == "__main__":
     agentbaseui = AgentBaseUI(agentbase_path="data/agentbase.csv", platforms_path="data/platforms.csv")
-    agentbaseui.show_header()
-    agentbaseui.show_agents()
-    agentbaseui.show_platforms()
+    agentbaseui.header_panel()
+    agentbaseui.retrieval_panel()
+    agentbaseui.info_panel()