9ba2e0d6c1
<!-- .github/pull_request_template.md --> ## Description <!-- Provide a clear description of the changes in this PR --> ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Introduced enhanced visualization capabilities that let users launch a dedicated server for visual displays. - **Documentation** - Updated several interactive notebooks to include execution outputs and expanded explanatory content for better user guidance. - **Style** - Refined formatting and layout across notebooks to ensure consistent presentation and improved readability. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Igor Ilic <30923996+dexters1@users.noreply.github.com>
258 lines
8.2 KiB
Text
258 lines
8.2 KiB
Text
{
|
||
"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",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"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"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Set environment variables"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 2,
|
||
"metadata": {
|
||
"ExecuteTime": {
|
||
"end_time": "2025-01-15T10:43:57.893763Z",
|
||
"start_time": "2025-01-15T10:43:57.891332Z"
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"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\""
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": "## Input texts with temporal information"
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 3,
|
||
"metadata": {
|
||
"ExecuteTime": {
|
||
"end_time": "2025-01-15T10:43:57.928664Z",
|
||
"start_time": "2025-01-15T10:43:57.927105Z"
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"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",
|
||
"]"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": "## Running graphiti + transforming its graph into cognee's core system (graph transformation + vector embeddings)"
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 4,
|
||
"metadata": {
|
||
"ExecuteTime": {
|
||
"end_time": "2025-01-15T10:44:25.008501Z",
|
||
"start_time": "2025-01-15T10:43:57.932240Z"
|
||
}
|
||
},
|
||
"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"
|
||
]
|
||
}
|
||
],
|
||
"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()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": "## Retrieving and generating answer from graphiti graph with cognee retriever"
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 5,
|
||
"metadata": {
|
||
"ExecuteTime": {
|
||
"end_time": "2025-01-15T10:44:27.844438Z",
|
||
"start_time": "2025-01-15T10:44:25.013325Z"
|
||
}
|
||
},
|
||
"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"
|
||
]
|
||
}
|
||
],
|
||
"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 = {\"question\": query, \"context\": context}\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}\")"
|
||
]
|
||
}
|
||
],
|
||
"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
|
||
}
|