From d2c2cc8e61831c65d929de41c782858123e7187e Mon Sep 17 00:00:00 2001
From: natsukium <tomoya.otabi@gmail.com>
Date: Sat, 13 Jul 2024 10:15:49 +0900
Subject: [PATCH] python312Packages.graphrag: init at 0.1.1

---
 .../python-modules/graphrag/default.nix       | 122 ++++++++++++++++++
 pkgs/top-level/python-packages.nix            |   2 +
 2 files changed, 124 insertions(+)
 create mode 100644 pkgs/development/python-modules/graphrag/default.nix

diff --git a/pkgs/development/python-modules/graphrag/default.nix b/pkgs/development/python-modules/graphrag/default.nix
new file mode 100644
index 000000000000..204b7a747597
--- /dev/null
+++ b/pkgs/development/python-modules/graphrag/default.nix
@@ -0,0 +1,122 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+  poetry-core,
+  poetry-dynamic-versioning,
+  aiofiles,
+  aiolimiter,
+  azure-identity,
+  azure-search-documents,
+  azure-storage-blob,
+  datashaper,
+  devtools,
+  environs,
+  fastparquet,
+  graspologic,
+  lancedb,
+  networkx,
+  nltk,
+  numba,
+  numpy,
+  openai,
+  pyaml-env,
+  pydantic,
+  python-dotenv,
+  pyyaml,
+  rich,
+  scipy,
+  swifter,
+  tenacity,
+  textual,
+  tiktoken,
+  typing-extensions,
+  uvloop,
+  nbformat,
+  pytest-asyncio,
+  pytestCheckHook,
+}:
+
+buildPythonPackage rec {
+  pname = "graphrag";
+  version = "0.1.1";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "microsoft";
+    repo = "graphrag";
+    rev = "refs/tags/v${version}";
+    hash = "sha256-hIAQOIqm9S9AtssE6UxcXfaIbSt3+506ueMrlathNaQ=";
+  };
+
+  build-system = [
+    poetry-core
+    poetry-dynamic-versioning
+  ];
+
+  pythonRelaxDeps = [
+    "aiofiles"
+    "azure-identity"
+    "scipy"
+    "tiktoken"
+  ];
+
+  dependencies = [
+    aiofiles
+    aiolimiter
+    azure-identity
+    azure-search-documents
+    azure-storage-blob
+    datashaper
+    devtools
+    environs
+    fastparquet
+    graspologic
+    lancedb
+    networkx
+    nltk
+    numba
+    numpy
+    openai
+    pyaml-env
+    pydantic
+    python-dotenv
+    pyyaml
+    rich
+    scipy
+    swifter
+    tenacity
+    textual
+    tiktoken
+    typing-extensions
+    uvloop
+  ];
+
+  env.NUMBA_CACHE_DIR = "$TMPDIR";
+
+  pythonImportsCheck = [ "graphrag" ];
+
+  nativeCheckInputs = [
+    nbformat
+    pytest-asyncio
+    pytestCheckHook
+  ];
+
+  pytestFlagsArray = [ "tests/unit" ];
+
+  disabledTests = [
+    # touch the network
+    "test_child"
+    "test_dotprefix"
+    "test_find"
+    "test_run_extract_entities_multiple_documents"
+    "test_run_extract_entities_single_document"
+  ];
+
+  meta = {
+    description = "Modular graph-based Retrieval-Augmented Generation (RAG) system";
+    homepage = "https://github.com/microsoft/graphrag";
+    license = lib.licenses.mit;
+    maintainers = with lib.maintainers; [ natsukium ];
+  };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 828ccc82ecb3..b5d5fe1f0c2d 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -5218,6 +5218,8 @@ self: super: with self; {
 
   graphql-subscription-manager = callPackage ../development/python-modules/graphql-subscription-manager { };
 
+  graphrag = callPackage ../development/python-modules/graphrag { };
+
   graph-tool = callPackage ../development/python-modules/graph-tool { };
 
   graphtage = callPackage ../development/python-modules/graphtage { };