notebooks/05_create_audio_video_embeddings.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Video and Audio Content Analysis with Amazon Bedrock and Amazon Aurora PostgreSQL pgvector\n",
    "\n",
    "This notebook demonstrates how to process video and audio content using [Amazon Bedrock](https://aws.amazon.com/bedrock/) to invoke [Amazon Titan Multimodal Embeddings G1 model](https://docs.aws.amazon.com/bedrock/latest/userguide/titan-multiemb-models.html) for generating multimodal embeddings, [Amazon Transcribe](https://aws.amazon.com/transcribe/) for converting speech to text, and [Amazon Aurora PostgreSQL](https://docs.aws.amazon.com/AmazonRDS/latest/AuroraUserGuide/data-api.html) with pgvector for efficient vector storage and similarity search, you will build an app that understands both visual and audio content, enabling natural language queries to find specific moments in videos.\n",
    "\n",
    "> Create Amazon Aurora PostgreSQL with this [Amazon CDK Stack](https://github.com/build-on-aws/langchain-embeddings/tree/main/create-aurora-pgvector)\n",
    "\n",
    "![Diagram](data/video-embedding.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### Processing Flow\n",
    "\n",
    "The pipeline processes videos through these steps:\n",
    "\n",
    "1. A video file is processed using Python code in a Jupyter notebook, utilizing the boto3 SDK to interact with AWS services.\n",
    "\n",
    "2. The audio stream is extracted and sent to Amazon Transcribe for speech-to-text conversion.\n",
    "\n",
    "3. Simultaneously, the video is processed to extract key frames, which are stored in an Amazon S3 bucket.\n",
    "\n",
    "4. The extracted frames are processed through Amazon Bedrock's Titan embedding model to generate multimodal vectors that represent the visual content.\n",
    "\n",
    "5. Finally, all the processed data (transcriptions, frame data, and vectors) is stored in Amazon Aurora Serverless PostgreSQL with pgvector extension, enabling vector-based searches through standard RDS API calls.\n",
    "\n",
    "![Diagram](data/diagram_video.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 💰 Cost to complete: \n",
    "- [Amazon Bedrock Pricing](https://aws.amazon.com/bedrock/pricing/)\n",
    "- [Amazon S3 Pricing](https://aws.amazon.com/s3/pricing/)\n",
    "- [Amazon Aurora Pricing](https://aws.amazon.com/rds/aurora/pricing/)\n",
    "- [Amazon Transcribe Pricing](https://aws.amazon.com/transcribe/pricing/)\n",
    "\n",
    "### Configuration\n",
    "- [AWS SDK for Python ](https://docs.aws.amazon.com/AmazonS3/latest/userguide/UsingTheBotoAPI.html)\n",
    "- [Configure AWS credentials](https://docs.aws.amazon.com/braket/latest/developerguide/braket-using-boto3.html) \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install boto3\n",
    "#!pip install json\n",
    "#!pip install base64\n",
    "#!pip install uuid\n",
    "# or install requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import json\n",
    "import os\n",
    "from PIL import Image as PILImage\n",
    "import random\n",
    "\n",
    "_region_name = \"us-west-2\"\n",
    "ssm = boto3.client(service_name=\"ssm\", region_name=_region_name)\n",
    "\n",
    "# Default model settings\n",
    "default_model_id = os.environ.get(\"DEFAULT_MODEL_ID\", \"amazon.titan-embed-image-v1\")\n",
    "default_embedding_dimmesion = os.environ.get(\"DEFAULT_EMBEDDING_DIMENSION\", \"1024\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Database Interface (AuroraPostgres Class)\n",
    "\n",
    "An `AuroraPostgres` class that interacts with Amazon Aurora PostgreSQL [using RDS Data API](https://docs.aws.amazon.com/AmazonRDS/latest/AuroraUserGuide/data-api.html)\n",
    "\n",
    "Code: [aurora_service.py](create_audio_video_helper/aurora_service.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from create_audio_video_helper import AuroraPostgres"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Video Content Processing\n",
    "\n",
    "A `VideoProcessor` class uses the [ffmpeg libavcodec library](https://ffmpeg.org/) to proccess the audio and create frames. \n",
    "\n",
    "The class is set to process frames every 1 second, you can modify this by changing the FPS value in command.\n",
    "\n",
    "Code: [video_processor.py](create_audio_video_helper/video_processor.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from create_audio_video_helper import VideoProcessor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Video Download and Processing\n",
    "\n",
    "Code: [video_manager.py](create_audio_video_helper/video_manager.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from create_audio_video_helper import VideoManager\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Audio Processing with Amazon Transcribe\n",
    "\n",
    "The `AudioProcessing` class extracts the audio track from the video file using [Amazon Transcribe StartTranscriptionJob API](https://docs.aws.amazon.com/transcribe/latest/APIReference/API_StartTranscriptionJob.html), converting speech into accurate text transcripts.\n",
    "With `IdentifyMultipleLanguages` as True, Transcribe uses [Amazon Comprehend](https://aws.amazon.com/comprehend/)to identify the language in the audio, If you know the language of your media file, specify it using the `LanguageCode` parameter.  \n",
    "\n",
    "`ShowSpeakerLabels` parameter as `True` enables speaker partitioning (diarization) in the transcription output. Speaker partitioning labels the speech from individual speakers in the media file and include `MaxSpeakerLabels` to specify the maximum number of speakers, in this case is 10. \n",
    "\n",
    "Code: [audio_processor.py](create_audio_video_helper/audio_processing.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from create_audio_video_helper import AudioProcessing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Embedding Generation\n",
    "\n",
    "Generate Embeddings for each extracted frame. Embeddins are created with the Amazon Titan Multimodal Embeddings G1 model using Amazon Bedrock Invoke Model API. \n",
    "\n",
    "Code: [embedding_generation.py](create_audio_video_helper/embedding_generation.py)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from create_audio_video_helper import EmbeddingGeneration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Select Key Frames\n",
    "\n",
    "The app uses the `CompareFrame` class to identifies significant visual changes by detecting when frame similarity falls below a defined threshold, in this case 0.8. This comparison leverages Cosine Similarity, calculating the cosine of the angle between frame vectors. The similarity score ranges from -1 to 1, with higher values indicating greater visual similarity between frames.\n",
    "\n",
    "Code: [compare_frames.py](create_audio_video_helper/compare_frames.py)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from create_audio_video_helper import CompareFrames"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration\n",
    "The system uses environment variables and AWS Systems Manager Parameter Store for configuration:\n",
    "\n",
    "**DEFAULT_MODEL_ID:** Bedrock model ID (default: \"amazon.titan-embed-image-v1\")\n",
    "\n",
    "**DEFAULT_EMBEDDING_DIMENSION:** Embedding dimension (default: \"1024\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def get_ssm_parameter(name):\n",
    "    response = ssm.get_parameter(Name=name, WithDecryption=True)\n",
    "    return response[\"Parameter\"][\"Value\"]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get Data from environment variables, never share secrets!\n",
    "\n",
    "cluster_arn = get_ssm_parameter(\"/videopgvector/cluster_arn\")\n",
    "credentials_arn = get_ssm_parameter(\"/videopgvector/secret_arn\")\n",
    "table_name = get_ssm_parameter(\"/videopgvector/video_table_name\")\n",
    "default_database_name = \"kbdata\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize Aurora PostgreSQL client\n",
    "aurora = AuroraPostgres(cluster_arn, default_database_name, credentials_arn,_region_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify Aurora Cluster conectivity:\n",
    "aurora.execute_statement(\"select count(*) from bedrock_integration.knowledge_bases\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Upload Video to Amazon S3 bucket and Obtain s3_uri\n",
    "\n",
    "This code shows how to upload a video from the `tmp` folder to an S3 bucket and obtain the S3 URI needed for further processing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from create_audio_video_helper.video_s3_uploader import UploadVideoS3\n",
    "# Configure the parameters\n",
    "video_path = \"tmp/video.mp4\"  # Path to the video in the tmp folder\n",
    "bucket_name = \"you-bucket-1234\"     # Name of your S3 bucket\n",
    "uploadvideo = UploadVideoS3(bucket_name)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# You can also specify a custom path in S3 (optional)\n",
    "s3_key = \"videos/sample_video.mp4\"\n",
    "\n",
    "# Subir el video y obtener el S3 URI\n",
    "s3_uri = uploadvideo.upload_video_to_s3(video_path, s3_key)\n",
    "print(f\"S3 URI: {s3_uri}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the file\n",
    "# Create directory if it doesn't exist\n",
    "\n",
    "tmp_path                    = \"./tmp\"\n",
    "\n",
    "#s3_uri = \"s3://you-bucket-1234/videos/you-video.mp4\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "videomanager = VideoManager(s3_uri,_region_name)\n",
    "\n",
    "bucket, prefix, fileName, extension, file  = videomanager.parse_location(s3_uri)\n",
    "\n",
    "local_path              = f\"{tmp_path}/{file}\"\n",
    "location                = f\"{prefix}/{file}\"\n",
    "output_dir              = f\"{tmp_path}/{fileName}\"\n",
    "\n",
    "\n",
    "os.makedirs(os.path.dirname(local_path), exist_ok=True)\n",
    "print(f\"dowloading s3://{bucket}/{prefix}/{file} to {local_path}\")\n",
    "result = videomanager.download_file(bucket,location, local_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify ffmpeg is installed \n",
    "videoprocessor = VideoProcessor()\n",
    "videoprocessor.ffmpeg_check() ## Check if ffmpeg is installed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Process the media file\n",
    "This part involves:\n",
    "1. For visual content:\n",
    "\n",
    "![Diagram](data/frames_processing.png)\n",
    "\n",
    "2. Transcribing the audio to text using Amazon Transcribe\n",
    "\n",
    "![Diagram](data/audio_processing.png)\n",
    "\n",
    "This notebook assumes you have a valid media file in s3://path/to/video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ✅ Start Amazon Transcribe Job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_processing = AudioProcessing(_region_name,videomanager)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job_name = audio_processing.transcribe(s3_uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ✅  Extract Key Frames with ffmpeg and Amazon Bedrock with Titan Multimodal Embeddings Model\n",
    "\n",
    "![Diagram](data/extract_frames.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "files = videoprocessor.extract_frames(local_path, output_dir, every=1) # 1 frame per second"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Text and Image embeddings \n",
    "\n",
    "![Diagram](data/get_images_embeddings.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_generation = EmbeddingGeneration(videomanager,_region_name,default_model_id,default_embedding_dimmesion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate embeddings for all extracted frames (1 per second)\n",
    "embed_1024 = embedding_generation.get_images_embeddings(files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compareframes = CompareFrames()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get only different frames by calculating cosine similarity sequentially\n",
    "selected_frames = compareframes.filter_relevant_frames(embed_1024, difference_threshold=0.8) # frame is skipped if is similar to previous \n",
    "\n",
    "print (f\"from {len(embed_1024)} frames to {len(selected_frames)} relevant frames:\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ✅  Check the transcription Job and process text results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#job_name = \"XXXX\" # For existing jobs put the job name here\n",
    "transcript_url =audio_processing.wait_transcription_complete(job_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transcripts, duration = audio_processing.process_transcript(transcript_url, max_chars_per_segment=1000)\n",
    "print (f\"Duration:{duration}s\")\n",
    "for seg, speaker, text in transcripts[:2]:\n",
    "    print (f\"sec: {seg}\\n{speaker}:\\n   {text}\\n\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "selected_frames_files = [(sf, files[sf])for sf in selected_frames]\n",
    "selected_frames_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text_embeddings = embedding_generation.create_text_embeddings(transcripts, transcript_url)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The list of embeddings for text should look like this: \n",
    "\n",
    "![Diagram](data/images_embeddings.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print (\"Text Embeddings:\\n\")\n",
    "for te in text_embeddings:\n",
    "    print(f\"Chunk:{te.get('chunks')[:50]}, embedding(3): {te.get('embedding')[:3]}, metadata: {te.get('metadata')} \")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frames_embeddings = embedding_generation.create_frames_embeddings(selected_frames_files, s3_uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The list of embeddings for image should look like this:\n",
    "\n",
    "![Diagram](data/text_embedding.png)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print (\"\\nImage Embeddings:\\n\")\n",
    "for fe in frames_embeddings:\n",
    "    print(f\"Source:{fe.get('source')}, embedding(3): {fe.get('embedding')[:3]}, metadata: {fe.get('metadata')} \")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Insert to Vector Database Aurora PostgreSQL (pgvector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aurora.execute_statement(\"select count(*) from bedrock_integration.knowledge_bases\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optionally clean the table\n",
    "aurora.execute_statement(\"delete from bedrock_integration.knowledge_bases\")\n",
    "aurora.execute_statement(\"select count(*) from bedrock_integration.knowledge_bases\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Insert text embeddings into Aurora PostgreSQL\n",
    "if text_embeddings:\n",
    "    aurora.insert(text_embeddings)\n",
    "    print(f\"Inserted {len(text_embeddings)} text embeddings\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Insert frame embeddings into Aurora PostgreSQL\n",
    "if frames_embeddings:\n",
    "    aurora.insert(frames_embeddings)\n",
    "    print(f\"Inserted {len(frames_embeddings)} frame embeddings\")\n",
    "\n",
    "\n",
    "aurora.execute_statement(\"select count(*) from bedrock_integration.knowledge_bases\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Similarity Search\n",
    "\n",
    "Implements functions for:\n",
    "- `retrieve()`: Performs similarity searches in the database and displays results\n",
    "- `aurora.similarity_search()`: Executes the vector similarity search in the database\n",
    "- `get_embeddings()`: Generates embeddings for the search query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display\n",
    "\n",
    "def retrieve(search_query, how=\"cosine\", k=5):\n",
    "    search_vector = embedding_generation.get_embeddings(search_query)\n",
    "    \n",
    "    result = aurora.similarity_search(search_vector,how=how, k=k)\n",
    "    rows = json.loads(result.get(\"formattedRecords\"))\n",
    "    for row in rows:\n",
    "        metric = \"similarity\" if how == \"cosine\" else \"distance\"\n",
    "        metric_value = row.get(metric)\n",
    "        if row.get(\"content_type\") == \"text\":\n",
    "            print (f\"text:\\n{row.get('chunks')}\\n{metric}:{metric_value}\\nmetadata:{row.get('metadata')}\\n\")\n",
    "        if row.get(\"content_type\") == \"image\":\n",
    "            img = PILImage.open(row.get('source'))            \n",
    "            print (f\"Image:\\n{row.get('source')}\\n{metric}:{metric_value}\\nmetadata:{row.get('metadata')}\\n\")\n",
    "            display(img)\n",
    "        del row[\"embedding\"]\n",
    "        del row[\"id\"]\n",
    "\n",
    "    return rows"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I tested the notebook with my AWS re:Invent 2024 sesion [AI self-service support with knowledge retrieval using PostgreSQL](https://www.youtube.com/watch?v=fpi3awGakyg?trk=fccf147c-636d-45bf-bf0a-7ab087d5691a&sc_channel=video). \n",
    "\n",
    "I ask for Aurora and it brings me images and texts where it mentions:\n",
    "\n",
    "![Diagram](data/cosine.png)\n",
    "\n",
    "```bash\n",
    "text:\n",
    "memory . A place where all the information is stored and can easily be retrievable , and that's where the vector database comes in . This is the the first building block . And a vector database stores and retrieves data in the form of vector embeddeds or mathematical representations . This allows us to find similarities between data rather than relying on the exact keyword match that is what usually happens up to today . This is essential for systems like retrieval ofmented generation or RAC , which combines external knowledge with the AI response to deliver those accurate and context aware response . And by the way , I think yesterday we announced the re-rank API for RAC . So now your rack applications , you can score and it will prioritize those documents that have the most accurate information . So at the end will be even faster and cheaper building rack . We're gonna use Amazon Aurora postgrade SQL with vector support that will give us a scalable and fully managed solution for our AI tasks .\n",
    "similarity:0.5754164493071239\n",
    "metadata:{\"speaker\":\"spk_0\",\"second\":321}\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "search_query = \"<your-query>\"\n",
    "docs = retrieve(search_query, how=\"cosine\", k=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "search_query = \"<your-query>\"\n",
    "docs = retrieve(search_query, how=\"l2\", k=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Search using images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, I send this image for the search:\n",
    "![Diagram](data/image_sends.png)\n",
    "\n",
    "This is one of the answers that the query gave me:\n",
    "![Diagram](data/image_response.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "one_image = random.choice(files)\n",
    "print(one_image)\n",
    "display(PILImage.open(one_image))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = retrieve(videomanager.read_image_from_local(one_image), how=\"cosine\", k=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RAG Implementation\n",
    "\n",
    "Finally, the notebook implements a complete RAG system:\n",
    "- `CustomMultimodalRetriever`: A custom retriever class that extends BaseRetriever\n",
    "- `_get_relevant_documents()`: Core retrieval method that finds similar content\n",
    "- `image_content_block()`: Formats image content for LLM consumption\n",
    "- `text_content_block()`: Formats text content for LLM consumption\n",
    "- `parse_docs_for_context()`: Processes documents for context (text and images)\n",
    "- `ThinkingLLM`: Uses an LLM to answer questions based on retrieved content\n",
    "\n",
    "> Based on https://github.com/langchain-ai/langchain/blob/master/docs/docs/how_to/custom_retriever.ipynb\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from langchain_core.callbacks import CallbackManagerForRetrieverRun\n",
    "from langchain_core.documents import Document\n",
    "from langchain_core.retrievers import BaseRetriever\n",
    "\n",
    "class CustomMultimodalRetriever(BaseRetriever):\n",
    "    \"\"\"A retriever that contains the top k documents that contain the user query.\n",
    "    query could be text or image_bytes\n",
    "    \"\"\"\n",
    "    k: int\n",
    "    \"\"\"Number of top results to return\"\"\"\n",
    "    how: str\n",
    "    \"\"\"How to calculate the similarity between the query and the documents.\"\"\"\n",
    "\n",
    "    def _get_relevant_documents(\n",
    "        self, query: str, *, run_manager: CallbackManagerForRetrieverRun\n",
    "    ) -> List[Document]:\n",
    "        \"\"\"Sync implementations for retriever.\"\"\"\n",
    "        search_vector = embedding_generation.get_embeddings(query)\n",
    "        result = aurora.similarity_search(search_vector, how=self.how, k=self.k)\n",
    "        rows = json.loads(result.get(\"formattedRecords\"))\n",
    "\n",
    "        matching_documents = []\n",
    "\n",
    "        for row in rows:\n",
    "            document_kwargs = dict(\n",
    "                metadata=dict(**json.loads(row.get(\"metadata\")), content_type = row.get(\"content_type\"), source=row.get(\"sourceurl\")))\n",
    "            \n",
    "            if self.how == \"cosine\":\n",
    "                document_kwargs[\"similarity\"] = row.get(\"similarity\")\n",
    "            elif self.how == \"l2\":\n",
    "                document_kwargs[\"distance\"] = row.get(\"distance\")\n",
    "\n",
    "            if row.get(\"content_type\") == \"text\":\n",
    "                matching_documents.append( Document( page_content=row.get(\"chunks\"), **document_kwargs ))\n",
    "            if row.get(\"content_type\") == \"image\":\n",
    "                matching_documents.append( Document( page_content=row.get(\"source\"),**document_kwargs ))\n",
    "\n",
    "        return matching_documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = CustomMultimodalRetriever(how=\"cosine\", k=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"<your-query>\"\n",
    "docs = retriever.invoke(query)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building the RAG "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Dict\n",
    "bedrock_runtime = boto3.client(service_name=\"bedrock-runtime\", region_name=_region_name)\n",
    "\n",
    "\n",
    "budget_tokens = 0\n",
    "max_tokens = 1024\n",
    "conversation: List[Dict] = []\n",
    "reasoning_config = {\"thinking\": {\"type\": \"enabled\", \"budget_tokens\": budget_tokens}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def image_content_block(image_file):\n",
    "    image_bytes = videomanager.read_image_from_local(image_file)\n",
    "    extension = image_file.split('.')[-1]\n",
    "    print (f\"Including Image :{image_file}\")\n",
    "    if extension == 'jpg':\n",
    "        extension = 'jpeg'\n",
    "    \n",
    "    block = { \"image\": { \"format\": extension, \"source\": { \"bytes\": image_bytes}}}\n",
    "    return block\n",
    "\n",
    "def text_content_block(text):\n",
    "    return { \"text\": text }\n",
    "\n",
    "def parse_docs_for_context(docs):\n",
    "    blocks = []\n",
    "    for doc in docs:\n",
    "        if doc.metadata.get('content_type') == \"image\":\n",
    "            blocks.append(image_content_block(doc.page_content))\n",
    "        else:\n",
    "            blocks.append(text_content_block(doc.page_content))\n",
    "    return blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def answer(model_id,system_prompt,content) -> str:\n",
    "    \"\"\"Get completion from Claude model based on conversation history.\n",
    "\n",
    "    Returns:\n",
    "        str: Model completion text\n",
    "    \"\"\"\n",
    "\n",
    "    # Invoke model\n",
    "    kwargs = dict(\n",
    "        modelId=model_id,\n",
    "        inferenceConfig=dict(maxTokens=max_tokens),\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": content,\n",
    "            }\n",
    "        ],\n",
    "\n",
    "    )\n",
    "\n",
    "    kwargs[\"system\"] = [{\"text\": system_prompt}]\n",
    "\n",
    "    response = bedrock_runtime.converse(**kwargs)\n",
    "    \n",
    "    return response.get(\"output\",{}).get(\"message\",{}).get(\"content\", [])\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parsed_docs = parse_docs_for_context(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "system_prompt = \"\"\"Answer the user's questions based on the below context. If the context has an image, indicate that it can be reviewed for further feedback.\n",
    "If the context doesn't contain any relevant information to the question, don't make something up and just say \"I don't know\". (IF YOU MAKE SOMETHING UP BY YOUR OWN YOU WILL BE FIRED). For each statement in your response provide a [n] where n is the document number that provides the response. \"\"\"\n",
    "model_id = \"us.amazon.nova-pro-v1:0\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"<your-query>\"\n",
    "docs = retriever.invoke(query)\n",
    "parsed_docs = parse_docs_for_context(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_response = answer(model_id,system_prompt,[text_content_block(f\"question:{query}\\n\\nDocs:\\n\"), *parsed_docs])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(llm_response[0].get(\"text\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"<your-query>\"\n",
    "docs = retriever.invoke(query)\n",
    "parsed_docs = parse_docs_for_context(docs)\n",
    "llm_response = answer(model_id,system_prompt,[text_content_block(f\"question:{query}\\n\\nDocs:\\n\"), *parsed_docs])\n",
    "print(llm_response[0].get(\"text\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}