pyairbyte_notebooks/rag_with_fb_marketing_milvus_lite.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook demonstrates simple RAG (Retrieval-Augmented Generation) pipeline with Facebook Marketing, Milvus Lite and PyAirbyte."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prerequisites\n",
    "- [PyAirbyte](https://airbyte.com/product/pyairbyte)\n",
    "  \n",
    "  PyAirbyte is an open-source that packages Airbyte connectors and makes them available in Python. We will connect to  \n",
    "  `source-facebook-marketing`, and retrive its streams\n",
    "- [Milvus Lite](https://milvus.io/docs/milvus_lite.md)\n",
    "  \n",
    "  Milvus Lite is the lightweight version of [Milvus](https://github.com/milvus-io/milvus) that enables vector emdeddings and similarity search\n",
    "  into the Python application.\n",
    "- OpenAI API Key\n",
    "  \n",
    "  Go to the [API Keys page](https://platform.openai.com/api-keys) to create the new secret key."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Install and set dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install airbyte pymilvus openai milvus-model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import airbyte as ab\n",
    "from openai import OpenAI\n",
    "from pymilvus import MilvusClient\n",
    "\n",
    "\n",
    "# in production, you might want to avoid putting the key here.\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-proj-****\"\n",
    "openai_client = OpenAI()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Set the source"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Connect `source-facebook-marketing` to fetch streams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source = ab.get_source(\n",
    "    \"source-facebook-marketing\",\n",
    "    config={\n",
    "        \"start_date\": \"2024-06-01T00:00:00Z\",\n",
    "        \"account_id\": \"account\",\n",
    "        \"access_token\": \"token\"\n",
    "    },\n",
    "    install_if_missing=True,\n",
    ")\n",
    "source.check()\n",
    "source.select_all_streams()\n",
    "result = source.read()\n",
    "\n",
    "for name, records in result.streams.items():\n",
    "    print(f\"Streams of {name} has {len(records)} records.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Milvus Lite & Text Embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"type\": \"DEBUG\", \"message\": \"Created new connection using: 5a7d992d451b41db831d254213b64892\", \"data\": {}}\n"
     ]
    }
   ],
   "source": [
    "milvus_client = MilvusClient(\"./milvus_source_fake.db\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will create the `milvus_source_fake.db` if this is the first initialization. There are some [limitations](https://milvus.io/docs/milvus_lite.md#Limits), but this quick setup for local development should be enought to test the pipeline."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's focused to get the `products` data. We will keep it simple by just getting the relevant fields:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = []\n",
    "for record in result.streams[\"products\"]:\n",
    "    make = record[\"make\"]\n",
    "    model = record[\"model\"]\n",
    "    year = record[\"year\"]\n",
    "\n",
    "    text = f\"{make} {model} {year}\"\n",
    "    data.append(text)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Mazda MX-5 2023',\n",
       " 'Mercedes-Benz C-Class 2023',\n",
       " 'Honda Accord Crosstour 2023',\n",
       " 'GMC Jimmy 2023',\n",
       " 'Infiniti FX 2023']"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymilvus import model\n",
    "\n",
    "openai_ef = model.dense.OpenAIEmbeddingFunction(\n",
    "    model_name='text-embedding-3-large', # Specify the model name\n",
    "    api_key=os.environ[\"OPENAI_API_KEY\"], # Provide your OpenAI API key\n",
    "    dimensions=512 # Set the embedding dimensionality\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedded_data = openai_ef.encode_documents(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"type\": \"DEBUG\", \"message\": \"Successfully created collection: products\", \"data\": {}}\n",
      "{\"type\": \"DEBUG\", \"message\": \"Successfully created an index on collection: products\", \"data\": {}}\n"
     ]
    }
   ],
   "source": [
    "milvus_client.create_collection(\"products\", dimension=512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedded_docs = []\n",
    "for _id, embedded_text in enumerate(embedded_data):\n",
    "    embedded_docs.append({\"id\": _id+1, \"vector\": embedded_text, \"text\": data[_id]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'insert_count': 50,\n",
       " 'ids': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50],\n",
       " 'cost': 0}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "milvus_client.insert(collection_name=\"products\", data=embedded_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Inspect the search results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Give list of products from Suzuki\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "search_res = milvus_client.search(\n",
    "    collection_name=\"products\",\n",
    "    data=[\n",
    "        openai_ef.encode_documents([question])[0]\n",
    "    ],  # Use the `emb_text` function to convert the question to an embedding vector\n",
    "    limit=3,  # Return top 3 results\n",
    "    search_params={\"metric_type\": \"COSINE\", \"params\": {}},  # Inner product distance\n",
    "    output_fields=[\"text\"],  # Return the text field\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[\n",
      "    [\n",
      "        \"Suzuki SJ 410 2023\",\n",
      "        0.5219288468360901\n",
      "    ],\n",
      "    [\n",
      "        \"Isuzu VehiCROSS 2023\",\n",
      "        0.38782158493995667\n",
      "    ],\n",
      "    [\n",
      "        \"Jaguar S-Type 2023\",\n",
      "        0.35628464818000793\n",
      "    ]\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "retrieved_lines_with_distances = [\n",
    "    (res[\"entity\"][\"text\"], res[\"distance\"]) for res in search_res[0]\n",
    "]\n",
    "print(json.dumps(retrieved_lines_with_distances, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Use OpenAI ChatGPT to get the RAG response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define the system and user prompts for the Language Model. This prompt is assembled with the retrieved documents from Milvus.\n",
    "\n",
    "We also use OpenAI ChatGPT to generate a response based on the prompts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "context = \"\\n\".join(\n",
    "    [line_with_distance[0] for line_with_distance in retrieved_lines_with_distances]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "SYSTEM_PROMPT = \"\"\"\n",
    "Human: You are an AI assistant. You are able to find answers to the questions from the contextual passage snippets provided.\n",
    "\"\"\"\n",
    "USER_PROMPT = f\"\"\"\n",
    "Use the following pieces of information enclosed in <context> tags to provide an answer to the question enclosed in <question> tags.\n",
    "<context>\n",
    "{context}\n",
    "</context>\n",
    "<question>\n",
    "{question}\n",
    "</question>\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<answer>\n",
      "Suzuki SJ 410 2023\n",
      "</answer>\n"
     ]
    }
   ],
   "source": [
    "response = openai_client.chat.completions.create(\n",
    "    model=\"gpt-3.5-turbo\",\n",
    "    messages=[\n",
    "        {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
    "        {\"role\": \"user\", \"content\": USER_PROMPT},\n",
    "    ],\n",
    ")\n",
    "print(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Summary\n",
    "\n",
    "This shows how easy to build RAG pipeline in Python for quick local development which helps us to speed our development iterations. All within the comport of Python environment."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pyairbyte-hackathon",
   "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.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}