notebooks/retriever_api_usage.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5b5cfe48-22f4-4526-91f3-b555c45b406d",
   "metadata": {},
   "source": [
    "### Retriever API Usage\n",
    "- Replace IPADDRESS `http://localhost:8081` with the actual server URL where the API is hosted.\n",
    "- Please run the [ingestion notebook](./ingestion_api_usage.ipynb) as a prerequisite to using this notebook.\n",
    "- Ensure the server is running before executing the notebook.\n",
    "\n",
    "You can now execute each cell in sequence to test the API.\n",
    "#### 1. Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc28c582-c0ac-4f04-8c97-433e79be8366",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install requests"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c541f87b-c9b7-46dc-8d1e-7fb6eb885374",
   "metadata": {},
   "source": [
    "#### 2. Setup Base Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fcab5bd-60d5-4354-a3f7-bf8039d78389",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import json\n",
    "from typing import Dict, Any\n",
    "\n",
    "IPADDRESS = \"localhost\" #Replace this with the correct IP address\n",
    "RAG_PORT = \"8081\"\n",
    "BASE_URL = f\"http://{IPADDRESS}:{RAG_PORT}\"  # Replace with your server URL\n",
    "\n",
    "def print_response(response: requests.Response):\n",
    "    \"\"\"Helper to print API response.\"\"\"\n",
    "    print(f\"Status Code: {response.status_code}\")\n",
    "    try:\n",
    "        print(json.dumps(response.json(), indent=2))\n",
    "    except requests.JSONDecodeError:\n",
    "        print(response.text)\n",
    "    except Exception:\n",
    "        print(response.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6451651b-23b6-4854-a734-656c3fac253f",
   "metadata": {},
   "source": [
    "#### 3. Health Check Endpoint\n",
    "\n",
    "**Purpose:**\n",
    "This endpoint performs a health check on the server. It returns a 200 status code if the server is operational."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d753b52b-2728-4201-bf0e-0fbbf4d115bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# GET /health\n",
    "url = f\"{BASE_URL}/health\"\n",
    "response = requests.get(url)\n",
    "print_response(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "decf4e4b-78cd-4a6f-a97a-971c18dafcaf",
   "metadata": {},
   "source": [
    "#### 4. Generate Answer Endpoint\n",
    "\n",
    "**Purpose:**\n",
    "This endpoint generates an AI response to a given user message. The system message is specified in the [prompts.yaml](src/prompt.yaml) file. This API retrieves the relevant chunks related to the query from knowledge base, adds them as part of the LLM prompt and returns a streaming response. It supports parameters like temperature, top_p, knowledge base usage, and also generates based on the specified vector collection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a4fdcfd-c845-4cbf-b1ae-d53b2d631b52",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = f\"{BASE_URL}/generate\"\n",
    "payload = {\n",
    "    \"messages\": [\n",
    "        {\"role\": \"user\", \"content\": \"Who is the Largest EV Maker?\"}\n",
    "    ],\n",
    "    \"use_knowledge_base\": True,\n",
    "    \"temperature\": 0.5,\n",
    "    \"top_p\": 0.7,\n",
    "    \"max_tokens\": 200,\n",
    "    \"top_k\": 4,\n",
    "    \"collection_name\": \"nvidia_blogs\"\n",
    "}\n",
    "\n",
    "response = requests.post(url, json=payload)\n",
    "print_response(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcdc614b-846f-4a91-a09a-061d8e66e21c",
   "metadata": {},
   "source": [
    "#### 5. Document Search Endpoint\n",
    "\n",
    "**Purpose:**\n",
    "This endpoint searches for the most relevant documents in the vector store based on a query. You can specify the maximum number of documents to retrieve."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d0e2709-9de9-49ca-aa43-66182fe31d52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# POST /search\n",
    "url = f\"{BASE_URL}/search\"\n",
    "payload = {\n",
    "    \"query\": \"artificial intelligence trends\",\n",
    "    \"top_k\": 4,\n",
    "    \"collection_name\": \"nvidia_blogs\"\n",
    "}\n",
    "\n",
    "response = requests.post(url, json=payload)\n",
    "print_response(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}