cognee/notebooks/cognee_code_graph_demo.ipynb
2024-12-03 23:48:12 +01:00

136 lines
3.8 KiB
Text

{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ['GRAPHISTRY_USERNAME'] = input(\"Please enter your graphistry username\")\n",
"os.environ['GRAPHISTRY_PASSWORD'] = input(\"Please enter your graphistry password\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import pathlib\n",
"import cognee\n",
"from cognee.infrastructure.databases.relational import create_db_and_tables\n",
"\n",
"notebook_path = os.path.abspath(\"\")\n",
"data_directory_path = str(pathlib.Path(os.path.join(notebook_path, \".data_storage/code_graph\")).resolve())\n",
"cognee.config.data_root_directory(data_directory_path)\n",
"cognee_directory_path = str(pathlib.Path(os.path.join(notebook_path, \".cognee_system/code_graph\")).resolve())\n",
"cognee.config.system_root_directory(cognee_directory_path)\n",
"\n",
"await cognee.prune.prune_data()\n",
"await cognee.prune.prune_system(metadata = True)\n",
"\n",
"await create_db_and_tables()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from os import path\n",
"from pathlib import Path\n",
"from cognee.infrastructure.files.storage import LocalStorage\n",
"import git\n",
"\n",
"notebook_path = path.abspath(\"\")\n",
"repo_clone_location = path.join(notebook_path, \"data/graphrag\")\n",
"\n",
"LocalStorage.remove_all(repo_clone_location)\n",
"\n",
"git.Repo.clone_from(\n",
" \"git@github.com:microsoft/graphrag.git\",\n",
" Path(repo_clone_location),\n",
" branch = \"main\",\n",
" single_branch = True,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from cognee.tasks.repo_processor import enrich_dependency_graph, expand_dependency_graph, get_repo_file_dependencies\n",
"from cognee.tasks.storage import add_data_points\n",
"from cognee.modules.pipelines.tasks.Task import Task\n",
"\n",
"tasks = [\n",
" Task(get_repo_file_dependencies),\n",
" Task(add_data_points, task_config = { \"batch_size\": 50 }),\n",
" Task(enrich_dependency_graph, task_config = { \"batch_size\": 50 }),\n",
" Task(expand_dependency_graph, task_config = { \"batch_size\": 50 }),\n",
" Task(add_data_points, task_config = { \"batch_size\": 50 }),\n",
"]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from cognee.modules.pipelines import run_tasks\n",
"\n",
"notebook_path = os.path.abspath(\"\")\n",
"repo_clone_location = os.path.join(notebook_path, \"data/graphrag\")\n",
"\n",
"pipeline = run_tasks(tasks, repo_clone_location, \"code_graph_pipeline\")\n",
"\n",
"async for result in pipeline:\n",
" print(result)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from cognee.shared.utils import render_graph\n",
"await render_graph(None, include_nodes = True, include_labels = True)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"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.11.8"
}
},
"nbformat": 4,
"nbformat_minor": 2
}