{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Client Setup\n",
    "from dotenv import load_dotenv\n",
    "import voyageai\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "client = voyageai.Client()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Chunk by section\n",
    "import re\n",
    "\n",
    "\n",
    "def chunk_by_section(document_text):\n",
    "    pattern = r\"\\n## \"\n",
    "    return re.split(pattern, document_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embedding Generation\n",
    "def generate_embedding(chunks, model=\"voyage-3-large\", input_type=\"query\"):\n",
    "    is_list = isinstance(chunks, list)\n",
    "    input = chunks if is_list else [chunks]\n",
    "    result = client.embed(input, model=model, input_type=input_type)\n",
    "    return result.embeddings if is_list else result.embeddings[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "# VectorIndex implementation\n",
    "import math\n",
    "from typing import Optional, Any, List, Dict, Tuple\n",
    "\n",
    "\n",
    "class VectorIndex:\n",
    "    def __init__(\n",
    "        self,\n",
    "        distance_metric: str = \"cosine\",\n",
    "        embedding_fn=None,\n",
    "    ):\n",
    "        self.vectors: List[List[float]] = []\n",
    "        self.documents: List[Dict[str, Any]] = []\n",
    "        self._vector_dim: Optional[int] = None\n",
    "        if distance_metric not in [\"cosine\", \"euclidean\"]:\n",
    "            raise ValueError(\"distance_metric must be 'cosine' or 'euclidean'\")\n",
    "        self._distance_metric = distance_metric\n",
    "        self._embedding_fn = embedding_fn\n",
    "\n",
    "    def add_document(self, document: Dict[str, Any]):\n",
    "        if not self._embedding_fn:\n",
    "            raise ValueError(\n",
    "                \"Embedding function not provided during initialization.\"\n",
    "            )\n",
    "        if not isinstance(document, dict):\n",
    "            raise TypeError(\"Document must be a dictionary.\")\n",
    "        if \"content\" not in document:\n",
    "            raise ValueError(\n",
    "                \"Document dictionary must contain a 'content' key.\"\n",
    "            )\n",
    "\n",
    "        content = document[\"content\"]\n",
    "        if not isinstance(content, str):\n",
    "            raise TypeError(\"Document 'content' must be a string.\")\n",
    "\n",
    "        vector = self._embedding_fn(content)\n",
    "        self.add_vector(vector=vector, document=document)\n",
    "\n",
    "    def add_documents(self, documents: List[Dict[str, Any]]):\n",
    "        if not self._embedding_fn:\n",
    "            raise ValueError(\n",
    "                \"Embedding function not provided during initialization.\"\n",
    "            )\n",
    "\n",
    "        if not isinstance(documents, list):\n",
    "            raise TypeError(\"Documents must be a list of dictionaries.\")\n",
    "\n",
    "        if not documents:\n",
    "            return\n",
    "\n",
    "        contents = []\n",
    "        for i, doc in enumerate(documents):\n",
    "            if not isinstance(doc, dict):\n",
    "                raise TypeError(f\"Document at index {i} must be a dictionary.\")\n",
    "            if \"content\" not in doc:\n",
    "                raise ValueError(\n",
    "                    f\"Document at index {i} must contain a 'content' key.\"\n",
    "                )\n",
    "            if not isinstance(doc[\"content\"], str):\n",
    "                raise TypeError(\n",
    "                    f\"Document 'content' at index {i} must be a string.\"\n",
    "                )\n",
    "            contents.append(doc[\"content\"])\n",
    "\n",
    "        vectors = self._embedding_fn(contents)\n",
    "\n",
    "        for vector, document in zip(vectors, documents):\n",
    "            self.add_vector(vector=vector, document=document)\n",
    "\n",
    "    def search(\n",
    "        self, query: Any, k: int = 1\n",
    "    ) -> List[Tuple[Dict[str, Any], float]]:\n",
    "        if not self.vectors:\n",
    "            return []\n",
    "\n",
    "        if isinstance(query, str):\n",
    "            if not self._embedding_fn:\n",
    "                raise ValueError(\n",
    "                    \"Embedding function not provided for string query.\"\n",
    "                )\n",
    "            query_vector = self._embedding_fn(query)\n",
    "        elif isinstance(query, list) and all(\n",
    "            isinstance(x, (int, float)) for x in query\n",
    "        ):\n",
    "            query_vector = query\n",
    "        else:\n",
    "            raise TypeError(\n",
    "                \"Query must be either a string or a list of numbers.\"\n",
    "            )\n",
    "\n",
    "        if self._vector_dim is None:\n",
    "            return []\n",
    "\n",
    "        if len(query_vector) != self._vector_dim:\n",
    "            raise ValueError(\n",
    "                f\"Query vector dimension mismatch. Expected {self._vector_dim}, got {len(query_vector)}\"\n",
    "            )\n",
    "\n",
    "        if k <= 0:\n",
    "            raise ValueError(\"k must be a positive integer.\")\n",
    "\n",
    "        if self._distance_metric == \"cosine\":\n",
    "            dist_func = self._cosine_distance\n",
    "        else:\n",
    "            dist_func = self._euclidean_distance\n",
    "\n",
    "        distances = []\n",
    "        for i, stored_vector in enumerate(self.vectors):\n",
    "            distance = dist_func(query_vector, stored_vector)\n",
    "            distances.append((distance, self.documents[i]))\n",
    "\n",
    "        distances.sort(key=lambda item: item[0])\n",
    "\n",
    "        return [(doc, dist) for dist, doc in distances[:k]]\n",
    "\n",
    "    def add_vector(self, vector, document: Dict[str, Any]):\n",
    "        if not isinstance(vector, list) or not all(\n",
    "            isinstance(x, (int, float)) for x in vector\n",
    "        ):\n",
    "            raise TypeError(\"Vector must be a list of numbers.\")\n",
    "        if not isinstance(document, dict):\n",
    "            raise TypeError(\"Document must be a dictionary.\")\n",
    "        if \"content\" not in document:\n",
    "            raise ValueError(\n",
    "                \"Document dictionary must contain a 'content' key.\"\n",
    "            )\n",
    "\n",
    "        if not self.vectors:\n",
    "            self._vector_dim = len(vector)\n",
    "        elif len(vector) != self._vector_dim:\n",
    "            raise ValueError(\n",
    "                f\"Inconsistent vector dimension. Expected {self._vector_dim}, got {len(vector)}\"\n",
    "            )\n",
    "\n",
    "        self.vectors.append(list(vector))\n",
    "        self.documents.append(document)\n",
    "\n",
    "    def _euclidean_distance(\n",
    "        self, vec1: List[float], vec2: List[float]\n",
    "    ) -> float:\n",
    "        if len(vec1) != len(vec2):\n",
    "            raise ValueError(\"Vectors must have the same dimension\")\n",
    "        return math.sqrt(sum((p - q) ** 2 for p, q in zip(vec1, vec2)))\n",
    "\n",
    "    def _dot_product(self, vec1: List[float], vec2: List[float]) -> float:\n",
    "        if len(vec1) != len(vec2):\n",
    "            raise ValueError(\"Vectors must have the same dimension\")\n",
    "        return sum(p * q for p, q in zip(vec1, vec2))\n",
    "\n",
    "    def _magnitude(self, vec: List[float]) -> float:\n",
    "        return math.sqrt(sum(x * x for x in vec))\n",
    "\n",
    "    def _cosine_distance(self, vec1: List[float], vec2: List[float]) -> float:\n",
    "        if len(vec1) != len(vec2):\n",
    "            raise ValueError(\"Vectors must have the same dimension\")\n",
    "\n",
    "        mag1 = self._magnitude(vec1)\n",
    "        mag2 = self._magnitude(vec2)\n",
    "\n",
    "        if mag1 == 0 and mag2 == 0:\n",
    "            return 0.0\n",
    "        elif mag1 == 0 or mag2 == 0:\n",
    "            return 1.0\n",
    "\n",
    "        dot_prod = self._dot_product(vec1, vec2)\n",
    "        cosine_similarity = dot_prod / (mag1 * mag2)\n",
    "        cosine_similarity = max(-1.0, min(1.0, cosine_similarity))\n",
    "\n",
    "        return 1.0 - cosine_similarity\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.vectors)\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        has_embed_fn = \"Yes\" if self._embedding_fn else \"No\"\n",
    "        return f\"VectorIndex(count={len(self)}, dim={self._vector_dim}, metric='{self._distance_metric}', has_embedding_fn='{has_embed_fn}')\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# BM25 implementation\n",
    "from collections import Counter\n",
    "from typing import Callable, Any, List, Dict, Tuple\n",
    "\n",
    "\n",
    "class BM25Index:\n",
    "    def __init__(\n",
    "        self,\n",
    "        k1: float = 1.5,\n",
    "        b: float = 0.75,\n",
    "        tokenizer: Optional[Callable[[str], List[str]]] = None,\n",
    "    ):\n",
    "        self.documents: List[Dict[str, Any]] = []\n",
    "        self._corpus_tokens: List[List[str]] = []\n",
    "        self._doc_len: List[int] = []\n",
    "        self._doc_freqs: Dict[str, int] = {}\n",
    "        self._avg_doc_len: float = 0.0\n",
    "        self._idf: Dict[str, float] = {}\n",
    "        self._index_built: bool = False\n",
    "\n",
    "        self.k1 = k1\n",
    "        self.b = b\n",
    "        self._tokenizer = tokenizer if tokenizer else self._default_tokenizer\n",
    "\n",
    "    def _default_tokenizer(self, text: str) -> List[str]:\n",
    "        text = text.lower()\n",
    "        tokens = re.split(r\"\\W+\", text)\n",
    "        return [token for token in tokens if token]\n",
    "\n",
    "    def _update_stats_add(self, doc_tokens: List[str]):\n",
    "        self._doc_len.append(len(doc_tokens))\n",
    "\n",
    "        seen_in_doc = set()\n",
    "        for token in doc_tokens:\n",
    "            if token not in seen_in_doc:\n",
    "                self._doc_freqs[token] = self._doc_freqs.get(token, 0) + 1\n",
    "                seen_in_doc.add(token)\n",
    "\n",
    "        self._index_built = False\n",
    "\n",
    "    def _calculate_idf(self):\n",
    "        N = len(self.documents)\n",
    "        self._idf = {}\n",
    "        for term, freq in self._doc_freqs.items():\n",
    "            idf_score = math.log(((N - freq + 0.5) / (freq + 0.5)) + 1)\n",
    "            self._idf[term] = idf_score\n",
    "\n",
    "    def _build_index(self):\n",
    "        if not self.documents:\n",
    "            self._avg_doc_len = 0.0\n",
    "            self._idf = {}\n",
    "            self._index_built = True\n",
    "            return\n",
    "\n",
    "        self._avg_doc_len = sum(self._doc_len) / len(self.documents)\n",
    "        self._calculate_idf()\n",
    "        self._index_built = True\n",
    "\n",
    "    def add_document(self, document: Dict[str, Any]):\n",
    "        if not isinstance(document, dict):\n",
    "            raise TypeError(\"Document must be a dictionary.\")\n",
    "        if \"content\" not in document:\n",
    "            raise ValueError(\n",
    "                \"Document dictionary must contain a 'content' key.\"\n",
    "            )\n",
    "\n",
    "        content = document.get(\"content\", \"\")\n",
    "        if not isinstance(content, str):\n",
    "            raise TypeError(\"Document 'content' must be a string.\")\n",
    "\n",
    "        doc_tokens = self._tokenizer(content)\n",
    "\n",
    "        self.documents.append(document)\n",
    "        self._corpus_tokens.append(doc_tokens)\n",
    "        self._update_stats_add(doc_tokens)\n",
    "\n",
    "    def add_documents(self, documents: List[Dict[str, Any]]):\n",
    "        if not isinstance(documents, list):\n",
    "            raise TypeError(\"Documents must be a list of dictionaries.\")\n",
    "\n",
    "        if not documents:\n",
    "            return\n",
    "\n",
    "        for i, doc in enumerate(documents):\n",
    "            if not isinstance(doc, dict):\n",
    "                raise TypeError(f\"Document at index {i} must be a dictionary.\")\n",
    "            if \"content\" not in doc:\n",
    "                raise ValueError(\n",
    "                    f\"Document at index {i} must contain a 'content' key.\"\n",
    "                )\n",
    "            if not isinstance(doc[\"content\"], str):\n",
    "                raise TypeError(\n",
    "                    f\"Document 'content' at index {i} must be a string.\"\n",
    "                )\n",
    "\n",
    "            content = doc[\"content\"]\n",
    "            doc_tokens = self._tokenizer(content)\n",
    "\n",
    "            self.documents.append(doc)\n",
    "            self._corpus_tokens.append(doc_tokens)\n",
    "            self._update_stats_add(doc_tokens)\n",
    "\n",
    "        self._index_built = False\n",
    "\n",
    "    def _compute_bm25_score(\n",
    "        self, query_tokens: List[str], doc_index: int\n",
    "    ) -> float:\n",
    "        score = 0.0\n",
    "        doc_term_counts = Counter(self._corpus_tokens[doc_index])\n",
    "        doc_length = self._doc_len[doc_index]\n",
    "\n",
    "        for token in query_tokens:\n",
    "            if token not in self._idf:\n",
    "                continue\n",
    "\n",
    "            idf = self._idf[token]\n",
    "            term_freq = doc_term_counts.get(token, 0)\n",
    "\n",
    "            numerator = idf * term_freq * (self.k1 + 1)\n",
    "            denominator = term_freq + self.k1 * (\n",
    "                1 - self.b + self.b * (doc_length / self._avg_doc_len)\n",
    "            )\n",
    "            score += numerator / (denominator + 1e-9)\n",
    "\n",
    "        return score\n",
    "\n",
    "    def search(\n",
    "        self,\n",
    "        query: Any,\n",
    "        k: int = 1,\n",
    "        score_normalization_factor: float = 0.1,\n",
    "    ) -> List[Tuple[Dict[str, Any], float]]:\n",
    "        if not self.documents:\n",
    "            return []\n",
    "\n",
    "        if isinstance(query, str):\n",
    "            query_text = query\n",
    "        else:\n",
    "            raise TypeError(\"Query must be a string for BM25Index.\")\n",
    "\n",
    "        if k <= 0:\n",
    "            raise ValueError(\"k must be a positive integer.\")\n",
    "\n",
    "        if not self._index_built:\n",
    "            self._build_index()\n",
    "\n",
    "        if self._avg_doc_len == 0:\n",
    "            return []\n",
    "\n",
    "        query_tokens = self._tokenizer(query_text)\n",
    "        if not query_tokens:\n",
    "            return []\n",
    "\n",
    "        raw_scores = []\n",
    "        for i in range(len(self.documents)):\n",
    "            raw_score = self._compute_bm25_score(query_tokens, i)\n",
    "            if raw_score > 1e-9:\n",
    "                raw_scores.append((raw_score, self.documents[i]))\n",
    "\n",
    "        raw_scores.sort(key=lambda item: item[0], reverse=True)\n",
    "\n",
    "        normalized_results = []\n",
    "        for raw_score, doc in raw_scores[:k]:\n",
    "            normalized_score = math.exp(-score_normalization_factor * raw_score)\n",
    "            normalized_results.append((doc, normalized_score))\n",
    "\n",
    "        normalized_results.sort(key=lambda item: item[1])\n",
    "\n",
    "        return normalized_results\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.documents)\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        return f\"BM25VectorStore(count={len(self)}, k1={self.k1}, b={self.b}, index_built={self._index_built})\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retriever implementation\n",
    "from typing import Any, List, Dict, Tuple, Protocol\n",
    "\n",
    "\n",
    "class SearchIndex(Protocol):\n",
    "    def add_document(self, document: Dict[str, Any]) -> None: ...\n",
    "\n",
    "    # Added the 'add_documents' method to avoid rate limiting errors from VoyageAI\n",
    "    def add_documents(self, documents: List[Dict[str, Any]]) -> None: ...\n",
    "\n",
    "    def search(\n",
    "        self, query: Any, k: int = 1\n",
    "    ) -> List[Tuple[Dict[str, Any], float]]: ...\n",
    "\n",
    "\n",
    "class Retriever:\n",
    "    def __init__(self, *indexes: SearchIndex):\n",
    "        if len(indexes) == 0:\n",
    "            raise ValueError(\"At least one index must be provided\")\n",
    "        self._indexes = list(indexes)\n",
    "\n",
    "    def add_document(self, document: Dict[str, Any]):\n",
    "        for index in self._indexes:\n",
    "            index.add_document(document)\n",
    "\n",
    "    # Added the 'add_documents' method to avoid rate limiting errors from VoyageAI\n",
    "    def add_documents(self, documents: List[Dict[str, Any]]):\n",
    "        for index in self._indexes:\n",
    "            index.add_documents(documents)\n",
    "\n",
    "    def search(\n",
    "        self, query_text: str, k: int = 1, k_rrf: int = 60\n",
    "    ) -> List[Tuple[Dict[str, Any], float]]:\n",
    "        if not isinstance(query_text, str):\n",
    "            raise TypeError(\"Query text must be a string.\")\n",
    "        if k <= 0:\n",
    "            raise ValueError(\"k must be a positive integer.\")\n",
    "        if k_rrf < 0:\n",
    "            raise ValueError(\"k_rrf must be non-negative.\")\n",
    "\n",
    "        all_results = [\n",
    "            index.search(query_text, k=k * 5) for index in self._indexes\n",
    "        ]\n",
    "\n",
    "        doc_ranks = {}\n",
    "        for idx, results in enumerate(all_results):\n",
    "            for rank, (doc, _) in enumerate(results):\n",
    "                doc_id = id(doc)\n",
    "                if doc_id not in doc_ranks:\n",
    "                    doc_ranks[doc_id] = {\n",
    "                        \"doc_obj\": doc,\n",
    "                        \"ranks\": [float(\"inf\")] * len(self._indexes),\n",
    "                    }\n",
    "                doc_ranks[doc_id][\"ranks\"][idx] = rank + 1\n",
    "\n",
    "        def calc_rrf_score(ranks: List[float]) -> float:\n",
    "            return sum(1.0 / (k_rrf + r) for r in ranks if r != float(\"inf\"))\n",
    "\n",
    "        scored_docs: List[Tuple[Dict[str, Any], float]] = [\n",
    "            (ranks[\"doc_obj\"], calc_rrf_score(ranks[\"ranks\"]))\n",
    "            for ranks in doc_ranks.values()\n",
    "        ]\n",
    "\n",
    "        filtered_docs = [\n",
    "            (doc, score) for doc, score in scored_docs if score > 0\n",
    "        ]\n",
    "        filtered_docs.sort(key=lambda x: x[1], reverse=True)\n",
    "\n",
    "        return filtered_docs[:k]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Chunk source text by section\n",
    "with open(\"./report.md\", \"r\") as f:\n",
    "    text = f.read()\n",
    "\n",
    "chunks = chunk_by_section(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a vector index, a bm25 index, then use them to create a Retriever\n",
    "vector_index = VectorIndex(embedding_fn=generate_embedding)\n",
    "bm25_index = BM25Index()\n",
    "\n",
    "retriever = Retriever(bm25_index, vector_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add all chunks to the retriever, which internally passes them along to both indexes\n",
    "# Note: converted to a bulk operation to avoid rate limiting errors from VoyageAI\n",
    "retriever.add_documents([{\"content\": chunk} for chunk in chunks])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}