From 57ae93b3cd0c6061a5192dfaa5dee35eb48b0520 Mon Sep 17 00:00:00 2001
From: Vasilije <8619304+Vasilije1990@users.noreply.github.com>
Date: Tue, 17 Dec 2024 18:53:28 +0100
Subject: [PATCH] Update get_cognify_router.py

---
 .../v1/cognify/routers/get_cognify_router.py  | 22 ++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/cognee/api/v1/cognify/routers/get_cognify_router.py b/cognee/api/v1/cognify/routers/get_cognify_router.py
index a56dbd7a5..1a53a7338 100644
--- a/cognee/api/v1/cognify/routers/get_cognify_router.py
+++ b/cognee/api/v1/cognify/routers/get_cognify_router.py
@@ -1,6 +1,6 @@
 from fastapi import APIRouter
-from typing import List, Optional
-from pydantic import BaseModel
+from typing import List, Optional, Any
+from pydantic import BaseModel, create_model
 from cognee.modules.users.models import User
 from fastapi.responses import JSONResponse
 from cognee.modules.users.methods import get_authenticated_user
@@ -9,7 +9,14 @@ from fastapi import Depends
 
 class CognifyPayloadDTO(BaseModel):
     datasets: List[str]
-    graph_model: Optional[BaseModel] = None
+    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()})
 
 def get_cognify_router() -> APIRouter:
     router = APIRouter()
@@ -18,6 +25,15 @@ def get_cognify_router() -> APIRouter:
     async def cognify(payload: CognifyPayloadDTO, user: User = Depends(get_authenticated_user)):
         """ 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 = None
         try:
             await cognee_cognify(payload.datasets, user, payload.graph_model)
         except Exception as error: