Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Add graph_model parameter to cognify api call #377

Open
wants to merge 3 commits into
base: dev
Choose a base branch
from
Open

feat: Add graph_model parameter to cognify api call #377

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
57ae93b
Update get_cognify_router.py
Vasilije1990 Dec 17, 2024
a42ab3e
fix: Resolve syntax issue with cognify router
dexters1 Dec 18, 2024
a34d010
fix: instantiate KnowledgeGraph
borisarzentar Dec 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 19 additions & 3 deletions cognee/api/v1/cognify/routers/get_cognify_router.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -9,7 +9,14 @@

class CognifyPayloadDTO(BaseModel):
datasets: List[str]
graph_model: Optional[BaseModel] = None
graph_model: Optional[Any] = None
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Consider using a more specific type than Any

Using Optional[Any] reduces type safety and could lead to runtime errors. Consider using a Union type with specific expected types or creating a base interface that all graph models must implement.

-    graph_model: Optional[Any] = None
+    graph_model: Optional[Union[dict, BaseModel]] = None
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
graph_model: Optional[Any] = None
graph_model: Optional[Union[dict, BaseModel]] = None

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why changing to Any? We assume that graph model will be a pydantic model.



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()})
Comment on lines +17 to +21
Copy link
Contributor

@coderabbitai coderabbitai bot Dec 17, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Add schema validation and improve type mapping

The current implementation has several potential issues:

  1. No validation of the input schema structure
  2. Assumes all fields are required (...)
  3. Simple type inference might not handle complex nested types correctly

Consider this improved implementation:

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_field_type(value: Any) -> tuple:
+        if isinstance(value, dict):
+            return (dict, ...)
+        if isinstance(value, list):
+            return (List[Any], ...)
+        return (type(value), None)  # Make fields optional by default
+
+    if not isinstance(json_schema, dict):
+        raise ValueError("JSON schema must be a dictionary")
+
+    return create_model(
+        name,
+        **{k: get_field_type(v) for k, v in json_schema.items()}
+    )
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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 json_to_pydantic_model(name: str, json_schema: dict) -> BaseModel:
"""
Create a Pydantic model on the fly from JSON.
"""
def get_field_type(value: Any) -> tuple:
if isinstance(value, dict):
return (dict, ...)
if isinstance(value, list):
return (List[Any], ...)
return (type(value), None) # Make fields optional by default
if not isinstance(json_schema, dict):
raise ValueError("JSON schema must be a dictionary")
return create_model(
name,
**{k: get_field_type(v) for k, v in json_schema.items()}
)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Worth merging?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@borisarzentar, yes, the suggested improvements provide additional schema validation and handle nested data structures more robustly. By making fields optional by default, they also reduce the chances of unexpected errors. Unless you foresee any compatibility issues, it seems worth merging.


def get_cognify_router() -> APIRouter:
router = APIRouter()
Expand All @@ -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)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please remove the print statement

else:
graph_model_instance = None
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Fix control flow and remove debug code

Several issues need attention:

  1. Remove the print statement
  2. Fix indentation of the else clause
  3. Add error handling for model creation
  4. Restructure the nested try-except blocks

Apply this fix:

        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:
+            else:
+                graph_model_instance = None
+            
+            await cognee_cognify(payload.datasets, user, graph_model_instance)
+        except ValueError as error:
+            return JSONResponse(
+                status_code=400,
+                content={"error": f"Invalid graph model schema: {str(error)}"}
+            )
+        except Exception as error:
            return JSONResponse(
                status_code=409,
                content={"error": str(error)}
            )
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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:
# 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)
else:
graph_model_instance = None
await cognee_cognify(payload.datasets, user, graph_model_instance)
except ValueError as error:
return JSONResponse(
status_code=400,
content={"error": f"Invalid graph model schema: {str(error)}"}
)
except Exception as error:
return JSONResponse(
status_code=409,
content={"error": str(error)}
)

try:
await cognee_cognify(payload.datasets, user, payload.graph_model)
except Exception as error:
Expand Down