{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"GRAPHISTRY_USERNAME\"] = input(\"Please enter your graphistry username\")\n",
    "os.environ[\"GRAPHISTRY_PASSWORD\"] = input(\"Please enter your graphistry password\")\n",
    "os.environ[\"OPENAI_API_KEY\"] = input(\"Please enter your OpenAI API key\")"
   ]
  },
  {
   "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(\n",
    "    pathlib.Path(os.path.join(notebook_path, \".data_storage/code_graph\")).resolve()\n",
    ")\n",
    "cognee.config.data_root_directory(data_directory_path)\n",
    "cognee_directory_path = str(\n",
    "    pathlib.Path(os.path.join(notebook_path, \".cognee_system/code_graph\")).resolve()\n",
    ")\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 (\n",
    "    get_repo_file_dependencies,\n",
    ")\n",
    "from cognee.tasks.storage import add_data_points\n",
    "from cognee.modules.pipelines.tasks.Task import Task\n",
    "\n",
    "detailed_extraction = True\n",
    "\n",
    "tasks = [\n",
    "    Task(get_repo_file_dependencies, detailed_extraction=detailed_extraction),\n",
    "    Task(add_data_points, task_config={\"batch_size\": 100 if detailed_extraction else 500}),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cognee.modules.pipelines import run_tasks\n",
    "from uuid import uuid5, NAMESPACE_OID\n",
    "\n",
    "pipeline = run_tasks(tasks, uuid5(NAMESPACE_OID, repo_clone_location), 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",
    "\n",
    "await render_graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Let's check the evaluations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cognee import search, SearchType\n",
    "\n",
    "results = await search(query_type=SearchType.CODE, query_text=\"def create_graphrag_config\")\n",
    "\n",
    "print(results)\n"
   ]
  }
 ],
 "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
}