gmakstutis/graphiti
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

dedupe updates

Browse source
This commit is contained in:
prestonrasmussen 2025-03-21 12:10:56 -04:00
parent a95c046cbb
commit 948a0057fb
4 changed files with 2773 additions and 3364 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

802
tests/evals/data/LongMemEval_Snippetization.ipynb
View file

@ -2,7 +2,7 @@
"cells": [ "cells": [
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 2, "execution_count": null,
"metadata": { "metadata": {
"colab": { "colab": {
"base_uri": "https://localhost:8080/" "base_uri": "https://localhost:8080/"
@ -10,59 +10,7 @@
"id": "i9i9uUZ3pWQE", "id": "i9i9uUZ3pWQE",
"outputId": "84404bb8-5841-4f2f-dd87-7f909c6e95aa" "outputId": "84404bb8-5841-4f2f-dd87-7f909c6e95aa"
}, },
"outputs": [ "outputs": [],
{
"name": "stdout",
"output_type": "stream",
"text": [
"Collecting gdown\n",
" Using cached gdown-5.2.0-py3-none-any.whl.metadata (5.8 kB)\n",
"Requirement already satisfied: pandas in ./.venv/lib/python3.13/site-packages (2.2.3)\n",
"Collecting beautifulsoup4 (from gdown)\n",
" Using cached beautifulsoup4-4.12.3-py3-none-any.whl.metadata (3.8 kB)\n",
"Collecting filelock (from gdown)\n",
" Using cached filelock-3.16.1-py3-none-any.whl.metadata (2.9 kB)\n",
"Collecting requests[socks] (from gdown)\n",
" Using cached requests-2.32.3-py3-none-any.whl.metadata (4.6 kB)\n",
"Collecting tqdm (from gdown)\n",
" Using cached tqdm-4.67.1-py3-none-any.whl.metadata (57 kB)\n",
"Requirement already satisfied: numpy>=1.26.0 in ./.venv/lib/python3.13/site-packages (from pandas) (2.2.1)\n",
"Requirement already satisfied: python-dateutil>=2.8.2 in ./.venv/lib/python3.13/site-packages (from pandas) (2.9.0.post0)\n",
"Requirement already satisfied: pytz>=2020.1 in ./.venv/lib/python3.13/site-packages (from pandas) (2024.2)\n",
"Requirement already satisfied: tzdata>=2022.7 in ./.venv/lib/python3.13/site-packages (from pandas) (2024.2)\n",
"Requirement already satisfied: six>=1.5 in ./.venv/lib/python3.13/site-packages (from python-dateutil>=2.8.2->pandas) (1.17.0)\n",
"Collecting soupsieve>1.2 (from beautifulsoup4->gdown)\n",
" Using cached soupsieve-2.6-py3-none-any.whl.metadata (4.6 kB)\n",
"Collecting charset-normalizer<4,>=2 (from requests[socks]->gdown)\n",
" Downloading charset_normalizer-3.4.1-cp313-cp313-macosx_10_13_universal2.whl.metadata (35 kB)\n",
"Collecting idna<4,>=2.5 (from requests[socks]->gdown)\n",
" Using cached idna-3.10-py3-none-any.whl.metadata (10 kB)\n",
"Collecting urllib3<3,>=1.21.1 (from requests[socks]->gdown)\n",
" Downloading urllib3-2.3.0-py3-none-any.whl.metadata (6.5 kB)\n",
"Collecting certifi>=2017.4.17 (from requests[socks]->gdown)\n",
" Downloading certifi-2024.12.14-py3-none-any.whl.metadata (2.3 kB)\n",
"Collecting PySocks!=1.5.7,>=1.5.6 (from requests[socks]->gdown)\n",
" Using cached PySocks-1.7.1-py3-none-any.whl.metadata (13 kB)\n",
"Using cached gdown-5.2.0-py3-none-any.whl (18 kB)\n",
"Using cached beautifulsoup4-4.12.3-py3-none-any.whl (147 kB)\n",
"Using cached filelock-3.16.1-py3-none-any.whl (16 kB)\n",
"Using cached tqdm-4.67.1-py3-none-any.whl (78 kB)\n",
"Downloading certifi-2024.12.14-py3-none-any.whl (164 kB)\n",
"Downloading charset_normalizer-3.4.1-cp313-cp313-macosx_10_13_universal2.whl (195 kB)\n",
"Using cached idna-3.10-py3-none-any.whl (70 kB)\n",
"Using cached PySocks-1.7.1-py3-none-any.whl (16 kB)\n",
"Using cached soupsieve-2.6-py3-none-any.whl (36 kB)\n",
"Downloading urllib3-2.3.0-py3-none-any.whl (128 kB)\n",
"Using cached requests-2.32.3-py3-none-any.whl (64 kB)\n",
"Installing collected packages: urllib3, tqdm, soupsieve, PySocks, idna, filelock, charset-normalizer, certifi, requests, beautifulsoup4, gdown\n",
"Successfully installed PySocks-1.7.1 beautifulsoup4-4.12.3 certifi-2024.12.14 charset-normalizer-3.4.1 filelock-3.16.1 gdown-5.2.0 idna-3.10 requests-2.32.3 soupsieve-2.6 tqdm-4.67.1 urllib3-2.3.0\n",
"\n",
"\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.3.1\u001b[0m\n",
"\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
"Note: you may need to restart the kernel to use updated packages.\n"
]
}
],
"source": [ "source": [
"######## Installations - BE SURE TO MAKE YOUR OWN LOCAL VENV FIRST\n", "######## Installations - BE SURE TO MAKE YOUR OWN LOCAL VENV FIRST\n",
"\n", "\n",
@ -71,7 +19,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 3, "execution_count": null,
"metadata": { "metadata": {
"id": "KgzveXyAp35v" "id": "KgzveXyAp35v"
}, },
@ -79,18 +27,19 @@
"source": [ "source": [
"######## Imports\n", "######## Imports\n",
"\n", "\n",
"import pandas as pd\n", "import csv\n",
"import gdown\n",
"import tarfile\n",
"import os\n",
"import json\n", "import json\n",
"import os\n",
"import tarfile\n",
"from datetime import datetime\n", "from datetime import datetime\n",
"import csv" "\n",
"import gdown\n",
"import pandas as pd"
] ]
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 4, "execution_count": null,
"metadata": { "metadata": {
"colab": { "colab": {
"base_uri": "https://localhost:8080/" "base_uri": "https://localhost:8080/"
@ -98,16 +47,7 @@
"id": "iwcBxgxgqES8", "id": "iwcBxgxgqES8",
"outputId": "4af8de17-57bd-4857-f49c-2ee5d39ed248" "outputId": "4af8de17-57bd-4857-f49c-2ee5d39ed248"
}, },
"outputs": [ "outputs": [],
{
"name": "stdout",
"output_type": "stream",
"text": [
"'./longmemeval_data/longmemeval_data.tar.gz' already exists, skipping download.\n",
"'longmemeval_oracle.json' already exists, so skipping extraction.\n"
]
}
],
"source": [ "source": [
"######## Download the eval dataset from the official Google Drive source\n", "######## Download the eval dataset from the official Google Drive source\n",
"\n", "\n",
@ -136,175 +76,24 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 5, "execution_count": null,
"metadata": { "metadata": {
"id": "_RjEZnk5v530" "id": "_RjEZnk5v530"
}, },
"outputs": [ "outputs": [],
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>question_id</th>\n",
" <th>question_type</th>\n",
" <th>question</th>\n",
" <th>answer</th>\n",
" <th>question_date</th>\n",
" <th>haystack_dates</th>\n",
" <th>haystack_session_ids</th>\n",
" <th>haystack_sessions</th>\n",
" <th>answer_session_ids</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>0</th>\n",
" <td>gpt4_2655b836</td>\n",
" <td>temporal-reasoning</td>\n",
" <td>What was the first issue I had with my new car...</td>\n",
" <td>GPS system not functioning correctly</td>\n",
" <td>2023/04/10 (Mon) 23:07</td>\n",
" <td>[2023/04/10 (Mon) 17:50, 2023/04/10 (Mon) 14:4...</td>\n",
" <td>[answer_4be1b6b4_2, answer_4be1b6b4_3, answer_...</td>\n",
" <td>[[{'role': 'user', 'content': 'I'm thinking of...</td>\n",
" <td>[answer_4be1b6b4_2, answer_4be1b6b4_3, answer_...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1</th>\n",
" <td>gpt4_2487a7cb</td>\n",
" <td>temporal-reasoning</td>\n",
" <td>Which event did I attend first, the 'Effective...</td>\n",
" <td>'Data Analysis using Python' webinar</td>\n",
" <td>2023/05/28 (Sun) 06:47</td>\n",
" <td>[2023/05/28 (Sun) 21:04, 2023/05/28 (Sun) 07:17]</td>\n",
" <td>[answer_1c6b85ea_1, answer_1c6b85ea_2]</td>\n",
" <td>[[{'role': 'user', 'content': 'I'm trying to g...</td>\n",
" <td>[answer_1c6b85ea_1, answer_1c6b85ea_2]</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2</th>\n",
" <td>gpt4_76048e76</td>\n",
" <td>temporal-reasoning</td>\n",
" <td>Which vehicle did I take care of first in Febr...</td>\n",
" <td>bike</td>\n",
" <td>2023/03/10 (Fri) 23:15</td>\n",
" <td>[2023/03/10 (Fri) 22:50, 2023/03/10 (Fri) 08:11]</td>\n",
" <td>[answer_b535969f_2, answer_b535969f_1]</td>\n",
" <td>[[{'role': 'user', 'content': 'I'm thinking of...</td>\n",
" <td>[answer_b535969f_2, answer_b535969f_1]</td>\n",
" </tr>\n",
" <tr>\n",
" <th>3</th>\n",
" <td>gpt4_2312f94c</td>\n",
" <td>temporal-reasoning</td>\n",
" <td>Which device did I got first, the Samsung Gala...</td>\n",
" <td>Samsung Galaxy S22</td>\n",
" <td>2023/03/15 (Wed) 03:53</td>\n",
" <td>[2023/03/15 (Wed) 00:56, 2023/03/15 (Wed) 10:31]</td>\n",
" <td>[answer_5328c3c2_1, answer_5328c3c2_2]</td>\n",
" <td>[[{'role': 'user', 'content': 'I'm planning a ...</td>\n",
" <td>[answer_5328c3c2_1, answer_5328c3c2_2]</td>\n",
" </tr>\n",
" <tr>\n",
" <th>4</th>\n",
" <td>0bb5a684</td>\n",
" <td>temporal-reasoning</td>\n",
" <td>How many days before the team meeting I was pr...</td>\n",
" <td>7 days. 8 days (including the last day) is als...</td>\n",
" <td>2023/01/13 (Fri) 19:39</td>\n",
" <td>[2023/01/13 (Fri) 18:07, 2023/01/13 (Fri) 21:38]</td>\n",
" <td>[answer_e936197f_1, answer_e936197f_2]</td>\n",
" <td>[[{'role': 'user', 'content': 'I'm preparing f...</td>\n",
" <td>[answer_e936197f_1, answer_e936197f_2]</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" question_id question_type \\\n",
"0 gpt4_2655b836 temporal-reasoning \n",
"1 gpt4_2487a7cb temporal-reasoning \n",
"2 gpt4_76048e76 temporal-reasoning \n",
"3 gpt4_2312f94c temporal-reasoning \n",
"4 0bb5a684 temporal-reasoning \n",
"\n",
" question \\\n",
"0 What was the first issue I had with my new car... \n",
"1 Which event did I attend first, the 'Effective... \n",
"2 Which vehicle did I take care of first in Febr... \n",
"3 Which device did I got first, the Samsung Gala... \n",
"4 How many days before the team meeting I was pr... \n",
"\n",
" answer question_date \\\n",
"0 GPS system not functioning correctly 2023/04/10 (Mon) 23:07 \n",
"1 'Data Analysis using Python' webinar 2023/05/28 (Sun) 06:47 \n",
"2 bike 2023/03/10 (Fri) 23:15 \n",
"3 Samsung Galaxy S22 2023/03/15 (Wed) 03:53 \n",
"4 7 days. 8 days (including the last day) is als... 2023/01/13 (Fri) 19:39 \n",
"\n",
" haystack_dates \\\n",
"0 [2023/04/10 (Mon) 17:50, 2023/04/10 (Mon) 14:4... \n",
"1 [2023/05/28 (Sun) 21:04, 2023/05/28 (Sun) 07:17] \n",
"2 [2023/03/10 (Fri) 22:50, 2023/03/10 (Fri) 08:11] \n",
"3 [2023/03/15 (Wed) 00:56, 2023/03/15 (Wed) 10:31] \n",
"4 [2023/01/13 (Fri) 18:07, 2023/01/13 (Fri) 21:38] \n",
"\n",
" haystack_session_ids \\\n",
"0 [answer_4be1b6b4_2, answer_4be1b6b4_3, answer_... \n",
"1 [answer_1c6b85ea_1, answer_1c6b85ea_2] \n",
"2 [answer_b535969f_2, answer_b535969f_1] \n",
"3 [answer_5328c3c2_1, answer_5328c3c2_2] \n",
"4 [answer_e936197f_1, answer_e936197f_2] \n",
"\n",
" haystack_sessions \\\n",
"0 [[{'role': 'user', 'content': 'I'm thinking of... \n",
"1 [[{'role': 'user', 'content': 'I'm trying to g... \n",
"2 [[{'role': 'user', 'content': 'I'm thinking of... \n",
"3 [[{'role': 'user', 'content': 'I'm planning a ... \n",
"4 [[{'role': 'user', 'content': 'I'm preparing f... \n",
"\n",
" answer_session_ids \n",
"0 [answer_4be1b6b4_2, answer_4be1b6b4_3, answer_... \n",
"1 [answer_1c6b85ea_1, answer_1c6b85ea_2] \n",
"2 [answer_b535969f_2, answer_b535969f_1] \n",
"3 [answer_5328c3c2_1, answer_5328c3c2_2] \n",
"4 [answer_e936197f_1, answer_e936197f_2] "
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [ "source": [
"######## Load the eval dataset\n", "######## Load the eval dataset\n",
"\n", "\n",
"lme_dataset_option = os.path.join(folder_path, 'longmemeval_oracle.json') # Can be _oracle, _s, or _m\n", "lme_dataset_option = os.path.join(\n",
" folder_path, 'longmemeval_oracle.json'\n",
") # Can be _oracle, _s, or _m\n",
"lme_dataset_df = pd.read_json(lme_dataset_option)\n", "lme_dataset_df = pd.read_json(lme_dataset_option)\n",
"lme_dataset_df.head()" "lme_dataset_df.head()"
] ]
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 13, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@ -321,7 +110,6 @@
"\n", "\n",
" all_snippets = []\n", " all_snippets = []\n",
" for index, row in lme_dataset_df.iterrows():\n", " for index, row in lme_dataset_df.iterrows():\n",
"\n",
" question_id = row['question_id']\n", " question_id = row['question_id']\n",
"\n", "\n",
" # Extract the haystack_sessions and dates\n", " # Extract the haystack_sessions and dates\n",
@ -330,58 +118,61 @@
"\n", "\n",
" # Combine into list of dictionaries\n", " # Combine into list of dictionaries\n",
" sessions_data = [\n", " sessions_data = [\n",
" {\n", " {'session': session, 'date': datetime.strptime(date, '%Y/%m/%d (%a) %H:%M')}\n",
" \"session\": session,\n",
" \"date\": datetime.strptime(date, \"%Y/%m/%d (%a) %H:%M\")\n",
" } \n",
" for session, date in zip(sessions, session_dates)\n", " for session, date in zip(sessions, session_dates)\n",
" ]\n", " ]\n",
"\n", "\n",
" # Sort by date from earliest to latest\n", " # Sort by date from earliest to latest\n",
" sessions_data.sort(key=lambda x: x[\"date\"])\n", " sessions_data.sort(key=lambda x: x['date'])\n",
"\n",
"\n", "\n",
" all_snippets_this_session = []\n", " all_snippets_this_session = []\n",
"\n", "\n",
" message_index_across_sessions = 0\n", " message_index_across_sessions = 0\n",
" for session_index, session_and_date in enumerate(sessions_data):\n", " for session_index, session_and_date in enumerate(sessions_data):\n",
" for message_index_within_session, message in enumerate(session_and_date[\"session\"]):\n", " for message_index_within_session, message in enumerate(session_and_date['session']):\n",
" \n", " num_previous_messages = min(\n",
" num_previous_messages = min(max_num_previous_messages, message_index_across_sessions)\n", " max_num_previous_messages, message_index_across_sessions\n",
" previous_snippets = all_snippets_this_session[message_index_across_sessions-num_previous_messages:]\n", " )\n",
" previous_messages_only = [{\"role\":previous_snippet[\"message\"][\"role\"], \"content\":previous_snippet[\"message\"][\"content\"]} for previous_snippet in previous_snippets]\n", " previous_snippets = all_snippets_this_session[\n",
" message_index_across_sessions - num_previous_messages:\n",
" ]\n",
" previous_messages_only = [\n",
" {\n",
" 'role': previous_snippet['message']['role'],\n",
" 'content': previous_snippet['message']['content'],\n",
" }\n",
" for previous_snippet in previous_snippets\n",
" ]\n",
"\n", "\n",
" snippet = {\n", " snippet = {\n",
" \"question_id\": question_id,\n", " 'question_id': question_id,\n",
" \"question_type\": row[\"question_type\"],\n", " 'question_type': row['question_type'],\n",
" \"multisession_index\": index,\n", " 'multisession_index': index,\n",
" \"session_index\": session_index,\n", " 'session_index': session_index,\n",
" \"message_index_within_session\": message_index_within_session,\n", " 'message_index_within_session': message_index_within_session,\n",
" \"message_index_across_sessions\": message_index_across_sessions,\n", " 'message_index_across_sessions': message_index_across_sessions,\n",
" \"session_date\": session_and_date[\"date\"],\n", " 'session_date': session_and_date['date'],\n",
" \"message\": message,\n", " 'message': message,\n",
" \"previous_messages\": previous_messages_only,\n", " 'previous_messages': previous_messages_only,\n",
" \"num_previous_messages\": num_previous_messages,\n", " 'num_previous_messages': num_previous_messages,\n",
" }\n", " }\n",
"\n", "\n",
" if lme_filename == \"longmemeval_oracle.json\":\n", " if lme_filename == 'longmemeval_oracle.json':\n",
" snippet[\"message_has_answer\"] = message[\"has_answer\"]\n", " snippet['message_has_answer'] = message['has_answer']\n",
"\n", "\n",
" all_snippets_this_session.append(snippet)\n", " all_snippets_this_session.append(snippet)\n",
" message_index_across_sessions += 1\n", " message_index_across_sessions += 1\n",
"\n", "\n",
" all_snippets.extend(all_snippets_this_session)\n", " all_snippets.extend(all_snippets_this_session)\n",
"\n", "\n",
" \n", " snippetized_folder = os.path.join(folder_path, 'snippetized_data')\n",
"\n",
" snippetized_folder = os.path.join(folder_path, \"snippetized_data\")\n",
" if not os.path.exists(snippetized_folder):\n", " if not os.path.exists(snippetized_folder):\n",
" os.makedirs(snippetized_folder)\n", " os.makedirs(snippetized_folder)\n",
"\n", "\n",
" filename = lme_filename.replace(\".json\", \"_snippetized.csv\")\n", " filename = lme_filename.replace('.json', '_snippetized.csv')\n",
" filepath = os.path.join(snippetized_folder, filename)\n", " filepath = os.path.join(snippetized_folder, filename)\n",
"\n", "\n",
" with open(filepath, \"w\", newline=\"\") as csvfile:\n", " with open(filepath, 'w', newline='') as csvfile:\n",
" writer = csv.DictWriter(csvfile, fieldnames=all_snippets[0].keys())\n", " writer = csv.DictWriter(csvfile, fieldnames=all_snippets[0].keys())\n",
" writer.writeheader()\n", " writer.writeheader()\n",
" for snippet in all_snippets:\n", " for snippet in all_snippets:\n",
@ -390,32 +181,35 @@
" writer.writerow(snippet)\n", " writer.writerow(snippet)\n",
"\n", "\n",
"\n", "\n",
"\n",
"\n",
"def snippetize_and_check(lme_filename):\n", "def snippetize_and_check(lme_filename):\n",
" folder_path = './longmemeval_data/snippetized_data'\n", " folder_path = './longmemeval_data/snippetized_data'\n",
" file_path = os.path.join(folder_path, lme_filename.replace(\".json\", \"_snippetized.csv\"))\n", " file_path = os.path.join(folder_path, lme_filename.replace('.json', '_snippetized.csv'))\n",
" if not os.path.exists(file_path):\n", " if not os.path.exists(file_path):\n",
" print(f\"Snippetizing {lme_filename}...\")\n", " print(f'Snippetizing {lme_filename}...')\n",
" snippetize_lme_dataset(lme_filename)\n", " snippetize_lme_dataset(lme_filename)\n",
" else:\n", " else:\n",
" print(f\"Skipping snippetization for {lme_filename} because it already exists.\")\n", " print(f'Skipping snippetization for {lme_filename} because it already exists.')\n",
"\n", "\n",
" # Check first few rows of the csv\n", " # Check first few rows of the csv\n",
" df = pd.read_csv(file_path)\n", " df = pd.read_csv(file_path)\n",
" display(df.head(10))\n" " display(df.head(10))"
] ]
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 14, "execution_count": 8,
"metadata": {}, "metadata": {
"ExecuteTime": {
"end_time": "2025-03-20T17:17:33.985547Z",
"start_time": "2025-03-20T17:17:33.569496Z"
}
},
"outputs": [ "outputs": [
{ {
"name": "stdout", "name": "stdout",
"output_type": "stream", "output_type": "stream",
"text": [ "text": [
"Snippetizing longmemeval_oracle.json...\n" "Skipping snippetization for longmemeval_oracle.json because it already exists.\n"
] ]
}, },
{ {
@ -664,485 +458,27 @@
} }
], ],
"source": [ "source": [
"lme_filename = \"longmemeval_oracle.json\"\n", "lme_filename = 'longmemeval_oracle.json'\n",
"snippetize_and_check(lme_filename)" "snippetize_and_check(lme_filename)"
] ]
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 15, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [ "outputs": [],
{
"name": "stdout",
"output_type": "stream",
"text": [
"Snippetizing longmemeval_s.json...\n"
]
},
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>question_id</th>\n",
" <th>question_type</th>\n",
" <th>multisession_index</th>\n",
" <th>session_index</th>\n",
" <th>message_index_within_session</th>\n",
" <th>message_index_across_sessions</th>\n",
" <th>session_date</th>\n",
" <th>message</th>\n",
" <th>previous_messages</th>\n",
" <th>num_previous_messages</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>0</th>\n",
" <td>e47becba</td>\n",
" <td>single-session-user</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>2023-05-20 02:21:00</td>\n",
" <td>{\"role\": \"user\", \"content\": \"The farmer needs ...</td>\n",
" <td>[]</td>\n",
" <td>0</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1</th>\n",
" <td>e47becba</td>\n",
" <td>single-session-user</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>1</td>\n",
" <td>2023-05-20 02:21:00</td>\n",
" <td>{\"role\": \"assistant\", \"content\": \"To solve thi...</td>\n",
" <td>[{\"role\": \"user\", \"content\": \"The farmer needs...</td>\n",
" <td>1</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2</th>\n",
" <td>e47becba</td>\n",
" <td>single-session-user</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>0</td>\n",
" <td>2</td>\n",
" <td>2023-05-20 02:57:00</td>\n",
" <td>{\"role\": \"user\", \"content\": \"I'm trying to sta...</td>\n",
" <td>[{\"role\": \"user\", \"content\": \"The farmer needs...</td>\n",
" <td>2</td>\n",
" </tr>\n",
" <tr>\n",
" <th>3</th>\n",
" <td>e47becba</td>\n",
" <td>single-session-user</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>1</td>\n",
" <td>3</td>\n",
" <td>2023-05-20 02:57:00</td>\n",
" <td>{\"role\": \"assistant\", \"content\": \"Congratulati...</td>\n",
" <td>[{\"role\": \"user\", \"content\": \"The farmer needs...</td>\n",
" <td>3</td>\n",
" </tr>\n",
" <tr>\n",
" <th>4</th>\n",
" <td>e47becba</td>\n",
" <td>single-session-user</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>2</td>\n",
" <td>4</td>\n",
" <td>2023-05-20 02:57:00</td>\n",
" <td>{\"role\": \"user\", \"content\": \"I've been doing s...</td>\n",
" <td>[{\"role\": \"user\", \"content\": \"The farmer needs...</td>\n",
" <td>4</td>\n",
" </tr>\n",
" <tr>\n",
" <th>5</th>\n",
" <td>e47becba</td>\n",
" <td>single-session-user</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>3</td>\n",
" <td>5</td>\n",
" <td>2023-05-20 02:57:00</td>\n",
" <td>{\"role\": \"assistant\", \"content\": \"Yoga is an e...</td>\n",
" <td>[{\"role\": \"user\", \"content\": \"The farmer needs...</td>\n",
" <td>5</td>\n",
" </tr>\n",
" <tr>\n",
" <th>6</th>\n",
" <td>e47becba</td>\n",
" <td>single-session-user</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>4</td>\n",
" <td>6</td>\n",
" <td>2023-05-20 02:57:00</td>\n",
" <td>{\"role\": \"user\", \"content\": \"That's really hel...</td>\n",
" <td>[{\"role\": \"assistant\", \"content\": \"To solve th...</td>\n",
" <td>5</td>\n",
" </tr>\n",
" <tr>\n",
" <th>7</th>\n",
" <td>e47becba</td>\n",
" <td>single-session-user</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>5</td>\n",
" <td>7</td>\n",
" <td>2023-05-20 02:57:00</td>\n",
" <td>{\"role\": \"assistant\", \"content\": \"Foam rolling...</td>\n",
" <td>[{\"role\": \"user\", \"content\": \"I'm trying to st...</td>\n",
" <td>5</td>\n",
" </tr>\n",
" <tr>\n",
" <th>8</th>\n",
" <td>e47becba</td>\n",
" <td>single-session-user</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>6</td>\n",
" <td>8</td>\n",
" <td>2023-05-20 02:57:00</td>\n",
" <td>{\"role\": \"user\", \"content\": \"I've also been tr...</td>\n",
" <td>[{\"role\": \"assistant\", \"content\": \"Congratulat...</td>\n",
" <td>5</td>\n",
" </tr>\n",
" <tr>\n",
" <th>9</th>\n",
" <td>e47becba</td>\n",
" <td>single-session-user</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>7</td>\n",
" <td>9</td>\n",
" <td>2023-05-20 02:57:00</td>\n",
" <td>{\"role\": \"assistant\", \"content\": \"Monitoring y...</td>\n",
" <td>[{\"role\": \"user\", \"content\": \"I've been doing ...</td>\n",
" <td>5</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" question_id question_type multisession_index session_index \\\n",
"0 e47becba single-session-user 0 0 \n",
"1 e47becba single-session-user 0 0 \n",
"2 e47becba single-session-user 0 1 \n",
"3 e47becba single-session-user 0 1 \n",
"4 e47becba single-session-user 0 1 \n",
"5 e47becba single-session-user 0 1 \n",
"6 e47becba single-session-user 0 1 \n",
"7 e47becba single-session-user 0 1 \n",
"8 e47becba single-session-user 0 1 \n",
"9 e47becba single-session-user 0 1 \n",
"\n",
" message_index_within_session message_index_across_sessions \\\n",
"0 0 0 \n",
"1 1 1 \n",
"2 0 2 \n",
"3 1 3 \n",
"4 2 4 \n",
"5 3 5 \n",
"6 4 6 \n",
"7 5 7 \n",
"8 6 8 \n",
"9 7 9 \n",
"\n",
" session_date message \\\n",
"0 2023-05-20 02:21:00 {\"role\": \"user\", \"content\": \"The farmer needs ... \n",
"1 2023-05-20 02:21:00 {\"role\": \"assistant\", \"content\": \"To solve thi... \n",
"2 2023-05-20 02:57:00 {\"role\": \"user\", \"content\": \"I'm trying to sta... \n",
"3 2023-05-20 02:57:00 {\"role\": \"assistant\", \"content\": \"Congratulati... \n",
"4 2023-05-20 02:57:00 {\"role\": \"user\", \"content\": \"I've been doing s... \n",
"5 2023-05-20 02:57:00 {\"role\": \"assistant\", \"content\": \"Yoga is an e... \n",
"6 2023-05-20 02:57:00 {\"role\": \"user\", \"content\": \"That's really hel... \n",
"7 2023-05-20 02:57:00 {\"role\": \"assistant\", \"content\": \"Foam rolling... \n",
"8 2023-05-20 02:57:00 {\"role\": \"user\", \"content\": \"I've also been tr... \n",
"9 2023-05-20 02:57:00 {\"role\": \"assistant\", \"content\": \"Monitoring y... \n",
"\n",
" previous_messages num_previous_messages \n",
"0 [] 0 \n",
"1 [{\"role\": \"user\", \"content\": \"The farmer needs... 1 \n",
"2 [{\"role\": \"user\", \"content\": \"The farmer needs... 2 \n",
"3 [{\"role\": \"user\", \"content\": \"The farmer needs... 3 \n",
"4 [{\"role\": \"user\", \"content\": \"The farmer needs... 4 \n",
"5 [{\"role\": \"user\", \"content\": \"The farmer needs... 5 \n",
"6 [{\"role\": \"assistant\", \"content\": \"To solve th... 5 \n",
"7 [{\"role\": \"user\", \"content\": \"I'm trying to st... 5 \n",
"8 [{\"role\": \"assistant\", \"content\": \"Congratulat... 5 \n",
"9 [{\"role\": \"user\", \"content\": \"I've been doing ... 5 "
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [ "source": [
"lme_filename = \"longmemeval_s.json\"\n", "lme_filename = 'longmemeval_s.json'\n",
"snippetize_and_check(lme_filename)" "snippetize_and_check(lme_filename)"
] ]
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 16, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [ "outputs": [],
{
"name": "stdout",
"output_type": "stream",
"text": [
"Snippetizing longmemeval_m.json...\n"
]
},
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>question_id</th>\n",
" <th>question_type</th>\n",
" <th>multisession_index</th>\n",
" <th>session_index</th>\n",
" <th>message_index_within_session</th>\n",
" <th>message_index_across_sessions</th>\n",
" <th>session_date</th>\n",
" <th>message</th>\n",
" <th>previous_messages</th>\n",
" <th>num_previous_messages</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>0</th>\n",
" <td>7161e7e2</td>\n",
" <td>single-session-assistant</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>2023-05-20 00:04:00</td>\n",
" <td>{\"role\": \"user\", \"content\": \"Can you provide t...</td>\n",
" <td>[]</td>\n",
" <td>0</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1</th>\n",
" <td>7161e7e2</td>\n",
" <td>single-session-assistant</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>1</td>\n",
" <td>2023-05-20 00:04:00</td>\n",
" <td>{\"role\": \"assistant\", \"content\": \"Yes, here ar...</td>\n",
" <td>[{\"role\": \"user\", \"content\": \"Can you provide ...</td>\n",
" <td>1</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2</th>\n",
" <td>7161e7e2</td>\n",
" <td>single-session-assistant</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>2</td>\n",
" <td>2</td>\n",
" <td>2023-05-20 00:04:00</td>\n",
" <td>{\"role\": \"user\", \"content\": \"Wow, it's amazing...</td>\n",
" <td>[{\"role\": \"user\", \"content\": \"Can you provide ...</td>\n",
" <td>2</td>\n",
" </tr>\n",
" <tr>\n",
" <th>3</th>\n",
" <td>7161e7e2</td>\n",
" <td>single-session-assistant</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>3</td>\n",
" <td>3</td>\n",
" <td>2023-05-20 00:04:00</td>\n",
" <td>{\"role\": \"assistant\", \"content\": \"Yes, achievi...</td>\n",
" <td>[{\"role\": \"user\", \"content\": \"Can you provide ...</td>\n",
" <td>3</td>\n",
" </tr>\n",
" <tr>\n",
" <th>4</th>\n",
" <td>7161e7e2</td>\n",
" <td>single-session-assistant</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>4</td>\n",
" <td>4</td>\n",
" <td>2023-05-20 00:04:00</td>\n",
" <td>{\"role\": \"user\", \"content\": \"It's crazy to thi...</td>\n",
" <td>[{\"role\": \"user\", \"content\": \"Can you provide ...</td>\n",
" <td>4</td>\n",
" </tr>\n",
" <tr>\n",
" <th>5</th>\n",
" <td>7161e7e2</td>\n",
" <td>single-session-assistant</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>5</td>\n",
" <td>5</td>\n",
" <td>2023-05-20 00:04:00</td>\n",
" <td>{\"role\": \"assistant\", \"content\": \"Astronauts u...</td>\n",
" <td>[{\"role\": \"user\", \"content\": \"Can you provide ...</td>\n",
" <td>5</td>\n",
" </tr>\n",
" <tr>\n",
" <th>6</th>\n",
" <td>7161e7e2</td>\n",
" <td>single-session-assistant</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>0</td>\n",
" <td>6</td>\n",
" <td>2023-05-20 00:16:00</td>\n",
" <td>{\"role\": \"user\", \"content\": \"What about aliena...</td>\n",
" <td>[{\"role\": \"assistant\", \"content\": \"Yes, here a...</td>\n",
" <td>5</td>\n",
" </tr>\n",
" <tr>\n",
" <th>7</th>\n",
" <td>7161e7e2</td>\n",
" <td>single-session-assistant</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>1</td>\n",
" <td>7</td>\n",
" <td>2023-05-20 00:16:00</td>\n",
" <td>{\"role\": \"assistant\", \"content\": \"Alienation i...</td>\n",
" <td>[{\"role\": \"user\", \"content\": \"Wow, it's amazin...</td>\n",
" <td>5</td>\n",
" </tr>\n",
" <tr>\n",
" <th>8</th>\n",
" <td>7161e7e2</td>\n",
" <td>single-session-assistant</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>2</td>\n",
" <td>8</td>\n",
" <td>2023-05-20 00:16:00</td>\n",
" <td>{\"role\": \"user\", \"content\": \"Why you did not m...</td>\n",
" <td>[{\"role\": \"assistant\", \"content\": \"Yes, achiev...</td>\n",
" <td>5</td>\n",
" </tr>\n",
" <tr>\n",
" <th>9</th>\n",
" <td>7161e7e2</td>\n",
" <td>single-session-assistant</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>3</td>\n",
" <td>9</td>\n",
" <td>2023-05-20 00:16:00</td>\n",
" <td>{\"role\": \"assistant\", \"content\": \"I apologize ...</td>\n",
" <td>[{\"role\": \"user\", \"content\": \"It's crazy to th...</td>\n",
" <td>5</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" question_id question_type multisession_index session_index \\\n",
"0 7161e7e2 single-session-assistant 0 0 \n",
"1 7161e7e2 single-session-assistant 0 0 \n",
"2 7161e7e2 single-session-assistant 0 0 \n",
"3 7161e7e2 single-session-assistant 0 0 \n",
"4 7161e7e2 single-session-assistant 0 0 \n",
"5 7161e7e2 single-session-assistant 0 0 \n",
"6 7161e7e2 single-session-assistant 0 1 \n",
"7 7161e7e2 single-session-assistant 0 1 \n",
"8 7161e7e2 single-session-assistant 0 1 \n",
"9 7161e7e2 single-session-assistant 0 1 \n",
"\n",
" message_index_within_session message_index_across_sessions \\\n",
"0 0 0 \n",
"1 1 1 \n",
"2 2 2 \n",
"3 3 3 \n",
"4 4 4 \n",
"5 5 5 \n",
"6 0 6 \n",
"7 1 7 \n",
"8 2 8 \n",
"9 3 9 \n",
"\n",
" session_date message \\\n",
"0 2023-05-20 00:04:00 {\"role\": \"user\", \"content\": \"Can you provide t... \n",
"1 2023-05-20 00:04:00 {\"role\": \"assistant\", \"content\": \"Yes, here ar... \n",
"2 2023-05-20 00:04:00 {\"role\": \"user\", \"content\": \"Wow, it's amazing... \n",
"3 2023-05-20 00:04:00 {\"role\": \"assistant\", \"content\": \"Yes, achievi... \n",
"4 2023-05-20 00:04:00 {\"role\": \"user\", \"content\": \"It's crazy to thi... \n",
"5 2023-05-20 00:04:00 {\"role\": \"assistant\", \"content\": \"Astronauts u... \n",
"6 2023-05-20 00:16:00 {\"role\": \"user\", \"content\": \"What about aliena... \n",
"7 2023-05-20 00:16:00 {\"role\": \"assistant\", \"content\": \"Alienation i... \n",
"8 2023-05-20 00:16:00 {\"role\": \"user\", \"content\": \"Why you did not m... \n",
"9 2023-05-20 00:16:00 {\"role\": \"assistant\", \"content\": \"I apologize ... \n",
"\n",
" previous_messages num_previous_messages \n",
"0 [] 0 \n",
"1 [{\"role\": \"user\", \"content\": \"Can you provide ... 1 \n",
"2 [{\"role\": \"user\", \"content\": \"Can you provide ... 2 \n",
"3 [{\"role\": \"user\", \"content\": \"Can you provide ... 3 \n",
"4 [{\"role\": \"user\", \"content\": \"Can you provide ... 4 \n",
"5 [{\"role\": \"user\", \"content\": \"Can you provide ... 5 \n",
"6 [{\"role\": \"assistant\", \"content\": \"Yes, here a... 5 \n",
"7 [{\"role\": \"user\", \"content\": \"Wow, it's amazin... 5 \n",
"8 [{\"role\": \"assistant\", \"content\": \"Yes, achiev... 5 \n",
"9 [{\"role\": \"user\", \"content\": \"It's crazy to th... 5 "
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [ "source": [
"lme_filename = \"longmemeval_m.json\"\n", "lme_filename = 'longmemeval_m.json'\n",
"snippetize_and_check(lme_filename)" "snippetize_and_check(lme_filename)"
] ]
} }
@ -1152,9 +488,9 @@
"provenance": [] "provenance": []
}, },
"kernelspec": { "kernelspec": {
"display_name": "Python (.venv)", "display_name": "Python 3 (ipykernel)",
"language": "python", "language": "python",
"name": ".venv" "name": "python3"
}, },
"language_info": { "language_info": {
"codemirror_mode": { "codemirror_mode": {

205
tests/evals/data/LongMemEval_mini_dataset_loading.ipynb
View file

@ -4,6 +4,10 @@
"cell_type": "code", "cell_type": "code",
"execution_count": 1, "execution_count": 1,
"metadata": { "metadata": {
"ExecuteTime": {
"end_time": "2025-03-20T16:16:43.665690Z",
"start_time": "2025-03-20T16:16:42.923920Z"
},
"colab": { "colab": {
"base_uri": "https://localhost:8080/" "base_uri": "https://localhost:8080/"
}, },
@ -15,35 +19,35 @@
"name": "stdout", "name": "stdout",
"output_type": "stream", "output_type": "stream",
"text": [ "text": [
"Requirement already satisfied: pandas in ./.venv/lib/python3.13/site-packages (2.2.3)\n", "Requirement already satisfied: pandas in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (2.2.2)\r\n",
"Requirement already satisfied: graphiti-core in ./.venv/lib/python3.13/site-packages (0.5.1)\n", "Requirement already satisfied: graphiti-core in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (0.4.0)\r\n",
"Requirement already satisfied: numpy>=1.26.0 in ./.venv/lib/python3.13/site-packages (from pandas) (2.2.1)\n", "Requirement already satisfied: numpy>=1.26.0 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from pandas) (2.1.3)\r\n",
"Requirement already satisfied: python-dateutil>=2.8.2 in ./.venv/lib/python3.13/site-packages (from pandas) (2.9.0.post0)\n", "Requirement already satisfied: python-dateutil>=2.8.2 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from pandas) (2.9.0.post0)\r\n",
"Requirement already satisfied: pytz>=2020.1 in ./.venv/lib/python3.13/site-packages (from pandas) (2024.2)\n", "Requirement already satisfied: pytz>=2020.1 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from pandas) (2024.2)\r\n",
"Requirement already satisfied: tzdata>=2022.7 in ./.venv/lib/python3.13/site-packages (from pandas) (2024.2)\n", "Requirement already satisfied: tzdata>=2022.7 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from pandas) (2024.1)\r\n",
"Requirement already satisfied: diskcache<6.0.0,>=5.6.3 in ./.venv/lib/python3.13/site-packages (from graphiti-core) (5.6.3)\n", "Requirement already satisfied: diskcache<6.0.0,>=5.6.3 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from graphiti-core) (5.6.3)\r\n",
"Requirement already satisfied: neo4j<6.0.0,>=5.23.0 in ./.venv/lib/python3.13/site-packages (from graphiti-core) (5.27.0)\n", "Requirement already satisfied: neo4j<6.0.0,>=5.23.0 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from graphiti-core) (5.26.0)\r\n",
"Requirement already satisfied: openai<2.0.0,>=1.53.0 in ./.venv/lib/python3.13/site-packages (from graphiti-core) (1.59.3)\n", "Requirement already satisfied: openai<2.0.0,>=1.53.0 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from graphiti-core) (1.54.3)\r\n",
"Requirement already satisfied: pydantic<3.0.0,>=2.8.2 in ./.venv/lib/python3.13/site-packages (from graphiti-core) (2.10.4)\n", "Requirement already satisfied: pydantic<3.0.0,>=2.8.2 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from graphiti-core) (2.9.2)\r\n",
"Requirement already satisfied: python-dotenv<2.0.0,>=1.0.1 in ./.venv/lib/python3.13/site-packages (from graphiti-core) (1.0.1)\n", "Requirement already satisfied: python-dotenv<2.0.0,>=1.0.1 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from graphiti-core) (1.0.1)\r\n",
"Requirement already satisfied: tenacity==9.0.0 in ./.venv/lib/python3.13/site-packages (from graphiti-core) (9.0.0)\n", "Requirement already satisfied: tenacity==9.0.0 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from graphiti-core) (9.0.0)\r\n",
"Requirement already satisfied: anyio<5,>=3.5.0 in ./.venv/lib/python3.13/site-packages (from openai<2.0.0,>=1.53.0->graphiti-core) (4.7.0)\n", "Requirement already satisfied: anyio<5,>=3.5.0 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from openai<2.0.0,>=1.53.0->graphiti-core) (4.6.2.post1)\r\n",
"Requirement already satisfied: distro<2,>=1.7.0 in ./.venv/lib/python3.13/site-packages (from openai<2.0.0,>=1.53.0->graphiti-core) (1.9.0)\n", "Requirement already satisfied: distro<2,>=1.7.0 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from openai<2.0.0,>=1.53.0->graphiti-core) (1.9.0)\r\n",
"Requirement already satisfied: httpx<1,>=0.23.0 in ./.venv/lib/python3.13/site-packages (from openai<2.0.0,>=1.53.0->graphiti-core) (0.28.1)\n", "Requirement already satisfied: httpx<1,>=0.23.0 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from openai<2.0.0,>=1.53.0->graphiti-core) (0.27.2)\r\n",
"Requirement already satisfied: jiter<1,>=0.4.0 in ./.venv/lib/python3.13/site-packages (from openai<2.0.0,>=1.53.0->graphiti-core) (0.8.2)\n", "Requirement already satisfied: jiter<1,>=0.4.0 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from openai<2.0.0,>=1.53.0->graphiti-core) (0.7.0)\r\n",
"Requirement already satisfied: sniffio in ./.venv/lib/python3.13/site-packages (from openai<2.0.0,>=1.53.0->graphiti-core) (1.3.1)\n", "Requirement already satisfied: sniffio in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from openai<2.0.0,>=1.53.0->graphiti-core) (1.3.1)\r\n",
"Requirement already satisfied: tqdm>4 in ./.venv/lib/python3.13/site-packages (from openai<2.0.0,>=1.53.0->graphiti-core) (4.67.1)\n", "Requirement already satisfied: tqdm>4 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from openai<2.0.0,>=1.53.0->graphiti-core) (4.67.0)\r\n",
"Requirement already satisfied: typing-extensions<5,>=4.11 in ./.venv/lib/python3.13/site-packages (from openai<2.0.0,>=1.53.0->graphiti-core) (4.12.2)\n", "Requirement already satisfied: typing-extensions<5,>=4.11 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from openai<2.0.0,>=1.53.0->graphiti-core) (4.12.2)\r\n",
"Requirement already satisfied: annotated-types>=0.6.0 in ./.venv/lib/python3.13/site-packages (from pydantic<3.0.0,>=2.8.2->graphiti-core) (0.7.0)\n", "Requirement already satisfied: annotated-types>=0.6.0 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from pydantic<3.0.0,>=2.8.2->graphiti-core) (0.7.0)\r\n",
"Requirement already satisfied: pydantic-core==2.27.2 in ./.venv/lib/python3.13/site-packages (from pydantic<3.0.0,>=2.8.2->graphiti-core) (2.27.2)\n", "Requirement already satisfied: pydantic-core==2.23.4 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from pydantic<3.0.0,>=2.8.2->graphiti-core) (2.23.4)\r\n",
"Requirement already satisfied: six>=1.5 in ./.venv/lib/python3.13/site-packages (from python-dateutil>=2.8.2->pandas) (1.17.0)\n", "Requirement already satisfied: six>=1.5 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from python-dateutil>=2.8.2->pandas) (1.16.0)\r\n",
"Requirement already satisfied: idna>=2.8 in ./.venv/lib/python3.13/site-packages (from anyio<5,>=3.5.0->openai<2.0.0,>=1.53.0->graphiti-core) (3.10)\n", "Requirement already satisfied: idna>=2.8 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from anyio<5,>=3.5.0->openai<2.0.0,>=1.53.0->graphiti-core) (3.10)\r\n",
"Requirement already satisfied: certifi in ./.venv/lib/python3.13/site-packages (from httpx<1,>=0.23.0->openai<2.0.0,>=1.53.0->graphiti-core) (2024.12.14)\n", "Requirement already satisfied: certifi in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from httpx<1,>=0.23.0->openai<2.0.0,>=1.53.0->graphiti-core) (2024.8.30)\r\n",
"Requirement already satisfied: httpcore==1.* in ./.venv/lib/python3.13/site-packages (from httpx<1,>=0.23.0->openai<2.0.0,>=1.53.0->graphiti-core) (1.0.7)\n", "Requirement already satisfied: httpcore==1.* in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from httpx<1,>=0.23.0->openai<2.0.0,>=1.53.0->graphiti-core) (1.0.6)\r\n",
"Requirement already satisfied: h11<0.15,>=0.13 in ./.venv/lib/python3.13/site-packages (from httpcore==1.*->httpx<1,>=0.23.0->openai<2.0.0,>=1.53.0->graphiti-core) (0.14.0)\n", "Requirement already satisfied: h11<0.15,>=0.13 in /Users/prestonrasmussen/Library/Caches/pypoetry/virtualenvs/graphiti-core-XzHUgKi9-py3.12/lib/python3.12/site-packages (from httpcore==1.*->httpx<1,>=0.23.0->openai<2.0.0,>=1.53.0->graphiti-core) (0.14.0)\r\n",
"\n", "\r\n",
"\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.3.1\u001b[0m\n", "\u001B[1m[\u001B[0m\u001B[34;49mnotice\u001B[0m\u001B[1;39;49m]\u001B[0m\u001B[39;49m A new release of pip is available: \u001B[0m\u001B[31;49m24.1\u001B[0m\u001B[39;49m -> \u001B[0m\u001B[32;49m25.0.1\u001B[0m\r\n",
"\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n", "\u001B[1m[\u001B[0m\u001B[34;49mnotice\u001B[0m\u001B[1;39;49m]\u001B[0m\u001B[39;49m To update, run: \u001B[0m\u001B[32;49mpip install --upgrade pip\u001B[0m\r\n",
"Note: you may need to restart the kernel to use updated packages.\n" "Note: you may need to restart the kernel to use updated packages.\n"
] ]
} }
@ -58,18 +62,24 @@
"cell_type": "code", "cell_type": "code",
"execution_count": 2, "execution_count": 2,
"metadata": { "metadata": {
"ExecuteTime": {
"end_time": "2025-03-20T16:16:48.236627Z",
"start_time": "2025-03-20T16:16:47.667113Z"
},
"id": "KgzveXyAp35v" "id": "KgzveXyAp35v"
}, },
"outputs": [], "outputs": [],
"source": [ "source": [
"######## Imports\n", "######## Imports\n",
"\n", "\n",
"import pandas as pd\n",
"import os\n",
"import json\n",
"import csv\n", "import csv\n",
"import json\n",
"import os\n",
"from datetime import datetime\n", "from datetime import datetime\n",
"from utils import ingest_and_label_minidataset, dict_to_entity, entity_to_dict\n", "\n",
"import pandas as pd\n",
"from utils import dict_to_entity, entity_to_dict, ingest_and_label_minidataset\n",
"\n",
"from graphiti_core.llm_client import OpenAIClient\n", "from graphiti_core.llm_client import OpenAIClient\n",
"from graphiti_core.llm_client.config import LLMConfig\n", "from graphiti_core.llm_client.config import LLMConfig\n",
"from graphiti_core.nodes import EntityNode" "from graphiti_core.nodes import EntityNode"
@ -77,8 +87,12 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 3, "execution_count": 4,
"metadata": { "metadata": {
"ExecuteTime": {
"end_time": "2025-03-20T17:18:04.504351Z",
"start_time": "2025-03-20T17:18:04.080796Z"
},
"id": "_RjEZnk5v530" "id": "_RjEZnk5v530"
}, },
"outputs": [ "outputs": [
@ -228,7 +242,7 @@
"4 False " "4 False "
] ]
}, },
"execution_count": 3, "execution_count": 4,
"metadata": {}, "metadata": {},
"output_type": "execute_result" "output_type": "execute_result"
} }
@ -289,19 +303,26 @@
" return df.drop('previous_messages', axis=1)\n", " return df.drop('previous_messages', axis=1)\n",
"\n", "\n",
"\n", "\n",
"\n",
"def make_messages_readable(df):\n", "def make_messages_readable(df):\n",
" \"\"\"\n", " \"\"\"\n",
" Makes the messages more readable.\n", " Makes the messages more readable.\n",
" \"\"\"\n", " \"\"\"\n",
" for i in range(MAX_NUM_PREVIOUS_MESSAGES):\n", " for i in range(MAX_NUM_PREVIOUS_MESSAGES):\n",
" df[f'previous_message_{i + 1}'] = df[f'previous_message_{i + 1}'].apply(\n", " df[f'previous_message_{i + 1}'] = df[f'previous_message_{i + 1}'].apply(\n",
" lambda x: \"|\"*10 + f\" {x['role']} \" + \"|\"*10 + \"\\n\\n\" + f\"{x['content']}\" if x is not None else None\n", " lambda x: '|' * 10 + f\" {x['role']} \" + '|' * 10 + '\\n\\n' + f\"{x['content']}\"\n",
" if x is not None\n",
" else None\n",
" )\n", " )\n",
"\n", "\n",
" df['message'] = df.apply(\n", " df['message'] = df.apply(\n",
" lambda row: \"|\"*10 + f\" {row['message_role']} \" + \"|\"*10 + \"\\n\\n\" + f\"{row['message']}\" if row['message'] is not None else None,\n", " lambda row: '|' * 10\n",
" axis=1\n", " + f\" {row['message_role']} \"\n",
" + '|' * 10\n",
" + '\\n\\n'\n",
" + f\"{row['message']}\"\n",
" if row['message'] is not None\n",
" else None,\n",
" axis=1,\n",
" )\n", " )\n",
" return df\n", " return df\n",
"\n", "\n",
@ -310,40 +331,56 @@
" \"\"\"\n", " \"\"\"\n",
" Orders the columns in the way we want them.\n", " Orders the columns in the way we want them.\n",
" \"\"\"\n", " \"\"\"\n",
" df = df[['question_id', 'question_type', 'multisession_index', 'session_index', 'message_index_within_session', 'message_index_across_sessions', 'session_date', 'message_role', 'num_previous_messages', 'message_has_answer', 'previous_message_1', 'previous_message_2', 'previous_message_3', 'previous_message_4', 'previous_message_5', 'message']]\n", " df = df[\n",
" [\n",
" 'question_id',\n",
" 'question_type',\n",
" 'multisession_index',\n",
" 'session_index',\n",
" 'message_index_within_session',\n",
" 'message_index_across_sessions',\n",
" 'session_date',\n",
" 'message_role',\n",
" 'num_previous_messages',\n",
" 'message_has_answer',\n",
" 'previous_message_1',\n",
" 'previous_message_2',\n",
" 'previous_message_3',\n",
" 'previous_message_4',\n",
" 'previous_message_5',\n",
" 'message',\n",
" ]\n",
" ]\n",
" return df\n", " return df\n",
"\n", "\n",
"\n", "\n",
"\n",
"def insert_answer_columns(df, num_prompt_instructions):\n", "def insert_answer_columns(df, num_prompt_instructions):\n",
"\n",
" for prompt_instruction_index in range(num_prompt_instructions, 0, -1):\n", " for prompt_instruction_index in range(num_prompt_instructions, 0, -1):\n",
"\n",
" for i in range(MAX_NUM_PREVIOUS_MESSAGES, 0, -1):\n", " for i in range(MAX_NUM_PREVIOUS_MESSAGES, 0, -1):\n",
" # Insert after each previous message\n", " # Insert after each previous message\n",
" column_tag = f\"({prompt_instruction_index}.{i})\"\n", " column_tag = f'({prompt_instruction_index}.{i})'\n",
" df.insert(\n", " df.insert(\n",
" loc=df.columns.get_loc(f'previous_message_{i}') + 1,\n", " loc=df.columns.get_loc(f'previous_message_{i}') + 1,\n",
" column=f'Answer to Prompt Instruction {prompt_instruction_index} {column_tag}',\n", " column=f'Answer to Prompt Instruction {prompt_instruction_index} {column_tag}',\n",
" value=''\n", " value='',\n",
" )\n", " )\n",
" df.insert(\n", " df.insert(\n",
" loc=df.columns.get_loc(f'previous_message_{i}') + 2,\n", " loc=df.columns.get_loc(f'previous_message_{i}') + 2,\n",
" column=f'Done? {column_tag}',\n", " column=f'Done? {column_tag}',\n",
" value=''\n", " value='',\n",
" )\n", " )\n",
"\n", "\n",
" column_tag = f\"({prompt_instruction_index}.{MAX_NUM_PREVIOUS_MESSAGES+1})\"\n", " column_tag = f'({prompt_instruction_index}.{MAX_NUM_PREVIOUS_MESSAGES + 1})'\n",
" # Insert after each previous message\n", " # Insert after each previous message\n",
" df.insert(\n", " df.insert(\n",
" loc=df.columns.get_loc(f'message') + 1,\n", " loc=df.columns.get_loc(f'message') + 1,\n",
" column=f'Answer to Prompt Instruction {prompt_instruction_index} {column_tag}',\n", " column=f'Answer to Prompt Instruction {prompt_instruction_index} {column_tag}',\n",
" value=''\n", " value='',\n",
" )\n", " )\n",
" df.insert(\n", " df.insert(\n",
" loc=df.columns.get_loc(f'message') + 2,\n", " loc=df.columns.get_loc(f'message') + 2,\n",
" column=f'Done? {column_tag}',\n", " column=f'Done? {column_tag}',\n",
" value=''\n", " value='',\n",
" )\n", " )\n",
"\n", "\n",
"\n", "\n",
@ -352,30 +389,26 @@
" Inserts default answers for the first round of prompt instructions.\n", " Inserts default answers for the first round of prompt instructions.\n",
" \"\"\"\n", " \"\"\"\n",
" for i in range(MAX_NUM_PREVIOUS_MESSAGES, 0, -1):\n", " for i in range(MAX_NUM_PREVIOUS_MESSAGES, 0, -1):\n",
" column_tag = f\"(1.{i})\"\n", " column_tag = f'(1.{i})'\n",
" answer_col = f'Answer to Prompt Instruction 1 {column_tag}'\n", " answer_col = f'Answer to Prompt Instruction 1 {column_tag}'\n",
" msg_col = f'previous_message_{i}'\n", " msg_col = f'previous_message_{i}'\n",
"\n", "\n",
" # Set default value based on role from previous message\n", " # Set default value based on role from previous message\n",
" df[answer_col] = df[msg_col].apply(\n", " df[answer_col] = df[msg_col].apply(lambda x: f\"[${x['role']}$, ]\" if x is not None else '')\n",
" lambda x: f\"[${x['role']}$, ]\" if x is not None else \"\"\n",
" )\n",
"\n", "\n",
" # Handle the final message\n", " # Handle the final message\n",
" column_tag = f\"(1.{MAX_NUM_PREVIOUS_MESSAGES+1})\"\n", " column_tag = f'(1.{MAX_NUM_PREVIOUS_MESSAGES + 1})'\n",
" answer_col = f'Answer to Prompt Instruction 1 {column_tag}'\n", " answer_col = f'Answer to Prompt Instruction 1 {column_tag}'\n",
"\n", "\n",
" # Set default value based on role from current message\n", " # Set default value based on role from current message\n",
" df[answer_col] = df.apply(\n", " df[answer_col] = df.apply(lambda row: f\"[${row['message_role']}$, ]\", axis=1)\n",
" lambda row: f\"[${row['message_role']}$, ]\",\n", "\n",
" axis=1\n",
" )\n",
"\n", "\n",
"def insert_example_row(df, num_prompt_instructions):\n", "def insert_example_row(df, num_prompt_instructions):\n",
" \"\"\"\n", " \"\"\"\n",
" Inserts an example row at the top of the dataframe with 'EXAMPLE' as values.\n", " Inserts an example row at the top of the dataframe with 'EXAMPLE' as values.\n",
" \"\"\"\n", " \"\"\"\n",
" example_row = {col: \"EXAMPLE\" for col in df.columns}\n", " example_row = {col: 'EXAMPLE' for col in df.columns}\n",
" # for i in range(2):\n", " # for i in range(2):\n",
" # for j in range(num_prompt_instructions):\n", " # for j in range(num_prompt_instructions):\n",
" # example_row[f\"Done? ({j+1}.{i+1})\"] = \"x\"\n", " # example_row[f\"Done? ({j+1}.{i+1})\"] = \"x\"\n",
@ -385,13 +418,12 @@
" return df\n", " return df\n",
"\n", "\n",
"\n", "\n",
"\n",
"def transform_eval_minidataset(df):\n", "def transform_eval_minidataset(df):\n",
" \"\"\"\n", " \"\"\"\n",
" Transforms the eval mini dataset so that there is a row for every message in previous messages.\n", " Transforms the eval mini dataset so that there is a row for every message in previous messages.\n",
" \"\"\"\n", " \"\"\"\n",
"\n", "\n",
" df = df.rename(columns={col: f\"snippet_{col}\" for col in df.columns})\n", " df = df.rename(columns={col: f'snippet_{col}' for col in df.columns})\n",
"\n", "\n",
" ### Add new columns\n", " ### Add new columns\n",
" df.reset_index(drop=True, inplace=True)\n", " df.reset_index(drop=True, inplace=True)\n",
@ -418,7 +450,14 @@
" transformed_df = pd.DataFrame(transformed_rows)\n", " transformed_df = pd.DataFrame(transformed_rows)\n",
"\n", "\n",
" transformed_rows = []\n", " transformed_rows = []\n",
" task_names = [\"extract_nodes\", \"dedupe_nodes\", \"extract_edges\", \"dedupe_edges\", \"extract_edge_dates\", \"edge_invalidation\"]\n", " task_names = [\n",
" 'extract_nodes',\n",
" 'dedupe_nodes',\n",
" 'extract_edges',\n",
" 'dedupe_edges',\n",
" 'extract_edge_dates',\n",
" 'edge_invalidation',\n",
" ]\n",
" for _, row in transformed_df.iterrows():\n", " for _, row in transformed_df.iterrows():\n",
" for task_index, task_name in enumerate(task_names):\n", " for task_index, task_name in enumerate(task_names):\n",
" new_row = row.copy()\n", " new_row = row.copy()\n",
@ -429,12 +468,23 @@
" transformed_df = pd.DataFrame(transformed_rows)\n", " transformed_df = pd.DataFrame(transformed_rows)\n",
"\n", "\n",
" # Reorder columns\n", " # Reorder columns\n",
" transformed_df = transformed_df[['snippet_index', 'message_index_within_snippet', 'task_index', 'task_name', 'snippet_message', 'snippet_previous_messages', 'input_message', 'input_previous_messages']]#, 'input_extracted_nodes', 'input_existing_relevant_nodes', 'input_extracted_edges', 'input_existing_relevant_edges', 'output_zep', 'output_gpt4o', 'output_human']]\n", " transformed_df = transformed_df[\n",
" [\n",
" 'snippet_index',\n",
" 'message_index_within_snippet',\n",
" 'task_index',\n",
" 'task_name',\n",
" 'snippet_message',\n",
" 'snippet_previous_messages',\n",
" 'input_message',\n",
" 'input_previous_messages',\n",
" ]\n",
" ] # , 'input_extracted_nodes', 'input_existing_relevant_nodes', 'input_extracted_edges', 'input_existing_relevant_edges', 'output_zep', 'output_gpt4o', 'output_human']]\n",
"\n", "\n",
" # Ensure to reset the indices to be sequential\n", " # Ensure to reset the indices to be sequential\n",
" transformed_df.reset_index(drop=True, inplace=True)\n", " transformed_df.reset_index(drop=True, inplace=True)\n",
"\n", "\n",
" return transformed_df\n" " return transformed_df"
] ]
}, },
{ {
@ -446,7 +496,7 @@
"# Set pandas settings to display all columns and have max width of columns\n", "# Set pandas settings to display all columns and have max width of columns\n",
"pd.set_option('display.max_columns', 100)\n", "pd.set_option('display.max_columns', 100)\n",
"pd.set_option('display.max_rows', 100)\n", "pd.set_option('display.max_rows', 100)\n",
"pd.set_option('display.max_colwidth', 120)\n" "pd.set_option('display.max_colwidth', 120)"
] ]
}, },
{ {
@ -1223,8 +1273,8 @@
"eval_minidataset = lme_dataset_df_filtered.copy()\n", "eval_minidataset = lme_dataset_df_filtered.copy()\n",
"eval_minidataset = transform_eval_minidataset(eval_minidataset)\n", "eval_minidataset = transform_eval_minidataset(eval_minidataset)\n",
"# Print the number of rows and columns\n", "# Print the number of rows and columns\n",
"print(f\"Number of rows: {len(eval_minidataset)}\")\n", "print(f'Number of rows: {len(eval_minidataset)}')\n",
"print(f\"Number of columns: {len(eval_minidataset.columns)}\")\n", "print(f'Number of columns: {len(eval_minidataset.columns)}')\n",
"eval_minidataset.head(100)" "eval_minidataset.head(100)"
] ]
}, },
@ -1992,11 +2042,13 @@
")\n", ")\n",
"llm_client = OpenAIClient(config=llm_config)\n", "llm_client = OpenAIClient(config=llm_config)\n",
"output_column_name = 'output_gpt4o_mini'\n", "output_column_name = 'output_gpt4o_mini'\n",
"eval_minidataset_labelled = await ingest_and_label_minidataset(llm_client, eval_minidataset, output_column_name)\n", "eval_minidataset_labelled = await ingest_and_label_minidataset(\n",
" llm_client, eval_minidataset, output_column_name\n",
")\n",
"\n", "\n",
"# Print the number of rows and columns\n", "# Print the number of rows and columns\n",
"print(f\"Number of rows: {len(eval_minidataset_labelled)}\")\n", "print(f'Number of rows: {len(eval_minidataset_labelled)}')\n",
"print(f\"Number of columns: {len(eval_minidataset_labelled.columns)}\")\n", "print(f'Number of columns: {len(eval_minidataset_labelled.columns)}')\n",
"eval_minidataset_labelled.head(100)" "eval_minidataset_labelled.head(100)"
] ]
}, },
@ -2036,17 +2088,16 @@
"source": [ "source": [
"# print the input message column for the first row\n", "# print the input message column for the first row\n",
"index = 2\n", "index = 2\n",
"print(\"Input Message:\")\n", "print('Input Message:')\n",
"print(eval_minidataset_labelled.iloc[index]['input_message'])\n", "print(eval_minidataset_labelled.iloc[index]['input_message'])\n",
"print(\"-\"*100)\n", "print('-' * 100)\n",
"cell_value = eval_minidataset_labelled.iloc[index][output_column_name]\n", "cell_value = eval_minidataset_labelled.iloc[index][output_column_name]\n",
"cell_value_dicts = json.loads(cell_value)\n", "cell_value_dicts = json.loads(cell_value)\n",
"for dict in cell_value_dicts:\n", "for dict in cell_value_dicts:\n",
" # Print only the 'fact' and 'name' values\n", " # Print only the 'fact' and 'name' values\n",
" print(f\"Fact: {dict.get('fact', 'N/A')}, Name: {dict.get('name', 'N/A')}\")\n", " print(f\"Fact: {dict.get('fact', 'N/A')}, Name: {dict.get('name', 'N/A')}\")\n",
" print(\"-\"*100)\n", " print('-' * 100)\n",
" print(\"\\n\")\n", " print('\\n')"
"\n"
] ]
}, },
{ {
@ -2102,9 +2153,9 @@
"provenance": [] "provenance": []
}, },
"kernelspec": { "kernelspec": {
"display_name": "Python (.venv)", "display_name": "Python 3 (ipykernel)",
"language": "python", "language": "python",
"name": ".venv" "name": "python3"
}, },
"language_info": { "language_info": {
"codemirror_mode": { "codemirror_mode": {

80
tests/evals/data/utils.py
View file

@ -1,7 +1,10 @@
import pandas as pd
from datetime import datetime, timedelta
import json import json
from graphiti_core.nodes import EpisodicNode, EpisodeType, EntityNode from datetime import datetime, timedelta
import pandas as pd
from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
def create_episodes_from_messages(input_message, input_previous_messages): def create_episodes_from_messages(input_message, input_previous_messages):
""" """
@ -11,28 +14,31 @@ def create_episodes_from_messages(input_message, input_previous_messages):
current_time = datetime.now() current_time = datetime.now()
# Create the current episode # Create the current episode
role = input_message["role"] role = input_message['role']
content = input_message["content"] content = input_message['content']
message_content = f"{role}: {content}" message_content = f'{role}: {content}'
episode = EpisodicNode( episode = EpisodicNode(
name="", name='',
group_id="", group_id='',
source=EpisodeType.message, source=EpisodeType.message,
type=EpisodeType.message, type=EpisodeType.message,
source_description="", source_description='',
content=message_content, content=message_content,
valid_at=current_time, valid_at=current_time,
) )
# Create previous episodes # Create previous episodes
num_previous_messages = len(input_previous_messages) num_previous_messages = len(input_previous_messages)
previous_times = [current_time - timedelta(minutes=num_previous_messages - i) for i in range(num_previous_messages)] previous_times = [
current_time - timedelta(minutes=num_previous_messages - i)
for i in range(num_previous_messages)
]
previous_episodes = [ previous_episodes = [
EpisodicNode( EpisodicNode(
name="", name='',
group_id="", group_id='',
source=EpisodeType.message, source=EpisodeType.message,
source_description="", source_description='',
content=f"{message['role']}: {message['content']}", content=f"{message['role']}: {message['content']}",
valid_at=previous_time, valid_at=previous_time,
) )
@ -41,10 +47,14 @@ def create_episodes_from_messages(input_message, input_previous_messages):
return episode, previous_episodes return episode, previous_episodes
async def ingest_and_label_snippet(llm_client, snippet_df, output_column_name): async def ingest_and_label_snippet(llm_client, snippet_df, output_column_name):
# Import necessary functions # Import necessary functions
from graphiti_core.utils.maintenance.node_operations import extract_nodes, resolve_extracted_nodes
from graphiti_core.utils.maintenance.edge_operations import extract_edges from graphiti_core.utils.maintenance.edge_operations import extract_edges
from graphiti_core.utils.maintenance.node_operations import (
extract_nodes,
resolve_extracted_nodes,
)
# Loop through each unique message_index_within_snippet in sorted order # Loop through each unique message_index_within_snippet in sorted order
for message_index in sorted(snippet_df['message_index_within_snippet'].unique()): for message_index in sorted(snippet_df['message_index_within_snippet'].unique()):
@ -52,16 +62,24 @@ async def ingest_and_label_snippet(llm_client, snippet_df, output_column_name):
#### Process 'extract_nodes' task #### Process 'extract_nodes' task
extract_nodes_row = message_df[message_df['task_name'] == 'extract_nodes'] extract_nodes_row = message_df[message_df['task_name'] == 'extract_nodes']
assert len(extract_nodes_row) == 1, f"There should be exactly one row for 'extract_nodes' but there are {len(extract_nodes_row)}" assert (
len(extract_nodes_row) == 1
), f"There should be exactly one row for 'extract_nodes' but there are {len(extract_nodes_row)}"
input_message = json.loads(extract_nodes_row.iloc[0]['input_message']) input_message = json.loads(extract_nodes_row.iloc[0]['input_message'])
input_previous_messages = json.loads(extract_nodes_row.iloc[0]['input_previous_messages']) input_previous_messages = json.loads(extract_nodes_row.iloc[0]['input_previous_messages'])
episode, previous_episodes = create_episodes_from_messages(input_message, input_previous_messages) episode, previous_episodes = create_episodes_from_messages(
input_message, input_previous_messages
)
extracted_nodes = await extract_nodes(llm_client, episode, previous_episodes) extracted_nodes = await extract_nodes(llm_client, episode, previous_episodes)
snippet_df.at[extract_nodes_row.index[0], output_column_name] = json.dumps([entity_to_dict(node) for node in extracted_nodes]) snippet_df.at[extract_nodes_row.index[0], output_column_name] = json.dumps(
[entity_to_dict(node) for node in extracted_nodes]
)
#### Process 'dedupe_nodes' task #### Process 'dedupe_nodes' task
dedupe_nodes_row = message_df[message_df['task_name'] == 'dedupe_nodes'] dedupe_nodes_row = message_df[message_df['task_name'] == 'dedupe_nodes']
assert len(dedupe_nodes_row) == 1, "There should be exactly one row for 'dedupe_nodes' but there are {len(dedupe_nodes_row)}" assert (
len(dedupe_nodes_row) == 1
), "There should be exactly one row for 'dedupe_nodes' but there are {len(dedupe_nodes_row)}"
# Calculate existing nodes list # Calculate existing nodes list
existing_nodes = [] existing_nodes = []
@ -71,8 +89,8 @@ async def ingest_and_label_snippet(llm_client, snippet_df, output_column_name):
# Filter for previous messages with 'extract_nodes' task # Filter for previous messages with 'extract_nodes' task
prev_message_df = snippet_df[ prev_message_df = snippet_df[
(snippet_df['message_index_within_snippet'] == prev_message_index) & (snippet_df['message_index_within_snippet'] == prev_message_index)
(snippet_df['task_name'] == 'extract_nodes') & (snippet_df['task_name'] == 'extract_nodes')
] ]
# Retrieve and deserialize the nodes # Retrieve and deserialize the nodes
@ -82,12 +100,18 @@ async def ingest_and_label_snippet(llm_client, snippet_df, output_column_name):
existing_nodes.extend(nodes) existing_nodes.extend(nodes)
existing_nodes_lists = [existing_nodes for _ in range(len(extracted_nodes))] existing_nodes_lists = [existing_nodes for _ in range(len(extracted_nodes))]
resolved_nodes, uuid_map = await resolve_extracted_nodes(llm_client, extracted_nodes, existing_nodes_lists, episode, previous_episodes) resolved_nodes, uuid_map = await resolve_extracted_nodes(
snippet_df.at[dedupe_nodes_row.index[0], output_column_name] = json.dumps([entity_to_dict(node) for node in resolved_nodes]) llm_client, extracted_nodes, existing_nodes_lists, episode, previous_episodes
)
snippet_df.at[dedupe_nodes_row.index[0], output_column_name] = json.dumps(
[entity_to_dict(node) for node in resolved_nodes]
)
#### Process 'extract_edges' task #### Process 'extract_edges' task
extract_edges_row = message_df[message_df['task_name'] == 'extract_edges'] extract_edges_row = message_df[message_df['task_name'] == 'extract_edges']
assert len(extract_edges_row) == 1, f"There should be exactly one row for 'extract_edges' but there are {len(extract_edges_row)}" assert (
len(extract_edges_row) == 1
), f"There should be exactly one row for 'extract_edges' but there are {len(extract_edges_row)}"
extracted_edges = await extract_edges( extracted_edges = await extract_edges(
llm_client, llm_client,
episode, episode,
@ -95,7 +119,9 @@ async def ingest_and_label_snippet(llm_client, snippet_df, output_column_name):
previous_episodes, previous_episodes,
group_id='', group_id='',
) )
snippet_df.at[extract_edges_row.index[0], output_column_name] = json.dumps([entity_to_dict(edge) for edge in extracted_edges]) snippet_df.at[extract_edges_row.index[0], output_column_name] = json.dumps(
[entity_to_dict(edge) for edge in extracted_edges]
)
########## TODO: Complete the implementation of the below ########## TODO: Complete the implementation of the below
@ -133,7 +159,9 @@ async def ingest_and_label_minidataset(llm_client, minidataset_df, output_column
snippet_df = minidataset_df[minidataset_df['snippet_index'] == snippet_index] snippet_df = minidataset_df[minidataset_df['snippet_index'] == snippet_index]
# Pass the output column name to the ingest_and_label_snippet function # Pass the output column name to the ingest_and_label_snippet function
snippet_df_labelled = await ingest_and_label_snippet(llm_client, snippet_df, output_column_name) snippet_df_labelled = await ingest_and_label_snippet(
llm_client, snippet_df, output_column_name
)
if minidataset_labelled_df is None: if minidataset_labelled_df is None:
minidataset_labelled_df = snippet_df_labelled minidataset_labelled_df = snippet_df_labelled
@ -142,6 +170,7 @@ async def ingest_and_label_minidataset(llm_client, minidataset_df, output_column
return minidataset_labelled_df return minidataset_labelled_df
def entity_to_dict(entity): def entity_to_dict(entity):
""" """
Convert an entity object to a dictionary, handling datetime serialization. Convert an entity object to a dictionary, handling datetime serialization.
@ -152,6 +181,7 @@ def entity_to_dict(entity):
entity_dict[key] = value.isoformat() # Convert datetime to ISO 8601 string entity_dict[key] = value.isoformat() # Convert datetime to ISO 8601 string
return entity_dict return entity_dict
def dict_to_entity(entity_dict, entity_class): def dict_to_entity(entity_dict, entity_class):
""" """
Convert a dictionary back to an entity object, handling datetime deserialization. Convert a dictionary back to an entity object, handling datetime deserialization.

82
tests/evals/eval_extract_nodes.py
View file

@ -14,9 +14,9 @@ See the License for the specific language governing permissions and
limitations under the License. limitations under the License.
""" """
import os import csv # Add this import at the top of the file
import json import json
from tests.evals.utils import setup_logging, ingest_snippet import os
from datetime import datetime, timedelta from datetime import datetime, timedelta
import pytest import pytest
@ -24,17 +24,11 @@ from dotenv import load_dotenv
from graphiti_core.edges import EntityEdge, EpisodicEdge from graphiti_core.edges import EntityEdge, EpisodicEdge
from graphiti_core.graphiti import Graphiti from graphiti_core.graphiti import Graphiti
from graphiti_core.nodes import EntityNode, EpisodicNode
from graphiti_core.utils.maintenance.node_operations import extract_nodes
from graphiti_core.llm_client import OpenAIClient from graphiti_core.llm_client import OpenAIClient
from graphiti_core.llm_client.config import LLMConfig from graphiti_core.llm_client.config import LLMConfig
from graphiti_core.nodes import EpisodeType from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
from graphiti_core.utils.maintenance.node_operations import extract_nodes
import csv # Add this import at the top of the file from tests.evals.utils import ingest_snippet, setup_logging
############# EVERYTHING BELOW IS OUTDATED ############# EVERYTHING BELOW IS OUTDATED
@ -63,68 +57,67 @@ async def general_extract_nodes_test(llm_client, data_sample):
return hypothesis_node_names return hypothesis_node_names
def prepare_data_from_csv(data_file_name, question_id, session_idx, message_idx): def prepare_data_from_csv(data_file_name, question_id, session_idx, message_idx):
samples_csv_path = 'tests/evals/data/' + data_file_name + '.csv'
samples_csv_path = "tests/evals/data/" + data_file_name + ".csv"
# From CSV path, load everything # From CSV path, load everything
with open(samples_csv_path, 'r') as file: with open(samples_csv_path, 'r') as file:
csv_reader = csv.DictReader(file) csv_reader = csv.DictReader(file)
lme_samples = list(csv_reader) lme_samples = list(csv_reader)
data_samples = [] data_samples = []
# Loop through each row # Loop through each row
for row in lme_samples: for row in lme_samples:
### Prepare episode ### Prepare episode
current_time = datetime.now() current_time = datetime.now()
message = json.loads(row["message"]) message = json.loads(row['message'])
role = message["role"] role = message['role']
content = message["content"] content = message['content']
message_content = role + ": " + content message_content = role + ': ' + content
episode = EpisodicNode( episode = EpisodicNode(
name="", name='',
group_id="", group_id='',
source=EpisodeType.message, source=EpisodeType.message,
type=EpisodeType.message, type=EpisodeType.message,
source_description="", source_description='',
content=message_content, content=message_content,
valid_at=current_time, valid_at=current_time,
) )
### Prepare previous episodes ### Prepare previous episodes
previous_messages = json.loads(row["previous_messages"]) previous_messages = json.loads(row['previous_messages'])
num_previous_messages = len(previous_messages) num_previous_messages = len(previous_messages)
previous_times = [current_time - timedelta(minutes=num_previous_messages-i) for i in range(num_previous_messages)] previous_times = [
previous_episodes = [EpisodicNode( current_time - timedelta(minutes=num_previous_messages - i)
name="", for i in range(num_previous_messages)
group_id="", ]
previous_episodes = [
EpisodicNode(
name='',
group_id='',
source=EpisodeType.message, source=EpisodeType.message,
source_description="", source_description='',
content=message["role"] + ": " + message["content"], content=message['role'] + ': ' + message['content'],
valid_at=previous_time, valid_at=previous_time,
) for message, previous_time in zip(previous_messages, previous_times)] )
for message, previous_time in zip(previous_messages, previous_times)
]
### TODO: Prepare gold answer names ### TODO: Prepare gold answer names
### Add to data samples list ### Add to data samples list
data_samples.append({ data_samples.append(
"episode": episode, {
"previous_episodes": previous_episodes, 'episode': episode,
"gold_answer_names": [], 'previous_episodes': previous_episodes,
}) 'gold_answer_names': [],
}
)
return data_samples return data_samples
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_extract_nodes(): async def test_extract_nodes():
model_name = 'gpt-4o-mini' model_name = 'gpt-4o-mini'
@ -135,17 +128,16 @@ async def test_extract_nodes():
llm_client = OpenAIClient(config=llm_config) llm_client = OpenAIClient(config=llm_config)
data_file_name = 'output_short' data_file_name = 'output_short'
question_id = "gpt4_2655b836" question_id = 'gpt4_2655b836'
session_idx = 0 session_idx = 0
message_idx = 0 message_idx = 0
data_samples = prepare_data_from_csv(data_file_name, question_id, session_idx, message_idx) data_samples = prepare_data_from_csv(data_file_name, question_id, session_idx, message_idx)
for data_sample in data_samples: for data_sample in data_samples:
print(f"\n\nEpisode: {data_sample['episode']}") print(f"\n\nEpisode: {data_sample['episode']}")
print("*"*50) print('*' * 50)
print(f"Previous Episodes: {data_sample['previous_episodes']}") print(f"Previous Episodes: {data_sample['previous_episodes']}")
print("*"*50) print('*' * 50)
# print(f"Gold Answer Names: {gold_answer_names}") # print(f"Gold Answer Names: {gold_answer_names}")
await general_extract_nodes_test(llm_client, data_sample) await general_extract_nodes_test(llm_client, data_sample)