{ "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: \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 }