azure-ai-search-cohere-embed-v3-multimodal.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Cohere Embed Models with Azure AI Search (Text and Image Embeddings)\n",
    "\n",
    "This notebook demonstrates how to use the Cohere Embed model deployed on Azure AI to generate embeddings for both text and image data. These embeddings are then stored in Azure AI Search for vector similarity search.\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "- **Azure Account**: Ensure you have an active Azure account with access to Azure AI services and Azure Cognitive Search.\n",
    "- **Cohere Embed Model**: Deploy the Cohere Embed model in Azure AI.\n",
    "- **API Keys and Endpoints**: Obtain the necessary API keys and endpoints for Azure AI and Azure Cognitive Search.\n",
    "- **Sample Images**: Have a few sample images available for embedding.\n",
    "\n",
    "## Install Required Packages\n",
    "\n",
    "First, install the necessary Python packages.\n",
    "\n",
    "```python\n",
    "!pip install cohere azure-search-documents python-dotenv\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "[notice] A new release of pip is available: 24.2 -> 24.3.1\n",
      "[notice] To update, run: python.exe -m pip install --upgrade pip\n",
      "\n",
      "[notice] A new release of pip is available: 24.2 -> 24.3.1\n",
      "[notice] To update, run: python.exe -m pip install --upgrade pip\n"
     ]
    }
   ],
   "source": [
    "! pip install --quiet azure-search-documents==11.6.0b6\n",
    "! pip install --quiet cohere python-dotenv azure-identity tqdm requests tenacity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64\n",
    "import cohere\n",
    "import json\n",
    "import numpy as np\n",
    "import os\n",
    "import requests\n",
    "from azure.core.credentials import AzureKeyCredential\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from azure.search.documents import SearchClient\n",
    "from azure.search.documents.indexes import SearchIndexClient\n",
    "from azure.search.documents.indexes.models import (\n",
    "    AIStudioModelCatalogName,\n",
    "    AzureMachineLearningParameters,\n",
    "    AzureMachineLearningVectorizer,\n",
    "    HnswAlgorithmConfiguration,\n",
    "    HnswParameters,\n",
    "    SearchField,\n",
    "    SearchFieldDataType,\n",
    "    SearchIndex,\n",
    "    SearchableField,\n",
    "    SimpleField,\n",
    "    VectorEncodingFormat,\n",
    "    VectorSearch,\n",
    "    VectorSearchAlgorithmKind,\n",
    "    VectorSearchAlgorithmMetric,\n",
    "    VectorSearchProfile,\n",
    ")\n",
    "from azure.search.documents.models import (\n",
    "    VectorizableTextQuery,\n",
    "    VectorizedQuery,\n",
    "    VectorQuery,\n",
    ")\n",
    "from dotenv import load_dotenv\n",
    "from tenacity import retry, stop_after_attempt, wait_exponential\n",
    "from tqdm.notebook import tqdm\n",
    "from typing import Dict, List, Optional\n",
    "\n",
    "\n",
    "# Load environment variables\n",
    "load_dotenv()\n",
    "\n",
    "# Azure AI Studio Cohere Configuration\n",
    "AZURE_AI_STUDIO_COHERE_EMBED_KEY = os.getenv(\"AZURE_AI_STUDIO_COHERE_EMBED_KEY\")\n",
    "AZURE_AI_STUDIO_COHERE_EMBED_ENDPOINT = os.getenv(\"AZURE_AI_STUDIO_COHERE_EMBED_ENDPOINT\")\n",
    "AZURE_AI_STUDIO_COHERE_COMMAND_ENDPOINT= os.getenv(\"AZURE_AI_STUDIO_COHERE_COMMAND_ENDPOINT\")\n",
    "AZURE_AI_STUDIO_COHERE_COMMAND_KEY=os.getenv(\"AZURE_AI_STUDIO_COHERE_COMMAND_KEY\")\n",
    "\n",
    "# Azure Cognitive Search Configuration\n",
    "AZURE_SEARCH_SERVICE_ENDPOINT = os.getenv(\"AZURE_SEARCH_SERVICE_ENDPOINT\")\n",
    "AZURE_SEARCH_ADMIN_KEY = os.getenv(\"AZURE_SEARCH_ADMIN_KEY\")\n",
    "INDEX_NAME = \"multimodal-cohere-index\"\n",
    "\n",
    "# Cohere Model Information\n",
    "EMBEDDING_MODEL_NAME = \"embed-english-v3.0\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Authenticate Clients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Authenticate Azure Cognitive Search client\n",
    "azure_search_credential = AzureKeyCredential(AZURE_SEARCH_ADMIN_KEY)\n",
    "\n",
    "# Initialize Azure Search clients\n",
    "index_client = SearchIndexClient(\n",
    "    endpoint=AZURE_SEARCH_SERVICE_ENDPOINT,\n",
    "    credential=azure_search_credential\n",
    ")\n",
    "\n",
    "search_client = SearchClient(\n",
    "    endpoint=AZURE_SEARCH_SERVICE_ENDPOINT,\n",
    "    index_name=INDEX_NAME,\n",
    "    credential=azure_search_credential\n",
    ")\n",
    "\n",
    "# Initialize Cohere client\n",
    "cohere_client = cohere.Client(\n",
    "    api_key=AZURE_AI_STUDIO_COHERE_EMBED_KEY,\n",
    "    base_url=AZURE_AI_STUDIO_COHERE_EMBED_ENDPOINT\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample Data\n",
    "### Text Documents and Image URLs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = [\n",
    "    {\n",
    "        \"id\": \"1\",\n",
    "        \"caption\": \"The image showcases a peaceful rural scene with several cows leisurely grazing in a sunlit pasture. The golden light of early morning creates a warm ambiance, illuminating the cows as they feed on the lush grass. In the background, soft shadows and a hazy atmosphere enhance the tranquil and idyllic nature of this farm setting\",\n",
    "        \"imageUrl\": \"https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample1-bbe41ac5.png\"\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"2\",\n",
    "        \"caption\": \"The image captures a solitary surfer standing on a rocky outcrop, gazing out at the ocean waves. Clad in a wetsuit and holding a surfboard, the figure embodies a spirit of adventure and anticipation as they prepare to embrace the surf. With the soft blue sky and gentle waves in the background, the scene conveys a sense of tranquility and the thrill of the ocean.\",\n",
    "        \"imageUrl\": \"https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample2-72b3c1ca.png\"\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"3\",\n",
    "        \"caption\": \"The image showcases the ornate architecture of a grand building, crowned by intricate sculptures that convey a sense of justice and authority. Prominently featured is a figure representing Lady Justice, holding scales, symbolizing fairness and the rule of law. Set against a backdrop of a bright blue sky, the details of the sculpture and the building's design reflect a commitment to justice and civic pride.\",\n",
    "        \"imageUrl\": \"https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample3-e03062c2.png\"\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"4\",\n",
    "        \"caption\": \"The image features a focused baseball player poised at the plate, ready to swing his bat. Dressed in a gray uniform with a striking red helmet, he embodies determination and athleticism under the stadium lights. The blurred background of the baseball field highlights the intensity of the moment, capturing the excitement of the game.\",\n",
    "        \"imageUrl\": \"https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample4-0559774f.png\"\n",
    "    }\n",
    "]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functions to Encode Images and Generate Embeddings\n",
    "### Encode Image to Base64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode_image_to_base64(image_url: str) -> str:\n",
    "    \"\"\"\n",
    "    Convert image URL to base64 data URI format required by Cohere.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        response = requests.get(image_url, timeout=10)\n",
    "        response.raise_for_status()\n",
    "        content_type = response.headers['Content-Type']\n",
    "        base64_data = base64.b64encode(response.content).decode('utf-8')\n",
    "        return f\"data:{content_type};base64,{base64_data}\"\n",
    "    except Exception as e:\n",
    "        print(f\"Error encoding image {image_url}: {e}\")\n",
    "        return None\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate Text Embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_text_embedding(text: str, model=EMBEDDING_MODEL_NAME) -> List[float]:\n",
    "    response = cohere_client.embed(\n",
    "        texts=[text],\n",
    "        model=model,\n",
    "        input_type=\"search_document\",\n",
    "        truncate=\"NONE\"\n",
    "    )\n",
    "    return response.embeddings[0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate Image Embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_image_embedding(image_base64: str, model=EMBEDDING_MODEL_NAME) -> List[float]:\n",
    "    response = cohere_client.embed(\n",
    "        images=[image_base64],\n",
    "        model=model,\n",
    "        input_type=\"image\",\n",
    "        truncate=\"NONE\"\n",
    "    )\n",
    "    return response.embeddings[0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Process and Embed Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8b73a148af4f43b6a8fa8ab83471dd67",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Processing documents:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "processed_documents = []\n",
    "\n",
    "for doc in tqdm(documents, desc=\"Processing documents\"):\n",
    "    caption = doc[\"caption\"]\n",
    "    image_url = doc[\"imageUrl\"]\n",
    "    \n",
    "    # Generate text embedding\n",
    "    caption_embedding = generate_text_embedding(caption)\n",
    "    \n",
    "    # Encode image and generate image embedding\n",
    "    image_base64 = encode_image_to_base64(image_url)\n",
    "    if image_base64:\n",
    "        image_embedding = generate_image_embedding(image_base64)\n",
    "    else:\n",
    "        image_embedding = None\n",
    "    \n",
    "    # Add embeddings to document\n",
    "    doc[\"captionVector\"] = caption_embedding\n",
    "    doc[\"imageVector\"] = image_embedding\n",
    "    \n",
    "    processed_documents.append(doc)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an Azure AI Search Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Index multimodal-cohere-index created successfully\n",
      "\n",
      "Index fields:\n",
      "- id (Edm.String)\n",
      "- imageUrl (Edm.String)\n",
      "- caption (Edm.String)\n",
      "- imageVector (Collection(Edm.Single))\n",
      "- captionVector (Collection(Edm.Single))\n"
     ]
    }
   ],
   "source": [
    "# Define the fields for the index\n",
    "fields = [\n",
    "    SimpleField(name=\"id\", type=SearchFieldDataType.String, key=True),\n",
    "    SimpleField(\n",
    "        name=\"imageUrl\",\n",
    "        type=SearchFieldDataType.String,\n",
    "        retrievable=True,\n",
    "        filterable=True\n",
    "    ),\n",
    "    SimpleField(\n",
    "        name=\"caption\",\n",
    "        type=SearchFieldDataType.String,\n",
    "        searchable=True,\n",
    "        retrievable=True\n",
    "    ),\n",
    "    SearchField(\n",
    "        name=\"imageVector\",\n",
    "        type=SearchFieldDataType.Collection(SearchFieldDataType.Single),\n",
    "        searchable=True,\n",
    "        vector_search_dimensions=1024,  # Adjust if your model uses different dimensions\n",
    "        vector_search_profile_name=\"vector_profile\"\n",
    "    ),\n",
    "    SearchField(\n",
    "        name=\"captionVector\",\n",
    "        type=SearchFieldDataType.Collection(SearchFieldDataType.Single),\n",
    "        searchable=True,\n",
    "        vector_search_dimensions=1024,  # Adjust if your model uses different dimensions\n",
    "        vector_search_profile_name=\"vector_profile\"\n",
    "    )\n",
    "]\n",
    "\n",
    "# Define vector search configuration\n",
    "vector_search = VectorSearch(\n",
    "    algorithms=[\n",
    "        HnswAlgorithmConfiguration(\n",
    "            name=\"hnsw_config\",\n",
    "            kind=VectorSearchAlgorithmKind.HNSW,\n",
    "            parameters=HnswParameters(\n",
    "                m=4,\n",
    "                ef_construction=400,\n",
    "                ef_search=500,\n",
    "                metric=VectorSearchAlgorithmMetric.COSINE,\n",
    "            ),\n",
    "        )\n",
    "    ],\n",
    "    profiles=[\n",
    "        VectorSearchProfile(\n",
    "            name=\"vector_profile\", \n",
    "            algorithm_configuration_name=\"hnsw_config\",\n",
    "            vectorizer_name=\"vectorizer\",\n",
    "        )\n",
    "    ],\n",
    "    vectorizers=[\n",
    "        AzureMachineLearningVectorizer(\n",
    "            vectorizer_name=\"vectorizer\",\n",
    "            aml_parameters=AzureMachineLearningParameters(\n",
    "                scoring_uri=AZURE_AI_STUDIO_COHERE_EMBED_ENDPOINT,\n",
    "                authentication_key=AZURE_AI_STUDIO_COHERE_EMBED_KEY,\n",
    "                model_name=AIStudioModelCatalogName.COHERE_EMBED_V3_ENGLISH,\n",
    "            ),\n",
    "        )\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Create the index\n",
    "index = SearchIndex(\n",
    "    name=INDEX_NAME,\n",
    "    fields=fields,\n",
    "    vector_search=vector_search\n",
    ")\n",
    "\n",
    "# Create or update the index\n",
    "try:\n",
    "    result = index_client.create_or_update_index(index)\n",
    "    print(f\"Index {INDEX_NAME} created successfully\")\n",
    "except Exception as e:\n",
    "    print(f\"Error creating index: {str(e)}\")\n",
    "\n",
    "# Optional: Verify index fields\n",
    "try:\n",
    "    index_info = index_client.get_index(INDEX_NAME)\n",
    "    print(\"\\nIndex fields:\")\n",
    "    for field in index_info.fields:\n",
    "        print(f\"- {field.name} ({field.type})\")\n",
    "except Exception as e:\n",
    "    print(f\"Error retrieving index fields: {str(e)}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Upload Documents to the Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Documents uploaded successfully.\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    result = search_client.merge_or_upload_documents(documents=processed_documents)\n",
    "    if result[0].succeeded:\n",
    "        print(\"Documents uploaded successfully.\")\n",
    "    else:\n",
    "        print(\"Error uploading documents.\")\n",
    "except Exception as e:\n",
    "    print(f\"Error uploading documents: {e}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform multimodal vector search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Text to Text Vector Search\n",
    "This section performs a vector search to find similar text entries based on the query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example usage of each search function\n",
    "query_text = \"sports activities like skateboarding\"\n",
    "image_url = \"https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample5-14b26724.png\" ## Image of man on skateboard\n",
    "\n",
    "# Helper function to display results\n",
    "def display_results(results):\n",
    "    for result in results:\n",
    "        print(f\"Caption: {result['caption']}\")\n",
    "        print(f\"Score: {result['@search.score']}\")\n",
    "        print(f\"URL: {result['imageUrl']}\")\n",
    "        display(HTML(f'<img src=\"{result[\"imageUrl\"]}\" style=\"width:200px;\"/>'))\n",
    "        print(\"-\" * 50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Text to Text Vector Search Results:\n",
      "Caption: The image features a focused baseball player poised at the plate, ready to swing his bat. Dressed in a gray uniform with a striking red helmet, he embodies determination and athleticism under the stadium lights. The blurred background of the baseball field highlights the intensity of the moment, capturing the excitement of the game.\n",
      "Score: 0.5928263\n",
      "URL: https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample4-0559774f.png\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample4-0559774f.png\" style=\"width:200px;\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# Text to Text Vector Search\n",
    "def text_to_text_search(query_text):\n",
    "    # Generate text embedding\n",
    "    text_embedding = generate_text_embedding(query_text)\n",
    "    \n",
    "    # Define the text vector query\n",
    "    text_vector_query = VectorizedQuery(\n",
    "        vector=text_embedding,\n",
    "        k_nearest_neighbors=1,\n",
    "        fields=\"captionVector\"\n",
    "    )\n",
    "    \n",
    "    # Perform search\n",
    "    results = search_client.search(\n",
    "        search_text=None, vector_queries=[text_vector_query], top=1\n",
    "    )\n",
    "    display_results(results)\n",
    "\n",
    "# Example usage\n",
    "print(\"Text to Text Vector Search Results:\")\n",
    "text_to_text_search(query_text)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Text to Image Vector Search\n",
    "This section performs a vector search to find images related to the query text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Text to Image Vector Search Results:\n",
      "Caption: The image captures a solitary surfer standing on a rocky outcrop, gazing out at the ocean waves. Clad in a wetsuit and holding a surfboard, the figure embodies a spirit of adventure and anticipation as they prepare to embrace the surf. With the soft blue sky and gentle waves in the background, the scene conveys a sense of tranquility and the thrill of the ocean.\n",
      "Score: 0.612399\n",
      "URL: https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample2-72b3c1ca.png\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample2-72b3c1ca.png\" style=\"width:200px;\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# Text to Image Vector Search\n",
    "def text_to_image_search(query_text):\n",
    "    # Generate text embedding\n",
    "    text_embedding = generate_text_embedding(query_text)\n",
    "    \n",
    "    # Define the text-to-image vector query\n",
    "    text_to_image_query = VectorizedQuery(\n",
    "        vector=text_embedding,\n",
    "        k_nearest_neighbors=1,\n",
    "        fields=\"imageVector\"\n",
    "    )\n",
    "    \n",
    "    \n",
    "    # Perform search\n",
    "    results = search_client.search(\n",
    "        search_text=None, vector_queries=[text_to_image_query], top=1\n",
    "    )\n",
    "    display_results(results)\n",
    "\n",
    "# Example usage\n",
    "print(\"Text to Image Vector Search Results:\")\n",
    "text_to_image_search(query_text)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Image to Text Vector Search\n",
    "This section performs a vector search to find text entries related to a given image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image to Text Vector Search Results:\n",
      "Caption: The image captures a solitary surfer standing on a rocky outcrop, gazing out at the ocean waves. Clad in a wetsuit and holding a surfboard, the figure embodies a spirit of adventure and anticipation as they prepare to embrace the surf. With the soft blue sky and gentle waves in the background, the scene conveys a sense of tranquility and the thrill of the ocean.\n",
      "Score: 0.6095421\n",
      "URL: https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample2-72b3c1ca.png\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample2-72b3c1ca.png\" style=\"width:200px;\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# Image to Text Vector Search\n",
    "def image_to_text_search(image_url):\n",
    "    # Generate image embedding\n",
    "    image_base64 = encode_image_to_base64(image_url)\n",
    "    image_embedding = generate_image_embedding(image_base64)\n",
    "    \n",
    "    # Define the image-to-text vector query\n",
    "    image_to_text_query = VectorizedQuery(\n",
    "        vector=image_embedding,\n",
    "        k_nearest_neighbors=1,\n",
    "        fields=\"captionVector\"\n",
    "    )\n",
    "    \n",
    "    # Perform search\n",
    "    results = search_client.search(\n",
    "        search_text=None, vector_queries=[image_to_text_query], top=1\n",
    "    )\n",
    "    display_results(results)\n",
    "\n",
    "# Example usage\n",
    "print(\"Image to Text Vector Search Results:\")\n",
    "image_to_text_search(image_url)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Image to Image Vector Search\n",
    "This section performs a vector search to find similar images based on a given image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image to Image Vector Search Results:\n",
      "Caption: The image captures a solitary surfer standing on a rocky outcrop, gazing out at the ocean waves. Clad in a wetsuit and holding a surfboard, the figure embodies a spirit of adventure and anticipation as they prepare to embrace the surf. With the soft blue sky and gentle waves in the background, the scene conveys a sense of tranquility and the thrill of the ocean.\n",
      "Score: 0.7631374\n",
      "URL: https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample2-72b3c1ca.png\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample2-72b3c1ca.png\" style=\"width:200px;\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# Image to Image Vector Search\n",
    "def image_to_image_search(image_url):\n",
    "    # Generate image embedding\n",
    "    image_base64 = encode_image_to_base64(image_url)\n",
    "    image_embedding = generate_image_embedding(image_base64)\n",
    "    \n",
    "    # Define the image-to-image vector query\n",
    "    image_to_image_query = VectorizedQuery(\n",
    "        vector=image_embedding,\n",
    "        k_nearest_neighbors=1,\n",
    "        fields=\"imageVector\"\n",
    "    )\n",
    "    \n",
    "    # Perform search\n",
    "    results = search_client.search(\n",
    "        search_text=None, vector_queries=[image_to_image_query], top=1\n",
    "    )\n",
    "    display_results(results)\n",
    "\n",
    "# Example usage\n",
    "print(\"Image to Image Vector Search Results:\")\n",
    "image_to_image_search(image_url)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cross-Field Vector Search: Text Embedding Query\n",
    "This section performs a cross-field vector search using a text embedding to query both image and caption fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cross-Field Vector Search Results (Text Embedding Query):\n",
      "Caption: The image captures a solitary surfer standing on a rocky outcrop, gazing out at the ocean waves. Clad in a wetsuit and holding a surfboard, the figure embodies a spirit of adventure and anticipation as they prepare to embrace the surf. With the soft blue sky and gentle waves in the background, the scene conveys a sense of tranquility and the thrill of the ocean.\n",
      "Score: 0.01666666753590107\n",
      "URL: https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample2-72b3c1ca.png\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample2-72b3c1ca.png\" style=\"width:200px;\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------\n",
      "Caption: The image features a focused baseball player poised at the plate, ready to swing his bat. Dressed in a gray uniform with a striking red helmet, he embodies determination and athleticism under the stadium lights. The blurred background of the baseball field highlights the intensity of the moment, capturing the excitement of the game.\n",
      "Score: 0.01666666753590107\n",
      "URL: https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample4-0559774f.png\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample4-0559774f.png\" style=\"width:200px;\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# Cross-Field Vector Search: Text Embedding Query\n",
    "def text_embedding_cross_field_search(query_text):\n",
    "    # Generate text embedding\n",
    "    text_embedding = generate_text_embedding(query_text)\n",
    "    \n",
    "    # Define the vector query for both caption and image fields\n",
    "    cross_field_query = VectorizedQuery(\n",
    "        vector=text_embedding,\n",
    "        k_nearest_neighbors=1,\n",
    "        fields=\"imageVector, captionVector\"\n",
    "    )\n",
    "    \n",
    "    # Perform search\n",
    "    results = search_client.search(\n",
    "        search_text=None, vector_queries=[cross_field_query], top=3\n",
    "    )\n",
    "    display_results(results)\n",
    "\n",
    "# Example usage\n",
    "print(\"Cross-Field Vector Search Results (Text Embedding Query):\")\n",
    "text_embedding_cross_field_search(query_text)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cross-Field Vector Search: Image Embedding Query\n",
    "This section performs a cross-field vector search using an image embedding to query both image and caption fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cross-Field Vector Search Results (Image Embedding Query):\n",
      "Caption: The image captures a solitary surfer standing on a rocky outcrop, gazing out at the ocean waves. Clad in a wetsuit and holding a surfboard, the figure embodies a spirit of adventure and anticipation as they prepare to embrace the surf. With the soft blue sky and gentle waves in the background, the scene conveys a sense of tranquility and the thrill of the ocean.\n",
      "Score: 0.03333333507180214\n",
      "URL: https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample2-72b3c1ca.png\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample2-72b3c1ca.png\" style=\"width:200px;\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# Cross-Field Vector Search: Image Embedding Query\n",
    "def image_embedding_cross_field_search(image_url):\n",
    "    # Generate image embedding\n",
    "    image_base64 = encode_image_to_base64(image_url)\n",
    "    image_embedding = generate_image_embedding(image_base64)\n",
    "    \n",
    "    # Define the vector query for both caption and image fields\n",
    "    cross_field_query = VectorizedQuery(\n",
    "        vector=image_embedding,\n",
    "        k_nearest_neighbors=1,\n",
    "        fields=\"imageVector, captionVector\"\n",
    "    )\n",
    "    \n",
    "    # Perform search\n",
    "    results = search_client.search(\n",
    "        search_text=None, vector_queries=[cross_field_query], top=1\n",
    "    )\n",
    "    display_results(results)\n",
    "\n",
    "# Example usage\n",
    "print(\"Cross-Field Vector Search Results (Image Embedding Query):\")\n",
    "image_embedding_cross_field_search(image_url)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-Vector Search: Text and Image Query\n",
    "This section performs a multi-vector search using both text and image vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Text and Image Multi-Vector Query Results:\n",
      "Caption: The image captures a solitary surfer standing on a rocky outcrop, gazing out at the ocean waves. Clad in a wetsuit and holding a surfboard, the figure embodies a spirit of adventure and anticipation as they prepare to embrace the surf. With the soft blue sky and gentle waves in the background, the scene conveys a sense of tranquility and the thrill of the ocean.\n",
      "Score: 0.01666666753590107\n",
      "URL: https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample2-72b3c1ca.png\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample2-72b3c1ca.png\" style=\"width:200px;\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------\n",
      "Caption: The image features a focused baseball player poised at the plate, ready to swing his bat. Dressed in a gray uniform with a striking red helmet, he embodies determination and athleticism under the stadium lights. The blurred background of the baseball field highlights the intensity of the moment, capturing the excitement of the game.\n",
      "Score: 0.01666666753590107\n",
      "URL: https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample4-0559774f.png\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://portal.vision.cognitive.azure.com/dist/assets/ImageCaptioningSample4-0559774f.png\" style=\"width:200px;\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# Multi-Vector Search: Text and Image Query\n",
    "def text_and_image_query_multi_vector(query_text, image_url):\n",
    "    # Generate the text embedding\n",
    "    text_embedding = generate_text_embedding(query_text)\n",
    "    \n",
    "    # Encode the image and generate the image embedding\n",
    "    image_base64 = encode_image_to_base64(image_url)\n",
    "    image_embedding = generate_image_embedding(image_base64)\n",
    "\n",
    "    # Define the text vector query\n",
    "    text_vector_query = VectorizedQuery(\n",
    "        vector=text_embedding,\n",
    "        k_nearest_neighbors=1,\n",
    "        fields=\"captionVector\"\n",
    "    )\n",
    "    \n",
    "    # Define the image vector query\n",
    "    image_vector_query = VectorizedQuery(\n",
    "        vector=image_embedding,\n",
    "        k_nearest_neighbors=1,\n",
    "        fields=\"imageVector\"\n",
    "    )\n",
    "\n",
    "    # Perform the search with both vector queries\n",
    "    results = search_client.search(\n",
    "        search_text=None,\n",
    "        vector_queries=[text_vector_query, image_vector_query],\n",
    "        top=2\n",
    "    )\n",
    "    \n",
    "    # Display the results\n",
    "    display_results(results)\n",
    "\n",
    "# Example usage\n",
    "print(\"Text and Image Multi-Vector Query Results:\")\n",
    "text_and_image_query_multi_vector(query_text, image_url)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Perform RAG using Command R+"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Grounded Response: Other sports activities include baseball.\n",
      "\n",
      "Formatted Response with Citations:\n",
      "Other sports activities include **baseball.**^(doc_1)\n"
     ]
    }
   ],
   "source": [
    "# Initialize Cohere client for Command R+ chat\n",
    "co_chat = cohere.Client(\n",
    "    base_url=f\"{AZURE_AI_STUDIO_COHERE_COMMAND_ENDPOINT}/v1\", \n",
    "    api_key=AZURE_AI_STUDIO_COHERE_COMMAND_KEY\n",
    ")\n",
    "\n",
    "# Cross-Field Vector Search: Text Embedding Query\n",
    "def text_embedding_cross_field_search(query_text):\n",
    "    \"\"\"\n",
    "    Performs a cross-field vector search on both imageVector and captionVector fields.\n",
    "    \n",
    "    Args:\n",
    "        query_text (str): The text query for search.\n",
    "\n",
    "    Returns:\n",
    "        list: A list of search results.\n",
    "    \"\"\"\n",
    "    # Generate text embedding\n",
    "    text_embedding = generate_text_embedding(query_text)\n",
    "    \n",
    "    # Define the vector query for both caption and image fields\n",
    "    cross_field_query = VectorizedQuery(\n",
    "        vector=text_embedding,\n",
    "        k_nearest_neighbors=1,\n",
    "        fields=\"imageVector, captionVector\"\n",
    "    )\n",
    "    \n",
    "    # Perform search and return results\n",
    "    results = search_client.search(\n",
    "        search_text=None, vector_queries=[cross_field_query], top=3\n",
    "    )\n",
    "    return results\n",
    "\n",
    "# Chat function to ground the response based on search results\n",
    "def ask(query_text):\n",
    "    \"\"\"\n",
    "    Ask a question and return a grounded response from the Cohere chat endpoint.\n",
    "\n",
    "    Args:\n",
    "        query_text (str): The question to ask.\n",
    "\n",
    "    Returns:\n",
    "        str: The grounded response from the chatbot.\n",
    "    \"\"\"\n",
    "    # Retrieve search results using the pre-defined search function\n",
    "    search_results = text_embedding_cross_field_search(query_text)\n",
    "\n",
    "    # Prepare documents for the chat API from search results\n",
    "    documents = [{\"text\": result[\"caption\"]} for result in search_results]\n",
    "\n",
    "    # Get grounded response using Cohere Command R+ chat\n",
    "    chat_response = co_chat.chat(\n",
    "        message=query_text, \n",
    "        documents=documents,\n",
    "        max_tokens=100\n",
    "    )\n",
    "\n",
    "    return chat_response\n",
    "\n",
    "# Function to clean and format the grounded response with citations\n",
    "def pretty_text(text, citations):\n",
    "    \"\"\"\n",
    "    Format the text response with citations in a readable way.\n",
    "\n",
    "    Args:\n",
    "        text (str): The response text from the chat.\n",
    "        citations (list): List of citation objects with start, end, and document_ids.\n",
    "\n",
    "    Returns:\n",
    "        str: The formatted text with citations.\n",
    "    \"\"\"\n",
    "    # Sort citations by start position to avoid issues when altering text indices\n",
    "    sorted_citations = sorted(citations, key=lambda x: x.start, reverse=True)\n",
    "\n",
    "    # Process each citation in reverse order to prevent index shifting\n",
    "    for citation in sorted_citations:\n",
    "        doc_ids_str = \", \".join(citation.document_ids)\n",
    "        citation_text = text[citation.start : citation.end]\n",
    "        # Bold the citation text and add document ids as superscript\n",
    "        new_text = f\"**{citation_text}**^({doc_ids_str})\"\n",
    "        text = text[: citation.start] + new_text + text[citation.end :]\n",
    "\n",
    "    return text\n",
    "\n",
    "# Example usage\n",
    "query_text = \"sports activities like skateboarding\"\n",
    "response = ask(query_text)\n",
    "\n",
    "# Print raw response\n",
    "print(\"Grounded Response:\", response.text)\n",
    "\n",
    "# Clean the response with citations\n",
    "pretty_text_output = pretty_text(response.text, response.citations)\n",
    "print(\"\\nFormatted Response with Citations:\")\n",
    "print(pretty_text_output)\n"
   ]
  }
 ],
 "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}