cognee 0.4.1__py3-none-any.whl → 0.5.0.dev0__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 (135) hide show
  1. cognee/__init__.py +1 -0
  2. cognee/api/client.py +8 -0
  3. cognee/api/v1/add/routers/get_add_router.py +3 -1
  4. cognee/api/v1/cognify/routers/get_cognify_router.py +28 -1
  5. cognee/api/v1/ontologies/__init__.py +4 -0
  6. cognee/api/v1/ontologies/ontologies.py +183 -0
  7. cognee/api/v1/ontologies/routers/__init__.py +0 -0
  8. cognee/api/v1/ontologies/routers/get_ontology_router.py +107 -0
  9. cognee/api/v1/permissions/routers/get_permissions_router.py +41 -1
  10. cognee/cli/commands/cognify_command.py +8 -1
  11. cognee/cli/config.py +1 -1
  12. cognee/context_global_variables.py +41 -9
  13. cognee/infrastructure/databases/cache/config.py +3 -1
  14. cognee/infrastructure/databases/cache/fscache/FsCacheAdapter.py +151 -0
  15. cognee/infrastructure/databases/cache/get_cache_engine.py +20 -10
  16. cognee/infrastructure/databases/exceptions/exceptions.py +16 -0
  17. cognee/infrastructure/databases/graph/config.py +4 -0
  18. cognee/infrastructure/databases/graph/get_graph_engine.py +2 -0
  19. cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py +9 -0
  20. cognee/infrastructure/databases/utils/get_or_create_dataset_database.py +37 -3
  21. cognee/infrastructure/databases/vector/config.py +3 -0
  22. cognee/infrastructure/databases/vector/create_vector_engine.py +5 -1
  23. cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py +1 -4
  24. cognee/infrastructure/engine/models/Edge.py +13 -1
  25. cognee/infrastructure/files/utils/guess_file_type.py +4 -0
  26. cognee/infrastructure/llm/config.py +2 -0
  27. cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py +5 -2
  28. cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py +7 -1
  29. cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/generic_llm_api/adapter.py +7 -1
  30. cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/get_llm_client.py +8 -16
  31. cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/mistral/adapter.py +12 -2
  32. cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/ollama/adapter.py +13 -2
  33. cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py +5 -2
  34. cognee/infrastructure/loaders/LoaderEngine.py +1 -0
  35. cognee/infrastructure/loaders/core/__init__.py +2 -1
  36. cognee/infrastructure/loaders/core/csv_loader.py +93 -0
  37. cognee/infrastructure/loaders/core/text_loader.py +1 -2
  38. cognee/infrastructure/loaders/external/advanced_pdf_loader.py +0 -9
  39. cognee/infrastructure/loaders/supported_loaders.py +2 -1
  40. cognee/memify_pipelines/persist_sessions_in_knowledge_graph.py +55 -0
  41. cognee/modules/chunking/CsvChunker.py +35 -0
  42. cognee/modules/chunking/models/DocumentChunk.py +2 -1
  43. cognee/modules/chunking/text_chunker_with_overlap.py +124 -0
  44. cognee/modules/data/methods/__init__.py +1 -0
  45. cognee/modules/data/methods/create_dataset.py +4 -2
  46. cognee/modules/data/methods/get_dataset_ids.py +5 -1
  47. cognee/modules/data/methods/get_unique_data_id.py +68 -0
  48. cognee/modules/data/methods/get_unique_dataset_id.py +66 -4
  49. cognee/modules/data/models/Dataset.py +2 -0
  50. cognee/modules/data/processing/document_types/CsvDocument.py +33 -0
  51. cognee/modules/data/processing/document_types/__init__.py +1 -0
  52. cognee/modules/graph/cognee_graph/CogneeGraph.py +4 -2
  53. cognee/modules/graph/utils/expand_with_nodes_and_edges.py +19 -2
  54. cognee/modules/graph/utils/resolve_edges_to_text.py +48 -49
  55. cognee/modules/ingestion/identify.py +4 -4
  56. cognee/modules/notebooks/operations/run_in_local_sandbox.py +3 -0
  57. cognee/modules/ontology/rdf_xml/RDFLibOntologyResolver.py +55 -23
  58. cognee/modules/pipelines/operations/run_tasks_data_item.py +1 -1
  59. cognee/modules/retrieval/EntityCompletionRetriever.py +10 -3
  60. cognee/modules/retrieval/base_graph_retriever.py +7 -3
  61. cognee/modules/retrieval/base_retriever.py +7 -3
  62. cognee/modules/retrieval/completion_retriever.py +11 -4
  63. cognee/modules/retrieval/graph_completion_context_extension_retriever.py +6 -2
  64. cognee/modules/retrieval/graph_completion_cot_retriever.py +14 -51
  65. cognee/modules/retrieval/graph_completion_retriever.py +4 -1
  66. cognee/modules/retrieval/temporal_retriever.py +9 -2
  67. cognee/modules/retrieval/utils/brute_force_triplet_search.py +1 -1
  68. cognee/modules/retrieval/utils/completion.py +2 -22
  69. cognee/modules/run_custom_pipeline/__init__.py +1 -0
  70. cognee/modules/run_custom_pipeline/run_custom_pipeline.py +69 -0
  71. cognee/modules/search/methods/search.py +5 -3
  72. cognee/modules/users/methods/create_user.py +12 -27
  73. cognee/modules/users/methods/get_authenticated_user.py +2 -1
  74. cognee/modules/users/methods/get_default_user.py +4 -2
  75. cognee/modules/users/methods/get_user.py +1 -1
  76. cognee/modules/users/methods/get_user_by_email.py +1 -1
  77. cognee/modules/users/models/DatasetDatabase.py +9 -0
  78. cognee/modules/users/models/Tenant.py +6 -7
  79. cognee/modules/users/models/User.py +6 -5
  80. cognee/modules/users/models/UserTenant.py +12 -0
  81. cognee/modules/users/models/__init__.py +1 -0
  82. cognee/modules/users/permissions/methods/get_all_user_permission_datasets.py +13 -13
  83. cognee/modules/users/roles/methods/add_user_to_role.py +3 -1
  84. cognee/modules/users/tenants/methods/__init__.py +1 -0
  85. cognee/modules/users/tenants/methods/add_user_to_tenant.py +21 -12
  86. cognee/modules/users/tenants/methods/create_tenant.py +22 -8
  87. cognee/modules/users/tenants/methods/select_tenant.py +62 -0
  88. cognee/shared/logging_utils.py +2 -0
  89. cognee/tasks/chunks/__init__.py +1 -0
  90. cognee/tasks/chunks/chunk_by_row.py +94 -0
  91. cognee/tasks/documents/classify_documents.py +2 -0
  92. cognee/tasks/feedback/generate_improved_answers.py +3 -3
  93. cognee/tasks/ingestion/ingest_data.py +1 -1
  94. cognee/tasks/memify/__init__.py +2 -0
  95. cognee/tasks/memify/cognify_session.py +41 -0
  96. cognee/tasks/memify/extract_user_sessions.py +73 -0
  97. cognee/tasks/storage/index_data_points.py +33 -22
  98. cognee/tasks/storage/index_graph_edges.py +37 -57
  99. cognee/tests/integration/documents/CsvDocument_test.py +70 -0
  100. cognee/tests/tasks/entity_extraction/entity_extraction_test.py +1 -1
  101. cognee/tests/test_add_docling_document.py +2 -2
  102. cognee/tests/test_cognee_server_start.py +84 -1
  103. cognee/tests/test_conversation_history.py +45 -4
  104. cognee/tests/test_data/example_with_header.csv +3 -0
  105. cognee/tests/test_delete_bmw_example.py +60 -0
  106. cognee/tests/test_edge_ingestion.py +27 -0
  107. cognee/tests/test_feedback_enrichment.py +1 -1
  108. cognee/tests/test_library.py +6 -4
  109. cognee/tests/test_load.py +62 -0
  110. cognee/tests/test_multi_tenancy.py +165 -0
  111. cognee/tests/test_parallel_databases.py +2 -0
  112. cognee/tests/test_relational_db_migration.py +54 -2
  113. cognee/tests/test_search_db.py +7 -1
  114. cognee/tests/unit/api/test_conditional_authentication_endpoints.py +12 -3
  115. cognee/tests/unit/api/test_ontology_endpoint.py +264 -0
  116. cognee/tests/unit/infrastructure/databases/cache/test_cache_config.py +5 -0
  117. cognee/tests/unit/infrastructure/databases/test_index_data_points.py +27 -0
  118. cognee/tests/unit/infrastructure/databases/test_index_graph_edges.py +14 -16
  119. cognee/tests/unit/modules/chunking/test_text_chunker.py +248 -0
  120. cognee/tests/unit/modules/chunking/test_text_chunker_with_overlap.py +324 -0
  121. cognee/tests/unit/modules/memify_tasks/test_cognify_session.py +111 -0
  122. cognee/tests/unit/modules/memify_tasks/test_extract_user_sessions.py +175 -0
  123. cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py +0 -51
  124. cognee/tests/unit/modules/retrieval/rag_completion_retriever_test.py +1 -0
  125. cognee/tests/unit/modules/retrieval/structured_output_test.py +204 -0
  126. cognee/tests/unit/modules/retrieval/summaries_retriever_test.py +1 -1
  127. cognee/tests/unit/modules/retrieval/temporal_retriever_test.py +0 -1
  128. cognee/tests/unit/modules/users/test_conditional_authentication.py +0 -63
  129. cognee/tests/unit/processing/chunks/chunk_by_row_test.py +52 -0
  130. {cognee-0.4.1.dist-info → cognee-0.5.0.dev0.dist-info}/METADATA +88 -71
  131. {cognee-0.4.1.dist-info → cognee-0.5.0.dev0.dist-info}/RECORD +135 -104
  132. {cognee-0.4.1.dist-info → cognee-0.5.0.dev0.dist-info}/WHEEL +1 -1
  133. {cognee-0.4.1.dist-info → cognee-0.5.0.dev0.dist-info}/entry_points.txt +0 -1
  134. {cognee-0.4.1.dist-info → cognee-0.5.0.dev0.dist-info}/licenses/LICENSE +0 -0
  135. {cognee-0.4.1.dist-info → cognee-0.5.0.dev0.dist-info}/licenses/NOTICE.md +0 -0
@@ -1,3 +1,4 @@
1
+ import os
1
2
  import pytest
2
3
  from unittest.mock import patch, AsyncMock, MagicMock
3
4
  from uuid import uuid4
@@ -5,8 +6,6 @@ from fastapi.testclient import TestClient
5
6
  from types import SimpleNamespace
6
7
  import importlib
7
8
 
8
- from cognee.api.client import app
9
-
10
9
 
11
10
  # Fixtures for reuse across test classes
12
11
  @pytest.fixture
@@ -32,6 +31,10 @@ def mock_authenticated_user():
32
31
  )
33
32
 
34
33
 
34
+ # To turn off authentication we need to set the environment variable before importing the module
35
+ # Also both require_authentication and backend access control must be false
36
+ os.environ["REQUIRE_AUTHENTICATION"] = "false"
37
+ os.environ["ENABLE_BACKEND_ACCESS_CONTROL"] = "false"
35
38
  gau_mod = importlib.import_module("cognee.modules.users.methods.get_authenticated_user")
36
39
 
37
40
 
@@ -40,6 +43,8 @@ class TestConditionalAuthenticationEndpoints:
40
43
 
41
44
  @pytest.fixture
42
45
  def client(self):
46
+ from cognee.api.client import app
47
+
43
48
  """Create a test client."""
44
49
  return TestClient(app)
45
50
 
@@ -133,6 +138,8 @@ class TestConditionalAuthenticationBehavior:
133
138
 
134
139
  @pytest.fixture
135
140
  def client(self):
141
+ from cognee.api.client import app
142
+
136
143
  return TestClient(app)
137
144
 
138
145
  @pytest.mark.parametrize(
@@ -209,6 +216,8 @@ class TestConditionalAuthenticationErrorHandling:
209
216
 
210
217
  @pytest.fixture
211
218
  def client(self):
219
+ from cognee.api.client import app
220
+
212
221
  return TestClient(app)
213
222
 
214
223
  @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
@@ -232,7 +241,7 @@ class TestConditionalAuthenticationErrorHandling:
232
241
  # The exact error message may vary depending on the actual database connection
233
242
  # The important thing is that we get a 500 error when user creation fails
234
243
 
235
- def test_current_environment_configuration(self):
244
+ def test_current_environment_configuration(self, client):
236
245
  """Test that current environment configuration is working properly."""
237
246
  # This tests the actual module state without trying to change it
238
247
  from cognee.modules.users.methods.get_authenticated_user import (
@@ -0,0 +1,264 @@
1
+ import pytest
2
+ import uuid
3
+ from fastapi.testclient import TestClient
4
+ from unittest.mock import patch, Mock, AsyncMock
5
+ from types import SimpleNamespace
6
+ import importlib
7
+ from cognee.api.client import app
8
+
9
+ gau_mod = importlib.import_module("cognee.modules.users.methods.get_authenticated_user")
10
+
11
+
12
+ @pytest.fixture
13
+ def client():
14
+ return TestClient(app)
15
+
16
+
17
+ @pytest.fixture
18
+ def mock_user():
19
+ user = Mock()
20
+ user.id = "test-user-123"
21
+ return user
22
+
23
+
24
+ @pytest.fixture
25
+ def mock_default_user():
26
+ """Mock default user for testing."""
27
+ return SimpleNamespace(
28
+ id=uuid.uuid4(), email="default@example.com", is_active=True, tenant_id=uuid.uuid4()
29
+ )
30
+
31
+
32
+ @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
33
+ def test_upload_ontology_success(mock_get_default_user, client, mock_default_user):
34
+ """Test successful ontology upload"""
35
+ import json
36
+
37
+ mock_get_default_user.return_value = mock_default_user
38
+ ontology_content = (
39
+ b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
40
+ )
41
+ unique_key = f"test_ontology_{uuid.uuid4().hex[:8]}"
42
+
43
+ response = client.post(
44
+ "/api/v1/ontologies",
45
+ files=[("ontology_file", ("test.owl", ontology_content, "application/xml"))],
46
+ data={"ontology_key": json.dumps([unique_key]), "description": json.dumps(["Test"])},
47
+ )
48
+
49
+ assert response.status_code == 200
50
+ data = response.json()
51
+ assert data["uploaded_ontologies"][0]["ontology_key"] == unique_key
52
+ assert "uploaded_at" in data["uploaded_ontologies"][0]
53
+
54
+
55
+ @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
56
+ def test_upload_ontology_invalid_file(mock_get_default_user, client, mock_default_user):
57
+ """Test 400 response for non-.owl files"""
58
+ mock_get_default_user.return_value = mock_default_user
59
+ unique_key = f"test_ontology_{uuid.uuid4().hex[:8]}"
60
+ response = client.post(
61
+ "/api/v1/ontologies",
62
+ files={"ontology_file": ("test.txt", b"not xml")},
63
+ data={"ontology_key": unique_key},
64
+ )
65
+ assert response.status_code == 400
66
+
67
+
68
+ @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
69
+ def test_upload_ontology_missing_data(mock_get_default_user, client, mock_default_user):
70
+ """Test 400 response for missing file or key"""
71
+ import json
72
+
73
+ mock_get_default_user.return_value = mock_default_user
74
+ # Missing file
75
+ response = client.post("/api/v1/ontologies", data={"ontology_key": json.dumps(["test"])})
76
+ assert response.status_code == 400
77
+
78
+ # Missing key
79
+ response = client.post(
80
+ "/api/v1/ontologies", files=[("ontology_file", ("test.owl", b"xml", "application/xml"))]
81
+ )
82
+ assert response.status_code == 400
83
+
84
+
85
+ @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
86
+ def test_upload_ontology_unauthorized(mock_get_default_user, client, mock_default_user):
87
+ """Test behavior when default user is provided (no explicit authentication)"""
88
+ import json
89
+
90
+ unique_key = f"test_ontology_{uuid.uuid4().hex[:8]}"
91
+ mock_get_default_user.return_value = mock_default_user
92
+ response = client.post(
93
+ "/api/v1/ontologies",
94
+ files=[("ontology_file", ("test.owl", b"<rdf></rdf>", "application/xml"))],
95
+ data={"ontology_key": json.dumps([unique_key])},
96
+ )
97
+
98
+ # The current system provides a default user when no explicit authentication is given
99
+ # This test verifies the system works with conditional authentication
100
+ assert response.status_code == 200
101
+ data = response.json()
102
+ assert data["uploaded_ontologies"][0]["ontology_key"] == unique_key
103
+ assert "uploaded_at" in data["uploaded_ontologies"][0]
104
+
105
+
106
+ @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
107
+ def test_upload_multiple_ontologies(mock_get_default_user, client, mock_default_user):
108
+ """Test uploading multiple ontology files in single request"""
109
+ import io
110
+
111
+ # Create mock files
112
+ file1_content = b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
113
+ file2_content = b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
114
+
115
+ files = [
116
+ ("ontology_file", ("vehicles.owl", io.BytesIO(file1_content), "application/xml")),
117
+ ("ontology_file", ("manufacturers.owl", io.BytesIO(file2_content), "application/xml")),
118
+ ]
119
+ data = {
120
+ "ontology_key": '["vehicles", "manufacturers"]',
121
+ "descriptions": '["Base vehicles", "Car manufacturers"]',
122
+ }
123
+
124
+ response = client.post("/api/v1/ontologies", files=files, data=data)
125
+
126
+ assert response.status_code == 200
127
+ result = response.json()
128
+ assert "uploaded_ontologies" in result
129
+ assert len(result["uploaded_ontologies"]) == 2
130
+ assert result["uploaded_ontologies"][0]["ontology_key"] == "vehicles"
131
+ assert result["uploaded_ontologies"][1]["ontology_key"] == "manufacturers"
132
+
133
+
134
+ @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
135
+ def test_upload_endpoint_accepts_arrays(mock_get_default_user, client, mock_default_user):
136
+ """Test that upload endpoint accepts array parameters"""
137
+ import io
138
+ import json
139
+
140
+ file_content = b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
141
+
142
+ files = [("ontology_file", ("single.owl", io.BytesIO(file_content), "application/xml"))]
143
+ data = {
144
+ "ontology_key": json.dumps(["single_key"]),
145
+ "descriptions": json.dumps(["Single ontology"]),
146
+ }
147
+
148
+ response = client.post("/api/v1/ontologies", files=files, data=data)
149
+
150
+ assert response.status_code == 200
151
+ result = response.json()
152
+ assert result["uploaded_ontologies"][0]["ontology_key"] == "single_key"
153
+
154
+
155
+ @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
156
+ def test_cognify_with_multiple_ontologies(mock_get_default_user, client, mock_default_user):
157
+ """Test cognify endpoint accepts multiple ontology keys"""
158
+ payload = {
159
+ "datasets": ["test_dataset"],
160
+ "ontology_key": ["ontology1", "ontology2"], # Array instead of string
161
+ "run_in_background": False,
162
+ }
163
+
164
+ response = client.post("/api/v1/cognify", json=payload)
165
+
166
+ # Should not fail due to ontology_key type
167
+ assert response.status_code in [200, 400, 409] # May fail for other reasons, not type
168
+
169
+
170
+ @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
171
+ def test_complete_multifile_workflow(mock_get_default_user, client, mock_default_user):
172
+ """Test complete workflow: upload multiple ontologies → cognify with multiple keys"""
173
+ import io
174
+ import json
175
+
176
+ # Step 1: Upload multiple ontologies
177
+ file1_content = b"""<?xml version="1.0"?>
178
+ <rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#"
179
+ xmlns:owl="http://www.w3.org/2002/07/owl#">
180
+ <owl:Class rdf:ID="Vehicle"/>
181
+ </rdf:RDF>"""
182
+
183
+ file2_content = b"""<?xml version="1.0"?>
184
+ <rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#"
185
+ xmlns:owl="http://www.w3.org/2002/07/owl#">
186
+ <owl:Class rdf:ID="Manufacturer"/>
187
+ </rdf:RDF>"""
188
+
189
+ files = [
190
+ ("ontology_file", ("vehicles.owl", io.BytesIO(file1_content), "application/xml")),
191
+ ("ontology_file", ("manufacturers.owl", io.BytesIO(file2_content), "application/xml")),
192
+ ]
193
+ data = {
194
+ "ontology_key": json.dumps(["vehicles", "manufacturers"]),
195
+ "descriptions": json.dumps(["Vehicle ontology", "Manufacturer ontology"]),
196
+ }
197
+
198
+ upload_response = client.post("/api/v1/ontologies", files=files, data=data)
199
+ assert upload_response.status_code == 200
200
+
201
+ # Step 2: Verify ontologies are listed
202
+ list_response = client.get("/api/v1/ontologies")
203
+ assert list_response.status_code == 200
204
+ ontologies = list_response.json()
205
+ assert "vehicles" in ontologies
206
+ assert "manufacturers" in ontologies
207
+
208
+ # Step 3: Test cognify with multiple ontologies
209
+ cognify_payload = {
210
+ "datasets": ["test_dataset"],
211
+ "ontology_key": ["vehicles", "manufacturers"],
212
+ "run_in_background": False,
213
+ }
214
+
215
+ cognify_response = client.post("/api/v1/cognify", json=cognify_payload)
216
+ # Should not fail due to ontology handling (may fail for dataset reasons)
217
+ assert cognify_response.status_code != 400 # Not a validation error
218
+
219
+
220
+ @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
221
+ def test_multifile_error_handling(mock_get_default_user, client, mock_default_user):
222
+ """Test error handling for invalid multifile uploads"""
223
+ import io
224
+ import json
225
+
226
+ # Test mismatched array lengths
227
+ file_content = b"<rdf:RDF></rdf:RDF>"
228
+ files = [("ontology_file", ("test.owl", io.BytesIO(file_content), "application/xml"))]
229
+ data = {
230
+ "ontology_key": json.dumps(["key1", "key2"]), # 2 keys, 1 file
231
+ "descriptions": json.dumps(["desc1"]),
232
+ }
233
+
234
+ response = client.post("/api/v1/ontologies", files=files, data=data)
235
+ assert response.status_code == 400
236
+ assert "Number of keys must match number of files" in response.json()["error"]
237
+
238
+ # Test duplicate keys
239
+ files = [
240
+ ("ontology_file", ("test1.owl", io.BytesIO(file_content), "application/xml")),
241
+ ("ontology_file", ("test2.owl", io.BytesIO(file_content), "application/xml")),
242
+ ]
243
+ data = {
244
+ "ontology_key": json.dumps(["duplicate", "duplicate"]),
245
+ "descriptions": json.dumps(["desc1", "desc2"]),
246
+ }
247
+
248
+ response = client.post("/api/v1/ontologies", files=files, data=data)
249
+ assert response.status_code == 400
250
+ assert "Duplicate ontology keys not allowed" in response.json()["error"]
251
+
252
+
253
+ @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
254
+ def test_cognify_missing_ontology_key(mock_get_default_user, client, mock_default_user):
255
+ """Test cognify with non-existent ontology key"""
256
+ payload = {
257
+ "datasets": ["test_dataset"],
258
+ "ontology_key": ["nonexistent_key"],
259
+ "run_in_background": False,
260
+ }
261
+
262
+ response = client.post("/api/v1/cognify", json=payload)
263
+ assert response.status_code == 409
264
+ assert "Ontology key 'nonexistent_key' not found" in response.json()["error"]
@@ -8,6 +8,7 @@ def test_cache_config_defaults():
8
8
  """Test that CacheConfig has the correct default values."""
9
9
  config = CacheConfig()
10
10
 
11
+ assert config.cache_backend == "fs"
11
12
  assert config.caching is False
12
13
  assert config.shared_kuzu_lock is False
13
14
  assert config.cache_host == "localhost"
@@ -19,6 +20,7 @@ def test_cache_config_defaults():
19
20
  def test_cache_config_custom_values():
20
21
  """Test that CacheConfig accepts custom values."""
21
22
  config = CacheConfig(
23
+ cache_backend="redis",
22
24
  caching=True,
23
25
  shared_kuzu_lock=True,
24
26
  cache_host="redis.example.com",
@@ -27,6 +29,7 @@ def test_cache_config_custom_values():
27
29
  agentic_lock_timeout=180,
28
30
  )
29
31
 
32
+ assert config.cache_backend == "redis"
30
33
  assert config.caching is True
31
34
  assert config.shared_kuzu_lock is True
32
35
  assert config.cache_host == "redis.example.com"
@@ -38,6 +41,7 @@ def test_cache_config_custom_values():
38
41
  def test_cache_config_to_dict():
39
42
  """Test the to_dict method returns all configuration values."""
40
43
  config = CacheConfig(
44
+ cache_backend="fs",
41
45
  caching=True,
42
46
  shared_kuzu_lock=True,
43
47
  cache_host="test-host",
@@ -49,6 +53,7 @@ def test_cache_config_to_dict():
49
53
  config_dict = config.to_dict()
50
54
 
51
55
  assert config_dict == {
56
+ "cache_backend": "fs",
52
57
  "caching": True,
53
58
  "shared_kuzu_lock": True,
54
59
  "cache_host": "test-host",
@@ -0,0 +1,27 @@
1
+ import pytest
2
+ from unittest.mock import AsyncMock, patch, MagicMock
3
+ from cognee.tasks.storage.index_data_points import index_data_points
4
+ from cognee.infrastructure.engine import DataPoint
5
+
6
+
7
+ class TestDataPoint(DataPoint):
8
+ name: str
9
+ metadata: dict = {"index_fields": ["name"]}
10
+
11
+
12
+ @pytest.mark.asyncio
13
+ async def test_index_data_points_calls_vector_engine():
14
+ """Test that index_data_points creates vector index and indexes data."""
15
+ data_points = [TestDataPoint(name="test1")]
16
+
17
+ mock_vector_engine = AsyncMock()
18
+ mock_vector_engine.embedding_engine.get_batch_size = MagicMock(return_value=100)
19
+
20
+ with patch.dict(
21
+ index_data_points.__globals__,
22
+ {"get_vector_engine": lambda: mock_vector_engine},
23
+ ):
24
+ await index_data_points(data_points)
25
+
26
+ assert mock_vector_engine.create_vector_index.await_count >= 1
27
+ assert mock_vector_engine.index_data_points.await_count >= 1
@@ -5,8 +5,7 @@ from cognee.tasks.storage.index_graph_edges import index_graph_edges
5
5
 
6
6
  @pytest.mark.asyncio
7
7
  async def test_index_graph_edges_success():
8
- """Test that index_graph_edges uses the index datapoints and creates vector index."""
9
- # Create the mocks for the graph and vector engines.
8
+ """Test that index_graph_edges retrieves edges and delegates to index_data_points."""
10
9
  mock_graph_engine = AsyncMock()
11
10
  mock_graph_engine.get_graph_data.return_value = (
12
11
  None,
@@ -15,26 +14,23 @@ async def test_index_graph_edges_success():
15
14
  [{"relationship_name": "rel2"}],
16
15
  ],
17
16
  )
18
- mock_vector_engine = AsyncMock()
19
- mock_vector_engine.embedding_engine.get_batch_size = MagicMock(return_value=100)
17
+ mock_index_data_points = AsyncMock()
20
18
 
21
- # Patch the globals of the function so that when it does:
22
- # vector_engine = get_vector_engine()
23
- # graph_engine = await get_graph_engine()
24
- # it uses the mocked versions.
25
19
  with patch.dict(
26
20
  index_graph_edges.__globals__,
27
21
  {
28
22
  "get_graph_engine": AsyncMock(return_value=mock_graph_engine),
29
- "get_vector_engine": lambda: mock_vector_engine,
23
+ "index_data_points": mock_index_data_points,
30
24
  },
31
25
  ):
32
26
  await index_graph_edges()
33
27
 
34
- # Assertions on the mock calls.
35
28
  mock_graph_engine.get_graph_data.assert_awaited_once()
36
- assert mock_vector_engine.create_vector_index.await_count == 1
37
- assert mock_vector_engine.index_data_points.await_count == 1
29
+ mock_index_data_points.assert_awaited_once()
30
+
31
+ call_args = mock_index_data_points.call_args[0][0]
32
+ assert len(call_args) == 2
33
+ assert all(hasattr(item, "relationship_name") for item in call_args)
38
34
 
39
35
 
40
36
  @pytest.mark.asyncio
@@ -42,20 +38,22 @@ async def test_index_graph_edges_no_relationships():
42
38
  """Test that index_graph_edges handles empty relationships correctly."""
43
39
  mock_graph_engine = AsyncMock()
44
40
  mock_graph_engine.get_graph_data.return_value = (None, [])
45
- mock_vector_engine = AsyncMock()
41
+ mock_index_data_points = AsyncMock()
46
42
 
47
43
  with patch.dict(
48
44
  index_graph_edges.__globals__,
49
45
  {
50
46
  "get_graph_engine": AsyncMock(return_value=mock_graph_engine),
51
- "get_vector_engine": lambda: mock_vector_engine,
47
+ "index_data_points": mock_index_data_points,
52
48
  },
53
49
  ):
54
50
  await index_graph_edges()
55
51
 
56
52
  mock_graph_engine.get_graph_data.assert_awaited_once()
57
- mock_vector_engine.create_vector_index.assert_not_awaited()
58
- mock_vector_engine.index_data_points.assert_not_awaited()
53
+ mock_index_data_points.assert_awaited_once()
54
+
55
+ call_args = mock_index_data_points.call_args[0][0]
56
+ assert len(call_args) == 0
59
57
 
60
58
 
61
59
  @pytest.mark.asyncio