{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1f115a5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "822970e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "from elasticsearch import Elasticsearch\n",
    "\n",
    "os.environ[\"PROD_ES_URL\"] = \"https://7fbd80c274354d87ba108f274b50b3f8.us-east-2.aws.elastic-cloud.com:443\"\n",
    "os.environ[\"PROD_API_KEY\"] = \"V3hWYmZJOEJtM3UwVVlDWlBCbi06XzRSMklOZGdURkNLY091Xy1faTFWUQ==\"\n",
    "\n",
    "# Create client\n",
    "es = Elasticsearch(\n",
    "    os.environ[\"PROD_ES_URL\"],\n",
    "    api_key=os.environ[\"PROD_API_KEY\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "92ba4e14",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== local_clip ===\n",
      "Error fetching index local_clip: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Content-Type, Accept]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n",
      "\n",
      "=== prod_playlist_v3 ===\n",
      "Error fetching index prod_playlist_v3: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Content-Type, Accept]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n",
      "\n",
      "=== prod_playlist_v2 ===\n",
      "Error fetching index prod_playlist_v2: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Content-Type, Accept]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n",
      "\n",
      "=== local_user ===\n",
      "Error fetching index local_user: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Accept, Content-Type]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n",
      "\n",
      "=== prod_genre ===\n",
      "Error fetching index prod_genre: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Content-Type, Accept]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n",
      "\n",
      "=== metrics-endpoint.metadata_current_default ===\n",
      "Error fetching index metrics-endpoint.metadata_current_default: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Content-Type, Accept]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n",
      "\n",
      "=== prod_hook ===\n",
      "Error fetching index prod_hook: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Content-Type, Accept]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n",
      "\n",
      "=== prod_user_v3 ===\n",
      "Error fetching index prod_user_v3: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Content-Type, Accept]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n",
      "\n",
      "=== prod_clip_v3 ===\n",
      "Error fetching index prod_clip_v3: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Content-Type, Accept]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n",
      "\n",
      "=== prod_clip_v4 ===\n",
      "Error fetching index prod_clip_v4: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Accept, Content-Type]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n",
      "\n",
      "=== local_persona ===\n",
      "Error fetching index local_persona: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Accept, Content-Type]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n",
      "\n",
      "=== prod_clip_copy_warm ===\n",
      "Error fetching index prod_clip_copy_warm: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Accept, Content-Type]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n",
      "\n",
      "=== prod_persona ===\n",
      "Error fetching index prod_persona: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Content-Type, Accept]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n",
      "\n",
      "=== prod_playlist ===\n",
      "Error fetching index prod_playlist: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Content-Type, Accept]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n",
      "\n",
      "=== prod_clip ===\n",
      "Error fetching index prod_clip: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Accept, Content-Type]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)\n"
     ]
    }
   ],
   "source": [
    "# 1) Get all indices\n",
    "all_indices = es.cat.indices(format=\"json\")\n",
    "# Keep only indices that don't start with \".\"\n",
    "indices = [idx[\"index\"] for idx in all_indices if not idx[\"index\"].startswith(\".\")]\n",
    "\n",
    "# Dict to store samples\n",
    "samples = {}\n",
    "\n",
    "# 2) Loop through indices\n",
    "for index in indices:\n",
    "    print(f\"\\n=== {index} ===\")\n",
    "    \n",
    "    # Fetch 5 documents\n",
    "    try:\n",
    "        response = es.search(\n",
    "            index=index,\n",
    "            body={\"query\": {\"match_all\": {}}, \"size\": 5}\n",
    "        )\n",
    "        # Extract hits\n",
    "        hits = response[\"hits\"][\"hits\"]\n",
    "        # Convert to DataFrame\n",
    "        df = pd.json_normalize([hit[\"_source\"] for hit in hits])\n",
    "        \n",
    "        # Print column names\n",
    "        for col in df.columns:\n",
    "            print(f\"- {col}\")\n",
    "        \n",
    "        # Store sample\n",
    "        samples[index] = df\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"Error fetching index {index}: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5e3640d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Example: show sample keys and preview\n",
    "# print(\"\\nAvailable sample indices:\", list(samples.keys()))\n",
    "# for idx, df in samples.items():\n",
    "#     print(f\"\\nSample from {idx}:\")\n",
    "#     print(df.head(2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c0a5c5b9",
   "metadata": {},
   "outputs": [
    {
     "ename": "BadRequestError",
     "evalue": "BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Content-Type, Accept]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mBadRequestError\u001b[0m                           Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 23\u001b[0m\n\u001b[1;32m      2\u001b[0m query \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m      3\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msize\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;241m10\u001b[39m,\n\u001b[1;32m      4\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m_source\u001b[39m\u001b[38;5;124m\"\u001b[39m: [\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mid\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtitle\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mplay_count\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcreated_at\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     19\u001b[0m     ]\n\u001b[1;32m     20\u001b[0m }\n\u001b[1;32m     22\u001b[0m \u001b[38;5;66;03m# Run search\u001b[39;00m\n\u001b[0;32m---> 23\u001b[0m response \u001b[38;5;241m=\u001b[39m \u001b[43mes\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msearch\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     24\u001b[0m \u001b[43m    \u001b[49m\u001b[43mindex\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mprod_clip_v3\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     25\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbody\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mquery\u001b[49m\n\u001b[1;32m     26\u001b[0m \u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env_dev/lib/python3.10/site-packages/elasticsearch/_sync/client/utils.py:415\u001b[0m, in \u001b[0;36m_rewrite_parameters.<locals>.wrapper.<locals>.wrapped\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    412\u001b[0m         \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m:\n\u001b[1;32m    413\u001b[0m             \u001b[38;5;28;01mpass\u001b[39;00m\n\u001b[0;32m--> 415\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mapi\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env_dev/lib/python3.10/site-packages/elasticsearch/_sync/client/__init__.py:4916\u001b[0m, in \u001b[0;36mElasticsearch.search\u001b[0;34m(self, index, aggregations, aggs, allow_no_indices, allow_partial_search_results, analyze_wildcard, analyzer, batched_reduce_size, ccs_minimize_roundtrips, collapse, default_operator, df, docvalue_fields, error_trace, expand_wildcards, explain, ext, fields, filter_path, force_synthetic_source, from_, highlight, human, ignore_throttled, ignore_unavailable, include_named_queries_score, indices_boost, knn, lenient, max_concurrent_shard_requests, min_score, pit, post_filter, pre_filter_shard_size, preference, pretty, profile, q, query, rank, request_cache, rescore, rest_total_hits_as_int, retriever, routing, runtime_mappings, script_fields, scroll, search_after, search_type, seq_no_primary_term, size, slice, sort, source, source_excludes, source_includes, stats, stored_fields, suggest, suggest_field, suggest_mode, suggest_size, suggest_text, terminate_after, timeout, track_scores, track_total_hits, typed_keys, version, body)\u001b[0m\n\u001b[1;32m   4914\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m __body \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m   4915\u001b[0m     __headers[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcontent-type\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mapplication/json\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m-> 4916\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mperform_request\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# type: ignore[return-value]\u001b[39;49;00m\n\u001b[1;32m   4917\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mPOST\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4918\u001b[0m \u001b[43m    \u001b[49m\u001b[43m__path\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4919\u001b[0m \u001b[43m    \u001b[49m\u001b[43mparams\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m__query\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4920\u001b[0m \u001b[43m    \u001b[49m\u001b[43mheaders\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m__headers\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4921\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbody\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m__body\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4922\u001b[0m \u001b[43m    \u001b[49m\u001b[43mendpoint_id\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43msearch\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4923\u001b[0m \u001b[43m    \u001b[49m\u001b[43mpath_parts\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m__path_parts\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4924\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env_dev/lib/python3.10/site-packages/elasticsearch/_sync/client/_base.py:271\u001b[0m, in \u001b[0;36mBaseClient.perform_request\u001b[0;34m(self, method, path, params, headers, body, endpoint_id, path_parts)\u001b[0m\n\u001b[1;32m    255\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mperform_request\u001b[39m(\n\u001b[1;32m    256\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m    257\u001b[0m     method: \u001b[38;5;28mstr\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    264\u001b[0m     path_parts: Optional[Mapping[\u001b[38;5;28mstr\u001b[39m, Any]] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    265\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m ApiResponse[Any]:\n\u001b[1;32m    266\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_otel\u001b[38;5;241m.\u001b[39mspan(\n\u001b[1;32m    267\u001b[0m         method,\n\u001b[1;32m    268\u001b[0m         endpoint_id\u001b[38;5;241m=\u001b[39mendpoint_id,\n\u001b[1;32m    269\u001b[0m         path_parts\u001b[38;5;241m=\u001b[39mpath_parts \u001b[38;5;129;01mor\u001b[39;00m {},\n\u001b[1;32m    270\u001b[0m     ) \u001b[38;5;28;01mas\u001b[39;00m otel_span:\n\u001b[0;32m--> 271\u001b[0m         response \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_perform_request\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    272\u001b[0m \u001b[43m            \u001b[49m\u001b[43mmethod\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    273\u001b[0m \u001b[43m            \u001b[49m\u001b[43mpath\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    274\u001b[0m \u001b[43m            \u001b[49m\u001b[43mparams\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    275\u001b[0m \u001b[43m            \u001b[49m\u001b[43mheaders\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mheaders\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    276\u001b[0m \u001b[43m            \u001b[49m\u001b[43mbody\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbody\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    277\u001b[0m \u001b[43m            \u001b[49m\u001b[43motel_span\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43motel_span\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    278\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    279\u001b[0m         otel_span\u001b[38;5;241m.\u001b[39mset_elastic_cloud_metadata(response\u001b[38;5;241m.\u001b[39mmeta\u001b[38;5;241m.\u001b[39mheaders)\n\u001b[1;32m    280\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m response\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env_dev/lib/python3.10/site-packages/elasticsearch/_sync/client/_base.py:351\u001b[0m, in \u001b[0;36mBaseClient._perform_request\u001b[0;34m(self, method, path, params, headers, body, otel_span)\u001b[0m\n\u001b[1;32m    348\u001b[0m         \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mValueError\u001b[39;00m, \u001b[38;5;167;01mKeyError\u001b[39;00m, \u001b[38;5;167;01mTypeError\u001b[39;00m):\n\u001b[1;32m    349\u001b[0m             \u001b[38;5;28;01mpass\u001b[39;00m\n\u001b[0;32m--> 351\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m HTTP_EXCEPTIONS\u001b[38;5;241m.\u001b[39mget(meta\u001b[38;5;241m.\u001b[39mstatus, ApiError)(\n\u001b[1;32m    352\u001b[0m         message\u001b[38;5;241m=\u001b[39mmessage, meta\u001b[38;5;241m=\u001b[39mmeta, body\u001b[38;5;241m=\u001b[39mresp_body\n\u001b[1;32m    353\u001b[0m     )\n\u001b[1;32m    355\u001b[0m \u001b[38;5;66;03m# 'X-Elastic-Product: Elasticsearch' should be on every 2XX response.\u001b[39;00m\n\u001b[1;32m    356\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_verified_elasticsearch:\n\u001b[1;32m    357\u001b[0m     \u001b[38;5;66;03m# If the header is set we mark the server as verified.\u001b[39;00m\n",
      "\u001b[0;31mBadRequestError\u001b[0m: BadRequestError(400, 'media_type_header_exception', 'Invalid media-type value on headers [Content-Type, Accept]', Accept version must be either version 8 or 7, but found 9. Accept=application/vnd.elasticsearch+json; compatible-with=9)"
     ]
    }
   ],
   "source": [
    "# Search for title \"can it be real\", sorted by play_count desc\n",
    "query = {\n",
    "    \"size\": 10,\n",
    "    \"_source\": [\"id\", \"title\", \"play_count\", \"created_at\"],\n",
    "    \"query\": {\n",
    "        \"bool\": {\n",
    "            \"must\": {\n",
    "                \"match_phrase\": {\n",
    "                    \"title\": \"can it be real\"\n",
    "                }\n",
    "            },\n",
    "            \"filter\": [\n",
    "                {\"term\": {\"is_public\": True}}\n",
    "            ]\n",
    "        }\n",
    "    },\n",
    "    \"sort\": [\n",
    "        {\"play_count\": {\"order\": \"desc\"}}\n",
    "    ]\n",
    "}\n",
    "\n",
    "# Run search\n",
    "response = es.search(\n",
    "    index=\"prod_clip_v3\",\n",
    "    body=query\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d180f3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display results\n",
    "for hit in response[\"hits\"][\"hits\"]:\n",
    "    print(f\"ID: {hit['_source'].get('id')}\")\n",
    "    print(f\"Title: {hit['_source'].get('title')}\")\n",
    "    print(f\"Plays: {hit['_source'].get('play_count')}\")\n",
    "    print(f\"Created At: {hit['_source'].get('created_at')}\")\n",
    "    print(\"---\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41271cf7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "535747f7",
   "metadata": {},
   "source": [
    "### 'Good to me' similarity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c60afd31",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fetch the doc by ID\n",
    "doc = es.get(\n",
    "    index=\"prod_clip_v3\",\n",
    "    id=\"3e96d1ee-d6ed-41c7-bc07-89ca7d65e0e7\",\n",
    "    _source_includes=[\"song_vector_v2\"]\n",
    ")\n",
    "\n",
    "# Extract the vector\n",
    "query_vector = doc[\"_source\"][\"song_vector_v2\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5166574a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve top 1000 candidates\n",
    "knn_response = es.knn_search(\n",
    "    index=\"prod_clip_v3\",\n",
    "    knn={\n",
    "        \"field\": \"song_vector_v2\",\n",
    "        \"query_vector\": query_vector,\n",
    "        \"k\": 1000,\n",
    "        \"num_candidates\": 2000\n",
    "    },\n",
    "    source=[\"id\", \"title\", \"play_count\"]\n",
    ")\n",
    "\n",
    "# Extract candidate IDs\n",
    "candidate_ids = [hit[\"_id\"] for hit in knn_response[\"hits\"][\"hits\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a527f450",
   "metadata": {},
   "outputs": [],
   "source": [
    "rerank_query = {\n",
    "    \"size\": 10,\n",
    "    \"_source\": [\"id\", \"title\", \"play_count\"],\n",
    "    \"query\": {\n",
    "        \"script_score\": {\n",
    "            \"query\": {\n",
    "                \"ids\": {\"values\": candidate_ids}\n",
    "            },\n",
    "            \"script\": {\n",
    "                \"source\": \"\"\"\n",
    "                    double popularity = 0;\n",
    "                    if (doc.containsKey('play_count') && !doc['play_count'].empty) {\n",
    "                        popularity = Math.log(1 + doc['play_count'].value);\n",
    "                    }\n",
    "                    return popularity;\n",
    "                \"\"\"\n",
    "            }\n",
    "        }\n",
    "    },\n",
    "    \"sort\": [\n",
    "        {\"_score\": {\"order\": \"desc\"}}\n",
    "    ]\n",
    "}\n",
    "\n",
    "response = es.search(\n",
    "    index=\"prod_clip_v3\",\n",
    "    body=rerank_query\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d88b64ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "for hit in response[\"hits\"][\"hits\"]:\n",
    "    print(f\"ID: {hit['_source'].get('id')}\")\n",
    "    print(f\"Title: {hit['_source'].get('title')}\")\n",
    "    print(f\"Play count: {hit['_source'].get('play_count')}\")\n",
    "    print(f\"Rerank score: {hit['_score']}\")\n",
    "    print(\"---\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbd3bf7c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b524bc9f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "588f2539",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env_dev",
   "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.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}