feat: added indexing configuration

data/agentbase.csv

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f8ac2555a8cacd01b1e319abe638b5da638451b67cd247350acd3fe46e71599e
-size 13923827
+oid sha256:8e02bc6fc96153821bccb4ebd290a9ba848bbf50a0e05e9d21e4f7f1b46c05dc
+size 13917037

data/agentbase_mini.csv

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cc54651a1cb6d275ea43b663e324efc3181bd0dfd2d9eb1df04dff81cd6c9ff9
-size 439127
+oid sha256:f13942b04f3f6ae6f8eafdc7eb35d73f725288c39939bfb13af42ef789d1197c
+size 438503

data/embeddings/{BAAI_bge-large-en-v1.5_88f95605539a0823.npz → BAAI_bge-large-en-v1.5_18d58ad4370dc14d.npz}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:75f3cfb24da3dba8c17590bbe7d59958f2bcc31d6cb6547e20d791d045453393
-size 37193250
+oid sha256:8204e3f290f8f7b2e1ce2fd42728e82ebe2c5d9b4eb73eb2a6578d6efe0e6411
+size 37193225

data/embeddings/BAAI_bge-large-en-v1.5_8c6dc4a78211c0a6.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1e757ca5e98cf0243dc147e093f1b4ca663eb7dc3b03cc6cffc5d5f1c503da2
+size 37191771

data/embeddings/{BAAI_bge-large-en-v1.5_6380768bd184d185.npz → BAAI_bge-large-en-v1.5_fb62913d86e36ab8.npz}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cffc46d72e08a7acf60b72a587896bc581f5129ecba306f059db31b435b4cdb3
-size 1144547
+oid sha256:7c7a12acaacfac522fccf7db985f6d105f13427caeecff1cec3133742edb3ed0
+size 1144274

interface.py

@@ -15,7 +15,7 @@ from retrieval.models.sentence_bert import DenseRetriever
 from retrieval.utils import load_queries
 
 
-def keyword_filter(df, query, columns, top_k=100) -> List[Tuple[str, float]]:
+def keyword_filter(df, query, columns=["agent_name", "agent_description"], top_k=100) -> List[Tuple[str, float]]:
     """
     Simple keyword-based boolean filter across specified columns.
     """
@@ -39,11 +39,22 @@ class AgentBaseUI:
         self.platforms_df = pd.read_csv(platforms_path)
 
         # initialise retrievers
-        self.retrieval_models = ["bm25", "bge-large-en-v1.5", "keyword"]
+        self.retrieval_models = ["bge-large-en-v1.5", "bm25", "keyword"]
         self.selected_model = "bge-large-en-v1.5"
-        self.columns = ["agent_name", "agent_description"] # experimental
-        self.bm25 = BM25Retriever(agentbase_path, columns=self.columns)
-        self.bge = DenseRetriever("BAAI/bge-large-en-v1.5", agentbase_path, columns=self.columns)
+        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.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):
         st.set_page_config(page_title="AgentBase", layout="wide")
@@ -59,15 +70,13 @@ class AgentBaseUI:
             if suggestion_cols[i].button(suggestion):
                 st.session_state.query = suggestion
 
-        col1, col2 = st.columns([3, 1])
+        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 col2:
-            self.selected_model = st.selectbox(
-                "",
-                self.retrieval_models,
-                index=0
-            )
+            self.selected_model = st.selectbox("", self.retrieval_models, index=0)
+        with col3:
+            self.indexing_config = st.selectbox("", self.visible_configs, index=0)
 
     def show_agents(self):
         st.header("Agents")
@@ -110,11 +119,11 @@ class AgentBaseUI:
         Default maximum top_k of 100
         """
         if self.selected_model == 'keyword':
-            return keyword_filter(self.agents_df, query, self.columns, top_k)
+            return keyword_filter(self.agents_df, query, top_k)
         elif self.selected_model == 'bm25':
-            res = self.bm25.retrieve(query, top_k)
+            res = self.bm25s[self.visible_configs_to_indexing[self.indexing_config]].retrieve(query, top_k)
         elif self.selected_model == 'bge-large-en-v1.5':
-            res = self.bge.retrieve(query, top_k)
+            res = self.bges[self.visible_configs_to_indexing[self.indexing_config]].retrieve(query, top_k)
         else:
             raise ValueError(f"Selected model must be one of {self.retrieval_models}")
 

retrieval/base.py

@@ -1,5 +1,8 @@
 from abc import ABC, abstractmethod
-from typing import List, Tuple, Dict
+from typing import List, Tuple, Dict, Literal
+from functools import partial
+
+from .utils import load_documents, agentbase_indexing
 
 
 class BaseRetriever(ABC):
@@ -7,12 +10,16 @@ class BaseRetriever(ABC):
     Abstract base class for AgentBase retrievers.
     """
     
-    def __init__(self, db_path: str, columns: List[str]):
+    def __init__(self, db_path: str, index_config: Literal["naive", "v1"]):
         self.db_path = db_path
-        self.columns = columns
         self.agent_ids = []
         self.documents = []
-    
+        self.index_config = index_config
+        self.indexing_func = {
+            "naive": partial(load_documents, self.db_path),
+            "v1": partial(agentbase_indexing, self.db_path),
+        }
+
     @abstractmethod
     def build_index(self) -> None:
         """Build retrieval index from database."""

retrieval/models/bm25.py

@@ -1,17 +1,17 @@
-from typing import List, Tuple, Dict
+from typing import List, Tuple, Dict, Literal
 
 import numpy as np
 from rank_bm25 import BM25Okapi
 
 from ..base import BaseRetriever
-from ..utils import load_documents, tokenise
+from ..utils import tokenise
 
 
 class BM25Retriever(BaseRetriever):
     """BM25 sparse retrieval for AgentBase."""
     
-    def __init__(self, db_path: str, columns: List[str], **bm25_params):
-        super().__init__(db_path, columns)
+    def __init__(self, db_path: str, index_config: Literal["naive", "v1"], **bm25_params):
+        super().__init__(db_path, index_config)
         self.bm25_params = bm25_params
         self.index = None
         self.build_index()
@@ -20,7 +20,7 @@ class BM25Retriever(BaseRetriever):
         """
         Load documents and build BM25 index.
         """
-        self.agent_ids, self.documents = load_documents(self.db_path, self.columns)
+        self.agent_ids, self.documents = self.indexing_func[self.index_config]()
         tokenised_docs = [tokenise(doc) for doc in self.documents]
         self.index = BM25Okapi(tokenised_docs, **self.bm25_params)
     
@@ -30,6 +30,5 @@ class BM25Retriever(BaseRetriever):
         """
         tokenized_query = tokenise(query)
         scores = self.index.get_scores(tokenized_query)
-        
         top_indices = np.argsort(scores)[-top_k:][::-1]
         return [(self.agent_ids[idx], float(scores[idx])) for idx in top_indices]

retrieval/models/sentence_bert.py

@@ -1,20 +1,20 @@
-from typing import List, Tuple, Dict
+from typing import List, Tuple, Dict, Literal
 from pathlib import Path
 import hashlib
 import json
 
 import numpy as np
+import pandas as pd
 from sentence_transformers import SentenceTransformer
 
 from ..base import BaseRetriever
-from ..utils import load_documents
 
 
 class DenseRetriever(BaseRetriever):
     """Dense retrieval using sentence transformers."""
     
-    def __init__(self, model_name: str, db_path: str, columns: List[str]):
-        super().__init__(db_path, columns)
+    def __init__(self, model_name: str, db_path: str, index_config: Literal["naive", "v1"]):
+        super().__init__(db_path, index_config)
         self.model_name = model_name
         self.model = SentenceTransformer(model_name)
         self.corpus_embeddings = None
@@ -23,11 +23,10 @@ class DenseRetriever(BaseRetriever):
 
     def _default_embeddings_path(self) -> str:
         model_safe = self.model_name.replace("/", "_")
-        agent_ids, _ = load_documents(self.db_path, self.columns)
         payload = {
             "model": self.model_name,
-            "columns": self.columns,
-            "rows": len(agent_ids)
+            "columns": self.index_config,
+            "rows": len(pd.read_csv(self.db_path))
         }
         # stable SHA-256 hash based on payload
         hash_str = hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()[:16]
@@ -44,7 +43,7 @@ class DenseRetriever(BaseRetriever):
             embeddings=self.corpus_embeddings,
             agent_ids=np.array(self.agent_ids),
             model_name=self.model_name,
-            columns=self.columns
+            index_config=self.index_config
         )
 
     def load_index(self):
@@ -57,17 +56,17 @@ class DenseRetriever(BaseRetriever):
         
         # verify metadata
         stored_model = str(data['model_name'])
-        stored_cols = data['columns'].tolist()
+        stored_index_config = data['index_config'].tolist()
         if stored_model != self.model_name:
             print(f"WARNING: Loaded embeddings from {stored_model}, but using {self.model_name}")
-        if stored_cols != self.columns:
-            raise ValueError(f"Column mismatch! Stored: {stored_cols}, Expected: {self.columns}")
+        if stored_index_config != self.index_config:
+            raise ValueError(f"Index configuration mismatch! Stored: {stored_index_config}, Expected: {self.index_config}")
 
     def build_index(self):
         """
         Build your embeddings.
         """
-        self.agent_ids, self.corpus = load_documents(self.db_path, self.columns)
+        self.agent_ids, self.corpus = self.indexing_func[self.index_config]()
         self.corpus_embeddings = self.model.encode(
             self.corpus, 
             show_progress_bar=True,
@@ -77,8 +76,9 @@ class DenseRetriever(BaseRetriever):
     
     def retrieve(self, query: str, top_k: int = 10):
         """
-        Retrieve using cosine similarity.
+        Retrieve using dot product.
         """
+        # NOTE: .encode() internally normalises our vectors
         query_embedding = self.model.encode([query], convert_to_numpy=True)[0]
         scores = np.dot(self.corpus_embeddings, query_embedding)
         top_indices = np.argsort(scores)[-top_k:][::-1]

retrieval/utils.py

@@ -1,10 +1,36 @@
-from typing import Dict, List
+from typing import Dict, List, Tuple
 import json
 
 import pandas as pd
 
 
-def load_documents(db_path, columns=["agent_name", "agent_description"]):
+def agentbase_indexing(db_path: str) -> Tuple[pd.DataFrame, List[str]]:
+    """
+    Another indexing configuration for AgentBase dense models.
+        1. concatenate and embed all columns, except agent_id (redundant) and misc (not to go over max_seq_length)
+        2. handle missing/null values
+        3. prioritise important fields (see field semantics)
+
+    :returns: ids and prepared documents
+    """
+    agents_df = pd.read_csv(db_path)
+    agent_ids = agents_df["agent_id"]
+    agents_df.drop(columns=["agent_id", "misc"], inplace=True)
+    columns = agents_df.columns
+
+    high_priority_cols = ["agent_name", "agent_description", "agent_category"]
+    columns = high_priority_cols + [col for col in agents_df.columns if col not in high_priority_cols]
+    documents = agents_df.apply(
+        lambda row: ' '.join([f"{row[col]}" for col in columns if pd.notna(row[col])]), axis=1
+    ).tolist()
+    return agent_ids, documents
+
+
+def load_documents(db_path: str, columns=["agent_name", "agent_description"]) -> Tuple[pd.DataFrame, List[str]]:
+    """
+    Loads documents (for sparse and dense models) by concatenating all column fields
+    :returns: ids and prepared documents
+    """
     agents_df = pd.read_csv(db_path)
     agent_ids = agents_df["agent_id"] # keep agent IDs (mapping back after retrieval)
     documents = agents_df[columns].astype(str).agg(' '.join, axis=1).tolist()
@@ -15,7 +41,7 @@ def tokenise(doc: str) -> List[str]:
     return doc.lower().split()
 
 
-def load_queries(queries_path) -> Dict[str, str]:
+def load_queries(queries_path: str) -> Dict[str, str]:
     with open(queries_path) as json_file:
         data = json.load(json_file)
     return data