cognee 0.2.3.dev0__py3-none-any.whl → 0.2.4__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (179) hide show
  1. cognee/__main__.py +4 -0
  2. cognee/api/v1/add/add.py +18 -11
  3. cognee/api/v1/cognify/code_graph_pipeline.py +7 -1
  4. cognee/api/v1/cognify/cognify.py +22 -115
  5. cognee/api/v1/cognify/routers/get_cognify_router.py +11 -3
  6. cognee/api/v1/config/config.py +5 -13
  7. cognee/api/v1/datasets/routers/get_datasets_router.py +2 -2
  8. cognee/api/v1/delete/delete.py +1 -1
  9. cognee/api/v1/exceptions/__init__.py +13 -0
  10. cognee/api/v1/{delete → exceptions}/exceptions.py +15 -12
  11. cognee/api/v1/responses/default_tools.py +4 -0
  12. cognee/api/v1/responses/dispatch_function.py +6 -1
  13. cognee/api/v1/responses/models.py +1 -1
  14. cognee/api/v1/search/search.py +6 -7
  15. cognee/cli/__init__.py +10 -0
  16. cognee/cli/_cognee.py +180 -0
  17. cognee/cli/commands/__init__.py +1 -0
  18. cognee/cli/commands/add_command.py +80 -0
  19. cognee/cli/commands/cognify_command.py +128 -0
  20. cognee/cli/commands/config_command.py +225 -0
  21. cognee/cli/commands/delete_command.py +80 -0
  22. cognee/cli/commands/search_command.py +149 -0
  23. cognee/cli/config.py +33 -0
  24. cognee/cli/debug.py +21 -0
  25. cognee/cli/echo.py +45 -0
  26. cognee/cli/exceptions.py +23 -0
  27. cognee/cli/minimal_cli.py +97 -0
  28. cognee/cli/reference.py +26 -0
  29. cognee/cli/suppress_logging.py +12 -0
  30. cognee/eval_framework/corpus_builder/corpus_builder_executor.py +2 -2
  31. cognee/eval_framework/eval_config.py +1 -1
  32. cognee/exceptions/__init__.py +5 -5
  33. cognee/exceptions/exceptions.py +37 -17
  34. cognee/infrastructure/data/exceptions/__init__.py +7 -0
  35. cognee/infrastructure/data/exceptions/exceptions.py +22 -0
  36. cognee/infrastructure/data/utils/extract_keywords.py +3 -3
  37. cognee/infrastructure/databases/exceptions/__init__.py +3 -0
  38. cognee/infrastructure/databases/exceptions/exceptions.py +57 -9
  39. cognee/infrastructure/databases/graph/get_graph_engine.py +4 -9
  40. cognee/infrastructure/databases/graph/kuzu/adapter.py +64 -2
  41. cognee/infrastructure/databases/graph/neo4j_driver/adapter.py +49 -0
  42. cognee/infrastructure/databases/graph/neptune_driver/exceptions.py +15 -10
  43. cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py +2 -2
  44. cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py +4 -5
  45. cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py +2 -2
  46. cognee/infrastructure/databases/vector/embeddings/FastembedEmbeddingEngine.py +5 -3
  47. cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py +17 -8
  48. cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py +5 -5
  49. cognee/infrastructure/databases/vector/embeddings/config.py +2 -2
  50. cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py +6 -6
  51. cognee/infrastructure/databases/vector/exceptions/exceptions.py +3 -3
  52. cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py +2 -2
  53. cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py +4 -3
  54. cognee/infrastructure/files/utils/get_data_file_path.py +14 -9
  55. cognee/infrastructure/files/utils/get_file_metadata.py +2 -1
  56. cognee/infrastructure/llm/LLMGateway.py +14 -5
  57. cognee/infrastructure/llm/config.py +5 -5
  58. cognee/infrastructure/llm/exceptions.py +30 -2
  59. cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/knowledge_graph/extract_content_graph.py +16 -5
  60. cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/extract_content_graph.py +19 -15
  61. cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py +5 -5
  62. cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py +6 -6
  63. cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/generic_llm_api/adapter.py +2 -2
  64. cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/get_llm_client.py +24 -15
  65. cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/ollama/adapter.py +6 -4
  66. cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py +9 -7
  67. cognee/infrastructure/llm/tokenizer/Gemini/adapter.py +2 -2
  68. cognee/infrastructure/llm/tokenizer/HuggingFace/adapter.py +3 -3
  69. cognee/infrastructure/llm/tokenizer/Mistral/adapter.py +3 -3
  70. cognee/infrastructure/llm/tokenizer/TikToken/adapter.py +6 -6
  71. cognee/infrastructure/llm/utils.py +7 -7
  72. cognee/modules/data/exceptions/exceptions.py +18 -5
  73. cognee/modules/data/methods/__init__.py +2 -0
  74. cognee/modules/data/methods/create_authorized_dataset.py +19 -0
  75. cognee/modules/data/methods/delete_data.py +2 -4
  76. cognee/modules/data/methods/get_authorized_dataset.py +11 -5
  77. cognee/modules/data/methods/get_authorized_dataset_by_name.py +16 -0
  78. cognee/modules/data/methods/load_or_create_datasets.py +2 -20
  79. cognee/modules/data/processing/document_types/exceptions/exceptions.py +2 -2
  80. cognee/modules/graph/cognee_graph/CogneeGraph.py +6 -4
  81. cognee/modules/graph/cognee_graph/CogneeGraphElements.py +5 -10
  82. cognee/modules/graph/exceptions/__init__.py +2 -0
  83. cognee/modules/graph/exceptions/exceptions.py +25 -3
  84. cognee/modules/graph/methods/get_formatted_graph_data.py +3 -2
  85. cognee/modules/ingestion/exceptions/exceptions.py +2 -2
  86. cognee/modules/ontology/exceptions/exceptions.py +4 -4
  87. cognee/modules/pipelines/__init__.py +1 -1
  88. cognee/modules/pipelines/exceptions/exceptions.py +2 -2
  89. cognee/modules/pipelines/exceptions/tasks.py +18 -0
  90. cognee/modules/pipelines/layers/__init__.py +1 -0
  91. cognee/modules/pipelines/layers/check_pipeline_run_qualification.py +59 -0
  92. cognee/modules/pipelines/layers/pipeline_execution_mode.py +127 -0
  93. cognee/modules/pipelines/layers/reset_dataset_pipeline_run_status.py +12 -0
  94. cognee/modules/pipelines/layers/resolve_authorized_user_dataset.py +34 -0
  95. cognee/modules/pipelines/layers/resolve_authorized_user_datasets.py +55 -0
  96. cognee/modules/pipelines/layers/setup_and_check_environment.py +41 -0
  97. cognee/modules/pipelines/layers/validate_pipeline_tasks.py +20 -0
  98. cognee/modules/pipelines/methods/__init__.py +2 -0
  99. cognee/modules/pipelines/methods/get_pipeline_runs_by_dataset.py +34 -0
  100. cognee/modules/pipelines/methods/reset_pipeline_run_status.py +16 -0
  101. cognee/modules/pipelines/operations/__init__.py +0 -1
  102. cognee/modules/pipelines/operations/log_pipeline_run_initiated.py +1 -1
  103. cognee/modules/pipelines/operations/pipeline.py +23 -138
  104. cognee/modules/retrieval/base_feedback.py +11 -0
  105. cognee/modules/retrieval/cypher_search_retriever.py +1 -9
  106. cognee/modules/retrieval/exceptions/exceptions.py +12 -6
  107. cognee/modules/retrieval/graph_completion_context_extension_retriever.py +9 -2
  108. cognee/modules/retrieval/graph_completion_cot_retriever.py +13 -6
  109. cognee/modules/retrieval/graph_completion_retriever.py +89 -5
  110. cognee/modules/retrieval/graph_summary_completion_retriever.py +2 -0
  111. cognee/modules/retrieval/natural_language_retriever.py +0 -4
  112. cognee/modules/retrieval/user_qa_feedback.py +83 -0
  113. cognee/modules/retrieval/utils/extract_uuid_from_node.py +18 -0
  114. cognee/modules/retrieval/utils/models.py +40 -0
  115. cognee/modules/search/exceptions/__init__.py +7 -0
  116. cognee/modules/search/exceptions/exceptions.py +15 -0
  117. cognee/modules/search/methods/search.py +47 -7
  118. cognee/modules/search/types/SearchType.py +1 -0
  119. cognee/modules/settings/get_settings.py +2 -2
  120. cognee/modules/users/exceptions/exceptions.py +6 -6
  121. cognee/shared/CodeGraphEntities.py +1 -0
  122. cognee/shared/exceptions/exceptions.py +2 -2
  123. cognee/shared/logging_utils.py +142 -31
  124. cognee/shared/utils.py +0 -1
  125. cognee/tasks/completion/exceptions/exceptions.py +3 -3
  126. cognee/tasks/documents/classify_documents.py +4 -0
  127. cognee/tasks/documents/exceptions/__init__.py +11 -0
  128. cognee/tasks/documents/exceptions/exceptions.py +36 -0
  129. cognee/tasks/documents/extract_chunks_from_documents.py +8 -2
  130. cognee/tasks/graph/exceptions/__init__.py +12 -0
  131. cognee/tasks/graph/exceptions/exceptions.py +41 -0
  132. cognee/tasks/graph/extract_graph_from_data.py +34 -2
  133. cognee/tasks/ingestion/exceptions/__init__.py +8 -0
  134. cognee/tasks/ingestion/exceptions/exceptions.py +12 -0
  135. cognee/tasks/ingestion/resolve_data_directories.py +5 -0
  136. cognee/tasks/repo_processor/get_local_dependencies.py +2 -0
  137. cognee/tasks/repo_processor/get_repo_file_dependencies.py +120 -48
  138. cognee/tasks/storage/add_data_points.py +41 -3
  139. cognee/tasks/storage/exceptions/__init__.py +9 -0
  140. cognee/tasks/storage/exceptions/exceptions.py +13 -0
  141. cognee/tasks/storage/index_data_points.py +1 -1
  142. cognee/tasks/summarization/exceptions/__init__.py +9 -0
  143. cognee/tasks/summarization/exceptions/exceptions.py +14 -0
  144. cognee/tasks/summarization/summarize_text.py +8 -1
  145. cognee/tests/integration/cli/__init__.py +3 -0
  146. cognee/tests/integration/cli/test_cli_integration.py +331 -0
  147. cognee/tests/integration/documents/PdfDocument_test.py +2 -2
  148. cognee/tests/integration/documents/TextDocument_test.py +2 -4
  149. cognee/tests/integration/documents/UnstructuredDocument_test.py +5 -8
  150. cognee/tests/test_delete_by_id.py +1 -1
  151. cognee/tests/{test_deletion.py → test_delete_hard.py} +0 -37
  152. cognee/tests/test_delete_soft.py +85 -0
  153. cognee/tests/test_kuzu.py +2 -2
  154. cognee/tests/test_neo4j.py +2 -2
  155. cognee/tests/test_search_db.py +126 -7
  156. cognee/tests/unit/cli/__init__.py +3 -0
  157. cognee/tests/unit/cli/test_cli_commands.py +483 -0
  158. cognee/tests/unit/cli/test_cli_edge_cases.py +625 -0
  159. cognee/tests/unit/cli/test_cli_main.py +173 -0
  160. cognee/tests/unit/cli/test_cli_runner.py +62 -0
  161. cognee/tests/unit/cli/test_cli_utils.py +127 -0
  162. cognee/tests/unit/modules/graph/cognee_graph_elements_test.py +5 -5
  163. cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py +3 -3
  164. cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py +3 -3
  165. cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py +3 -3
  166. cognee/tests/unit/modules/search/search_methods_test.py +4 -2
  167. {cognee-0.2.3.dev0.dist-info → cognee-0.2.4.dist-info}/METADATA +7 -5
  168. {cognee-0.2.3.dev0.dist-info → cognee-0.2.4.dist-info}/RECORD +172 -121
  169. cognee-0.2.4.dist-info/entry_points.txt +2 -0
  170. cognee/infrastructure/databases/exceptions/EmbeddingException.py +0 -20
  171. cognee/infrastructure/databases/graph/networkx/__init__.py +0 -0
  172. cognee/infrastructure/databases/graph/networkx/adapter.py +0 -1017
  173. cognee/infrastructure/pipeline/models/Operation.py +0 -60
  174. cognee/infrastructure/pipeline/models/__init__.py +0 -0
  175. cognee/notebooks/github_analysis_step_by_step.ipynb +0 -37
  176. cognee/tests/tasks/descriptive_metrics/networkx_metrics_test.py +0 -7
  177. {cognee-0.2.3.dev0.dist-info → cognee-0.2.4.dist-info}/WHEEL +0 -0
  178. {cognee-0.2.3.dev0.dist-info → cognee-0.2.4.dist-info}/licenses/LICENSE +0 -0
  179. {cognee-0.2.3.dev0.dist-info → cognee-0.2.4.dist-info}/licenses/NOTICE.md +0 -0
@@ -0,0 +1,331 @@
1
+ """
2
+ Integration tests for CLI commands that test end-to-end functionality.
3
+ """
4
+
5
+ import tempfile
6
+ import os
7
+ import sys
8
+ import subprocess
9
+ from pathlib import Path
10
+ from unittest.mock import patch, MagicMock
11
+
12
+
13
+ class TestCliIntegration:
14
+ """Integration tests for CLI commands"""
15
+
16
+ def test_cli_help(self):
17
+ """Test that CLI help works"""
18
+ result = subprocess.run(
19
+ [sys.executable, "-m", "cognee.cli._cognee", "--help"],
20
+ capture_output=True,
21
+ text=True,
22
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
23
+ )
24
+
25
+ assert result.returncode == 0
26
+ assert "cognee" in result.stdout.lower()
27
+ assert "available commands" in result.stdout.lower()
28
+
29
+ def test_cli_version(self):
30
+ """Test that CLI version works"""
31
+ result = subprocess.run(
32
+ [sys.executable, "-m", "cognee.cli._cognee", "--version"],
33
+ capture_output=True,
34
+ text=True,
35
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
36
+ )
37
+
38
+ assert result.returncode == 0
39
+ assert "cognee" in result.stdout.lower()
40
+
41
+ def test_command_help(self):
42
+ """Test that individual command help works"""
43
+ commands = ["add", "search", "cognify", "delete", "config"]
44
+
45
+ for command in commands:
46
+ result = subprocess.run(
47
+ [sys.executable, "-m", "cognee.cli._cognee", command, "--help"],
48
+ capture_output=True,
49
+ text=True,
50
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
51
+ )
52
+
53
+ assert result.returncode == 0, f"Command {command} help failed"
54
+ assert command in result.stdout.lower()
55
+
56
+ def test_invalid_command(self):
57
+ """Test that invalid commands are handled properly"""
58
+ result = subprocess.run(
59
+ [sys.executable, "-m", "cognee.cli._cognee", "invalid_command"],
60
+ capture_output=True,
61
+ text=True,
62
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
63
+ )
64
+
65
+ assert result.returncode != 0
66
+
67
+ @patch("cognee.add")
68
+ def test_add_command_integration(self, mock_add):
69
+ """Test add command integration"""
70
+ mock_add.return_value = None
71
+
72
+ with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
73
+ f.write("Test content for CLI integration")
74
+ temp_file = f.name
75
+
76
+ try:
77
+ result = subprocess.run(
78
+ [sys.executable, "-m", "cognee.cli._cognee", "add", temp_file],
79
+ capture_output=True,
80
+ text=True,
81
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
82
+ )
83
+
84
+ # Note: This might fail due to dependencies, but we're testing the CLI structure
85
+ # The important thing is that it doesn't crash with argument parsing errors
86
+ # Allow litellm logging worker cancellation errors as they're expected during process shutdown
87
+ stderr_lower = result.stderr.lower()
88
+ has_error = "error" in stderr_lower
89
+ has_expected_failure = "failed to add data" in stderr_lower
90
+ has_litellm_cancellation = (
91
+ "loggingworker cancelled" in stderr_lower or "cancellederror" in stderr_lower
92
+ )
93
+
94
+ assert not has_error or has_expected_failure or has_litellm_cancellation
95
+
96
+ finally:
97
+ os.unlink(temp_file)
98
+
99
+ def test_config_subcommands(self):
100
+ """Test config subcommands help"""
101
+ subcommands = ["get", "set", "list", "unset", "reset"]
102
+
103
+ for subcommand in subcommands:
104
+ result = subprocess.run(
105
+ [sys.executable, "-m", "cognee.cli._cognee", "config", subcommand, "--help"],
106
+ capture_output=True,
107
+ text=True,
108
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
109
+ )
110
+
111
+ assert result.returncode == 0, f"Config {subcommand} help failed"
112
+
113
+ def test_search_command_missing_query(self):
114
+ """Test search command fails when query is missing"""
115
+ result = subprocess.run(
116
+ [sys.executable, "-m", "cognee.cli._cognee", "search"],
117
+ capture_output=True,
118
+ text=True,
119
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
120
+ )
121
+
122
+ assert result.returncode != 0
123
+ assert "required" in result.stderr.lower() or "error" in result.stderr.lower()
124
+
125
+ def test_delete_command_no_target(self):
126
+ """Test delete command with no target specified"""
127
+ result = subprocess.run(
128
+ [sys.executable, "-m", "cognee.cli._cognee", "delete"],
129
+ capture_output=True,
130
+ text=True,
131
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
132
+ )
133
+
134
+ # Should run but show error message about missing target
135
+ # Return code might be 0 since the command handles this gracefully
136
+ assert (
137
+ "specify what to delete" in result.stdout.lower()
138
+ or "specify what to delete" in result.stderr.lower()
139
+ )
140
+
141
+
142
+ class TestCliArgumentParsing:
143
+ """Test CLI argument parsing edge cases"""
144
+
145
+ def test_add_multiple_files(self):
146
+ """Test add command with multiple file arguments"""
147
+ with tempfile.TemporaryDirectory() as temp_dir:
148
+ file1 = os.path.join(temp_dir, "file1.txt")
149
+ file2 = os.path.join(temp_dir, "file2.txt")
150
+
151
+ with open(file1, "w") as f:
152
+ f.write("Content 1")
153
+ with open(file2, "w") as f:
154
+ f.write("Content 2")
155
+
156
+ result = subprocess.run(
157
+ [
158
+ sys.executable,
159
+ "-m",
160
+ "cognee.cli._cognee",
161
+ "add",
162
+ file1,
163
+ file2,
164
+ "--dataset-name",
165
+ "test",
166
+ ],
167
+ capture_output=True,
168
+ text=True,
169
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
170
+ )
171
+
172
+ # Test that argument parsing works (regardless of actual execution)
173
+ assert (
174
+ "argument" not in result.stderr.lower() or "failed to add" in result.stderr.lower()
175
+ )
176
+
177
+ def test_search_with_all_options(self):
178
+ """Test search command with all possible options"""
179
+ result = subprocess.run(
180
+ [
181
+ sys.executable,
182
+ "-m",
183
+ "cognee.cli._cognee",
184
+ "search",
185
+ "test query",
186
+ "--query-type",
187
+ "CHUNKS",
188
+ "--datasets",
189
+ "dataset1",
190
+ "dataset2",
191
+ "--top-k",
192
+ "5",
193
+ "--output-format",
194
+ "json",
195
+ ],
196
+ capture_output=True,
197
+ text=True,
198
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
199
+ )
200
+
201
+ # Should not have argument parsing errors
202
+ assert "unrecognized arguments" not in result.stderr.lower()
203
+ assert "invalid choice" not in result.stderr.lower()
204
+
205
+ def test_cognify_with_all_options(self):
206
+ """Test cognify command with all possible options"""
207
+ result = subprocess.run(
208
+ [
209
+ sys.executable,
210
+ "-m",
211
+ "cognee.cli._cognee",
212
+ "cognify",
213
+ "--datasets",
214
+ "dataset1",
215
+ "dataset2",
216
+ "--chunk-size",
217
+ "1024",
218
+ "--chunker",
219
+ "TextChunker",
220
+ "--background",
221
+ "--verbose",
222
+ ],
223
+ capture_output=True,
224
+ text=True,
225
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
226
+ )
227
+
228
+ # Should not have argument parsing errors
229
+ assert "unrecognized arguments" not in result.stderr.lower()
230
+ assert "invalid choice" not in result.stderr.lower()
231
+
232
+ def test_config_set_command(self):
233
+ """Test config set command argument parsing"""
234
+ result = subprocess.run(
235
+ [sys.executable, "-m", "cognee.cli._cognee", "config", "set", "test_key", "test_value"],
236
+ capture_output=True,
237
+ text=True,
238
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
239
+ )
240
+
241
+ # Should not have argument parsing errors
242
+ assert "unrecognized arguments" not in result.stderr.lower()
243
+ assert "required" not in result.stderr.lower() or "failed to set" in result.stderr.lower()
244
+
245
+ def test_delete_with_force(self):
246
+ """Test delete command with force flag"""
247
+ result = subprocess.run(
248
+ [
249
+ sys.executable,
250
+ "-m",
251
+ "cognee.cli._cognee",
252
+ "delete",
253
+ "--dataset-name",
254
+ "test_dataset",
255
+ "--force",
256
+ ],
257
+ capture_output=True,
258
+ text=True,
259
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
260
+ )
261
+
262
+ # Should not have argument parsing errors
263
+ assert "unrecognized arguments" not in result.stderr.lower()
264
+
265
+
266
+ class TestCliErrorHandling:
267
+ """Test CLI error handling and edge cases"""
268
+
269
+ def test_debug_mode_flag(self):
270
+ """Test that debug flag is accepted"""
271
+ result = subprocess.run(
272
+ [sys.executable, "-m", "cognee.cli._cognee", "--debug", "search", "test query"],
273
+ capture_output=True,
274
+ text=True,
275
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
276
+ )
277
+
278
+ # Should not have argument parsing errors for debug flag
279
+ assert "unrecognized arguments" not in result.stderr.lower()
280
+
281
+ def test_invalid_search_type(self):
282
+ """Test invalid search type handling"""
283
+ result = subprocess.run(
284
+ [
285
+ sys.executable,
286
+ "-m",
287
+ "cognee.cli._cognee",
288
+ "search",
289
+ "test query",
290
+ "--query-type",
291
+ "INVALID_TYPE",
292
+ ],
293
+ capture_output=True,
294
+ text=True,
295
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
296
+ )
297
+
298
+ assert result.returncode != 0
299
+ assert "invalid choice" in result.stderr.lower()
300
+
301
+ def test_invalid_chunker(self):
302
+ """Test invalid chunker handling"""
303
+ result = subprocess.run(
304
+ [sys.executable, "-m", "cognee.cli._cognee", "cognify", "--chunker", "InvalidChunker"],
305
+ capture_output=True,
306
+ text=True,
307
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
308
+ )
309
+
310
+ assert result.returncode != 0
311
+ assert "invalid choice" in result.stderr.lower()
312
+
313
+ def test_invalid_output_format(self):
314
+ """Test invalid output format handling"""
315
+ result = subprocess.run(
316
+ [
317
+ sys.executable,
318
+ "-m",
319
+ "cognee.cli._cognee",
320
+ "search",
321
+ "test query",
322
+ "--output-format",
323
+ "invalid",
324
+ ],
325
+ capture_output=True,
326
+ text=True,
327
+ cwd=Path(__file__).parent.parent.parent, # Go to project root
328
+ )
329
+
330
+ assert result.returncode != 0
331
+ assert "invalid choice" in result.stderr.lower()
@@ -2,6 +2,7 @@ import os
2
2
  import sys
3
3
  import uuid
4
4
  import pytest
5
+ import pathlib
5
6
  from unittest.mock import patch
6
7
 
7
8
  from cognee.modules.chunking.TextChunker import TextChunker
@@ -24,8 +25,7 @@ GROUND_TRUTH = [
24
25
  @pytest.mark.asyncio
25
26
  async def test_PdfDocument(mock_engine):
26
27
  test_file_path = os.path.join(
27
- os.sep,
28
- *(os.path.dirname(__file__).split(os.sep)[:-2]),
28
+ pathlib.Path(__file__).parent.parent.parent,
29
29
  "test_data",
30
30
  "artificial-intelligence.pdf",
31
31
  )
@@ -2,6 +2,7 @@ import os
2
2
  import sys
3
3
  import uuid
4
4
  import pytest
5
+ import pathlib
5
6
  from unittest.mock import patch
6
7
 
7
8
  from cognee.modules.chunking.TextChunker import TextChunker
@@ -34,10 +35,7 @@ GROUND_TRUTH = {
34
35
  @pytest.mark.asyncio
35
36
  async def test_TextDocument(mock_engine, input_file, chunk_size):
36
37
  test_file_path = os.path.join(
37
- os.sep,
38
- *(os.path.dirname(__file__).split(os.sep)[:-2]),
39
- "test_data",
40
- input_file,
38
+ pathlib.Path(__file__).parent.parent.parent, "test_data", input_file
41
39
  )
42
40
  document = TextDocument(
43
41
  id=uuid.uuid4(),
@@ -2,6 +2,7 @@ import os
2
2
  import sys
3
3
  import uuid
4
4
  import pytest
5
+ import pathlib
5
6
  from unittest.mock import patch
6
7
 
7
8
  from cognee.modules.chunking.TextChunker import TextChunker
@@ -18,29 +19,25 @@ chunk_by_sentence_module = sys.modules.get("cognee.tasks.chunks.chunk_by_sentenc
18
19
  async def test_UnstructuredDocument(mock_engine):
19
20
  # Define file paths of test data
20
21
  pptx_file_path = os.path.join(
21
- os.sep,
22
- *(os.path.dirname(__file__).split(os.sep)[:-2]),
22
+ pathlib.Path(__file__).parent.parent.parent,
23
23
  "test_data",
24
24
  "example.pptx",
25
25
  )
26
26
 
27
27
  docx_file_path = os.path.join(
28
- os.sep,
29
- *(os.path.dirname(__file__).split(os.sep)[:-2]),
28
+ pathlib.Path(__file__).parent.parent.parent,
30
29
  "test_data",
31
30
  "example.docx",
32
31
  )
33
32
 
34
33
  csv_file_path = os.path.join(
35
- os.sep,
36
- *(os.path.dirname(__file__).split(os.sep)[:-2]),
34
+ pathlib.Path(__file__).parent.parent.parent,
37
35
  "test_data",
38
36
  "example.csv",
39
37
  )
40
38
 
41
39
  xlsx_file_path = os.path.join(
42
- os.sep,
43
- *(os.path.dirname(__file__).split(os.sep)[:-2]),
40
+ pathlib.Path(__file__).parent.parent.parent,
44
41
  "test_data",
45
42
  "example.xlsx",
46
43
  )
@@ -7,7 +7,7 @@ from cognee.shared.logging_utils import get_logger
7
7
  from cognee.modules.users.methods import get_default_user, create_user
8
8
  from cognee.modules.users.permissions.methods import authorized_give_permission_on_datasets
9
9
  from cognee.modules.data.methods import get_dataset_data, get_datasets_by_name
10
- from cognee.api.v1.delete.exceptions import DocumentNotFoundError, DatasetNotFoundError
10
+ from cognee.api.v1.exceptions import DocumentNotFoundError, DatasetNotFoundError
11
11
 
12
12
  logger = get_logger()
13
13
 
@@ -45,8 +45,6 @@ async def main():
45
45
  Each of these car manufacturer contributes to Germany's reputation as a leader in the global automotive industry, showcasing a blend of innovation, performance, and design excellence.
46
46
  """
47
47
 
48
- ################### HARD DELETE
49
-
50
48
  # Add documents and get dataset information
51
49
  add_result = await cognee.add(
52
50
  [
@@ -80,41 +78,6 @@ async def main():
80
78
 
81
79
  assert len(nodes) == 0 and len(edges) == 0, "Document is not deleted with hard delete."
82
80
 
83
- ################### SOFT DELETE
84
-
85
- # Add documents and get dataset information
86
- add_result = await cognee.add(
87
- [
88
- pdf_document,
89
- txt_document,
90
- text_document_as_literal,
91
- unstructured_document,
92
- audio_document,
93
- image_document,
94
- ]
95
- )
96
- dataset_id = add_result.dataset_id
97
-
98
- await cognee.cognify()
99
-
100
- from cognee.infrastructure.databases.graph import get_graph_engine
101
-
102
- graph_engine = await get_graph_engine()
103
- nodes, edges = await graph_engine.get_graph_data()
104
- assert len(nodes) > 10 and len(edges) > 10, "Graph database is not loaded."
105
-
106
- # Get the data IDs from the dataset
107
- dataset_data = await get_dataset_data(dataset_id)
108
- assert len(dataset_data) > 0, "Dataset should contain data"
109
-
110
- # Delete each document using its ID
111
- for data_item in dataset_data:
112
- await cognee.delete(data_item.id, dataset_id, mode="soft")
113
-
114
- nodes, edges = await graph_engine.get_graph_data()
115
-
116
- assert len(nodes) == 0 and len(edges) == 0, "Document is not deleted with soft delete."
117
-
118
81
 
119
82
  if __name__ == "__main__":
120
83
  import asyncio
@@ -0,0 +1,85 @@
1
+ import os
2
+ import shutil
3
+ import cognee
4
+ import pathlib
5
+ from cognee.shared.logging_utils import get_logger
6
+ from cognee.modules.data.methods import get_dataset_data
7
+
8
+ logger = get_logger()
9
+
10
+
11
+ async def main():
12
+ await cognee.prune.prune_data()
13
+ await cognee.prune.prune_system(metadata=True)
14
+
15
+ pdf_document = os.path.join(
16
+ pathlib.Path(__file__).parent, "test_data/artificial-intelligence.pdf"
17
+ )
18
+
19
+ txt_document = os.path.join(
20
+ pathlib.Path(__file__).parent, "test_data/Natural_language_processing_copy.txt"
21
+ )
22
+
23
+ audio_document = os.path.join(pathlib.Path(__file__).parent, "test_data/text_to_speech.mp3")
24
+
25
+ image_document = os.path.join(pathlib.Path(__file__).parent, "test_data/example.png")
26
+
27
+ unstructured_document = os.path.join(pathlib.Path(__file__).parent, "test_data/example.pptx")
28
+
29
+ text_document_as_literal = """
30
+ 1. Audi
31
+ Audi is known for its modern designs and advanced technology. Founded in the early 1900s, the brand has earned a reputation for precision engineering and innovation. With features like the Quattro all-wheel-drive system, Audi offers a range of vehicles from stylish sedans to high-performance sports cars.
32
+
33
+ 2. BMW
34
+ BMW, short for Bayerische Motoren Werke, is celebrated for its focus on performance and driving pleasure. The company's vehicles are designed to provide a dynamic and engaging driving experience, and their slogan, "The Ultimate Driving Machine," reflects that commitment. BMW produces a variety of cars that combine luxury with sporty performance.
35
+
36
+ 3. Mercedes-Benz
37
+ Mercedes-Benz is synonymous with luxury and quality. With a history dating back to the early 20th century, the brand is known for its elegant designs, innovative safety features, and high-quality engineering. Mercedes-Benz manufactures not only luxury sedans but also SUVs, sports cars, and commercial vehicles, catering to a wide range of needs.
38
+
39
+ 4. Porsche
40
+ Porsche is a name that stands for high-performance sports cars. Founded in 1931, the brand has become famous for models like the iconic Porsche 911. Porsche cars are celebrated for their speed, precision, and distinctive design, appealing to car enthusiasts who value both performance and style.
41
+
42
+ 5. Volkswagen
43
+ Volkswagen, which means "people's car" in German, was established with the idea of making affordable and reliable vehicles accessible to everyone. Over the years, Volkswagen has produced several iconic models, such as the Beetle and the Golf. Today, it remains one of the largest car manufacturers in the world, offering a wide range of vehicles that balance practicality with quality.
44
+
45
+ Each of these car manufacturer contributes to Germany's reputation as a leader in the global automotive industry, showcasing a blend of innovation, performance, and design excellence.
46
+ """
47
+
48
+ # Add documents and get dataset information
49
+ add_result = await cognee.add(
50
+ [
51
+ pdf_document,
52
+ txt_document,
53
+ text_document_as_literal,
54
+ unstructured_document,
55
+ audio_document,
56
+ image_document,
57
+ ]
58
+ )
59
+ dataset_id = add_result.dataset_id
60
+
61
+ await cognee.cognify()
62
+
63
+ from cognee.infrastructure.databases.graph import get_graph_engine
64
+
65
+ graph_engine = await get_graph_engine()
66
+ nodes, edges = await graph_engine.get_graph_data()
67
+ assert len(nodes) > 10 and len(edges) > 10, "Graph database is not loaded."
68
+
69
+ # Get the data IDs from the dataset
70
+ dataset_data = await get_dataset_data(dataset_id)
71
+ assert len(dataset_data) > 0, "Dataset should contain data"
72
+
73
+ # Delete each document using its ID
74
+ for data_item in dataset_data:
75
+ await cognee.delete(data_item.id, dataset_id, mode="soft")
76
+
77
+ nodes, edges = await graph_engine.get_graph_data()
78
+
79
+ assert len(nodes) == 0 and len(edges) == 0, "Document is not deleted with soft delete."
80
+
81
+
82
+ if __name__ == "__main__":
83
+ import asyncio
84
+
85
+ asyncio.run(main())
cognee/tests/test_kuzu.py CHANGED
@@ -94,12 +94,12 @@ async def main():
94
94
 
95
95
  await cognee.cognify([dataset_name])
96
96
 
97
- context_nonempty = await GraphCompletionRetriever(
97
+ context_nonempty, _ = await GraphCompletionRetriever(
98
98
  node_type=NodeSet,
99
99
  node_name=["first"],
100
100
  ).get_context("What is in the context?")
101
101
 
102
- context_empty = await GraphCompletionRetriever(
102
+ context_empty, _ = await GraphCompletionRetriever(
103
103
  node_type=NodeSet,
104
104
  node_name=["nonexistent"],
105
105
  ).get_context("What is in the context?")
@@ -98,12 +98,12 @@ async def main():
98
98
 
99
99
  await cognee.cognify([dataset_name])
100
100
 
101
- context_nonempty = await GraphCompletionRetriever(
101
+ context_nonempty, _ = await GraphCompletionRetriever(
102
102
  node_type=NodeSet,
103
103
  node_name=["first"],
104
104
  ).get_context("What is in the context?")
105
105
 
106
- context_empty = await GraphCompletionRetriever(
106
+ context_empty, _ = await GraphCompletionRetriever(
107
107
  node_type=NodeSet,
108
108
  node_name=["nonexistent"],
109
109
  ).get_context("What is in the context?")