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

feat: adds graphiti demo notebook

Browse source
This commit is contained in:
hajdul88 2025-01-15 11:45:06 +01:00
parent dd8a488003
commit 1db44de7de
2 changed files with 262 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
examples/python/graphiti_example.py
View file

@ -68,5 +68,5 @@ async def main():
if __name__ == "__main__":
setup_logging(logging.INFO)
setup_logging(logging.ERROR)
asyncio.run(main())

261
notebooks/cognee_graphiti_demo.ipynb Normal file
View file

@ -0,0 +1,261 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": "# Cognee Graphiti integration demo"
},
{
"cell_type": "markdown",
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"source": "First we import the necessary libaries"
},
{
"cell_type": "code",
"metadata": {},
"source": [
"import cognee\n",
"import logging\n",
"import warnings\n",
"from cognee.modules.pipelines import Task, run_tasks\n",
"from cognee.shared.utils import setup_logging\n",
"from cognee.tasks.temporal_awareness import build_graph_with_temporal_awareness\n",
"from cognee.infrastructure.databases.relational import (\n",
" create_db_and_tables as create_relational_db_and_tables,\n",
")\n",
"from cognee.tasks.temporal_awareness.index_graphiti_objects import (\n",
" index_and_transform_graphiti_nodes_and_edges,\n",
")\n",
"from cognee.modules.retrieval.brute_force_triplet_search import brute_force_triplet_search\n",
"from cognee.tasks.completion.graph_query_completion import retrieved_edges_to_string\n",
"from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt\n",
"from cognee.infrastructure.llm.get_llm_client import get_llm_client"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Set environment variables"
]
},
{
"cell_type": "code",
"metadata": {
"ExecuteTime": {
"end_time": "2025-01-15T10:43:57.893763Z",
"start_time": "2025-01-15T10:43:57.891332Z"
}
},
"source": [
"import os\n",
"\n",
"# We ignore warnigns for now\n",
"warnings.filterwarnings('ignore')\n",
"\n",
"# API key for cognee\n",
"if \"LLM_API_KEY\" not in os.environ:\n",
" os.environ[\"LLM_API_KEY\"] = \"\"\n",
"\n",
"# API key for graphiti\n",
"if \"OPENAI_API_KEY\" not in os.environ:\n",
" os.environ[\"OPENAI_API_KEY\"] = \"\"\n",
"\n",
"# Graphiti integration is only tested with neo4j + pgvector + postgres for now\n",
"GRAPH_DATABASE_PROVIDER=\"neo4j\"\n",
"GRAPH_DATABASE_URL=\"bolt://localhost:7687\"\n",
"GRAPH_DATABASE_USERNAME=\"neo4j\"\n",
"GRAPH_DATABASE_PASSWORD=\"pleaseletmein\"\n",
"\n",
"os.environ[\"VECTOR_DB_PROVIDER\"] = \"pgvector\"\n",
"\n",
"os.environ[\"DB_PROVIDER\"] = \"postgres\"\n",
"\n",
"os.environ[\"DB_NAME\"] = \"cognee_db\"\n",
"\n",
"os.environ[\"DB_HOST\"]=\"127.0.0.1\"\n",
"os.environ[\"DB_PORT\"]=\"5432\"\n",
"os.environ[\"DB_USERNAME\"]=\"cognee\"\n",
"os.environ[\"DB_PASSWORD\"]=\"cognee\""
],
"outputs": [],
"execution_count": 2
},
{
"cell_type": "markdown",
"metadata": {},
"source": "## Input texts with temporal information"
},
{
"cell_type": "code",
"metadata": {
"ExecuteTime": {
"end_time": "2025-01-15T10:43:57.928664Z",
"start_time": "2025-01-15T10:43:57.927105Z"
}
},
"source": [
"text_list = [\n",
" \"Kamala Harris is the Attorney General of California. She was previously \"\n",
" \"the district attorney for San Francisco.\",\n",
" \"As AG, Harris was in office from January 3, 2011 – January 3, 2017\",\n",
"]"
],
"outputs": [],
"execution_count": 3
},
{
"cell_type": "markdown",
"metadata": {},
"source": "## Running graphiti + transforming its graph into cognee's core system (graph transformation + vector embeddings)"
},
{
"cell_type": "code",
"metadata": {
"ExecuteTime": {
"end_time": "2025-01-15T10:44:25.008501Z",
"start_time": "2025-01-15T10:43:57.932240Z"
}
},
"source": [
"# 🔧 Setting Up Logging to Suppress Errors\n",
"setup_logging(logging.ERROR) # Keeping logs clean and focused\n",
"\n",
"# 🧹 Pruning Old Data and Metadata\n",
"await cognee.prune.prune_data() # Removing outdated data\n",
"await cognee.prune.prune_system(metadata=True)\n",
"\n",
"# 🏗️ Creating Relational Database and Tables\n",
"await create_relational_db_and_tables()\n",
"\n",
"# 📚 Adding Text Data to Cognee\n",
"for text in text_list:\n",
" await cognee.add(text)\n",
"\n",
"# 🕰️ Building Temporal-Aware Graphs\n",
"tasks = [\n",
" Task(build_graph_with_temporal_awareness, text_list=text_list),\n",
"]\n",
"\n",
"# 🚀 Running the Task Pipeline\n",
"pipeline = run_tasks(tasks)\n",
"\n",
"# 🌟 Processing Pipeline Results\n",
"async for result in pipeline:\n",
" print(f\"✅ Result Processed: {result}\")\n",
"\n",
"# 🔄 Indexing and Transforming Graph Data\n",
"await index_and_transform_graphiti_nodes_and_edges()"
],
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Database deleted successfully.\n",
"Database deleted successfully.\n",
"User d3b51a32-38e1-4fe5-8270-6dc1d6ebfdf0 has registered.\n",
"Pipeline file_load_from_filesystem load step completed in 0.10 seconds\n",
"1 load package(s) were loaded to destination sqlalchemy and into dataset public\n",
"The sqlalchemy destination used postgresql://cognee:***@127.0.0.1:5432/cognee_db location to store data\n",
"Load package 1736937839.7739599 is LOADED and contains no failed jobs\n",
"Pipeline file_load_from_filesystem load step completed in 0.06 seconds\n",
"1 load package(s) were loaded to destination sqlalchemy and into dataset public\n",
"The sqlalchemy destination used postgresql://cognee:***@127.0.0.1:5432/cognee_db location to store data\n",
"Load package 1736937841.8467042 is LOADED and contains no failed jobs\n",
"Graph database initialized.\n",
"Added text: Kamala Harris is the Attorney Gener...\n",
"Added text: As AG, Harris was in office from Ja...\n",
"✅ Result Processed: <graphiti_core.graphiti.Graphiti object at 0x326fe0ce0>\n"
]
}
],
"execution_count": 4
},
{
"metadata": {},
"cell_type": "markdown",
"source": "## Retrieving and generating answer from graphiti graph with cognee retriever"
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-01-15T10:44:27.844438Z",
"start_time": "2025-01-15T10:44:25.013325Z"
}
},
"cell_type": "code",
"source": [
"# Step 1: Formulating the Query 🔍\n",
"query = \"When was Kamala Harris in office?\"\n",
"\n",
"# Step 2: Searching for Relevant Triplets 📊\n",
"triplets = await brute_force_triplet_search(\n",
" query=query,\n",
" top_k=3,\n",
" collections=[\"graphitinode_content\", \"graphitinode_name\", \"graphitinode_summary\"],\n",
")\n",
"\n",
"# Step 3: Preparing the Context for the LLM\n",
"context = retrieved_edges_to_string(triplets)\n",
"\n",
"args = {\n",
" \"question\": query,\n",
" \"context\": context\n",
"}\n",
"\n",
"# Step 4: Generating Prompts ✍️\n",
"user_prompt = render_prompt(\"graph_context_for_question.txt\", args)\n",
"system_prompt = read_query_prompt(\"answer_simple_question_restricted.txt\")\n",
"\n",
"# Step 5: Interacting with the LLM 🤖\n",
"llm_client = get_llm_client()\n",
"computed_answer = await llm_client.acreate_structured_output(\n",
" text_input=user_prompt, # Input prompt for the user context\n",
" system_prompt=system_prompt, # System-level instructions for the model\n",
" response_model=str,\n",
")\n",
"\n",
"# Step 6: Displaying the Computed Answer ✨\n",
"print(f\"💡 Answer: {computed_answer}\")"
],
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"💡 Answer: Kamala Harris was in office as Attorney General of California from January 3, 2011, to January 3, 2017.\n"
]
}
],
"execution_count": 5
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"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.9.6"
}
},
"nbformat": 4,
"nbformat_minor": 2
}