{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "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": 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 = 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 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": [
    "with open(\"./report.md\", \"r\") as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Chunk the text by section"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2. Generate embeddings for each chunk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3. Create a vector store and add each embedding to it\n",
    "# Note: converted to a bulk operation to avoid rate limiting errors from VoyageAI\n",
    "store = VectorIndex()\n",
    "store.add_documents([{\"content\": chunk} for chunk in chunks])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 4. Some time later, a user will ask a question. Generate an embedding for it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5. Search the store with the embedding, find the 2 most relevant chunks"
   ]
  }
 ],
 "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}