examples/javelin_rag_embeddings_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Secure RAG Implementation with Azure OpenAI and Javelin\n",
    "\n",
    "This notebook demonstrates a secure Retrieval Augmented Generation (RAG) implementation using Azure OpenAI and Javelin for embeddings and LLM queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "pytest-httpx 0.22.0 requires httpx==0.24.*, but you have httpx 0.27.2 which is incompatible.\n",
      "javelin-sdk 18.5.15 requires httpx<0.25.0,>=0.24.0, but you have httpx 0.27.2 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.3.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install --quiet --upgrade langchain langchain-community langchain-chroma"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup and Dependencies\n",
    "\n",
    "First, let's import the required libraries and set up our environment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import bs4\n",
    "import dotenv\n",
    "from langchain import hub\n",
    "from langchain_chroma import Chroma\n",
    "from langchain_community.document_loaders import WebBaseLoader\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langchain_openai import AzureChatOpenAI\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from openai import AzureOpenAI\n",
    "\n",
    "dotenv.load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "\n",
    "Set up API keys and headers for Javelin and Azure OpenAI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# API Keys and Headers\n",
    "javelin_api_key = os.getenv(\"JAVELIN_API_KEY\")\n",
    "llm_api_key = os.getenv(\"JAVELIN_AZURE_OPENAI_API_KEY\")\n",
    "\n",
    "# Headers for LLM and embeddings\n",
    "javelin_headers_llm = {\"x-api-key\": javelin_api_key, \"x-javelin-route\": \"azureopenai\"}\n",
    "javelin_headers_embeddings = {\n",
    "    \"x-api-key\": javelin_api_key,\n",
    "    \"x-javelin-route\": \"azureopenaiembeddings\",\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize Azure OpenAI Clients\n",
    "\n",
    "Set up clients for embeddings and LLM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize Azure OpenAI client for embeddings\n",
    "azure_openai_client = AzureOpenAI(\n",
    "    api_key=llm_api_key,\n",
    "    base_url=\"https://api.javelin.live/query\",\n",
    "    default_headers=javelin_headers_embeddings,\n",
    "    api_version=\"2023-05-15\",\n",
    ")\n",
    "\n",
    "# Initialize LLM\n",
    "llm = AzureChatOpenAI(\n",
    "    api_key=llm_api_key,\n",
    "    azure_endpoint=\"https://api.javelin.live/query/azureopenai\",\n",
    "    azure_deployment=\"gpt35\",\n",
    "    openai_api_version=\"2024-02-15-preview\",\n",
    "    model_kwargs={\"extra_headers\": javelin_headers_llm}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample Data\n",
    "\n",
    "Define sample texts for testing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and prepare sample texts\n",
    "sample_texts = [\n",
    "    \"\"\"Authored by Shyam: Climate change is one of the most pressing global challenges of our time. \n",
    "    Rising temperatures, extreme weather events, and melting polar ice caps are \n",
    "    clear indicators of global warming. Greenhouse gas emissions from human activities \n",
    "    continue to be the primary driver of these environmental changes.\"\"\",\n",
    "    \n",
    "    \"\"\"Authored by Shyam: Renewable energy sources like solar, wind, and hydroelectric power are crucial \n",
    "    in combating climate change. These clean energy alternatives are becoming increasingly \n",
    "    cost-effective and efficient. Many countries are setting ambitious targets to transition \n",
    "    away from fossil fuels to reduce their carbon footprint.\"\"\",\n",
    "    \n",
    "    \"\"\"Authored by Shyam: Conservation efforts and sustainable practices play a vital role in environmental \n",
    "    protection. This includes protecting biodiversity, reducing deforestation, and \n",
    "    implementing sustainable agriculture methods. Individual actions like reducing waste, \n",
    "    recycling, and choosing eco-friendly products also contribute to environmental preservation.\n",
    "    \n",
    "    This article is authored by Shyam\"\"\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom Embeddings Class\n",
    "\n",
    "Create a custom embeddings class for Chroma:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomEmbeddings:\n",
    "    def __init__(self, client):\n",
    "        self.client = client\n",
    "    \n",
    "    def embed_documents(self, texts):\n",
    "        response = self.client.embeddings.create(\n",
    "            input=texts,\n",
    "            model=\"text-embedding-3-small\"\n",
    "        )\n",
    "        \n",
    "        return [item.embedding for item in response.data]\n",
    "    \n",
    "    def embed_query(self, text):\n",
    "        response = self.client.embeddings.create(\n",
    "            input=[text],\n",
    "            model=\"text-embedding-3-small\"\n",
    "        )\n",
    "        return response.data[0].embedding\n",
    "\n",
    "# Initialize custom embeddings\n",
    "custom_embeddings = CustomEmbeddings(azure_openai_client)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up Vector Store and RAG Chain\n",
    "\n",
    "Create the vector store and set up the RAG pipeline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create vector store with smaller chunk size\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)\n",
    "split_texts = text_splitter.split_text(\"\\n\\n\".join(sample_texts))\n",
    "\n",
    "vectorstore = Chroma.from_texts(\n",
    "    texts=split_texts,\n",
    "    embedding=custom_embeddings\n",
    ")\n",
    "\n",
    "# Set up retriever and prompt\n",
    "retriever = vectorstore.as_retriever()\n",
    "prompt = hub.pull(\"rlm/rag-prompt\")\n",
    "\n",
    "def format_docs(docs):\n",
    "    return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
    "\n",
    "# Create RAG chain\n",
    "rag_chain = (\n",
    "    {\"context\": retriever | format_docs, \"question\": RunnablePassthrough()}\n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test the RAG System\n",
    "\n",
    "Run test questions through the RAG system:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test questions\n",
    "questions = [\n",
    "    \"What are the main indicators of climate change?\",\n",
    "    \"How are renewable energy sources helping to address climate change?\",\n",
    "    \"What role do individual actions play in environmental conservation?\",\n",
    "    \"Who is the author of this article?\"\n",
    "]\n",
    "\n",
    "# Run questions through the RAG chain\n",
    "for question in questions:\n",
    "    print(f\"\\nQuestion: {question}\")\n",
    "    print(\"Answer:\", rag_chain.invoke(question))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}