codegen/text2sql/text2sql_gen_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fbf8056a-ab21-4f76-adce-af62ea2d2372",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 依赖安装"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e84906db-5ec4-4270-8c7a-e32f9f8c5460",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "###如果是国内，设置清华pip repo源\n",
    "!pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cf80354-7edd-4403-9c8b-97a4712f4dcc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install langchain[all]\n",
    "!pip install requests_aws4auth\n",
    "!pip install opensearch-py\n",
    "!pip install pydantic==1.10.0\n",
    "!pip install PyAthena[SQLAlchemy]==1.0.0\n",
    "!pip install PyAthena[JDBC]==1.0.0\n",
    "!pip install sqlalchemy-redshift\n",
    "!pip install redshift_connector\n",
    "!pip install pymysql\n",
    "!pip install langchain_experimental"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80e3fea8-7c8d-4118-9578-ce911ea0ca8b",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 初始化rds元数据，aos index\n",
    " * 如果已经有RDS实例，修改setup.sh，skip create db instance 步骤\n",
    " * 确保你的aws configure正确设置aksk及region\n",
    " * 确保网络在同一VPC且入站规则互联互通"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c271d3f0-0039-41f0-a140-fc2121737301",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!chmod 777 ./setup.sh\n",
    "!./setup.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da79c07e-365a-4f56-b83d-db53cf86f8aa",
   "metadata": {
    "tags": []
   },
   "source": [
    "* AOS domain creation "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "8b419706-b9b8-4da9-a828-dcfbacd411b0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vpc-my-domain-g63shn6r3volwzhs2gt7rzy7bq.us-west-2.es.amazonaws.com\n"
     ]
    }
   ],
   "source": [
    "#public公网aos创建\n",
    "#!aws opensearch create-domain --domain-name llm-rag-aos --engine-version OpenSearch_2.3  --ebs-options EBSEnabled=true,VolumeType=gp2,VolumeSize=10\n",
    "\n",
    "#vpc内aos创建\n",
    "!aws opensearch create-domain   --domain-name my-domain --engine-version OpenSearch_2.3 --cluster-config InstanceType=r6g.xlarge.search,InstanceCount=2 --ebs-options EBSEnabled=true,VolumeType=gp2,VolumeSize=10 --access-policies '{\"Version\":\"2012-10-17\",\"Statement\":[{\"Effect\":\"Allow\",\"Principal\":{\"AWS\":[\"*\"]},\"Action\":[\"es:*\"],\"Resource\":\"*\"}]}' --vpc-options SubnetIds=subnet-0135e88c8e8da7369,SecurityGroupIds=sg-0aa3d61256d687a0c\n",
    "!aws opensearch describe-domain --domain-name my-domain| jq -r '.DomainStatus.Endpoints.vpc'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6e5828a-55b8-4be8-85eb-b56f51c85db8",
   "metadata": {},
   "source": [
    "* metadata ingestion\n",
    "* 如用admin账户aos domain，修改上述脚本创建后的aos domain管理员账户为下述user password，并采用password auth认证方式创建index\n",
    "* 如果是没有开admin账户的aos domain，用下面的AWSV4SignerAuth认证方式创建index\n",
    "* 注意：如果开了aos的精细权限控制，则需要用账户密码或者AWSV4SignerAuth签名认证\n",
    "* 确保你的vpc网络和安全组与notebook instance互联互通\n",
    "* 修改aos_endpoint为上述创建的domain endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b216f11-a2b5-4f0b-b4de-060a1a64076a",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## initial sagemaker env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "19f05414-d44b-464e-a0c7-46c317378ebd",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sagemaker.config INFO - Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml\n",
      "sagemaker.config INFO - Not applying SDK defaults from location: /home/ec2-user/.config/sagemaker/config.yaml\n",
      "sagemaker.config INFO - Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml\n",
      "sagemaker.config INFO - Not applying SDK defaults from location: /home/ec2-user/.config/sagemaker/config.yaml\n",
      "sagemaker.config INFO - Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml\n",
      "sagemaker.config INFO - Not applying SDK defaults from location: /home/ec2-user/.config/sagemaker/config.yaml\n",
      "sagemaker role arn: arn:aws:iam::687912291502:role/service-role/AmazonSageMaker-ExecutionRole-20211013T113123\n",
      "sagemaker bucket: sagemaker-us-west-2-687912291502\n",
      "sagemaker session region: us-west-2\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sagemaker\n",
    "import boto3\n",
    "import json\n",
    "from typing import Dict\n",
    "from typing import Any, Dict, List, Optional\n",
    "\n",
    "sess = sagemaker.Session()\n",
    "# sagemaker session bucket -> used for uploading data, models and logs\n",
    "# sagemaker will automatically create this bucket if it not exists\n",
    "sagemaker_session_bucket=None\n",
    "if sagemaker_session_bucket is None and sess is not None:\n",
    "    # set to default bucket if a bucket name is not given\n",
    "    sagemaker_session_bucket = sess.default_bucket()\n",
    "\n",
    "try:\n",
    "    role = sagemaker.get_execution_role()\n",
    "except ValueError:\n",
    "    iam = boto3.client('iam')\n",
    "    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']\n",
    "\n",
    "sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)\n",
    "sm_client = boto3.client(\"sagemaker-runtime\")\n",
    "\n",
    "print(f\"sagemaker role arn: {role}\")\n",
    "print(f\"sagemaker bucket: {sess.default_bucket()}\")\n",
    "print(f\"sagemaker session region: {sess.boto_region_name}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0df98a7-0545-43b4-94fd-7b001532bd43",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth, helpers\n",
    "region='us-west-2'\n",
    "username=\"admin\"\n",
    "passwd=\"(OL>0p;/\"\n",
    "size=10\n",
    "\n",
    "auth = (username, passwd)\n",
    "\n",
    "credentials = boto3.Session().get_credentials()\n",
    "auth = AWSV4SignerAuth(credentials, region)\n",
    "index_name=\"prompt-optimal-index\"\n",
    "aos_endpoint=\"vpc-my-domain-g63shn6r3volwzhs2gt7rzy7bq.us-west-2.es.amazonaws.com\"\n",
    "\n",
    "schema={\n",
    "    \"settings\" : {\n",
    "        \"index\":{\n",
    "            \"number_of_shards\" : 5,\n",
    "            \"number_of_replicas\" : 0,\n",
    "            \"knn\": \"true\",\n",
    "            \"knn.algo_param.ef_search\": 32\n",
    "        }\n",
    "    },\n",
    "    \"mappings\": {\n",
    "        \"properties\": {\n",
    "            \"metadata_type\" : {\n",
    "                \"type\" : \"keyword\"\n",
    "            },\n",
    "            \"database_name\": {\n",
    "                \"type\" : \"keyword\"\n",
    "            },\n",
    "            \"table_name\": {\n",
    "               \"type\" : \"keyword\"\n",
    "            },\n",
    "            \"exactly_query_text\": {\n",
    "                \"type\": \"text\",\n",
    "                \"analyzer\": \"ik_max_word\",\n",
    "                \"search_analyzer\": \"ik_smart\"\n",
    "            },\n",
    "            \"exactly_query_embedding\": {\n",
    "                \"type\": \"knn_vector\",\n",
    "                \"dimension\": 1024,\n",
    "                \"method\": {\n",
    "                    \"name\": \"hnsw\",\n",
    "                    \"space_type\": \"l2\",\n",
    "                    \"engine\": \"faiss\",\n",
    "                    \"parameters\": {\n",
    "                        \"ef_construction\": 512,\n",
    "                        \"m\": 32\n",
    "                    }\n",
    "                }            \n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n",
    "search = OpenSearch(\n",
    "    hosts = [{'host': aos_endpoint, 'port': 443}],\n",
    "    ##http_auth = awsauth ,\n",
    "    http_auth = auth ,\n",
    "    use_ssl = True,\n",
    "    verify_certs = True,\n",
    "    connection_class = RequestsHttpConnection\n",
    ")\n",
    "search.indices.create(index=index_name, body=schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d3aed78-e14b-4d30-93d3-2871d6603bd8",
   "metadata": {
    "tags": []
   },
   "source": [
    "## intial lanchain lib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4aae9f94-a8df-49b5-b22b-f89e8d756b4d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.vectorstores import OpenSearchVectorSearch\n",
    "from langchain import PromptTemplate, SagemakerEndpoint\n",
    "from langchain.chains.question_answering import load_qa_chain\n",
    "from langchain.embeddings import SagemakerEndpointEmbeddings\n",
    "from langchain.llms.sagemaker_endpoint import ContentHandlerBase\n",
    "from langchain.llms.sagemaker_endpoint import LLMContentHandler\n",
    "from langchain.docstore.document import Document\n",
    "from langchain.memory import ConversationBufferWindowMemory,ConversationBufferMemory\n",
    "from langchain import LLMChain\n",
    "from typing import Dict\n",
    "\n",
    "\n",
    "## for chatglm\n",
    "class TextGenContentHandler(LLMContentHandler):\n",
    "    content_type = \"application/json\"\n",
    "    accepts = \"application/json\"\n",
    "\n",
    "    def transform_input(self, prompt: str, model_kwargs: Dict) -> bytes:\n",
    "        #input_str = json.dumps({prompt: prompt, **model_kwargs})\n",
    "        input_str = json.dumps({\n",
    "                \"ask\": prompt,\n",
    "                \"parameters\": model_kwargs\n",
    "            })\n",
    "        return input_str.encode('utf-8')\n",
    "    \n",
    "    def transform_output(self, output: bytes) -> str:\n",
    "        response_json = json.loads(output.read().decode(\"utf-8\"))\n",
    "        return response_json[\"answer\"]\n",
    "\n",
    "### for vicuna/llama\n",
    "class TextGenContentHandler2(LLMContentHandler):\n",
    "    content_type = \"application/json\"\n",
    "    accepts = \"application/json\"\n",
    "\n",
    "    def transform_input(self, prompt: str, model_kwargs: Dict) -> bytes:\n",
    "        input_str = json.dumps({\n",
    "                \"inputs\": prompt,\n",
    "                \"parameters\": model_kwargs\n",
    "            })\n",
    "        return input_str.encode('utf-8')\n",
    "    \n",
    "    def transform_output(self, output: bytes) -> str:\n",
    "        response_json = json.loads(output.read().decode(\"utf-8\"))\n",
    "        #print(response_json)\n",
    "        #sql_result=response_json[\"outputs\"].split(\"```sql\")[-1].split(\"```\")[0].split(\";\")[0].strip().replace(\"\\\\n\",\" \") + \";\"\n",
    "        sql_result=response_json[\"outputs\"]\n",
    "        return sql_result\n",
    "\n",
    "\n",
    "content_hander2=TextGenContentHandler2()\n",
    "\n",
    "parameters = {\n",
    "  #\"early_stopping\": True,\n",
    "  #\"length_penalty\": 2.0,\n",
    "  \"max_new_tokens\": 300,\n",
    "  #\"temperature\": 0.6,\n",
    "  #\"max_tokens\": 300,\n",
    "  #\"no_repeat_ngram_size\": 2,\n",
    "}\n",
    "sm_llm=SagemakerEndpoint(\n",
    "        #endpoint_name=\"chatglm-inference-0524-2023-06-01-07-11-27-379\",\n",
    "        endpoint_name=\"sqlcoder-2023-09-23-00-54-24-198-endpoint\",\n",
    "        #endpoint_name=\"codellame-2023-09-22-09-25-02-063-endpoint\",\n",
    "        region_name=\"us-west-2\", \n",
    "        model_kwargs=parameters,\n",
    "        content_handler=content_hander2,\n",
    "        #endpoint_kwargs={'CustomAttributes':'accept_eula=true'}\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa302e0a-b8ed-4fa0-934e-ab4ff0740316",
   "metadata": {
    "tags": []
   },
   "source": [
    "## major chain pipeline ################"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfa19652-7785-42f6-a74f-bc31ff4ca28a",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Opt1:直接使用 langchain SQLDatabaseChain \n",
    "* text2sql prompt见如下langchain sqldatabasaseChain 所示  \n",
    "* SqlDatabaseChain可以使用sagemaker endpoint llm，也可以使用BedRock\n",
    "* 此处使用langchain BedRock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "634c6070-af8c-40b7-b655-e6eb37eff456",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import re\n",
    "text2sq_prompt_template=\"\"\"\n",
    "You are a MySQL expert. Given an input question, first create a syntactically correct MySQL query to run, then look at the results of the query and return the answer to the input question.\n",
    "Unless the user specifies in the question a specific number of examples to obtain, query for at most 5 results using the LIMIT clause as per MySQL. You can order the results to return the most informative data in the database.\n",
    "Never query for all columns from a table. You must query only the columns that are needed to answer the question. Wrap each column name in backticks (`) to denote them as delimited identifiers.\n",
    "Pay attention to use only the column names you can see in the tables below. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.\n",
    "Pay attention to use CURDATE() function to get the current date, if the question involves \"today\".\n",
    "\n",
    "Use the following format:\n",
    "\n",
    "Question: Question here\n",
    "SQLQuery: SQL Query to run\n",
    "SQLResult: Result of the SQLQuery\n",
    "Answer: Final answer here\n",
    "\n",
    "Only use the following tables:\n",
    "\n",
    "CREATE TABLE ads_bi_quality_monitor_shipping_detail (\n",
    "\tshipping_order_code VARCHAR(100) COMMENT '派车单编码', \n",
    "\tlicense_plate VARCHAR(100) COMMENT '车牌号', \n",
    "\ttruck_type VARCHAR(50) COMMENT '车辆类型', \n",
    "\ttenant_id VARCHAR(100) COMMENT '租户编码', \n",
    "\ttenant_name VARCHAR(200) COMMENT '租户名称', \n",
    "\tfather_company_code VARCHAR(100) COMMENT '分子公司编码', \n",
    "\tfather_company_name VARCHAR(200) COMMENT '分子公司名称', \n",
    "\tfather_company_short_name VARCHAR(200) COMMENT '分子公司简称', \n",
    "\tstart_transport_time VARCHAR(50) COMMENT '运输出发时间', \n",
    "\tsigning_time VARCHAR(50) COMMENT '派车单签收时间', \n",
    "\tplan_start_time VARCHAR(50) COMMENT '派车单计划取货时间', \n",
    "\tplan_end_time VARCHAR(50) COMMENT '派车单计划送达时间', \n",
    "\tfrist_fence_time VARCHAR(50) COMMENT '第一次碰撞装货地电子围栏时间', \n",
    "\tleave_load_station_2km_time VARCHAR(50) COMMENT '离开最后一个装货地电子围栏2km的时间', \n",
    "\tfrist_arrive_unload_time VARCHAR(50) COMMENT '首次到达卸货点电子围栏时间', \n",
    "\ttransport_type VARCHAR(10) COMMENT '运输类型(干线/城配)', \n",
    "\tplan_mileage VARCHAR(50) COMMENT '计划里程', \n",
    "\tdriver_accept_time VARCHAR(50) COMMENT '司机接单时间', \n",
    "\tdriver_name VARCHAR(50) COMMENT '司机姓名', \n",
    "\tdriver_phone VARCHAR(50) COMMENT '司机电话', \n",
    "\tdriver_type VARCHAR(10) COMMENT '司机类型', \n",
    "\ttransport_tenant_id VARCHAR(100) COMMENT '运力承运商id', \n",
    "\ttransport_tenant_name VARCHAR(200) COMMENT '运力承运商名称', \n",
    "\twarm_area VARCHAR(50) COMMENT '温区信息', \n",
    "\twaybill_count INTEGER COMMENT '运单数量', \n",
    "\tgps_device_list VARCHAR(100) COMMENT 'gps设备', \n",
    "\tgps_report_dot_num INTEGER COMMENT 'gps上报点数', \n",
    "\ttemp_substandard_min DECIMAL(20, 8) COMMENT '温度不达标时长(分钟)', \n",
    "\ttemp_substandard_min_n12 DECIMAL(20, 8) COMMENT '温度不达标时长(分钟)_-12', \n",
    "\ttemp_substandard_min_n16 DECIMAL(20, 8) COMMENT '温度不达标时长(分钟)_-16', \n",
    "\tshipping_order_transport_min DECIMAL(20, 8) COMMENT '派车单运输总时长(分钟)', \n",
    "\tshipping_order_temp_standard_rate DECIMAL(20, 8) COMMENT '派车单温度达标率', \n",
    "\tshipping_order_temp_standard_rate_n12 DECIMAL(20, 8) COMMENT '派车单温度达标率_-12', \n",
    "\tshipping_order_temp_standard_rate_n16 DECIMAL(20, 8) COMMENT '派车单温度达标率_-16', \n",
    "\tdep_fence_match_num INTEGER COMMENT '出发地电子围栏匹配数量', \n",
    "\tdep_total_num INTEGER COMMENT '出发地总数量', \n",
    "\tdep_fence_match_rate DECIMAL(20, 8) COMMENT '出发地电子围栏匹配率', \n",
    "\tdes_fence_match_num INTEGER COMMENT '目的地电子围栏匹配数量', \n",
    "\tdes_total_num INTEGER COMMENT '目的地总数量', \n",
    "\tdes_fence_match_rate DECIMAL(20, 8) COMMENT '目的地电子围栏匹配率', \n",
    "\texception_shipping_order_type VARCHAR(10) COMMENT '异常派车单情况,0-非异常', \n",
    "\tsettlement_tenant_id VARCHAR(100) COMMENT '结算主体租户id', \n",
    "\tsettlement_tenant_name VARCHAR(100) COMMENT '结算主体租户名称', \n",
    "\ttruck_ownership VARCHAR(100) COMMENT '车辆所有权(自有/外请/临时)', \n",
    "\tload_finish_time VARCHAR(50) COMMENT '派车单点击装货完成时间', \n",
    "\tlast_leave_unload_time VARCHAR(50) COMMENT '最后离开卸货点时间', \n",
    "\ttemp_right_tag VARCHAR(20) COMMENT '温度是否合格(合格/不合格/不参与评估)', \n",
    "\twarm_area_type VARCHAR(20) COMMENT '温区类型。常温/冷链/自定义区间', \n",
    "\tsource_dt VARCHAR(20) COMMENT '派车单来源的增量表dt', \n",
    "\tcur_shipping_odr_cust TEXT CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci COMMENT '当前派车单对应的客户', \n",
    "\tis_app_operation VARCHAR(10) COMMENT '是否APP操作卡控', \n",
    "\tleave_load_station_time_app VARCHAR(50) COMMENT '离开第一个装货地时间_app', \n",
    "\tfrist_arrive_unload_time_app VARCHAR(50) COMMENT '到达第一个卸货点时间_app', \n",
    "\tlast_arrive_unload_time_app VARCHAR(50) COMMENT '到达最后一个卸货点时间_app', \n",
    "\ttemp_calc_start_time VARCHAR(50) COMMENT '温度计算开始时间', \n",
    "\ttemp_calc_end_time VARCHAR(50) COMMENT '温度计算结束时间', \n",
    "\texception_shipping_order_type_desc VARCHAR(100) COMMENT '异常派车单情况描述', \n",
    "\troot_shipping_odr_cust TEXT CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci COMMENT '根派车单对应的客户', \n",
    "\tis_use_custom_temp_range VARCHAR(10) COMMENT '是否使用自定义温度范围', \n",
    "\ttemp_eval_lowest INTEGER COMMENT '温度考核最低值', \n",
    "\ttemp_eval_highest INTEGER COMMENT '温度考核最高值', \n",
    "\tis_gps_cover VARCHAR(10) COMMENT 'GPS是否覆盖', \n",
    "\tshipping_create_time VARCHAR(50) COMMENT '派车单创建时间', \n",
    "\tgps_device_list_plan VARCHAR(100) COMMENT 'GPS设备_预估', \n",
    "\tis_gps_cover_plan VARCHAR(10) COMMENT 'GPS是否覆盖_预估', \n",
    "\tfirst_station_is_on_time VARCHAR(10) COMMENT '首店是否准时', \n",
    "\tis_many_warm VARCHAR(10) COMMENT '是否多温区', \n",
    "\tmany_temp_standard_rate DECIMAL(20, 8) COMMENT '多温区温度达标率', \n",
    "\tlc_standard_ratio DECIMAL(20, 8) COMMENT '冷藏温度达标率', \n",
    "\tld_standard_ratio DECIMAL(20, 8) COMMENT '冷冻温度达标率', \n",
    "\tprobe_warm_list TEXT CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci COMMENT '探头温区'\n",
    ")ENGINE=InnoDB DEFAULT CHARSET=utf8 COMMENT='BI看板_品质监控_派车单明细'\n",
    "\n",
    "Question: 最近一个月温度合格的派车单数量\n",
    "SQLQuery: \n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d5ef40b-89af-49fb-b98e-77ef9cd87796",
   "metadata": {},
   "source": [
    "* 定制langchain SqlDataBase\n",
    "* 可以实现对sql生成更个性化控制（如，去掉冗余信息，sql改写/优化等）\n",
    "* 此处示例对输出去掉SQLQuery的前缀，直接取sql语句"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "781f2be7-0a3e-4eac-a55b-00b182762326",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\"\"\"SQLAlchemy wrapper around a database.\"\"\"\n",
    "from __future__ import annotations\n",
    "\n",
    "import warnings\n",
    "from typing import Any, Iterable, List, Optional\n",
    "import sqlalchemy\n",
    "import re\n",
    "from langchain.utilities import SQLDatabase\n",
    "from langchain_experimental.sql import SQLDatabaseChain\n",
    "from langchain.chains import create_sql_query_chain\n",
    "from langchain.base_language import BaseLanguageModel\n",
    "from langchain.callbacks.manager import CallbackManagerForChainRun\n",
    "from sqlalchemy import MetaData, Table, create_engine, inspect, select, text\n",
    "from sqlalchemy.engine import Engine\n",
    "from sqlalchemy.exc import ProgrammingError, SQLAlchemyError\n",
    "from sqlalchemy.schema import CreateTable\n",
    "\n",
    "\n",
    "\n",
    "class CustomerizedSQLDatabaseChain(SQLDatabaseChain):\n",
    "    \n",
    "    def _call(\n",
    "        self,\n",
    "        inputs: Dict[str, Any],\n",
    "        run_manager: Optional[CallbackManagerForChainRun] = None,\n",
    "    ) -> Dict[str, Any]:\n",
    "        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()\n",
    "        input_text = f\"{inputs[self.input_key]}\\nSQLQuery:\"\n",
    "        _run_manager.on_text(input_text, verbose=self.verbose)\n",
    "        # If not present, then defaults to None which is all tables.\n",
    "        table_names_to_use = inputs.get(\"table_names_to_use\")\n",
    "        table_info = self.database.get_table_info(table_names=table_names_to_use)\n",
    "        llm_inputs = {\n",
    "            \"input\": input_text,\n",
    "            \"top_k\": str(self.top_k),\n",
    "            \"dialect\": self.database.dialect,\n",
    "            \"table_info\": table_info,\n",
    "            #\"stop\": [\"\\nSQLResult:\"],\n",
    "        }\n",
    "        intermediate_steps: List = []\n",
    "        try:\n",
    "            intermediate_steps.append(llm_inputs)  # input: sql generation\n",
    "            sql_cmd = self.llm_chain.predict(\n",
    "                callbacks=_run_manager.get_child(),\n",
    "                **llm_inputs,\n",
    "            ).strip()\n",
    "            #print(\"orginal sql_cmd==\"+sql_cmd)\n",
    "            if self.return_sql:\n",
    "                return {self.output_key: sql_cmd}\n",
    "            if not self.use_query_checker:\n",
    "                _run_manager.on_text(sql_cmd, color=\"green\", verbose=self.verbose)\n",
    "                intermediate_steps.append(\n",
    "                    sql_cmd\n",
    "                )  # output: sql generation (no checker)\n",
    "                #########定制bedRock模型输出##############\n",
    "                pattern = r\"SQL执行结果: (.*?)\\n\"\n",
    "                matches = re.findall(pattern, sql_cmd)\n",
    "                match = matches[1]\n",
    "                sql_cmd = match\n",
    "                #print(\"query sql==\"+sql_cmd) \n",
    "                \n",
    "                intermediate_steps.append({\"sql_cmd\": sql_cmd})  # input: sql exec\n",
    "                result = self.database.run(sql_cmd)\n",
    "                intermediate_steps.append(str(result))  # output: sql exec\n",
    "            else:\n",
    "                query_checker_prompt = self.query_checker_prompt or PromptTemplate(\n",
    "                    template=QUERY_CHECKER, input_variables=[\"query\", \"dialect\"]\n",
    "                )\n",
    "                query_checker_chain = LLMChain(\n",
    "                    llm=self.llm_chain.llm, prompt=query_checker_prompt\n",
    "                )\n",
    "                query_checker_inputs = {\n",
    "                    \"query\": sql_cmd,\n",
    "                    \"dialect\": self.database.dialect,\n",
    "                }\n",
    "                checked_sql_command: str = query_checker_chain.predict(\n",
    "                    callbacks=_run_manager.get_child(), **query_checker_inputs\n",
    "                ).strip()\n",
    "                intermediate_steps.append(\n",
    "                    checked_sql_command\n",
    "                )  # output: sql generation (checker)\n",
    "                _run_manager.on_text(\n",
    "                    checked_sql_command, color=\"green\", verbose=self.verbose\n",
    "                )\n",
    "                intermediate_steps.append(\n",
    "                    {\"sql_cmd\": checked_sql_command}\n",
    "                )  # input: sql exec\n",
    "                result = self.database.run(checked_sql_command)\n",
    "                intermediate_steps.append(str(result))  # output: sql exec\n",
    "                sql_cmd = checked_sql_command\n",
    "\n",
    "            _run_manager.on_text(\"\\nSQLResult: \", verbose=self.verbose)\n",
    "            _run_manager.on_text(result, color=\"yellow\", verbose=self.verbose)\n",
    "            # If return direct, we just set the final result equal to\n",
    "            # the result of the sql query result, otherwise try to get a human readable\n",
    "            # final answer\n",
    "            if self.return_direct:\n",
    "                final_result = result\n",
    "            else:\n",
    "                _run_manager.on_text(\"\\nAnswer:\", verbose=self.verbose)\n",
    "                input_text += f\"{sql_cmd}\\nSQLResult: {result}\\nAnswer:\"\n",
    "                llm_inputs[\"input\"] = input_text\n",
    "                intermediate_steps.append(llm_inputs)  # input: final answer\n",
    "                final_result = self.llm_chain.predict(\n",
    "                    callbacks=_run_manager.get_child(),\n",
    "                    **llm_inputs,\n",
    "                ).strip()\n",
    "                intermediate_steps.append(final_result)  # output: final answer\n",
    "                _run_manager.on_text(final_result, color=\"green\", verbose=self.verbose)\n",
    "            chain_result: Dict[str, Any] = {self.output_key: final_result}\n",
    "            if self.return_intermediate_steps:\n",
    "                chain_result[INTERMEDIATE_STEPS_KEY] = intermediate_steps\n",
    "            return chain_result\n",
    "        except Exception as exc:\n",
    "            # Append intermediate steps to exception, to aid in logging and later\n",
    "            # improvement of few shot prompt seeds\n",
    "            exc.intermediate_steps = intermediate_steps  # type: ignore\n",
    "            raise exc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6aee711d-dea0-4d05-8323-4469947e87de",
   "metadata": {},
   "source": [
    "* 初始化Bedrock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d21dbaa-d328-4ae2-a776-98e089667539",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from typing import Optional\n",
    "\n",
    "# External Dependencies:\n",
    "import boto3\n",
    "from botocore.config import Config\n",
    "from langchain.llms.bedrock import Bedrock\n",
    "\n",
    "def get_bedrock_client(\n",
    "    assumed_role: Optional[str] = None,\n",
    "    region: Optional[str] = None,\n",
    "    runtime: Optional[bool] = True,\n",
    "):\n",
    "  \n",
    "    if region is None:\n",
    "        target_region = os.environ.get(\"AWS_REGION\", os.environ.get(\"AWS_DEFAULT_REGION\"))\n",
    "    else:\n",
    "        target_region = region\n",
    "\n",
    "    print(f\"Create new client\\n  Using region: {target_region}\")\n",
    "    session_kwargs = {\"region_name\": target_region}\n",
    "    client_kwargs = {**session_kwargs}\n",
    "\n",
    "    profile_name = os.environ.get(\"AWS_PROFILE\")\n",
    "    if profile_name:\n",
    "        print(f\"  Using profile: {profile_name}\")\n",
    "        session_kwargs[\"profile_name\"] = profile_name\n",
    "\n",
    "    retry_config = Config(\n",
    "        region_name=target_region,\n",
    "        retries={\n",
    "            \"max_attempts\": 10,\n",
    "            \"mode\": \"standard\",\n",
    "        },\n",
    "    )\n",
    "    session = boto3.Session(**session_kwargs)\n",
    "\n",
    "    if assumed_role:\n",
    "        print(f\"  Using role: {assumed_role}\", end='')\n",
    "        sts = session.client(\"sts\")\n",
    "        response = sts.assume_role(\n",
    "            RoleArn=str(assumed_role),\n",
    "            RoleSessionName=\"langchain-llm-1\"\n",
    "        )\n",
    "        print(\" ... successful!\")\n",
    "        client_kwargs[\"aws_access_key_id\"] = response[\"Credentials\"][\"AccessKeyId\"]\n",
    "        client_kwargs[\"aws_secret_access_key\"] = response[\"Credentials\"][\"SecretAccessKey\"]\n",
    "        client_kwargs[\"aws_session_token\"] = response[\"Credentials\"][\"SessionToken\"]\n",
    "        \n",
    "\n",
    "    if runtime:\n",
    "        service_name='bedrock-runtime'\n",
    "    else:\n",
    "        service_name='bedrock'\n",
    "\n",
    "    client_kwargs[\"aws_access_key_id\"] = os.environ.get(\"AWS_ACCESS_KEY_ID\",\"\")\n",
    "    client_kwargs[\"aws_secret_access_key\"] = os.environ.get(\"AWS_SECRET_ACCESS_KEY\",\"\")\n",
    "    \n",
    "    bedrock_client = session.client(\n",
    "        service_name=service_name,\n",
    "        config=retry_config,\n",
    "        **client_kwargs\n",
    "    )\n",
    "\n",
    "    print(\"boto3 Bedrock client successfully created!\")\n",
    "    print(bedrock_client._endpoint)\n",
    "    return bedrock_client\n",
    "\n",
    "\n",
    "\n",
    "## for aksk bedrock\n",
    "def get_bedrock_aksk(secret_name='chatbot_bedrock', region_name = \"us-west-2\"):\n",
    "    # Create a Secrets Manager client\n",
    "    session = boto3.session.Session()\n",
    "    client = session.client(\n",
    "        service_name='secretsmanager',\n",
    "        region_name=region_name\n",
    "    )\n",
    "\n",
    "    try:\n",
    "        get_secret_value_response = client.get_secret_value(\n",
    "            SecretId=secret_name\n",
    "        )\n",
    "    except ClientError as e:\n",
    "        # For a list of exceptions thrown, see\n",
    "        # https://docs.aws.amazon.com/secretsmanager/latest/apireference/API_GetSecretValue.html\n",
    "        raise e\n",
    "\n",
    "    # Decrypts secret using the associated KMS key.\n",
    "    secret = json.loads(get_secret_value_response['SecretString'])\n",
    "    return secret['BEDROCK_ACCESS_KEY'],secret['BEDROCK_SECRET_KEY']\n",
    "\n",
    "ACCESS_KEY, SECRET_KEY=get_bedrock_aksk()\n",
    "\n",
    "#role based initial client#######\n",
    "os.environ[\"AWS_DEFAULT_REGION\"] = \"us-west-2\"  # E.g. \"us-east-1\"\n",
    "os.environ[\"AWS_PROFILE\"] = \"default\"\n",
    "#os.environ[\"BEDROCK_ASSUME_ROLE\"] = \"arn:aws:iam::687912291502:role/service-role/AmazonSageMaker-ExecutionRole-20211013T113123\"  # E.g. \"arn:aws:...\"\n",
    "os.environ[\"AWS_ACCESS_KEY_ID\"]=ACCESS_KEY\n",
    "os.environ[\"AWS_SECRET_ACCESS_KEY\"]=SECRET_KEY\n",
    "\n",
    "\n",
    "#新boto3 sdk只能session方式初始化bedrock\n",
    "boto3_bedrock = get_bedrock_client(\n",
    "    #assumed_role=os.environ.get(\"BEDROCK_ASSUME_ROLE\", None),\n",
    "    region=os.environ.get(\"AWS_DEFAULT_REGION\", None)\n",
    ")\n",
    "\n",
    "parameters_bedrock = {\n",
    "    \"max_tokens_to_sample\": 2048,\n",
    "    #\"temperature\": 0.5,\n",
    "    \"temperature\": 0,\n",
    "    #\"top_k\": 250,\n",
    "    #\"top_p\": 1,\n",
    "    \"stop_sequences\": [\"\\n\\nHuman\"],\n",
    "}\n",
    "\n",
    "\n",
    "bedrock_llm = Bedrock(model_id=\"anthropic.claude-v2\", \n",
    "                      client=boto3_bedrock, \n",
    "                      model_kwargs=parameters_bedrock)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "f6d4e53f-3e0a-4709-8be8-678ff542918a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' 对于Windows远程桌面连接,当连接断开后系统默认会进入锁屏状态,这是为了保护远程桌面不被其他人访问。但是如果需要避免锁屏,可以通过以下几种方法:\\n\\n1. 在远程桌面连接窗口底部,点击\"选项\" - \"远程\" - 勾选\"断开连接时不锁定计算机\"。这样断开连接后就不会锁屏了。\\n\\n2. 在注册表中设置相关键值也可以避免锁屏:\\n\\n(1)打开注册表编辑器,定位到HKEY_LOCAL_MACHINE\\\\SOFTWARE\\\\Microsoft\\\\Windows NT\\\\CurrentVersion\\\\Winlogon 键。\\n\\n(2)右击新建一个DWORD值,命名为\"DisableLockWorkstation\",并设置值为1。\\n\\n(3)重新启动计算机后设置才会生效。\\n\\n3. 使用第三方工具,例如DisableLockWorkstation,可以设置热键来禁用或启用锁屏功能。\\n\\n4. 通过组策略也可以设置远程桌面断开后不锁定屏幕。\\n\\n需要注意的是,关闭锁屏功能会降低安全性,如果不是非常必要,不建议这样操作。正确的做法是在断开连接前保存工作并妥善退出。'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bedrock_llm.predict(\"windows远程链接关闭后，如何不进入锁屏状态\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7fbcf0b-bc98-4513-9b84-9ef842444b5b",
   "metadata": {
    "tags": []
   },
   "source": [
    "* Langchain SqlDatabase chain支持多种数据源\n",
    "* 不同数据源的连接url及传参不一样\n",
    "* 此处使用rds mysql "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc28c0b2-3169-4d67-a083-041d31b507f5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## langchain agent demo test##########\n",
    "from langchain.utilities import SQLDatabase\n",
    "from langchain_experimental.sql import SQLDatabaseChain\n",
    "from langchain.chains import create_sql_query_chain\n",
    "import json\n",
    "import os\n",
    "\n",
    "os.environ[\"PGPASSWORD\"] = \"*******\"\n",
    "os.environ[\"LANGCHAIN_HANDLER\"] = \"langchain\"\n",
    "conn_str = \"awsathena+rest://{aws_access_key_id}:{aws_secret_access_key}@athena.{region_name}.amazonaws.com:443/\"\\\n",
    "           \"{schema_name}?s3_staging_dir={s3_staging_dir}\"\n",
    "conn_str = conn_str.format(\n",
    "    aws_access_key_id=\"**********\",\n",
    "    aws_secret_access_key=\"***********\",\n",
    "    region_name=\"cn-northwest-1\",\n",
    "    schema_name=\"specturmdb\",\n",
    "    s3_staging_dir=\"s3://tangqy-athenaoutput-us-west-2\")\n",
    "\n",
    "db = SQLDatabase.from_uri(\n",
    "    #conn_str,\n",
    "    #\"redshift+redshift_connector://admin@redshift-cluster-1.cp1kgq7oikv3.ap-southeast-1.redshift.amazonaws.com:5439/dev\",\n",
    "    \"mysql+pymysql://admin:admin12345678@database-us-west-2-demo.cluster-c1qvx9wzmmcz.us-west-2.rds.amazonaws.com/llm\",\n",
    "    #\"jdbc:awsathena://athena.us-west-2.amazonaws.com:443/tpcds_bin_partitioned_orc_300?s3_staging_dir=s3://tangqy-athenaoutput/&aws_credentials_provider_class=com.amazonaws.auth.DefaultAWSCredentialsProviderChain\",\n",
    "    include_tables=['dws_truck_portrait_index_sum_da','dws_ots_waybill_info_da','ads_bi_quality_monitor_shipping_detail','dim_pub_truck_info'], # we include only one table to save tokens in the prompt :)\n",
    "    #include_tables=[\"customer\"],\n",
    "    sample_rows_in_table_info=0)\n",
    "#chain = create_sql_query_chain(sm_llm, db)\n",
    "#response = chain.invoke({\"question\":\"最近一个月温度合格的派车单数量\"})\n",
    "#response = chain.invoke({\"question\":\"最近一个月下单最大的客户邮件地址\"})\n",
    "\n",
    "#print(response)\n",
    "#db_chain = CustomerizedSQLDatabaseChain(llm=llm, database=db, verbose=True, top_k=3)\n",
    "#db_chain.run(\"最近一个月温度合格的派车单数量\")\n",
    "#db_chain.run(\"I need to know the max sales customer's id in sales report\")\n",
    "### for bedrock#######\n",
    "db_chain = CustomerizedSQLDatabaseChain.from_llm(llm=bedrock_llm, db=db, verbose=False, return_sql=True)\n",
    "#db_chain.run(\"2023年7月派车单数量超过26次的4.2米车辆一共有多少辆?\")\n",
    "#db_chain.run(\"历史累计派车单数量、干线派车单数量、城配派车单数量\")\n",
    "db_chain.run(\"2022年的运输总量是多少吨?\")\n",
    "#db_chain.run(\"成都市的车辆资源累计有多少？\")\n",
    "#db_chain.run(\"I need to know the max sales customer's id in sales report\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70984638-0139-44c1-afe2-de6efb3216e9",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Opt2:元数据召回+langchain SQLDatabasechain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe2714ff-f4d0-4c1a-ab9f-7fb04de996a2",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### func for agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b48fc050-f656-46da-82d1-4dc88b553991",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import boto3\n",
    "import json\n",
    "import requests\n",
    "import time\n",
    "from collections import defaultdict\n",
    "from requests_aws4auth import AWS4Auth\n",
    "import os\n",
    "from opensearchpy import OpenSearch, RequestsHttpConnection\n",
    "from langchain.vectorstores import OpenSearchVectorSearch\n",
    "from langchain import PromptTemplate, SagemakerEndpoint\n",
    "from langchain.chains.question_answering import load_qa_chain\n",
    "from langchain.embeddings import SagemakerEndpointEmbeddings\n",
    "from langchain.llms.sagemaker_endpoint import ContentHandlerBase\n",
    "from langchain.docstore.document import Document\n",
    "from langchain.memory import ConversationBufferWindowMemory\n",
    "from langchain import LLMChain\n",
    "\n",
    "\n",
    "\n",
    "def aos_knn_search(client, field,q_embedding, index, size=1):\n",
    "    if not isinstance(client, OpenSearch):   \n",
    "        client = OpenSearch(\n",
    "            hosts=[{'host': aos_endpoint, 'port': 443}],\n",
    "            http_auth = pwdauth,\n",
    "            use_ssl=True,\n",
    "            verify_certs=True,\n",
    "            connection_class=RequestsHttpConnection\n",
    "        )\n",
    "    query = {\n",
    "        \"size\": size,\n",
    "        \"query\": {\n",
    "            \"knn\": {\n",
    "                field: {\n",
    "                    \"vector\": q_embedding,\n",
    "                    \"k\": size\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "    opensearch_knn_respose = []\n",
    "    query_response = client.search(\n",
    "        body=query,\n",
    "        index=index\n",
    "    )\n",
    "    opensearch_knn_respose = [{'idx':item['_source'].get('idx',1),'database_name':item['_source']['database_name'],'table_name':item['_source']['table_name'],'query_desc_text':item['_source']['query_desc_text'],\"score\":item[\"_score\"]}  for item in query_response[\"hits\"][\"hits\"]]\n",
    "    return opensearch_knn_respose\n",
    "\n",
    "def aos_knn_search_v2(client, field,q_embedding, index, size=1):\n",
    "    if not isinstance(client, OpenSearch):   \n",
    "        client = OpenSearch(\n",
    "            hosts=[{'host': aos_endpoint, 'port': 443}],\n",
    "            http_auth = auth,\n",
    "            use_ssl=True,\n",
    "            verify_certs=True,\n",
    "            connection_class=RequestsHttpConnection\n",
    "        )\n",
    "    query = {\n",
    "        \"size\": size,\n",
    "        \"query\": {\n",
    "            \"knn\": {\n",
    "                field: {\n",
    "                    \"vector\": q_embedding,\n",
    "                    \"k\": size\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "    opensearch_knn_respose = []\n",
    "    query_response = client.search(\n",
    "        body=query,\n",
    "        index=index\n",
    "    )\n",
    "    opensearch_knn_respose = [{'idx':item['_source'].get('idx',1),'database_name':item['_source']['database_name'],'table_name':item['_source']['table_name'],'exactly_query_text':item['_source']['exactly_query_text'],\"score\":item[\"_score\"]}  for item in query_response[\"hits\"][\"hits\"]]\n",
    "    return opensearch_knn_respose\n",
    "\n",
    "\n",
    "def aos_reverse_search(client, index_name, field, query_term, exactly_match=False, size=1):\n",
    "    \"\"\"\n",
    "    search opensearch with query.\n",
    "    :param host: AOS endpoint\n",
    "    :param index_name: Target Index Name\n",
    "    :param field: search field\n",
    "    :param query_term: query term\n",
    "    :return: aos response json\n",
    "    \"\"\"\n",
    "    if not isinstance(client, OpenSearch):   \n",
    "        client = OpenSearch(\n",
    "            hosts=[{'host': aos_endpoint, 'port': 443}],\n",
    "            http_auth = pwdauth,\n",
    "            use_ssl=True,\n",
    "            verify_certs=True,\n",
    "            connection_class=RequestsHttpConnection\n",
    "        )\n",
    "    query = None\n",
    "    if exactly_match:\n",
    "        query =  {\n",
    "            \"query\" : {\n",
    "                \"match_phrase\":{\n",
    "                    field: {\n",
    "                        \"query\": query_term,\n",
    "                        \"analyzer\": \"ik_smart\"\n",
    "                      }\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    else:\n",
    "        query = {\n",
    "            \"size\": size,\n",
    "            \"query\": {\n",
    "                \"query_string\": {\n",
    "                \"default_field\": \"exactly_query_text\",  \n",
    "                \"query\": query_term         \n",
    "              }\n",
    "            },\n",
    "           \"sort\": [{\n",
    "               \"_score\": {\n",
    "                   \"order\": \"desc\"\n",
    "               }\n",
    "           }]\n",
    "    }        \n",
    "    query_response = client.search(\n",
    "        body=query,\n",
    "        index=index_name\n",
    "    )\n",
    "    result_arr = [{'idx':item['_source'].get('idx',1),'database_name':item['_source']['database_name'],'table_name':item['_source']['table_name'],'exactly_query_text':item['_source']['exactly_query_text'],\"score\":item[\"_score\"]}  for item in query_response[\"hits\"][\"hits\"]]\n",
    "    return result_arr\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def get_vector_by_sm_endpoint(questions, sm_client, endpoint_name):\n",
    "    parameters = {\n",
    "    }\n",
    "\n",
    "    response_model = sm_client.invoke_endpoint(\n",
    "        EndpointName=endpoint_name,\n",
    "        Body=json.dumps(\n",
    "            {\n",
    "                \"inputs\": questions,\n",
    "                \"parameters\": parameters,\n",
    "                \"is_query\" : True,\n",
    "                \"instruction\" :  \"为这个句子生成表示以用于检索相关文章：\"\n",
    "            }\n",
    "        ),\n",
    "        ContentType=\"application/json\",\n",
    "    )\n",
    "    json_str = response_model['Body'].read().decode('utf8')\n",
    "    json_obj = json.loads(json_str)\n",
    "    embeddings = json_obj['sentence_embeddings']\n",
    "    return embeddings\n",
    "\n",
    "\n",
    "def k_nn_ingestion_by_aos(docs,index,hostname,username,passwd):\n",
    "    auth = (username, passwd)\n",
    "    search = OpenSearch(\n",
    "        hosts = [{'host': aos_endpoint, 'port': 443}],\n",
    "        ##http_auth = awsauth ,\n",
    "        http_auth = auth ,\n",
    "        use_ssl = True,\n",
    "        verify_certs = True,\n",
    "        connection_class = RequestsHttpConnection\n",
    "    )\n",
    "    for doc in docs:\n",
    "        query_desc_embedding = doc['query_desc_embedding']\n",
    "        database_name = doc['database_name']\n",
    "        table_name = doc['table_name']\n",
    "        query_desc_text = doc[\"query_desc_text\"]\n",
    "        document = { \"query_desc_embedding\": query_desc_embedding, 'database_name':database_name, \"table_name\": table_name,\"query_desc_text\":query_desc_text}\n",
    "        search.index(index=index, body=document)\n",
    "        \n",
    "def k_nn_ingestion_by_aos_v2(docs,index,hostname,username,passwd):\n",
    "    auth = (username, passwd)\n",
    "    search = OpenSearch(\n",
    "        hosts = [{'host': aos_endpoint, 'port': 443}],\n",
    "        ##http_auth = awsauth ,\n",
    "        http_auth = auth ,\n",
    "        use_ssl = True,\n",
    "        verify_certs = True,\n",
    "        connection_class = RequestsHttpConnection\n",
    "    )\n",
    "    for doc in docs:\n",
    "        exactly_query_embedding = doc['exactly_query_embedding']\n",
    "        database_name = doc['database_name']\n",
    "        table_name = doc['table_name']\n",
    "        exactly_query_text = doc[\"exactly_query_text\"]\n",
    "        document = { \"exactly_query_embedding\": exactly_query_embedding, 'database_name':database_name, \"table_name\": table_name,\"exactly_query_text\":exactly_query_text}\n",
    "        search.index(index=index, body=document)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccd35f03-b3ee-4f97-9bbc-cc12484173c7",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### 元数据ingestion 入aos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "149abfb0-b2c7-42da-8c3d-8a1dbc02ad70",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## data process\n",
    "all_querys = \"\"\"2023年7月派车单数量超过26次的4.2米车辆一共有多少辆\n",
    "请统计历史累计派车单数量、干线派车单数量、城配派车单数量。历史累计的意思是不限定时间范围\n",
    "奶茶品牌的站点数量和运输货品数量统计\n",
    "2022年的运输总量是多少吨？请注意：traff_weight的单位是千克，请把单位转换为吨\n",
    "车牌为'黑RG6696'的车辆的GPS最近定位上传时间、GPS最近定位省份、GPS最近定位城市、GPS最近定位区县、APP最近定位上传时间、APP最近定位省份、APP最近定位城市、APP最近定位区县。给出sql中字段名不要带上库名\n",
    "车牌归属城市为'成都'的车辆累计有多少？\n",
    "取货地城市名称为'北京市'的历史累计不重复的车牌有多少？请注意车牌号有可能有重复\n",
    "货主-行业列表相似于'西餐连锁'的一共有多少个品牌？多少个客户？\n",
    "品牌名称为'星巴克'的不重复的站点一共有多少个？\n",
    "查看租户简称为云南,车辆的车厢长为9.6米和15米的外廓车长、核定载重\"\"\"\n",
    "querys = all_querys.split(\"\\n\")\n",
    "\n",
    "all_tables = \"\"\"ads_bi_quality_monitor_shipping_detail\n",
    "dws_ots_waybill_info_da\n",
    "dws_station_portrait_index_sum_da\n",
    "dws_ots_waybill_info_da\n",
    "dws_truck_portrait_index_sum_da\n",
    "dws_truck_portrait_index_sum_da\n",
    "dws_ots_waybill_info_da\n",
    "ads_customer_portrait_index_sum_da\n",
    "dim_customer_enterprise_station_base_info\n",
    "dim_pub_truck_tenant,dim_pub_truck_info\"\"\"\n",
    "tables=all_tables.split(\"\\n\")\n",
    "\n",
    "all_dbs = \"\"\"llm\n",
    "llm\n",
    "llm\n",
    "llm\n",
    "llm\n",
    "llm\n",
    "llm\n",
    "llm\n",
    "llm\n",
    "llm\"\"\"\n",
    "dbs=all_dbs.split(\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "917d06d3-0556-430e-a10a-a237ddf41579",
   "metadata": {},
   "outputs": [],
   "source": [
    "index_name=\"prompt-optimal-index\"\n",
    "embedding_endpoint_name=\"bge-zh-15-2023-09-25-07-02-01-080-endpoint\"\n",
    "##########embedding by llm model##############\n",
    "sentense_vectors = []\n",
    "sentense_vectors=get_vector_by_sm_endpoint(querys,sm_client,embedding_endpoint_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa53fb15-ad46-4583-bbca-29a23d982b94",
   "metadata": {},
   "outputs": [],
   "source": [
    "docs=[]\n",
    "for index, sentence_vector in enumerate(sentense_vectors):\n",
    "    doc = {\n",
    "        \"metadata_type\":\"table\",\n",
    "        \"database_name\":dbs[index],\n",
    "        \"table_name\": tables[index],\n",
    "        \"exactly_query_text\":querys[index],\n",
    "        \"exactly_query_embedding\": sentence_vector\n",
    "          }\n",
    "    docs.append(doc)\n",
    "\n",
    "#print((doc[\"database_name\"]))\n",
    "#########ingestion into aos ###################\n",
    "k_nn_ingestion_by_aos_v2(docs,index_name,aos_endpoint,username,passwd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc16bf55-393e-47c5-a3e1-3db5ae0dea89",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "client = OpenSearch(\n",
    "            hosts=[{'host': aos_endpoint, 'port': 443}],\n",
    "            http_auth = pwdauth,\n",
    "            use_ssl=True,\n",
    "            verify_certs=True,\n",
    "            connection_class=RequestsHttpConnection\n",
    "        )\n",
    "\n",
    "query=\"上个月温度合格的派车单数量\"\n",
    "query_embedding = get_vector_by_sm_endpoint(query, sm_client, embedding_endpoint_name)\n",
    "rets=aos_knn_search_v2(client, \"exactly_query_embedding\",query_embedding[0],index_name,1)   \n",
    "print(rets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8526747e-ec52-4938-bec5-42ec26c76dcd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "query=\"上个月温度合格的派车单数量\"\n",
    "query_embedding = get_vector_by_sm_endpoint(query, sm_client, embedding_endpoint_name)\n",
    "rets=aos_knn_search_v2(client, \"exactly_query_embedding\",query_embedding[0],index_name,1)   \n",
    "print(rets)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16d7df82-5f33-4b9c-98c5-5fbd898081f8",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### e2e pipeline\n",
    "* 先使用reverse search召回表元数据\n",
    "* 如果reverse search召回为空，使用embedding向量召回\n",
    "* 将召回的表table name传入sqlDatabaseChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35e0ea32-0a32-4c2d-a301-e4e70e9c4c71",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.utilities import SQLDatabase\n",
    "from langchain_experimental.sql import SQLDatabaseChain\n",
    "from langchain.chains import create_sql_query_chain\n",
    "import json\n",
    "import os\n",
    "\n",
    "\n",
    "table_name = None\n",
    "query=\"2023年7月派车单数量超过26次的4.2米车辆一共有多少辆?\"\n",
    "aos_client = OpenSearch(\n",
    "            hosts=[{'host': aos_endpoint, 'port': 443}],\n",
    "            http_auth = auth,\n",
    "            use_ssl=True,\n",
    "            verify_certs=True,\n",
    "            connection_class=RequestsHttpConnection\n",
    "        )\n",
    "#### reverse 倒排召回 ############\n",
    "opensearch_query_response = aos_reverse_search(aos_client, index_name, \"exactly_query_text\", query)\n",
    "try:\n",
    "    table_name=opensearch_query_response[0][\"table_name\"].strip()\n",
    "except Exception as e:\n",
    "    print(e)\n",
    "    table_name=None\n",
    "\n",
    "#### reverse 向量召回 ############\n",
    "if table_name is None:\n",
    "    query_embedding = get_vector_by_sm_endpoint(query, sm_client, embedding_endpoint_name)\n",
    "    opensearch_query_response = aos_knn_search_v2(aos_client, \"exactly_query_embedding\",query_embedding[0], index_name, size=10)\n",
    "    try:\n",
    "        table_name = opensearch_query_response[0][\"table_name\"].strip()\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        table_name = None\n",
    "\n",
    "#####使用召回table name执行SqlDataBaseChain#######\n",
    "db = SQLDatabase.from_uri(\n",
    "    \"mysql+pymysql://admin:admin12345678@database-us-west-2-demo.cluster-c1qvx9wzmmcz.us-west-2.rds.amazonaws.com/llm\",\n",
    "    include_tables=[table_name], # we include only one table to save tokens in the prompt :)\n",
    "    sample_rows_in_table_info=0)\n",
    "\n",
    "db_chain = CustomerizedSQLDatabaseChain.from_llm(llm=bedrock_llm, db=db, verbose=False, return_sql=True)\n",
    "db_chain.run(query)\n",
    "#db_chain.run(\"历史累计派车单数量、干线派车单数量、城配派车单数量\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ece8bcef-d342-4578-956f-505125469c12",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p310",
   "language": "python",
   "name": "conda_pytorch_p310"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}