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