More cleanups

Author: hammerhead

topic/chatbot/table-augmented-generation/aws/cratedb_tag_inline_agent.ipynb

@@ -80,16 +80,10 @@
    "outputs": [],
    "source": [
     "# Install dependencies\n",
-    "!pip install pandas matplotlib sqlalchemy-cratedb python-dotenv ipython-sql\n",
-    "\n",
-    "# CrateDB Client\n",
-    "!pip install crate cratedb-mcp==0.0.0\n",
-    "\n",
     "# tqdm for nicer loading bars\n",
-    "!pip install tqdm\n",
-    "\n",
-    "# AWS Inline Agent library\n",
-    "!pip install /Users/niklasschmidtmer/Documents/GitHub/amazon-bedrock-agent-samples/src/InlineAgent/dist/inlineagent-0+untagged.276.g0d4f7be.dirty-py3-none-any.whl"
+    "!pip install -U pandas matplotlib ipython-sql tqdm \\\n",
+    "    crate cratedb-mcp==0.0.0 sqlalchemy-cratedb \\\n",
+    "    inlineagent-0+untagged.276.g0d4f7be.dirty-py3-none-any.whl"
    ]
   },
   {
@@ -137,16 +131,7 @@
    "execution_count": null,
    "id": "2585e16e",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "✅ Successfully connected to CrateDB!\n",
-      "Sample query result from sys.summits: Mont Blanc\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import os\n",
     "import sqlalchemy as sa\n",
@@ -193,15 +178,7 @@
    "execution_count": null,
    "id": "856582bb",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "✅ Table 'motor_readings' created (if not already existing).\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from sqlalchemy import text\n",
     "\n",
@@ -420,14 +397,17 @@
     "def fetch_table_schema(table_name):\n",
     "    \"Fetch the column names and data types for a given table from CrateDB's system catalog.\"\n",
     "\n",
-    "    query = \"\"\"\n",
-    "            SELECT column_name, data_type\n",
-    "            FROM information_schema.columns\n",
-    "            WHERE table_name = ?\n",
-    "            ORDER BY ordinal_position\n",
-    "            \"\"\"\n",
+    "    query = text(\n",
+    "        \"\"\"\n",
+    "        SELECT column_name, data_type\n",
+    "        FROM information_schema.columns\n",
+    "        WHERE table_name = :tbl\n",
+    "        ORDER BY ordinal_position\n",
+    "        \"\"\"\n",
+    "    )\n",
     "    try:\n",
-    "        df = pd.read_sql(query, con=engine, params=(table_name,))\n",
+    "        stmt = query.bindparams(tbl=table_name)\n",
+    "        df = pd.read_sql(stmt, con=engine)\n",
     "\n",
     "        schema_text = f\"Table: {table_name}\\nColumns:\\n\"\n",
     "        for _, row in df.iterrows():\n",
@@ -534,10 +514,15 @@
     "    # Prompt the LLM\n",
     "    response = prompt_llm(table_schema, question)\n",
     "\n",
+    "    # Uncomment the line below to see the complete LLM response\n",
+    "    # print(\"Debug: The LLM output was: \\n\", response)\n",
+    "\n",
     "    # Extract the actual SQL query from the response\n",
-    "    sql = re.search(\"(?<=`sql)([^`]*)(?=`)\", response).group(1).strip()\n",
+    "    match = re.search(\"(?<=`sql)([^`]*)(?=`)\", response)\n",
+    "    if match:\n",
+    "        return match.group(1).strip()\n",
     "\n",
-    "    return sql"
+    "    raise ValueError(\"Failed to extract SQL from LLM reply\", response)"
    ]
   },
   {
@@ -712,7 +697,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "cbbb38f2",
    "metadata": {},
    "outputs": [],
@@ -912,7 +897,7 @@
     "df_manuals = pd.DataFrame(manuals)\n",
     "\n",
     "# Store in CrateDB\n",
-    "df_manuals.to_sql(\"machine_manuals\", con=engine, if_exists=\"replace\", index=False)\n",
+    "df_manuals.to_sql(\"machine_manuals\", con=engine, if_exists=\"append\", index=False)\n",
     "print(f\"✅ Stored manuals for {len(df_manuals)} machines in 'machine_manuals'.\")\n",
     "_ = connection.execute(sa.text(\"REFRESH TABLE machine_manuals;\"))"
    ]