Update get_cognify_router.py

This commit is contained in:
Vasilije 2024-12-14 21:30:12 +01:00 committed by GitHub
parent 2f2aa81194
commit dc37646df3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -6,6 +6,7 @@ from fastapi.responses import JSONResponse
from cognee.modules.users.methods import get_authenticated_user from cognee.modules.users.methods import get_authenticated_user
from fastapi import Depends from fastapi import Depends
class CognifyPayloadDTO(BaseModel): class CognifyPayloadDTO(BaseModel):
datasets: List[str] datasets: List[str]
@ -13,11 +14,11 @@ def get_cognify_router() -> APIRouter:
router = APIRouter() router = APIRouter()
@router.post("/", response_model=None) @router.post("/", response_model=None)
async def cognify(payload: CognifyPayloadDTO, user: User = Depends(get_authenticated_user)): async def cognify(payload: CognifyPayloadDTO, user: User = Depends(get_authenticated_user), graph_model:BaseModel):
""" This endpoint is responsible for the cognitive processing of the content.""" """ This endpoint is responsible for the cognitive processing of the content."""
from cognee.api.v1.cognify.cognify_v2 import cognify as cognee_cognify from cognee.api.v1.cognify.cognify_v2 import cognify as cognee_cognify
try: try:
await cognee_cognify(payload.datasets, user) await cognee_cognify(payload.datasets, user, payload.graph_model)
except Exception as error: except Exception as error:
return JSONResponse( return JSONResponse(
status_code=409, status_code=409,