{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Client Setup\n",
    "import voyageai\n",
    "import json\n",
    "from dotenv import load_dotenv\n",
    "from anthropic import Anthropic\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "embedding_client = voyageai.Client()\n",
    "client = Anthropic()\n",
    "model = \"claude-3-7-sonnet-latest\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper functions\n",
    "from anthropic.types import Message\n",
    "\n",
    "\n",
    "def add_user_message(messages, message):\n",
    "    user_message = {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": message.content if isinstance(message, Message) else message,\n",
    "    }\n",
    "    messages.append(user_message)\n",
    "\n",
    "\n",
    "def add_assistant_message(messages, message):\n",
    "    assistant_message = {\n",
    "        \"role\": \"assistant\",\n",
    "        \"content\": message.content if isinstance(message, Message) else message,\n",
    "    }\n",
    "    messages.append(assistant_message)\n",
    "\n",
    "\n",
    "def chat(messages, system=None, temperature=1.0, stop_sequences=[], tools=None):\n",
    "    params = {\n",
    "        \"model\": model,\n",
    "        \"max_tokens\": 1000,\n",
    "        \"messages\": messages,\n",
    "        \"temperature\": temperature,\n",
    "        \"stop_sequences\": stop_sequences,\n",
    "    }\n",
    "\n",
    "    if tools:\n",
    "        params[\"tools\"] = tools\n",
    "\n",
    "    if system:\n",
    "        params[\"system\"] = system\n",
    "\n",
    "    message = client.messages.create(**params)\n",
    "    return message\n",
    "\n",
    "\n",
    "def text_from_message(message):\n",
    "    return \"\\n\".join(\n",
    "        [block.text for block in message.content if block.type == \"text\"]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "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 = embedding_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": null,
   "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": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retriever implementation\n",
    "\n",
    "from typing import Any, List, Dict, Tuple, Protocol, Callable, Optional\n",
    "import random\n",
    "import string\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__(\n",
    "        self,\n",
    "        *indexes: SearchIndex,\n",
    "        reranker_fn: Optional[\n",
    "            Callable[[List[Dict[str, Any]], str, int], List[str]]\n",
    "        ] = None,\n",
    "    ):\n",
    "        if len(indexes) == 0:\n",
    "            raise ValueError(\"At least one index must be provided\")\n",
    "        self._indexes = list(indexes)\n",
    "        self._reranker_fn = reranker_fn\n",
    "\n",
    "    def add_document(self, document: Dict[str, Any]):\n",
    "        if \"id\" not in document:\n",
    "            document[\"id\"] = \"\".join(\n",
    "                random.choices(string.ascii_letters + string.digits, k=4)\n",
    "            )\n",
    "\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",
    "        result = filtered_docs[:k]\n",
    "\n",
    "        if self._reranker_fn is not None:\n",
    "            docs_only = [doc for doc, _ in result]\n",
    "\n",
    "            for doc in docs_only:\n",
    "                if \"id\" not in doc:\n",
    "                    doc[\"id\"] = \"\".join(\n",
    "                        random.choices(\n",
    "                            string.ascii_letters + string.digits, k=4\n",
    "                        )\n",
    "                    )\n",
    "\n",
    "            doc_lookup = {doc[\"id\"]: doc for doc in docs_only}\n",
    "            reranked_ids = self._reranker_fn(docs_only, query_text, k)\n",
    "\n",
    "            new_result = []\n",
    "            original_scores = {id(doc): score for doc, score in result}\n",
    "\n",
    "            for doc_id in reranked_ids:\n",
    "                if doc_id in doc_lookup:\n",
    "                    doc = doc_lookup[doc_id]\n",
    "                    score = original_scores.get(id(doc), 0.0)\n",
    "                    new_result.append((doc, score))\n",
    "\n",
    "            result = new_result\n",
    "\n",
    "        return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reranker function\n",
    "def reranker_fn(docs, query_text, k):\n",
    "    joined_docs = \"\\n\".join(\n",
    "        [\n",
    "            f\"\"\"\n",
    "        <document>\n",
    "        <document_id>{doc[\"id\"]}</document_id>\n",
    "        <document_content>{doc[\"content\"]}</document_content>\n",
    "        </document>\n",
    "        \"\"\"\n",
    "            for doc in docs\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    prompt = f\"\"\"\n",
    "    You are about to be given a set of documents, along with an id of each.\n",
    "    Your task is to select the {k} most relevant documents to answer the user's question.\n",
    "\n",
    "    Here is the user's question:\n",
    "    <question>\n",
    "    {query_text}\n",
    "    </question>\n",
    "    \n",
    "    Here are the documents to select from:\n",
    "    <documents>\n",
    "    {joined_docs}\n",
    "    </documents>\n",
    "\n",
    "    Respond in the following format:\n",
    "    ```json\n",
    "    {{\n",
    "        \"document_ids\": str[] # List document ids, {k} elements long, sorted in order of decreasing relevance to the user's query.\n",
    "    }}\n",
    "    ```\n",
    "    \"\"\"\n",
    "\n",
    "    messages = []\n",
    "    add_user_message(messages, prompt)\n",
    "    add_assistant_message(messages, \"```json\")\n",
    "\n",
    "    result = chat(messages, stop_sequences=[\"```\"])\n",
    "\n",
    "    # Note: updated to use 'text_from_message' helper fn\n",
    "    return json.loads(text_from_message(result))[\"document_ids\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add context to a single chunk\n",
    "def add_context(text_chunk, source_text):\n",
    "    prompt = f\"\"\"\n",
    "    Write a short and succinct snippet of text to situate this chunk within the \n",
    "    overall source document for the purposes of improving search retrieval of the chunk. \n",
    "\n",
    "    Here is the original source document:\n",
    "    <document> \n",
    "    {source_text}\n",
    "    </document> \n",
    "\n",
    "    Here is the chunk we want to situate within the whole document:\n",
    "    <chunk> \n",
    "    {text_chunk}\n",
    "    </chunk>\n",
    "    \n",
    "    Answer only with the succinct context and nothing else. \n",
    "    \"\"\"\n",
    "\n",
    "    messages = []\n",
    "\n",
    "    add_user_message(messages, prompt)\n",
    "    result = chat(messages)\n",
    "\n",
    "    # Note: updated to use 'text_from_message' helper fn\n",
    "    return text_from_message(result) + \"\\n\" + text_chunk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read source document\n",
    "with open(\"./report.md\", \"r\") as f:\n",
    "    report_text = f.read()\n",
    "\n",
    "chunks = chunk_by_section(report_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "add_context(chunks[5], report_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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, reranker_fn=reranker_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add context to each chunk, then add to the retriever\n",
    "num_start_chunks = 2\n",
    "num_prev_chunks = 2\n",
    "contextualized_chunks = []\n",
    "\n",
    "for i, chunk in enumerate(chunks):\n",
    "    context_parts = []\n",
    "\n",
    "    # Initial set of chunks from the start of the doc\n",
    "    context_parts.extend(chunks[: min(num_start_chunks, len(chunks))])\n",
    "\n",
    "    # Additional chunks ahead of the current chunk we're contextualizing\n",
    "    start_idx = max(0, i - num_prev_chunks)\n",
    "    context_parts.extend(chunks[start_idx:i])\n",
    "\n",
    "    context = \"\\n\".join(context_parts)\n",
    "\n",
    "    contextualized_chunks.append(add_context(chunk, context))\n",
    "\n",
    "# Note: converted to a bulk operation to avoid rate limiting errors from VoyageAI\n",
    "retriever.add_documents([{\"content\": chunk} for chunk in contextualized_chunks])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = retriever.search(\"what did the eng team do with INC-2023-Q4-011?\", 2)\n",
    "\n",
    "for doc, score in results:\n",
    "    print(score, \"\\n\", doc[\"content\"][0:200], \"\\n---\\n\")"
   ]
  }
 ],
 "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
}