Last active
September 23, 2024 11:56
-
-
Save alonsosilvaallende/270a0ab5b276e708b80ee5aab9af522e to your computer and use it in GitHub Desktop.
LlamaIndex with existing Kuzu db
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| { | |
| "cells": [ | |
| { | |
| "cell_type": "markdown", | |
| "id": "362af633-644d-4656-acf5-9cb5d5f7e024", | |
| "metadata": {}, | |
| "source": [ | |
| "## Get the data" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 1, | |
| "id": "f5510434-fb17-4ef2-a24f-adef4ab17184", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:50.236387Z", | |
| "iopub.status.busy": "2024-09-23T11:56:50.235316Z", | |
| "iopub.status.idle": "2024-09-23T11:56:53.135286Z", | |
| "shell.execute_reply": "2024-09-23T11:56:53.134127Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:50.236333Z" | |
| } | |
| }, | |
| "outputs": [ | |
| { | |
| "data": { | |
| "text/html": [ | |
| "<div><style>\n", | |
| ".dataframe > thead > tr,\n", | |
| ".dataframe > tbody > tr {\n", | |
| " text-align: right;\n", | |
| " white-space: pre-wrap;\n", | |
| "}\n", | |
| "</style>\n", | |
| "<small>shape: (5, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>Song</th><th>Lyrics</th><th>Album</th><th>Artist</th></tr><tr><td>str</td><td>str</td><td>str</td><td>str</td></tr></thead><tbody><tr><td>"DON'T TREAD ON ME"</td><td>"Liberty or death\n", | |
| "what we so pr…</td><td>"THE BLACK ALBUM"</td><td>"Metallica"</td></tr><tr><td>"NOTHING ELSE MATTERS"</td><td>"So close no matter how far\n", | |
| "cou…</td><td>"THE BLACK ALBUM"</td><td>"Metallica"</td></tr><tr><td>"THROUGH THE NEVER"</td><td>"All that is\n", | |
| "was and will be\n", | |
| "un…</td><td>"THE BLACK ALBUM"</td><td>"Metallica"</td></tr><tr><td>"MY FRIEND OF MISERY"</td><td>"You just stood there screaming…</td><td>"THE BLACK ALBUM"</td><td>"Metallica"</td></tr><tr><td>"WHEREVER I MAY ROAM"</td><td>"...and the road becomes my bri…</td><td>"THE BLACK ALBUM"</td><td>"Metallica"</td></tr></tbody></table></div>" | |
| ], | |
| "text/plain": [ | |
| "shape: (5, 4)\n", | |
| "┌──────────────────────┬─────────────────────────────────┬─────────────────┬───────────┐\n", | |
| "│ Song ┆ Lyrics ┆ Album ┆ Artist │\n", | |
| "│ --- ┆ --- ┆ --- ┆ --- │\n", | |
| "│ str ┆ str ┆ str ┆ str │\n", | |
| "╞══════════════════════╪═════════════════════════════════╪═════════════════╪═══════════╡\n", | |
| "│ DON'T TREAD ON ME ┆ Liberty or death ┆ THE BLACK ALBUM ┆ Metallica │\n", | |
| "│ ┆ what we so pr… ┆ ┆ │\n", | |
| "│ NOTHING ELSE MATTERS ┆ So close no matter how far ┆ THE BLACK ALBUM ┆ Metallica │\n", | |
| "│ ┆ cou… ┆ ┆ │\n", | |
| "│ THROUGH THE NEVER ┆ All that is ┆ THE BLACK ALBUM ┆ Metallica │\n", | |
| "│ ┆ was and will be ┆ ┆ │\n", | |
| "│ ┆ un… ┆ ┆ │\n", | |
| "│ MY FRIEND OF MISERY ┆ You just stood there screaming… ┆ THE BLACK ALBUM ┆ Metallica │\n", | |
| "│ WHEREVER I MAY ROAM ┆ ...and the road becomes my bri… ┆ THE BLACK ALBUM ┆ Metallica │\n", | |
| "└──────────────────────┴─────────────────────────────────┴─────────────────┴───────────┘" | |
| ] | |
| }, | |
| "execution_count": 1, | |
| "metadata": {}, | |
| "output_type": "execute_result" | |
| } | |
| ], | |
| "source": [ | |
| "import polars as pl\n", | |
| "\n", | |
| "df = pl.read_csv(\n", | |
| " \"https://drive.google.com/uc?export=download&id=1uD3h7xYxr9EoZ0Ggoh99JtQXa3AxtxyU\"\n", | |
| ")\n", | |
| "\n", | |
| "df.head()" | |
| ] | |
| }, | |
| { | |
| "cell_type": "markdown", | |
| "id": "f1095f16-c182-472b-ab03-831686f2fd50", | |
| "metadata": {}, | |
| "source": [ | |
| "## Clean the data" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 2, | |
| "id": "2a9265be-cca8-486b-8835-196c745f46df", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:53.137388Z", | |
| "iopub.status.busy": "2024-09-23T11:56:53.136841Z", | |
| "iopub.status.idle": "2024-09-23T11:56:53.150451Z", | |
| "shell.execute_reply": "2024-09-23T11:56:53.149508Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:53.137344Z" | |
| } | |
| }, | |
| "outputs": [ | |
| { | |
| "data": { | |
| "text/html": [ | |
| "<div><style>\n", | |
| ".dataframe > thead > tr,\n", | |
| ".dataframe > tbody > tr {\n", | |
| " text-align: right;\n", | |
| " white-space: pre-wrap;\n", | |
| "}\n", | |
| "</style>\n", | |
| "<small>shape: (5, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>Song</th><th>Lyrics</th><th>Album</th><th>Artist</th></tr><tr><td>str</td><td>str</td><td>str</td><td>str</td></tr></thead><tbody><tr><td>"Don't Tread On Me"</td><td>"Liberty or death\n", | |
| "what we so pr…</td><td>"The Black Album"</td><td>"Metallica"</td></tr><tr><td>"Nothing Else Matters"</td><td>"So close no matter how far\n", | |
| "cou…</td><td>"The Black Album"</td><td>"Metallica"</td></tr><tr><td>"Through The Never"</td><td>"All that is\n", | |
| "was and will be\n", | |
| "un…</td><td>"The Black Album"</td><td>"Metallica"</td></tr><tr><td>"My Friend Of Misery"</td><td>"You just stood there screaming…</td><td>"The Black Album"</td><td>"Metallica"</td></tr><tr><td>"Wherever I May Roam"</td><td>"...and the road becomes my bri…</td><td>"The Black Album"</td><td>"Metallica"</td></tr></tbody></table></div>" | |
| ], | |
| "text/plain": [ | |
| "shape: (5, 4)\n", | |
| "┌──────────────────────┬─────────────────────────────────┬─────────────────┬───────────┐\n", | |
| "│ Song ┆ Lyrics ┆ Album ┆ Artist │\n", | |
| "│ --- ┆ --- ┆ --- ┆ --- │\n", | |
| "│ str ┆ str ┆ str ┆ str │\n", | |
| "╞══════════════════════╪═════════════════════════════════╪═════════════════╪═══════════╡\n", | |
| "│ Don't Tread On Me ┆ Liberty or death ┆ The Black Album ┆ Metallica │\n", | |
| "│ ┆ what we so pr… ┆ ┆ │\n", | |
| "│ Nothing Else Matters ┆ So close no matter how far ┆ The Black Album ┆ Metallica │\n", | |
| "│ ┆ cou… ┆ ┆ │\n", | |
| "│ Through The Never ┆ All that is ┆ The Black Album ┆ Metallica │\n", | |
| "│ ┆ was and will be ┆ ┆ │\n", | |
| "│ ┆ un… ┆ ┆ │\n", | |
| "│ My Friend Of Misery ┆ You just stood there screaming… ┆ The Black Album ┆ Metallica │\n", | |
| "│ Wherever I May Roam ┆ ...and the road becomes my bri… ┆ The Black Album ┆ Metallica │\n", | |
| "└──────────────────────┴─────────────────────────────────┴─────────────────┴───────────┘" | |
| ] | |
| }, | |
| "execution_count": 2, | |
| "metadata": {}, | |
| "output_type": "execute_result" | |
| } | |
| ], | |
| "source": [ | |
| "import string\n", | |
| "\n", | |
| "df = df.with_columns(\n", | |
| " pl.Series(\"Album\", [string.capwords(album) for album in df[\"Album\"]])\n", | |
| ")\n", | |
| "df = df.with_columns(pl.Series(\"Song\", [string.capwords(song) for song in df[\"Song\"]]))\n", | |
| "df = df.with_columns(pl.col(\"Lyrics\").fill_null(\"None\"))\n", | |
| "df.head()" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 3, | |
| "id": "84f7ab9a-dffb-4124-81e0-99e29b5c13d7", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:53.152594Z", | |
| "iopub.status.busy": "2024-09-23T11:56:53.151685Z", | |
| "iopub.status.idle": "2024-09-23T11:56:53.177325Z", | |
| "shell.execute_reply": "2024-09-23T11:56:53.175893Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:53.152562Z" | |
| } | |
| }, | |
| "outputs": [ | |
| { | |
| "data": { | |
| "text/html": [ | |
| "<div><style>\n", | |
| ".dataframe > thead > tr,\n", | |
| ".dataframe > tbody > tr {\n", | |
| " text-align: right;\n", | |
| " white-space: pre-wrap;\n", | |
| "}\n", | |
| "</style>\n", | |
| "<small>shape: (5, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>Song</th><th>Lyrics</th><th>Album</th><th>Artist</th><th>text</th></tr><tr><td>str</td><td>str</td><td>str</td><td>str</td><td>str</td></tr></thead><tbody><tr><td>"Don't Tread On Me"</td><td>"Liberty or death\n", | |
| "what we so pr…</td><td>"The Black Album"</td><td>"Metallica"</td><td>"# The Black Album: Don't Tread…</td></tr><tr><td>"Nothing Else Matters"</td><td>"So close no matter how far\n", | |
| "cou…</td><td>"The Black Album"</td><td>"Metallica"</td><td>"# The Black Album: Nothing Els…</td></tr><tr><td>"Through The Never"</td><td>"All that is\n", | |
| "was and will be\n", | |
| "un…</td><td>"The Black Album"</td><td>"Metallica"</td><td>"# The Black Album: Through The…</td></tr><tr><td>"My Friend Of Misery"</td><td>"You just stood there screaming…</td><td>"The Black Album"</td><td>"Metallica"</td><td>"# The Black Album: My Friend O…</td></tr><tr><td>"Wherever I May Roam"</td><td>"...and the road becomes my bri…</td><td>"The Black Album"</td><td>"Metallica"</td><td>"# The Black Album: Wherever I …</td></tr></tbody></table></div>" | |
| ], | |
| "text/plain": [ | |
| "shape: (5, 5)\n", | |
| "┌──────────────────────┬──────────────────────┬─────────────────┬───────────┬──────────────────────┐\n", | |
| "│ Song ┆ Lyrics ┆ Album ┆ Artist ┆ text │\n", | |
| "│ --- ┆ --- ┆ --- ┆ --- ┆ --- │\n", | |
| "│ str ┆ str ┆ str ┆ str ┆ str │\n", | |
| "╞══════════════════════╪══════════════════════╪═════════════════╪═══════════╪══════════════════════╡\n", | |
| "│ Don't Tread On Me ┆ Liberty or death ┆ The Black Album ┆ Metallica ┆ # The Black Album: │\n", | |
| "│ ┆ what we so pr… ┆ ┆ ┆ Don't Tread… │\n", | |
| "│ Nothing Else Matters ┆ So close no matter ┆ The Black Album ┆ Metallica ┆ # The Black Album: │\n", | |
| "│ ┆ how far ┆ ┆ ┆ Nothing Els… │\n", | |
| "│ ┆ cou… ┆ ┆ ┆ │\n", | |
| "│ Through The Never ┆ All that is ┆ The Black Album ┆ Metallica ┆ # The Black Album: │\n", | |
| "│ ┆ was and will be ┆ ┆ ┆ Through The… │\n", | |
| "│ ┆ un… ┆ ┆ ┆ │\n", | |
| "│ My Friend Of Misery ┆ You just stood there ┆ The Black Album ┆ Metallica ┆ # The Black Album: │\n", | |
| "│ ┆ screaming… ┆ ┆ ┆ My Friend O… │\n", | |
| "│ Wherever I May Roam ┆ ...and the road ┆ The Black Album ┆ Metallica ┆ # The Black Album: │\n", | |
| "│ ┆ becomes my bri… ┆ ┆ ┆ Wherever I … │\n", | |
| "└──────────────────────┴──────────────────────┴─────────────────┴───────────┴──────────────────────┘" | |
| ] | |
| }, | |
| "execution_count": 3, | |
| "metadata": {}, | |
| "output_type": "execute_result" | |
| } | |
| ], | |
| "source": [ | |
| "df = df.with_columns(\n", | |
| " text=pl.lit(\"# \")\n", | |
| " + pl.col(\"Album\")\n", | |
| " + pl.lit(\": \")\n", | |
| " + pl.col(\"Song\")\n", | |
| " + pl.lit(\"\\n\\n\")\n", | |
| " + pl.col(\"Lyrics\")\n", | |
| ")\n", | |
| "\n", | |
| "df.head()" | |
| ] | |
| }, | |
| { | |
| "cell_type": "markdown", | |
| "id": "6bc94fc5-a93e-4a7a-9326-6c6c99152b55", | |
| "metadata": {}, | |
| "source": [ | |
| "## Create Knowledge Graph" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 4, | |
| "id": "fde0b271-d7bf-459f-9a6d-7edd70b5c117", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:53.179484Z", | |
| "iopub.status.busy": "2024-09-23T11:56:53.179004Z", | |
| "iopub.status.idle": "2024-09-23T11:56:53.342379Z", | |
| "shell.execute_reply": "2024-09-23T11:56:53.341055Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:53.179437Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "# Initialize database\n", | |
| "import shutil\n", | |
| "import kuzu\n", | |
| "\n", | |
| "shutil.rmtree(\"test_kuzudb\", ignore_errors=True) # clean up\n", | |
| "db = kuzu.Database(\"test_kuzudb\")\n", | |
| "conn = kuzu.Connection(db)" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 5, | |
| "id": "d0be744d-8243-4ea4-afe3-8fd9daac1272", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:53.343442Z", | |
| "iopub.status.busy": "2024-09-23T11:56:53.343254Z", | |
| "iopub.status.idle": "2024-09-23T11:56:53.424335Z", | |
| "shell.execute_reply": "2024-09-23T11:56:53.423008Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:53.343426Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "# Create schema\n", | |
| "conn.execute(\"CREATE NODE TABLE Artist(name STRING, PRIMARY KEY (name))\")\n", | |
| "conn.execute(\"CREATE NODE TABLE Album(name STRING, PRIMARY KEY (name))\")\n", | |
| "conn.execute(\"CREATE NODE TABLE Song(ID SERIAL, name STRING, lyrics STRING, PRIMARY KEY(ID))\")\n", | |
| "conn.execute(\"CREATE REL TABLE InAlbum(FROM Song TO Album)\")\n", | |
| "conn.execute(\"CREATE REL TABLE FromArtist(FROM Album TO Artist)\");" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 6, | |
| "id": "7dea297e-64d8-49a9-bf76-3ab9912fa8c6", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:53.429068Z", | |
| "iopub.status.busy": "2024-09-23T11:56:53.428602Z", | |
| "iopub.status.idle": "2024-09-23T11:56:54.928527Z", | |
| "shell.execute_reply": "2024-09-23T11:56:54.926966Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:53.429020Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "# Insert nodes\n", | |
| "for artist in df[\"Artist\"].unique():\n", | |
| " conn.execute(f\"CREATE (artist:Artist {{name: '{artist}'}})\")\n", | |
| "\n", | |
| "for album in df[\"Album\"].unique():\n", | |
| " conn.execute(f\"\"\"CREATE (album:Album {{name: \"{album}\"}})\"\"\")\n", | |
| "\n", | |
| "for song, lyrics in df.select([\"Song\", \"text\"]).unique().rows():\n", | |
| " replaced_lyrics = lyrics.replace('\"', \"'\")\n", | |
| " conn.execute(\n", | |
| " f\"\"\"CREATE (song:Song {{name: \"{song}\", lyrics: \"{replaced_lyrics}\"}})\"\"\"\n", | |
| " )" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 7, | |
| "id": "0a088fff-f7de-4e9e-a5c7-7ad4e8639928", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:54.930970Z", | |
| "iopub.status.busy": "2024-09-23T11:56:54.930445Z", | |
| "iopub.status.idle": "2024-09-23T11:56:56.682977Z", | |
| "shell.execute_reply": "2024-09-23T11:56:56.681311Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:54.930920Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "# Insert edges\n", | |
| "for song, album, lyrics in df.select([\"Song\", \"Album\", \"text\"]).rows():\n", | |
| " replaced_lyrics = lyrics.replace('\"', \"'\")\n", | |
| " conn.execute(\n", | |
| " f\"\"\"\n", | |
| " MATCH (song:Song), (album:Album) WHERE song.name = \"{song}\" AND song.lyrics = \"{replaced_lyrics}\" AND album.name = \"{album}\"\n", | |
| " CREATE (song)-[:InAlbum]->(album)\n", | |
| " \"\"\"\n", | |
| " )\n", | |
| "\n", | |
| "for album, artist in df.select([\"Album\", \"Artist\"]).unique().rows():\n", | |
| " conn.execute(\n", | |
| " f\"\"\"\n", | |
| " MATCH (album:Album), (artist:Artist) WHERE album.name = \"{album}\" AND artist.name = \"{artist}\"\n", | |
| " CREATE (album)-[:FromArtist]->(artist)\n", | |
| " \"\"\"\n", | |
| " )" | |
| ] | |
| }, | |
| { | |
| "cell_type": "markdown", | |
| "id": "62434e8f-3def-453e-afb4-e4cbf4fcb7a9", | |
| "metadata": {}, | |
| "source": [ | |
| "## Verify things are working" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 8, | |
| "id": "efe9c577-fb7f-4fe7-a1a9-dbc3715fa0cf", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:56.685492Z", | |
| "iopub.status.busy": "2024-09-23T11:56:56.684971Z", | |
| "iopub.status.idle": "2024-09-23T11:56:56.871212Z", | |
| "shell.execute_reply": "2024-09-23T11:56:56.870716Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:56.685444Z" | |
| } | |
| }, | |
| "outputs": [ | |
| { | |
| "data": { | |
| "text/plain": [ | |
| "['Wherever I May Roam',\n", | |
| " 'Sad But True',\n", | |
| " 'Of Wolf And Man',\n", | |
| " \"Don't Tread On Me\",\n", | |
| " 'Enter Sandman',\n", | |
| " 'The God That Failed',\n", | |
| " 'Through The Never',\n", | |
| " 'Holier Than Thou',\n", | |
| " 'Nothing Else Matters',\n", | |
| " 'The Struggle Within',\n", | |
| " 'The Unforgiven',\n", | |
| " 'My Friend Of Misery']" | |
| ] | |
| }, | |
| "execution_count": 8, | |
| "metadata": {}, | |
| "output_type": "execute_result" | |
| } | |
| ], | |
| "source": [ | |
| "response = conn.execute(\n", | |
| " \"\"\"\n", | |
| " MATCH (a:Album {name: 'The Black Album'})<-[:InAlbum]-(s:Song) RETURN s.name\n", | |
| " \"\"\"\n", | |
| " )\n", | |
| "\n", | |
| "df_response = response.get_as_pl()\n", | |
| "\n", | |
| "df_response[\"s.name\"].to_list()" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 9, | |
| "id": "d8805d3b-db0a-40d3-a1aa-39e3062b59a2", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:56.872380Z", | |
| "iopub.status.busy": "2024-09-23T11:56:56.871959Z", | |
| "iopub.status.idle": "2024-09-23T11:56:56.875075Z", | |
| "shell.execute_reply": "2024-09-23T11:56:56.874616Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:56.872357Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "import nest_asyncio\n", | |
| "\n", | |
| "nest_asyncio.apply()" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 10, | |
| "id": "44160830-387e-4b58-a3e1-b1f0585cb67b", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:56.876083Z", | |
| "iopub.status.busy": "2024-09-23T11:56:56.875727Z", | |
| "iopub.status.idle": "2024-09-23T11:56:56.912546Z", | |
| "shell.execute_reply": "2024-09-23T11:56:56.911269Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:56.876064Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "from dotenv import find_dotenv, load_dotenv\n", | |
| "\n", | |
| "_ = load_dotenv(find_dotenv())" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 11, | |
| "id": "6a2413d1-0235-4382-a107-4218db45782b", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:56.914461Z", | |
| "iopub.status.busy": "2024-09-23T11:56:56.914018Z", | |
| "iopub.status.idle": "2024-09-23T11:56:56.920270Z", | |
| "shell.execute_reply": "2024-09-23T11:56:56.918990Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:56.914416Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "import os\n", | |
| "\n", | |
| "openai_api_key = os.getenv(\"OPENAI_API_KEY\")" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 12, | |
| "id": "5164d23b-15b3-4e10-a29d-58270e17078e", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:56.922175Z", | |
| "iopub.status.busy": "2024-09-23T11:56:56.921728Z", | |
| "iopub.status.idle": "2024-09-23T11:56:58.151578Z", | |
| "shell.execute_reply": "2024-09-23T11:56:58.150102Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:56.922130Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "from llama_index.graph_stores.kuzu import KuzuPropertyGraphStore\n", | |
| "\n", | |
| "graph_store = KuzuPropertyGraphStore(db)" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 13, | |
| "id": "4f09411f-45fc-4102-8ca7-8d2c71ea29ef", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:58.153946Z", | |
| "iopub.status.busy": "2024-09-23T11:56:58.153285Z", | |
| "iopub.status.idle": "2024-09-23T11:56:59.283391Z", | |
| "shell.execute_reply": "2024-09-23T11:56:59.282036Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:58.153895Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "from llama_index.core import PropertyGraphIndex\n", | |
| "from llama_index.embeddings.openai import OpenAIEmbedding\n", | |
| "from llama_index.llms.openai import OpenAI\n", | |
| "\n", | |
| "index = PropertyGraphIndex.from_existing(\n", | |
| " embed_model=OpenAIEmbedding(model_name=\"text-embedding-3-small\"),\n", | |
| " llm=OpenAI(model=\"gpt-4o-2024-08-06\", temperature=0.0),\n", | |
| " property_graph_store=graph_store,\n", | |
| ")" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 14, | |
| "id": "439c7edb-a9be-4935-95f3-af35fc70633f", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:59.284367Z", | |
| "iopub.status.busy": "2024-09-23T11:56:59.284197Z", | |
| "iopub.status.idle": "2024-09-23T11:56:59.288626Z", | |
| "shell.execute_reply": "2024-09-23T11:56:59.287929Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:59.284351Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "query_engine = index.as_query_engine(include_text=False)" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 15, | |
| "id": "14ed7c6c-bc8a-41fc-8424-f4d9a2ba8c09", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:56:59.289195Z", | |
| "iopub.status.busy": "2024-09-23T11:56:59.289057Z", | |
| "iopub.status.idle": "2024-09-23T11:57:00.539696Z", | |
| "shell.execute_reply": "2024-09-23T11:57:00.538108Z", | |
| "shell.execute_reply.started": "2024-09-23T11:56:59.289182Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "response = query_engine.query(\"How many songs does The Black Album have?\")" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 16, | |
| "id": "6d6eb8e0-dc68-43a9-9c6f-b0ce10ce4e23", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T11:57:00.541781Z", | |
| "iopub.status.busy": "2024-09-23T11:57:00.541400Z", | |
| "iopub.status.idle": "2024-09-23T11:57:00.548470Z", | |
| "shell.execute_reply": "2024-09-23T11:57:00.547511Z", | |
| "shell.execute_reply.started": "2024-09-23T11:57:00.541745Z" | |
| } | |
| }, | |
| "outputs": [ | |
| { | |
| "data": { | |
| "text/plain": [ | |
| "Response(response='Empty Response', source_nodes=[], metadata=None)" | |
| ] | |
| }, | |
| "execution_count": 16, | |
| "metadata": {}, | |
| "output_type": "execute_result" | |
| } | |
| ], | |
| "source": [ | |
| "response" | |
| ] | |
| } | |
| ], | |
| "metadata": { | |
| "kernelspec": { | |
| "display_name": "Python 3 (ipykernel)", | |
| "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.14" | |
| } | |
| }, | |
| "nbformat": 4, | |
| "nbformat_minor": 5 | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment