gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

chore: fixes ruff formatting

Browse source
This commit is contained in:
hajdul88 2025-11-06 09:48:01 +01:00
parent cc058b4482
commit 79bd2b2576
5 changed files with 44 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
cognee/api/v1/cognify/routers/get_cognify_router.py
View file

@ -42,8 +42,7 @@ class CognifyPayloadDTO(InDTO):
default="", description="Custom prompt for entity extraction and graph generation" default="", description="Custom prompt for entity extraction and graph generation"
) )
ontology_key: Optional[str] = Field( ontology_key: Optional[str] = Field(
default=None, default=None, description="Reference to previously uploaded ontology"
description="Reference to previously uploaded ontology"
) )
@ -123,16 +122,22 @@ def get_cognify_router() -> APIRouter:
if payload.ontology_key: if payload.ontology_key:
ontology_service = OntologyService() ontology_service = OntologyService()
try: try:
ontology_content = ontology_service.get_ontology_content(payload.ontology_key, user) ontology_content = ontology_service.get_ontology_content(
payload.ontology_key, user
)
from cognee.modules.ontology.ontology_config import Config from cognee.modules.ontology.ontology_config import Config
from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import (
RDFLibOntologyResolver,
)
from io import StringIO from io import StringIO
ontology_stream = StringIO(ontology_content) ontology_stream = StringIO(ontology_content)
config_to_use: Config = { config_to_use: Config = {
"ontology_config": { "ontology_config": {
"ontology_resolver": RDFLibOntologyResolver(ontology_file=ontology_stream) "ontology_resolver": RDFLibOntologyResolver(
ontology_file=ontology_stream
)
} }
} }
except ValueError as e: except ValueError as e:

2
cognee/api/v1/ontologies/__init__.py
View file

@ -1,4 +1,4 @@
from .ontologies import OntologyService from .ontologies import OntologyService
from .routers.get_ontology_router import get_ontology_router from .routers.get_ontology_router import get_ontology_router
__all__ = ["OntologyService", "get_ontology_router"] __all__ = ["OntologyService", "get_ontology_router"]

22
cognee/api/v1/ontologies/ontologies.py
View file

@ -6,6 +6,7 @@ from datetime import datetime, timezone
from typing import Optional from typing import Optional
from dataclasses import dataclass from dataclasses import dataclass
@dataclass @dataclass
class OntologyMetadata: class OntologyMetadata:
ontology_key: str ontology_key: str
@ -14,6 +15,7 @@ class OntologyMetadata:
uploaded_at: str uploaded_at: str
description: Optional[str] = None description: Optional[str] = None
class OntologyService: class OntologyService:
def __init__(self): def __init__(self):
pass pass
@ -33,18 +35,20 @@ class OntologyService:
def _load_metadata(self, user_dir: Path) -> dict: def _load_metadata(self, user_dir: Path) -> dict:
metadata_path = self._get_metadata_path(user_dir) metadata_path = self._get_metadata_path(user_dir)
if metadata_path.exists(): if metadata_path.exists():
with open(metadata_path, 'r') as f: with open(metadata_path, "r") as f:
return json.load(f) return json.load(f)
return {} return {}
def _save_metadata(self, user_dir: Path, metadata: dict): def _save_metadata(self, user_dir: Path, metadata: dict):
metadata_path = self._get_metadata_path(user_dir) metadata_path = self._get_metadata_path(user_dir)
with open(metadata_path, 'w') as f: with open(metadata_path, "w") as f:
json.dump(metadata, f, indent=2) json.dump(metadata, f, indent=2)
async def upload_ontology(self, ontology_key: str, file, user, description: Optional[str] = None) -> OntologyMetadata: async def upload_ontology(
self, ontology_key: str, file, user, description: Optional[str] = None
) -> OntologyMetadata:
# Validate file format # Validate file format
if not file.filename.lower().endswith('.owl'): if not file.filename.lower().endswith(".owl"):
raise ValueError("File must be in .owl format") raise ValueError("File must be in .owl format")
user_dir = self._get_user_dir(str(user.id)) user_dir = self._get_user_dir(str(user.id))
@ -61,7 +65,7 @@ class OntologyService:
# Save file # Save file
file_path = user_dir / f"{ontology_key}.owl" file_path = user_dir / f"{ontology_key}.owl"
with open(file_path, 'wb') as f: with open(file_path, "wb") as f:
f.write(content) f.write(content)
# Update metadata # Update metadata
@ -69,7 +73,7 @@ class OntologyService:
"filename": file.filename, "filename": file.filename,
"size_bytes": len(content), "size_bytes": len(content),
"uploaded_at": datetime.now(timezone.utc).isoformat(), "uploaded_at": datetime.now(timezone.utc).isoformat(),
"description": description "description": description,
} }
metadata[ontology_key] = ontology_metadata metadata[ontology_key] = ontology_metadata
self._save_metadata(user_dir, metadata) self._save_metadata(user_dir, metadata)
@ -79,7 +83,7 @@ class OntologyService:
filename=file.filename, filename=file.filename,
size_bytes=len(content), size_bytes=len(content),
uploaded_at=ontology_metadata["uploaded_at"], uploaded_at=ontology_metadata["uploaded_at"],
description=description description=description,
) )
def get_ontology_content(self, ontology_key: str, user) -> str: def get_ontology_content(self, ontology_key: str, user) -> str:
@ -93,9 +97,9 @@ class OntologyService:
if not file_path.exists(): if not file_path.exists():
raise ValueError(f"Ontology file for key '{ontology_key}' not found") raise ValueError(f"Ontology file for key '{ontology_key}' not found")
with open(file_path, 'r', encoding='utf-8') as f: with open(file_path, "r", encoding="utf-8") as f:
return f.read() return f.read()
def list_ontologies(self, user) -> dict: def list_ontologies(self, user) -> dict:
user_dir = self._get_user_dir(str(user.id)) user_dir = self._get_user_dir(str(user.id))
return self._load_metadata(user_dir) return self._load_metadata(user_dir)

11
cognee/api/v1/ontologies/routers/get_ontology_router.py
View file

@ -8,6 +8,7 @@ from cognee.shared.utils import send_telemetry
from cognee import __version__ as cognee_version from cognee import __version__ as cognee_version
from ..ontologies import OntologyService from ..ontologies import OntologyService
def get_ontology_router() -> APIRouter: def get_ontology_router() -> APIRouter:
router = APIRouter() router = APIRouter()
ontology_service = OntologyService() ontology_service = OntologyService()
@ -17,7 +18,7 @@ def get_ontology_router() -> APIRouter:
ontology_key: str = Form(...), ontology_key: str = Form(...),
ontology_file: UploadFile = File(...), ontology_file: UploadFile = File(...),
description: Optional[str] = Form(None), description: Optional[str] = Form(None),
user: User = Depends(get_authenticated_user) user: User = Depends(get_authenticated_user),
): ):
""" """
Upload an ontology file with a named key for later use in cognify operations. Upload an ontology file with a named key for later use in cognify operations.
@ -51,7 +52,7 @@ def get_ontology_router() -> APIRouter:
"ontology_key": result.ontology_key, "ontology_key": result.ontology_key,
"filename": result.filename, "filename": result.filename,
"size_bytes": result.size_bytes, "size_bytes": result.size_bytes,
"uploaded_at": result.uploaded_at "uploaded_at": result.uploaded_at,
} }
except ValueError as e: except ValueError as e:
return JSONResponse(status_code=400, content={"error": str(e)}) return JSONResponse(status_code=400, content={"error": str(e)})
@ -59,9 +60,7 @@ def get_ontology_router() -> APIRouter:
return JSONResponse(status_code=500, content={"error": str(e)}) return JSONResponse(status_code=500, content={"error": str(e)})
@router.get("", response_model=dict) @router.get("", response_model=dict)
async def list_ontologies( async def list_ontologies(user: User = Depends(get_authenticated_user)):
user: User = Depends(get_authenticated_user)
):
""" """
List all uploaded ontologies for the authenticated user. List all uploaded ontologies for the authenticated user.
@ -86,4 +85,4 @@ def get_ontology_router() -> APIRouter:
except Exception as e: except Exception as e:
return JSONResponse(status_code=500, content={"error": str(e)}) return JSONResponse(status_code=500, content={"error": str(e)})
return router return router

24
cognee/tests/test_ontology_endpoint.py
View file

@ -8,37 +8,40 @@ from cognee.api.client import app
gau_mod = importlib.import_module("cognee.modules.users.methods.get_authenticated_user") gau_mod = importlib.import_module("cognee.modules.users.methods.get_authenticated_user")
@pytest.fixture @pytest.fixture
def client(): def client():
return TestClient(app) return TestClient(app)
@pytest.fixture @pytest.fixture
def mock_user(): def mock_user():
user = Mock() user = Mock()
user.id = "test-user-123" user.id = "test-user-123"
return user return user
@pytest.fixture @pytest.fixture
def mock_default_user(): def mock_default_user():
"""Mock default user for testing.""" """Mock default user for testing."""
return SimpleNamespace( return SimpleNamespace(
id=uuid.uuid4(), id=uuid.uuid4(), email="default@example.com", is_active=True, tenant_id=uuid.uuid4()
email="default@example.com",
is_active=True,
tenant_id=uuid.uuid4()
) )
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock) @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
def test_upload_ontology_success(mock_get_default_user, client, mock_default_user): def test_upload_ontology_success(mock_get_default_user, client, mock_default_user):
"""Test successful ontology upload""" """Test successful ontology upload"""
mock_get_default_user.return_value = mock_default_user mock_get_default_user.return_value = mock_default_user
ontology_content = b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>" ontology_content = (
b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
)
unique_key = f"test_ontology_{uuid.uuid4().hex[:8]}" unique_key = f"test_ontology_{uuid.uuid4().hex[:8]}"
response = client.post( response = client.post(
"/api/v1/ontologies", "/api/v1/ontologies",
files={"ontology_file": ("test.owl", ontology_content)}, files={"ontology_file": ("test.owl", ontology_content)},
data={"ontology_key": unique_key, "description": "Test"} data={"ontology_key": unique_key, "description": "Test"},
) )
assert response.status_code == 200 assert response.status_code == 200
@ -46,6 +49,7 @@ def test_upload_ontology_success(mock_get_default_user, client, mock_default_use
assert data["ontology_key"] == unique_key assert data["ontology_key"] == unique_key
assert "uploaded_at" in data assert "uploaded_at" in data
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock) @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
def test_upload_ontology_invalid_file(mock_get_default_user, client, mock_default_user): def test_upload_ontology_invalid_file(mock_get_default_user, client, mock_default_user):
"""Test 400 response for non-.owl files""" """Test 400 response for non-.owl files"""
@ -54,10 +58,11 @@ def test_upload_ontology_invalid_file(mock_get_default_user, client, mock_defaul
response = client.post( response = client.post(
"/api/v1/ontologies", "/api/v1/ontologies",
files={"ontology_file": ("test.txt", b"not xml")}, files={"ontology_file": ("test.txt", b"not xml")},
data={"ontology_key": unique_key} data={"ontology_key": unique_key},
) )
assert response.status_code == 400 assert response.status_code == 400
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock) @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
def test_upload_ontology_missing_data(mock_get_default_user, client, mock_default_user): def test_upload_ontology_missing_data(mock_get_default_user, client, mock_default_user):
"""Test 400 response for missing file or key""" """Test 400 response for missing file or key"""
@ -70,6 +75,7 @@ def test_upload_ontology_missing_data(mock_get_default_user, client, mock_defaul
response = client.post("/api/v1/ontologies", files={"ontology_file": ("test.owl", b"xml")}) response = client.post("/api/v1/ontologies", files={"ontology_file": ("test.owl", b"xml")})
assert response.status_code == 400 assert response.status_code == 400
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock) @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
def test_upload_ontology_unauthorized(mock_get_default_user, client, mock_default_user): def test_upload_ontology_unauthorized(mock_get_default_user, client, mock_default_user):
"""Test behavior when default user is provided (no explicit authentication)""" """Test behavior when default user is provided (no explicit authentication)"""
@ -78,7 +84,7 @@ def test_upload_ontology_unauthorized(mock_get_default_user, client, mock_defaul
response = client.post( response = client.post(
"/api/v1/ontologies", "/api/v1/ontologies",
files={"ontology_file": ("test.owl", b"<rdf></rdf>")}, files={"ontology_file": ("test.owl", b"<rdf></rdf>")},
data={"ontology_key": unique_key} data={"ontology_key": unique_key},
) )
# The current system provides a default user when no explicit authentication is given # The current system provides a default user when no explicit authentication is given
@ -86,4 +92,4 @@ def test_upload_ontology_unauthorized(mock_get_default_user, client, mock_defaul
assert response.status_code == 200 assert response.status_code == 200
data = response.json() data = response.json()
assert data["ontology_key"] == unique_key assert data["ontology_key"] == unique_key
assert "uploaded_at" in data assert "uploaded_at" in data