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.
- cognee/__init__.py +1 -0
- cognee/api/client.py +8 -0
- cognee/api/v1/add/routers/get_add_router.py +3 -1
- cognee/api/v1/cognify/routers/get_cognify_router.py +28 -1
- cognee/api/v1/ontologies/__init__.py +4 -0
- cognee/api/v1/ontologies/ontologies.py +183 -0
- cognee/api/v1/ontologies/routers/__init__.py +0 -0
- cognee/api/v1/ontologies/routers/get_ontology_router.py +107 -0
- cognee/api/v1/permissions/routers/get_permissions_router.py +41 -1
- cognee/cli/commands/cognify_command.py +8 -1
- cognee/cli/config.py +1 -1
- cognee/context_global_variables.py +41 -9
- cognee/infrastructure/databases/cache/config.py +3 -1
- cognee/infrastructure/databases/cache/fscache/FsCacheAdapter.py +151 -0
- cognee/infrastructure/databases/cache/get_cache_engine.py +20 -10
- cognee/infrastructure/databases/exceptions/exceptions.py +16 -0
- cognee/infrastructure/databases/graph/config.py +4 -0
- cognee/infrastructure/databases/graph/get_graph_engine.py +2 -0
- cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py +9 -0
- cognee/infrastructure/databases/utils/get_or_create_dataset_database.py +37 -3
- cognee/infrastructure/databases/vector/config.py +3 -0
- cognee/infrastructure/databases/vector/create_vector_engine.py +5 -1
- cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py +1 -4
- cognee/infrastructure/engine/models/Edge.py +13 -1
- cognee/infrastructure/files/utils/guess_file_type.py +4 -0
- cognee/infrastructure/llm/config.py +2 -0
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py +5 -2
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py +7 -1
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/generic_llm_api/adapter.py +7 -1
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/get_llm_client.py +8 -16
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/mistral/adapter.py +12 -2
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/ollama/adapter.py +13 -2
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py +5 -2
- cognee/infrastructure/loaders/LoaderEngine.py +1 -0
- cognee/infrastructure/loaders/core/__init__.py +2 -1
- cognee/infrastructure/loaders/core/csv_loader.py +93 -0
- cognee/infrastructure/loaders/core/text_loader.py +1 -2
- cognee/infrastructure/loaders/external/advanced_pdf_loader.py +0 -9
- cognee/infrastructure/loaders/supported_loaders.py +2 -1
- cognee/memify_pipelines/persist_sessions_in_knowledge_graph.py +55 -0
- cognee/modules/chunking/CsvChunker.py +35 -0
- cognee/modules/chunking/models/DocumentChunk.py +2 -1
- cognee/modules/chunking/text_chunker_with_overlap.py +124 -0
- cognee/modules/data/methods/__init__.py +1 -0
- cognee/modules/data/methods/create_dataset.py +4 -2
- cognee/modules/data/methods/get_dataset_ids.py +5 -1
- cognee/modules/data/methods/get_unique_data_id.py +68 -0
- cognee/modules/data/methods/get_unique_dataset_id.py +66 -4
- cognee/modules/data/models/Dataset.py +2 -0
- cognee/modules/data/processing/document_types/CsvDocument.py +33 -0
- cognee/modules/data/processing/document_types/__init__.py +1 -0
- cognee/modules/graph/cognee_graph/CogneeGraph.py +4 -2
- cognee/modules/graph/utils/expand_with_nodes_and_edges.py +19 -2
- cognee/modules/graph/utils/resolve_edges_to_text.py +48 -49
- cognee/modules/ingestion/identify.py +4 -4
- cognee/modules/notebooks/operations/run_in_local_sandbox.py +3 -0
- cognee/modules/ontology/rdf_xml/RDFLibOntologyResolver.py +55 -23
- cognee/modules/pipelines/operations/run_tasks_data_item.py +1 -1
- cognee/modules/retrieval/EntityCompletionRetriever.py +10 -3
- cognee/modules/retrieval/base_graph_retriever.py +7 -3
- cognee/modules/retrieval/base_retriever.py +7 -3
- cognee/modules/retrieval/completion_retriever.py +11 -4
- cognee/modules/retrieval/graph_completion_context_extension_retriever.py +6 -2
- cognee/modules/retrieval/graph_completion_cot_retriever.py +14 -51
- cognee/modules/retrieval/graph_completion_retriever.py +4 -1
- cognee/modules/retrieval/temporal_retriever.py +9 -2
- cognee/modules/retrieval/utils/brute_force_triplet_search.py +1 -1
- cognee/modules/retrieval/utils/completion.py +2 -22
- cognee/modules/run_custom_pipeline/__init__.py +1 -0
- cognee/modules/run_custom_pipeline/run_custom_pipeline.py +69 -0
- cognee/modules/search/methods/search.py +5 -3
- cognee/modules/users/methods/create_user.py +12 -27
- cognee/modules/users/methods/get_authenticated_user.py +2 -1
- cognee/modules/users/methods/get_default_user.py +4 -2
- cognee/modules/users/methods/get_user.py +1 -1
- cognee/modules/users/methods/get_user_by_email.py +1 -1
- cognee/modules/users/models/DatasetDatabase.py +9 -0
- cognee/modules/users/models/Tenant.py +6 -7
- cognee/modules/users/models/User.py +6 -5
- cognee/modules/users/models/UserTenant.py +12 -0
- cognee/modules/users/models/__init__.py +1 -0
- cognee/modules/users/permissions/methods/get_all_user_permission_datasets.py +13 -13
- cognee/modules/users/roles/methods/add_user_to_role.py +3 -1
- cognee/modules/users/tenants/methods/__init__.py +1 -0
- cognee/modules/users/tenants/methods/add_user_to_tenant.py +21 -12
- cognee/modules/users/tenants/methods/create_tenant.py +22 -8
- cognee/modules/users/tenants/methods/select_tenant.py +62 -0
- cognee/shared/logging_utils.py +2 -0
- cognee/tasks/chunks/__init__.py +1 -0
- cognee/tasks/chunks/chunk_by_row.py +94 -0
- cognee/tasks/documents/classify_documents.py +2 -0
- cognee/tasks/feedback/generate_improved_answers.py +3 -3
- cognee/tasks/ingestion/ingest_data.py +1 -1
- cognee/tasks/memify/__init__.py +2 -0
- cognee/tasks/memify/cognify_session.py +41 -0
- cognee/tasks/memify/extract_user_sessions.py +73 -0
- cognee/tasks/storage/index_data_points.py +33 -22
- cognee/tasks/storage/index_graph_edges.py +37 -57
- cognee/tests/integration/documents/CsvDocument_test.py +70 -0
- cognee/tests/tasks/entity_extraction/entity_extraction_test.py +1 -1
- cognee/tests/test_add_docling_document.py +2 -2
- cognee/tests/test_cognee_server_start.py +84 -1
- cognee/tests/test_conversation_history.py +45 -4
- cognee/tests/test_data/example_with_header.csv +3 -0
- cognee/tests/test_delete_bmw_example.py +60 -0
- cognee/tests/test_edge_ingestion.py +27 -0
- cognee/tests/test_feedback_enrichment.py +1 -1
- cognee/tests/test_library.py +6 -4
- cognee/tests/test_load.py +62 -0
- cognee/tests/test_multi_tenancy.py +165 -0
- cognee/tests/test_parallel_databases.py +2 -0
- cognee/tests/test_relational_db_migration.py +54 -2
- cognee/tests/test_search_db.py +7 -1
- cognee/tests/unit/api/test_conditional_authentication_endpoints.py +12 -3
- cognee/tests/unit/api/test_ontology_endpoint.py +264 -0
- cognee/tests/unit/infrastructure/databases/cache/test_cache_config.py +5 -0
- cognee/tests/unit/infrastructure/databases/test_index_data_points.py +27 -0
- cognee/tests/unit/infrastructure/databases/test_index_graph_edges.py +14 -16
- cognee/tests/unit/modules/chunking/test_text_chunker.py +248 -0
- cognee/tests/unit/modules/chunking/test_text_chunker_with_overlap.py +324 -0
- cognee/tests/unit/modules/memify_tasks/test_cognify_session.py +111 -0
- cognee/tests/unit/modules/memify_tasks/test_extract_user_sessions.py +175 -0
- cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py +0 -51
- cognee/tests/unit/modules/retrieval/rag_completion_retriever_test.py +1 -0
- cognee/tests/unit/modules/retrieval/structured_output_test.py +204 -0
- cognee/tests/unit/modules/retrieval/summaries_retriever_test.py +1 -1
- cognee/tests/unit/modules/retrieval/temporal_retriever_test.py +0 -1
- cognee/tests/unit/modules/users/test_conditional_authentication.py +0 -63
- cognee/tests/unit/processing/chunks/chunk_by_row_test.py +52 -0
- {cognee-0.4.1.dist-info → cognee-0.5.0.dev0.dist-info}/METADATA +88 -71
- {cognee-0.4.1.dist-info → cognee-0.5.0.dev0.dist-info}/RECORD +135 -104
- {cognee-0.4.1.dist-info → cognee-0.5.0.dev0.dist-info}/WHEEL +1 -1
- {cognee-0.4.1.dist-info → cognee-0.5.0.dev0.dist-info}/entry_points.txt +0 -1
- {cognee-0.4.1.dist-info → cognee-0.5.0.dev0.dist-info}/licenses/LICENSE +0 -0
- {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
|
|
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
|
-
|
|
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
|
-
"
|
|
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
|
-
|
|
37
|
-
|
|
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
|
-
|
|
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
|
-
"
|
|
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
|
-
|
|
58
|
-
|
|
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
|