From f660127d2d25e0265cac1e865c8c962865551fe3 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Wed, 18 Dec 2024 14:24:14 +0100 Subject: [PATCH] refactor: Remove changes to cognify router Remove changes to cognify router Refactor COG-650 --- .../v1/cognify/routers/get_cognify_router.py | 26 +++---------------- 1 file changed, 4 insertions(+), 22 deletions(-) diff --git a/cognee/api/v1/cognify/routers/get_cognify_router.py b/cognee/api/v1/cognify/routers/get_cognify_router.py index 739bd1657..257ac994f 100644 --- a/cognee/api/v1/cognify/routers/get_cognify_router.py +++ b/cognee/api/v1/cognify/routers/get_cognify_router.py @@ -1,24 +1,15 @@ from fastapi import APIRouter -from typing import List, Optional, Any -from pydantic import BaseModel, create_model +from typing import List, Optional +from pydantic import BaseModel from cognee.modules.users.models import User from fastapi.responses import JSONResponse from cognee.modules.users.methods import get_authenticated_user from fastapi import Depends - from cognee.shared.data_models import KnowledgeGraph - class CognifyPayloadDTO(BaseModel): datasets: List[str] - graph_model: Optional[Any] = None - - -def json_to_pydantic_model(name: str, json_schema: dict) -> BaseModel: - """ - Create a Pydantic model on the fly from JSON. - """ - return create_model(name, **{k: (type(v), ...) for k, v in json_schema.items()}) + graph_model: Optional[BaseModel] = KnowledgeGraph def get_cognify_router() -> APIRouter: router = APIRouter() @@ -28,16 +19,7 @@ def get_cognify_router() -> APIRouter: """ This endpoint is responsible for the cognitive processing of the content.""" from cognee.api.v1.cognify.cognify_v2 import cognify as cognee_cognify try: - # Dynamic conversion of `graph_model` to Pydantic - if payload.graph_model: - graph_model_schema = payload.graph_model - GraphModelDynamic = json_to_pydantic_model("GraphModelDynamic", graph_model_schema) - graph_model_instance = GraphModelDynamic(**graph_model_schema) - print(graph_model_instance) - else: - graph_model_instance = KnowledgeGraph - - await cognee_cognify(payload.datasets, user, graph_model_instance) + await cognee_cognify(payload.datasets, user, payload.graph_model) except Exception as error: return JSONResponse( status_code=409,