feat: added 2 retrieval models

.gitignore

@@ -0,0 +1,208 @@
+*.DS_Store
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+#poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+#pdm.lock
+#pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+#pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Cursor
+#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
+#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
+#  refer to https://docs.cursor.com/context/ignore-files
+.cursorignore
+.cursorindexingignore
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/

data/agentbase.csv

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

platform/agentbase.db → data/agentbase_mini.csv

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

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

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

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

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

data/platforms.csv

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

data/queries.json

@@ -0,0 +1,5 @@
+{
+    "q:01": "Strategies to enhance the efficiency and responsiveness of our company's customer support services.",
+    "q:02": "Explain Newton’s first law of motion and its implications for objects in uniform motion.",
+    "q:03": "Design personalised fitness and nutrition programs tailored to individual goals and needs."
+}

platform/interface.py → interface.py

@@ -5,84 +5,78 @@ Author: Arastun Mammadli
 Date: [Current Date]
 """
 
+from typing import List, Tuple
 import streamlit as st
 import pandas as pd
 import numpy as np
-from pathlib import Path
 
-from database import SQLiteConnector
+from retrieval.models.bm25 import BM25Retriever
+from retrieval.models.sentence_bert import DenseRetriever
+from retrieval.utils import load_queries
 
 
-def retrieve_data(df, search_term, search_columns, method='baseline'):
-    """
-    Retrieve data from the loaded DataFrame based on the specified method.
-    """
-
-    if method == 'baseline': return keyword_filter(df, search_term, search_columns)
-
-def keyword_filter(df, keyword, columns):
+def keyword_filter(df, query, columns, top_k=100) -> List[Tuple[str, float]]:
     """
     Simple keyword-based boolean filter across specified columns.
     """
-
-    if not keyword: return df
+    if not query: return df
 
     mask = df[columns].astype(str).apply(
-        lambda col: col.str.contains(keyword, case=False, na=False)
+        lambda col: col.str.contains(query, case=False, na=False)
     ).any(axis=1)
-    return df[mask]
-
+    filtered_df = df[mask].head(top_k)
+    filtered_df["scores"] = 1    
+    return filtered_df
 
 class AgentBaseUI:
     """
     AgentBase Streamlit-based UI Components.
     """
 
-    def __init__(self, db_path):
-        # loading the data
-        self.db_connector = SQLiteConnector(db_path)
-        self.platforms_df, self.agents_df = self.db_connector.load_all_as_df()
+    def __init__(self, agentbase_path, platforms_path):
+        self.agents_df = pd.read_csv(agentbase_path)
+        self.agents_df['scores'] = 0.0
+        self.platforms_df = pd.read_csv(platforms_path)
+
+        # initialise retrievers
+        self.retrieval_models = ["bm25", "bge-large-en-v1.5", "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)
 
     def show_header(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")
-        self.search_term = st.text_input("Test our retrieval models", placeholder="Type to search...")
-
-    def show_platforms(self):
-        st.header("Platforms")
-        platforms_filtered = retrieve_data(
-            self.platforms_df, 
-            self.search_term, 
-            ['platform_id', 'platform_name', 'platform_url']
-        )
-        
-        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 →"),
-            }
-            st.dataframe(
-                platforms_filtered,
-                column_config=platform_config,
-                use_container_width=True,
-                hide_index=True
+        if "query" not in st.session_state:
+            st.session_state.query = ""
+
+        query_suggestions = list(load_queries("data/queries.json").values())
+        suggestion_cols = st.columns(len(query_suggestions))
+        for i, suggestion in enumerate(query_suggestions):
+            if suggestion_cols[i].button(suggestion):
+                st.session_state.query = suggestion
+
+        col1, col2 = st.columns([3, 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
             )
-        else:
-            st.info("No platforms match your search.")
 
     def show_agents(self):
         st.header("Agents")
-        agents_filtered = retrieve_data(
-            self.agents_df,
-            self.search_term,
-            ['agent_id', 'agent_name', 'agent_description', 'agent_category']
-        )
-        
-        if len(agents_filtered) > 0:
-            st.write(f"Showing {len(agents_filtered)} of {len(self.agents_df)} 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))
+
+        if len(filtered_df) > 0:
+            st.write(f"Showing {top_k} of {len(self.agents_df)} agents")
             # cropped display
             agent_config = {
                 "agent_url": st.column_config.LinkColumn("agent_url", display_text="Visit →"),
@@ -93,15 +87,16 @@ class AgentBaseUI:
             # option to show details
             if st.toggle("Show detailed view"):
                 st.dataframe(
-                    agents_filtered,
+                    filtered_df.head(top_k),
                     column_config=agent_config,
                     use_container_width=True,
                     hide_index=True
                 )
             else:
-                key_columns = ['agent_name', 'platform_id', 'agent_category', 'agent_accessibility', 'agent_url']
+                key_columns = ['agent_name', 'platform_name', 'agent_description', 'agent_url', 'scores']
+                if (filtered_df['scores'] == 0).all(): key_columns.remove("scores")
                 st.dataframe(
-                    agents_filtered[key_columns],
+                    filtered_df[key_columns].head(top_k),
                     column_config=agent_config,
                     use_container_width=True,
                     hide_index=True
@@ -109,11 +104,50 @@ class AgentBaseUI:
         else:
             st.info("No agents match your search.")
 
+    def retrieve_agents(self, query, top_k=100) -> pd.DataFrame:
+        """
+        Returns a filtered dataframe with updated scores.
+        Default maximum top_k of 100
+        """
+        if self.selected_model == 'keyword':
+            return keyword_filter(self.agents_df, query, self.columns, top_k)
+        elif self.selected_model == 'bm25':
+            res = self.bm25.retrieve(query, top_k)
+        elif self.selected_model == 'bge-large-en-v1.5':
+            res = self.bge.retrieve(query, top_k)
+        else:
+            raise ValueError(f"Selected model must be one of {self.retrieval_models}")
+
+        self.agents_df["scores"] = 0 # reset
+        agent_ids, _ = zip(*res)
+        filtered_df = self.agents_df.loc[self.agents_df.agent_id.isin(agent_ids)]
+        for index, row in filtered_df.iterrows():
+            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 →"),
+            }
+            st.dataframe(
+                platforms_filtered,
+                column_config=platform_config,
+                use_container_width=True,
+                hide_index=True
+            )
+        else:
+            st.info("No platforms match your search.")
+
 
 if __name__ == "__main__":
-    BASE_DIR = Path(__file__).parent
-    DB_PATH = BASE_DIR / "agentbase.db"
-    agentbaseui = AgentBaseUI(db_path=DB_PATH)
+    agentbaseui = AgentBaseUI(agentbase_path="data/agentbase.csv", platforms_path="data/platforms.csv")
     agentbaseui.show_header()
     agentbaseui.show_agents()
     agentbaseui.show_platforms()

platform/database.py

@@ -1,123 +0,0 @@
-"""
-AgentBase SQLite3 Database Connector.
-
-This module implements general-purpose functions to access, manipulate and query the AgentBase
-database platform.
-1. Schema Initialisation
-2. Database Connection and Query Execution
-3. Data Insertion (e.g., from temporary CSV files, TODO: updates through web scraping)
-
-Author: Arastun Mammadli
-Date: [Current Date]
-"""
-
-import json
-import pandas as pd
-import numpy as np
-import sqlite3
-
-
-class SQLiteConnector:
-    """
-    A class to handle SQLite database connections and operations.
-    """
-
-    def __init__(self, database_file):
-        self.conn = sqlite3.connect(database_file)
-        self.cursor = self.conn.cursor()
-        self.init_schema()
-
-    def init_schema(self):
-        """
-        Initialise the datbase schema (run once).
-        (1) Platform Entity, (2) Agent Entity, (3) Embeddings Table (optional)
-        """
-        self.conn.execute("""
-        CREATE TABLE IF NOT EXISTS platforms (
-            platform_id TEXT PRIMARY KEY,
-            platform_name TEXT,
-            platform_url TEXT
-        )
-        """)
-        self.conn.execute("""
-        CREATE TABLE IF NOT EXISTS agents (
-            agent_id TEXT PRIMARY KEY,
-            platform_id INTEGER,
-            agent_name TEXT,
-            platform_name TEXT,
-            agent_description TEXT,
-            agent_category TEXT,
-            agent_usage TEXT,  -- JSON string
-            usage_example TEXT,  -- JSON string
-            agent_url TEXT,
-            agent_accessibility TEXT,
-            update_time TEXT,
-            misc TEXT,
-            FOREIGN KEY (platform_id) REFERENCES platforms(platform_id)
-        )
-        """)
-        self.conn.execute("""
-        CREATE TABLE IF NOT EXISTS embeddings (
-            agent_id TEXT,
-            embedding_model TEXT,  -- e.g., 'all-MiniLM-L6-v2'
-            field TEXT,  -- 'description', 'combined', etc.
-            embedding BLOB,  -- numpy array as bytes
-            PRIMARY KEY (agent_id, embedding_model, field),
-            FOREIGN KEY (agent_id) REFERENCES agents(agent_id)
-        )
-        """)
-        self.conn.commit()
-
-    def insert_platform(self, platform_csv):
-        """
-        Insert a mock platform from a CSV file
-        """
-
-        df = pd.read_csv(platform_csv)
-        df.to_sql('platforms', self.conn, if_exists='replace', index=False)
-        self.conn.commit()
-    
-    def insert_agent(self, agent_csv):
-        """
-        Insert mock agents from a CSV file
-        """
-
-        df = pd.read_csv(agent_csv)
-        df.to_sql('agents', self.conn, if_exists='replace', index=False)
-        self.conn.commit()
-
-    def execute_query(self, query, params=None):
-        if params:
-            self.cursor.execute(query, params)
-        else:
-            self.cursor.execute(query)
-        self.conn.commit()
-        return self.cursor.fetchall()
-    
-    def load_all_as_df(self):
-        """
-        Loads all platforms and agents from database as a pandas dataframe.
-        """
-
-        platforms_df = pd.read_sql_query("SELECT * FROM platforms", self.conn)
-        agents_df = pd.read_sql_query("SELECT * FROM agents", self.conn)
-        return platforms_df, agents_df
-
-    def close(self):
-        self.conn.close()
-
-
-if __name__ == "__main__":
-    db_file = "agentbase.db"
-    platforms_csv = f"../data_collection/extracted_data/platforms.csv"
-    agents_csv = f"../data_collection/extracted_data/cleaned/agentbase_cleaned.csv"
-
-    # connect & insert
-    db_connector = SQLiteConnector(db_file)
-    db_connector.insert_platform(platforms_csv)
-    db_connector.insert_agent(agents_csv)
-
-    # test
-    # platforms, agents = db_connector.load_all_as_df()
-    # print(platforms)
-    db_connector.close()

retrieval/__init__.py

@@ -0,0 +1 @@
+

retrieval/base.py

@@ -0,0 +1,36 @@
+from abc import ABC, abstractmethod
+from typing import List, Tuple, Dict
+
+
+class BaseRetriever(ABC):
+    """
+    Abstract base class for AgentBase retrievers.
+    """
+    
+    def __init__(self, db_path: str, columns: List[str]):
+        self.db_path = db_path
+        self.columns = columns
+        self.agent_ids = []
+        self.documents = []
+    
+    @abstractmethod
+    def build_index(self) -> None:
+        """Build retrieval index from database."""
+        pass
+    
+    @abstractmethod
+    def retrieve(self, query: str, top_k: int = 10) -> List[Tuple[str, float]]:
+        """
+        Retrieve top-k agents for a single query.
+        
+        Returns:
+            List of (agent_id, score) tuples
+        """
+        pass
+    
+    def batch_retrieve(self, queries: Dict[str, str], top_k: int = 10) -> Dict[str, List[Tuple[str, float]]]:
+        """Retrieve for multiple queries (for evaluation)."""
+        results = {}
+        for qid, query in queries.items():
+            results[qid] = self.retrieve(query, top_k)
+        return results

retrieval/models/bm25.py

@@ -0,0 +1,35 @@
+from typing import List, Tuple, Dict
+
+import numpy as np
+from rank_bm25 import BM25Okapi
+
+from ..base import BaseRetriever
+from ..utils import load_documents, 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)
+        self.bm25_params = bm25_params
+        self.index = None
+        self.build_index()
+    
+    def build_index(self):
+        """
+        Load documents and build BM25 index.
+        """
+        self.agent_ids, self.documents = load_documents(self.db_path, self.columns)
+        tokenised_docs = [tokenise(doc) for doc in self.documents]
+        self.index = BM25Okapi(tokenised_docs, **self.bm25_params)
+    
+    def retrieve(self, query: str, top_k: int = 10):
+        """
+        Retrieve top-k agents using BM25.
+        """
+        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

@@ -0,0 +1,85 @@
+from typing import List, Tuple, Dict
+from pathlib import Path
+import hashlib
+import json
+
+import numpy as np
+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)
+        self.model_name = model_name
+        self.model = SentenceTransformer(model_name)
+        self.corpus_embeddings = None
+        self.embeddings_path = self._default_embeddings_path()
+        self.load_index() if self._embeddings_exist() else self.build_index()
+
+    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)
+        }
+        # stable SHA-256 hash based on payload
+        hash_str = hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()[:16]
+        return f"data/embeddings/{model_safe}_{hash_str}.npz"
+    
+    def _embeddings_exist(self) -> bool:
+        return Path(self.embeddings_path).exists()
+    
+    def _store_index(self):
+        path = Path(self.embeddings_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        np.savez_compressed(
+            path,
+            embeddings=self.corpus_embeddings,
+            agent_ids=np.array(self.agent_ids),
+            model_name=self.model_name,
+            columns=self.columns
+        )
+
+    def load_index(self):
+        """
+        Use precomputed embeddings.
+        """
+        data = np.load(self.embeddings_path, allow_pickle=True)
+        self.corpus_embeddings = data['embeddings']
+        self.agent_ids = data['agent_ids'].tolist()
+        
+        # verify metadata
+        stored_model = str(data['model_name'])
+        stored_cols = data['columns'].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}")
+
+    def build_index(self):
+        """
+        Build your embeddings.
+        """
+        self.agent_ids, self.corpus = load_documents(self.db_path, self.columns)
+        self.corpus_embeddings = self.model.encode(
+            self.corpus, 
+            show_progress_bar=True,
+            convert_to_numpy=True,
+        )
+        self._store_index() # avoid re-building
+    
+    def retrieve(self, query: str, top_k: int = 10):
+        """
+        Retrieve using cosine similarity.
+        """
+        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]
+        return [(self.agent_ids[idx], float(scores[idx])) for idx in top_indices]

retrieval/utils.py

@@ -0,0 +1,21 @@
+from typing import Dict, List
+import json
+
+import pandas as pd
+
+
+def load_documents(db_path, columns=["agent_name", "agent_description"]):
+    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()
+    return agent_ids, documents
+
+
+def tokenise(doc: str) -> List[str]:
+    return doc.lower().split()
+
+
+def load_queries(queries_path) -> Dict[str, str]:
+    with open(queries_path) as json_file:
+        data = json.load(json_file)
+    return data