cognee/cognee/api/v1/users/routers/get_visualize_router.py
Igor Ilic bcd418151a
fix: Secure api v2 (#1060)
<!-- .github/pull_request_template.md -->

## Description
<!-- Provide a clear description of the changes in this PR -->

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin.
2025-07-11 16:10:02 +02:00

35 lines
1.3 KiB
Python

from fastapi import APIRouter, Depends
from fastapi.responses import HTMLResponse, JSONResponse
from uuid import UUID
from cognee.shared.logging_utils import get_logger
from cognee.modules.users.methods import get_authenticated_user
from cognee.modules.data.methods import get_authorized_existing_datasets
from cognee.modules.users.models import User
from cognee.context_global_variables import set_database_global_context_variables
logger = get_logger()
def get_visualize_router() -> APIRouter:
router = APIRouter()
@router.get("", response_model=None)
async def visualize(dataset_id: UUID, user: User = Depends(get_authenticated_user)):
"""This endpoint is responsible for adding data to the graph."""
from cognee.api.v1.visualize import visualize_graph
try:
# Verify user has permission to read dataset
dataset = await get_authorized_existing_datasets([dataset_id], "read", user)
# Will only be used if ENABLE_BACKEND_ACCESS_CONTROL is set to True
await set_database_global_context_variables(dataset[0].id, dataset[0].owner_id)
html_visualization = await visualize_graph()
return HTMLResponse(html_visualization)
except Exception as error:
return JSONResponse(status_code=409, content={"error": str(error)})
return router