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"
)
ontology_key: Optional[str] = Field(
default=None,
description="Reference to previously uploaded ontology"
default=None, description="Reference to previously uploaded ontology"
)
@ -123,16 +122,22 @@ def get_cognify_router() -> APIRouter:
if payload.ontology_key:
ontology_service = OntologyService()
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.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver
from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import (
RDFLibOntologyResolver,
)
from io import StringIO
ontology_stream = StringIO(ontology_content)
config_to_use: Config = {
"ontology_config": {
"ontology_resolver": RDFLibOntologyResolver(ontology_file=ontology_stream)
"ontology_resolver": RDFLibOntologyResolver(
ontology_file=ontology_stream
)
}
}
except ValueError as e:

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

@ -1,4 +1,4 @@
from .ontologies import OntologyService
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 dataclasses import dataclass
@dataclass
class OntologyMetadata:
ontology_key: str
@ -14,6 +15,7 @@ class OntologyMetadata:
uploaded_at: str
description: Optional[str] = None
class OntologyService:
def __init__(self):
pass
@ -33,18 +35,20 @@ class OntologyService:
def _load_metadata(self, user_dir: Path) -> dict:
metadata_path = self._get_metadata_path(user_dir)
if metadata_path.exists():
with open(metadata_path, 'r') as f:
with open(metadata_path, "r") as f:
return json.load(f)
return {}
def _save_metadata(self, user_dir: Path, metadata: dict):
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)
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
if not file.filename.lower().endswith('.owl'):
if not file.filename.lower().endswith(".owl"):
raise ValueError("File must be in .owl format")
user_dir = self._get_user_dir(str(user.id))
@ -61,7 +65,7 @@ class OntologyService:
# Save file
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)
# Update metadata
@ -69,7 +73,7 @@ class OntologyService:
"filename": file.filename,
"size_bytes": len(content),
"uploaded_at": datetime.now(timezone.utc).isoformat(),
"description": description
"description": description,
}
metadata[ontology_key] = ontology_metadata
self._save_metadata(user_dir, metadata)
@ -79,7 +83,7 @@ class OntologyService:
filename=file.filename,
size_bytes=len(content),
uploaded_at=ontology_metadata["uploaded_at"],
description=description
description=description,
)
def get_ontology_content(self, ontology_key: str, user) -> str:
@ -93,9 +97,9 @@ class OntologyService:
if not file_path.exists():
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()
def list_ontologies(self, user) -> dict:
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 ..ontologies import OntologyService
def get_ontology_router() -> APIRouter:
router = APIRouter()
ontology_service = OntologyService()
@ -17,7 +18,7 @@ def get_ontology_router() -> APIRouter:
ontology_key: str = Form(...),
ontology_file: UploadFile = File(...),
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.
@ -51,7 +52,7 @@ def get_ontology_router() -> APIRouter:
"ontology_key": result.ontology_key,
"filename": result.filename,
"size_bytes": result.size_bytes,
"uploaded_at": result.uploaded_at
"uploaded_at": result.uploaded_at,
}
except ValueError as 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)})
@router.get("", response_model=dict)
async def list_ontologies(
user: User = Depends(get_authenticated_user)
):
async def list_ontologies(user: User = Depends(get_authenticated_user)):
"""
List all uploaded ontologies for the authenticated user.
@ -86,4 +85,4 @@ def get_ontology_router() -> APIRouter:
except Exception as 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")
@pytest.fixture
def client():
return TestClient(app)
@pytest.fixture
def mock_user():
user = Mock()
user.id = "test-user-123"
return user
@pytest.fixture
def mock_default_user():
"""Mock default user for testing."""
return SimpleNamespace(
id=uuid.uuid4(),
email="default@example.com",
is_active=True,
tenant_id=uuid.uuid4()
id=uuid.uuid4(), email="default@example.com", is_active=True, tenant_id=uuid.uuid4()
)
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
def test_upload_ontology_success(mock_get_default_user, client, mock_default_user):
"""Test successful ontology upload"""
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]}"
response = client.post(
"/api/v1/ontologies",
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
@ -46,6 +49,7 @@ def test_upload_ontology_success(mock_get_default_user, client, mock_default_use
assert data["ontology_key"] == unique_key
assert "uploaded_at" in data
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
def test_upload_ontology_invalid_file(mock_get_default_user, client, mock_default_user):
"""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(
"/api/v1/ontologies",
files={"ontology_file": ("test.txt", b"not xml")},
data={"ontology_key": unique_key}
data={"ontology_key": unique_key},
)
assert response.status_code == 400
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
def test_upload_ontology_missing_data(mock_get_default_user, client, mock_default_user):
"""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")})
assert response.status_code == 400
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
def test_upload_ontology_unauthorized(mock_get_default_user, client, mock_default_user):
"""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(
"/api/v1/ontologies",
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
@ -86,4 +92,4 @@ def test_upload_ontology_unauthorized(mock_get_default_user, client, mock_defaul
assert response.status_code == 200
data = response.json()
assert data["ontology_key"] == unique_key
assert "uploaded_at" in data
assert "uploaded_at" in data