{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b31c2849",
   "metadata": {},
   "source": [
    "# Utility Functions for Blog Post Loading and Processing\n",
    "\n",
    "This notebook contains utility functions for loading blog posts from the data directory, processing their metadata, and creating vector embeddings for use in the RAG system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "848b0a86",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from pathlib import Path\n",
    "from typing import List, Dict, Any, Optional\n",
    "\n",
    "from langchain_community.document_loaders import DirectoryLoader\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain.schema.document import Document\n",
    "from langchain_huggingface import HuggingFaceEmbeddings\n",
    "from langchain_community.vectorstores import Qdrant\n",
    "\n",
    "from IPython.display import Markdown, display\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# Load environment variables from .env file\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39e32435",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "\n",
    "Load configuration from environment variables or use defaults."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a6a5d6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuration with defaults\n",
    "DATA_DIR = os.environ.get(\"DATA_DIR\", \"data/\")\n",
    "VECTOR_STORAGE_PATH = os.environ.get(\"VECTOR_STORAGE_PATH\", \"./db/vectorstore_v3\")\n",
    "EMBEDDING_MODEL = os.environ.get(\"EMBEDDING_MODEL\", \"Snowflake/snowflake-arctic-embed-l\")\n",
    "QDRANT_COLLECTION = os.environ.get(\"QDRANT_COLLECTION\", \"thedataguy_documents\")\n",
    "BLOG_BASE_URL = os.environ.get(\"BLOG_BASE_URL\", \"https://thedataguy.pro/blog/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01454147",
   "metadata": {},
   "source": [
    "## Utility Functions\n",
    "\n",
    "These functions handle the loading, processing, and storing of blog posts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25792cd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_blog_posts(data_dir: str = DATA_DIR, \n",
    "                   glob_pattern: str = \"*.md\", \n",
    "                   recursive: bool = True, \n",
    "                   show_progress: bool = True) -> List[Document]:\n",
    "    \"\"\"\n",
    "    Load blog posts from the specified directory.\n",
    "    \n",
    "    Args:\n",
    "        data_dir: Directory containing the blog posts\n",
    "        glob_pattern: Pattern to match files\n",
    "        recursive: Whether to search subdirectories\n",
    "        show_progress: Whether to show a progress bar\n",
    "        \n",
    "    Returns:\n",
    "        List of Document objects containing the blog posts\n",
    "    \"\"\"\n",
    "    text_loader = DirectoryLoader(\n",
    "        data_dir, \n",
    "        glob=glob_pattern, \n",
    "        show_progress=show_progress,\n",
    "        recursive=recursive\n",
    "    )\n",
    "    \n",
    "    documents = text_loader.load()\n",
    "    print(f\"Loaded {len(documents)} documents from {data_dir}\")\n",
    "    return documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7ddba72",
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_document_metadata(documents: List[Document], \n",
    "                           data_dir_prefix: str = DATA_DIR,\n",
    "                           blog_base_url: str = BLOG_BASE_URL,\n",
    "                           remove_suffix: str = \"index.md\") -> List[Document]:\n",
    "    \"\"\"\n",
    "    Update the metadata of documents to include URL and other information.\n",
    "    \n",
    "    Args:\n",
    "        documents: List of Document objects to update\n",
    "        data_dir_prefix: Prefix to replace in source paths\n",
    "        blog_base_url: Base URL for the blog posts\n",
    "        remove_suffix: Suffix to remove from paths (like index.md)\n",
    "        \n",
    "    Returns:\n",
    "        Updated list of Document objects\n",
    "    \"\"\"\n",
    "    for doc in documents:\n",
    "        # Create URL from source path\n",
    "        doc.metadata[\"url\"] = doc.metadata[\"source\"].replace(data_dir_prefix, blog_base_url)\n",
    "        \n",
    "        # Remove index.md or other suffix if present\n",
    "        if remove_suffix and doc.metadata[\"url\"].endswith(remove_suffix):\n",
    "            doc.metadata[\"url\"] = doc.metadata[\"url\"][:-len(remove_suffix)]\n",
    "            \n",
    "        # Extract post title from the directory structure\n",
    "        path_parts = Path(doc.metadata[\"source\"]).parts\n",
    "        if len(path_parts) > 1:\n",
    "            # Use the directory name as post_slug\n",
    "            doc.metadata[\"post_slug\"] = path_parts[-2]\n",
    "            doc.metadata[\"post_title\"] = path_parts[-2].replace(\"-\", \" \").title()\n",
    "            \n",
    "        # Add document length as metadata\n",
    "        doc.metadata[\"content_length\"] = len(doc.page_content)\n",
    "    \n",
    "    return documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0dfe498",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_document_stats(documents: List[Document]) -> Dict[str, Any]:\n",
    "    \"\"\"\n",
    "    Get statistics about the documents.\n",
    "    \n",
    "    Args:\n",
    "        documents: List of Document objects\n",
    "        \n",
    "    Returns:\n",
    "        Dictionary with statistics\n",
    "    \"\"\"\n",
    "    stats = {\n",
    "        \"total_documents\": len(documents),\n",
    "        \"total_characters\": sum(len(doc.page_content) for doc in documents),\n",
    "        \"min_length\": min(len(doc.page_content) for doc in documents),\n",
    "        \"max_length\": max(len(doc.page_content) for doc in documents),\n",
    "        \"avg_length\": sum(len(doc.page_content) for doc in documents) / len(documents) if documents else 0,\n",
    "    }\n",
    "    \n",
    "    # Create a list of document info for analysis\n",
    "    doc_info = []\n",
    "    for doc in documents:\n",
    "        doc_info.append({\n",
    "            \"url\": doc.metadata.get(\"url\", \"\"),\n",
    "            \"source\": doc.metadata.get(\"source\", \"\"),\n",
    "            \"title\": doc.metadata.get(\"post_title\", \"\"),\n",
    "            \"text_length\": doc.metadata.get(\"content_length\", 0),\n",
    "        })\n",
    "    \n",
    "    stats[\"documents\"] = doc_info\n",
    "    return stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ae139c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_document_stats(stats: Dict[str, Any]):\n",
    "    \"\"\"\n",
    "    Display document statistics in a readable format.\n",
    "    \n",
    "    Args:\n",
    "        stats: Dictionary with statistics from get_document_stats\n",
    "    \"\"\"\n",
    "    print(f\"Total Documents: {stats['total_documents']}\")\n",
    "    print(f\"Total Characters: {stats['total_characters']}\")\n",
    "    print(f\"Min Length: {stats['min_length']} characters\")\n",
    "    print(f\"Max Length: {stats['max_length']} characters\")\n",
    "    print(f\"Average Length: {stats['avg_length']:.2f} characters\")\n",
    "    \n",
    "    # Display documents as a table\n",
    "    import pandas as pd\n",
    "    if stats[\"documents\"]:\n",
    "        df = pd.DataFrame(stats[\"documents\"])\n",
    "        display(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dcf66b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_documents(documents: List[Document], \n",
    "                   chunk_size: int = 1000, \n",
    "                   chunk_overlap: int = 200) -> List[Document]:\n",
    "    \"\"\"\n",
    "    Split documents into chunks for better embedding and retrieval.\n",
    "    \n",
    "    Args:\n",
    "        documents: List of Document objects to split\n",
    "        chunk_size: Size of each chunk in characters\n",
    "        chunk_overlap: Overlap between chunks in characters\n",
    "        \n",
    "    Returns:\n",
    "        List of split Document objects\n",
    "    \"\"\"\n",
    "    text_splitter = RecursiveCharacterTextSplitter(\n",
    "        chunk_size=chunk_size,\n",
    "        chunk_overlap=chunk_overlap,\n",
    "        length_function=len,\n",
    "    )\n",
    "    \n",
    "    split_docs = text_splitter.split_documents(documents)\n",
    "    print(f\"Split {len(documents)} documents into {len(split_docs)} chunks\")\n",
    "    return split_docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "527ad848",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_vector_store(documents: List[Document], \n",
    "                       storage_path: str = VECTOR_STORAGE_PATH,\n",
    "                       collection_name: str = QDRANT_COLLECTION,\n",
    "                       embedding_model: str = EMBEDDING_MODEL,\n",
    "                       force_recreate: bool = False) -> Qdrant:\n",
    "    \"\"\"\n",
    "    Create a vector store from documents.\n",
    "    \n",
    "    Args:\n",
    "        documents: List of Document objects to store\n",
    "        storage_path: Path to the vector store\n",
    "        collection_name: Name of the collection\n",
    "        embedding_model: Name of the embedding model\n",
    "        force_recreate: Whether to force recreation of the vector store\n",
    "        \n",
    "    Returns:\n",
    "        Qdrant vector store\n",
    "    \"\"\"\n",
    "    # Initialize the embedding model\n",
    "    embeddings = HuggingFaceEmbeddings(model_name=embedding_model)\n",
    "    \n",
    "    # Create the directory if it doesn't exist\n",
    "    storage_dir = Path(storage_path).parent\n",
    "    os.makedirs(storage_dir, exist_ok=True)\n",
    "    \n",
    "    # Check if vector store exists\n",
    "    vector_store_exists = Path(storage_path).exists() and not force_recreate\n",
    "    \n",
    "    if vector_store_exists:\n",
    "        print(f\"Loading existing vector store from {storage_path}\")\n",
    "        try:\n",
    "            vector_store = Qdrant(\n",
    "                path=storage_path,\n",
    "                embedding_function=embeddings,\n",
    "                collection_name=collection_name\n",
    "            )\n",
    "            return vector_store\n",
    "        except Exception as e:\n",
    "            print(f\"Error loading existing vector store: {e}\")\n",
    "            print(\"Creating new vector store...\")\n",
    "            force_recreate = True\n",
    "    \n",
    "    # Create new vector store\n",
    "    print(f\"Creating new vector store at {storage_path}\")\n",
    "    vector_store = Qdrant.from_documents(\n",
    "        documents=documents,\n",
    "        embedding=embeddings,\n",
    "        path=storage_path,\n",
    "        collection_name=collection_name,\n",
    "    )\n",
    "    \n",
    "    return vector_store"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c78f99fc",
   "metadata": {},
   "source": [
    "## Example Usage\n",
    "\n",
    "Here's how to use these utility functions for processing blog posts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "132d32c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_blog_posts(data_dir: str = DATA_DIR,\n",
    "                      create_embeddings: bool = True,\n",
    "                      force_recreate_embeddings: bool = False):\n",
    "    \"\"\"\n",
    "    Complete pipeline to process blog posts and optionally create vector embeddings.\n",
    "    \n",
    "    Args:\n",
    "        data_dir: Directory containing the blog posts\n",
    "        create_embeddings: Whether to create vector embeddings\n",
    "        force_recreate_embeddings: Whether to force recreation of embeddings\n",
    "        \n",
    "    Returns:\n",
    "        Dictionary with data and vector store (if created)\n",
    "    \"\"\"\n",
    "    # Load documents\n",
    "    documents = load_blog_posts(data_dir)\n",
    "    \n",
    "    # Update metadata\n",
    "    documents = update_document_metadata(documents)\n",
    "    \n",
    "    # Get and display stats\n",
    "    stats = get_document_stats(documents)\n",
    "    display_document_stats(stats)\n",
    "    \n",
    "    result = {\n",
    "        \"documents\": documents,\n",
    "        \"stats\": stats,\n",
    "        \"vector_store\": None\n",
    "    }\n",
    "    \n",
    "    # Create vector store if requested\n",
    "    if create_embeddings:\n",
    "        vector_store = create_vector_store(\n",
    "            documents, \n",
    "            force_recreate=force_recreate_embeddings\n",
    "        )\n",
    "        result[\"vector_store\"] = vector_store\n",
    "    \n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "266d4fb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example usage\n",
    "if __name__ == \"__main__\":\n",
    "    # Process blog posts without creating embeddings\n",
    "    result = process_blog_posts(create_embeddings=False)\n",
    "    \n",
    "    # Example: Access the documents\n",
    "    print(f\"\\nDocument example: {result['documents'][0].metadata}\")\n",
    "    \n",
    "    # Create embeddings if needed\n",
    "    # result = process_blog_posts(create_embeddings=True)\n",
    "    \n",
    "    # Retriever example\n",
    "    # retriever = result[\"vector_store\"].as_retriever()\n",
    "    # query = \"What is RAGAS?\"\n",
    "    # docs = retriever.invoke(query, k=2)\n",
    "    # print(f\"\\nRetrieved {len(docs)} documents for query: {query}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22132649",
   "metadata": {},
   "source": [
    "## Function for Loading Existing Vector Store\n",
    "\n",
    "This function can be used to load an existing vector store without reprocessing all blog posts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c24e0c02",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_vector_store(storage_path: str = VECTOR_STORAGE_PATH,\n",
    "                     collection_name: str = QDRANT_COLLECTION,\n",
    "                     embedding_model: str = EMBEDDING_MODEL) -> Optional[Qdrant]:\n",
    "    \"\"\"\n",
    "    Load an existing vector store.\n",
    "    \n",
    "    Args:\n",
    "        storage_path: Path to the vector store\n",
    "        collection_name: Name of the collection\n",
    "        embedding_model: Name of the embedding model\n",
    "        \n",
    "    Returns:\n",
    "        Qdrant vector store or None if it doesn't exist\n",
    "    \"\"\"\n",
    "    # Initialize the embedding model\n",
    "    embeddings = HuggingFaceEmbeddings(model_name=embedding_model)\n",
    "    \n",
    "    # Check if vector store exists\n",
    "    if not Path(storage_path).exists():\n",
    "        print(f\"Vector store not found at {storage_path}\")\n",
    "        return None\n",
    "    \n",
    "    try:\n",
    "        vector_store = Qdrant(\n",
    "            path=storage_path,\n",
    "            embedding_function=embeddings,\n",
    "            collection_name=collection_name\n",
    "        )\n",
    "        print(f\"Loaded vector store from {storage_path}\")\n",
    "        return vector_store\n",
    "    except Exception as e:\n",
    "        print(f\"Error loading vector store: {e}\")\n",
    "        return None"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}