1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5 poetry-core,
6 poetry-dynamic-versioning,
7 aiofiles,
8 aiolimiter,
9 azure-identity,
10 azure-search-documents,
11 azure-storage-blob,
12 datashaper,
13 devtools,
14 environs,
15 fnllm,
16 graspologic,
17 json-repair,
18 lancedb,
19 matplotlib,
20 networkx,
21 nltk,
22 numpy,
23 openai,
24 pandas,
25 pyaml-env,
26 pyarrow,
27 pydantic,
28 python-dotenv,
29 pyyaml,
30 rich,
31 tenacity,
32 tiktoken,
33 typer,
34 typing-extensions,
35 umap-learn,
36 nbformat,
37 pytest-asyncio,
38 pytestCheckHook,
39}:
40
41buildPythonPackage rec {
42 pname = "graphrag";
43 version = "2.4.0";
44 pyproject = true;
45
46 src = fetchFromGitHub {
47 owner = "microsoft";
48 repo = "graphrag";
49 tag = "v.${version}";
50 hash = "sha256-a8t6Nl9W/Cr7eueAvJ3dbz5G0oIhddqFMIm7HeZ8N9A=";
51 };
52
53 build-system = [
54 poetry-core
55 poetry-dynamic-versioning
56 ];
57
58 pythonRelaxDeps = true;
59
60 dependencies = [
61 aiofiles
62 aiolimiter
63 azure-identity
64 azure-search-documents
65 azure-storage-blob
66 datashaper
67 devtools
68 environs
69 fnllm
70 graspologic
71 json-repair
72 lancedb
73 matplotlib
74 networkx
75 nltk
76 numpy
77 openai
78 pandas
79 pyaml-env
80 pyarrow
81 pydantic
82 python-dotenv
83 pyyaml
84 rich
85 tenacity
86 tiktoken
87 typer
88 typing-extensions
89 umap-learn
90 ];
91
92 env.NUMBA_CACHE_DIR = "$TMPDIR";
93
94 pythonImportsCheck = [ "graphrag" ];
95
96 nativeCheckInputs = [
97 nbformat
98 pytest-asyncio
99 pytestCheckHook
100 ];
101
102 enabledTestPaths = [ "tests/unit" ];
103
104 disabledTests = [
105 # touch the network
106 "test_child"
107 "test_dotprefix"
108 "test_find"
109 "test_run_extract_entities_multiple_documents"
110 "test_run_extract_entities_single_document"
111 "test_sort_context"
112 "test_sort_context_max_tokens"
113 ];
114
115 meta = {
116 description = "Modular graph-based Retrieval-Augmented Generation (RAG) system";
117 homepage = "https://github.com/microsoft/graphrag";
118 changelog = "https://github.com/microsoft/graphrag/blob/${src.tag}/CHANGELOG.md";
119 license = lib.licenses.mit;
120 maintainers = with lib.maintainers; [ natsukium ];
121 };
122}