feat: add return type hints to API endpoint functions

cognee/api/v1/delete/delete.py

@@ -1,6 +1,7 @@
 from uuid import UUID
 from sqlalchemy import select
 from sqlalchemy.sql import delete as sql_delete
+from typing import Dict, Any
 
 from cognee.infrastructure.engine import DataPoint
 from cognee.infrastructure.databases.graph import get_graph_engine
@@ -30,7 +31,7 @@ async def delete(
     dataset_id: UUID,
     mode: str = "soft",
     user: User = None,
-):
+) -> dict:
     """Delete data by its ID from the specified dataset.
 
     Args:
@@ -92,7 +93,7 @@ async def delete(
     return await delete_single_document(data_id, dataset.id, mode)
 
 
-async def delete_single_document(data_id: str, dataset_id: UUID = None, mode: str = "soft"):
+async def delete_single_document(data_id: str, dataset_id: UUID = None, mode: str = "soft") -> dict:
     """Delete a single document by its content hash."""
 
     # Delete from graph database
@@ -215,7 +216,7 @@ async def delete_single_document(data_id: str, dataset_id: UUID = None, mode: st
     }
 
 
-async def delete_document_subgraph(document_id: str, mode: str = "soft"):
+async def delete_document_subgraph(document_id: str, mode: str = "soft") -> dict:
     """Delete a document and all its related nodes in the correct order."""
     graph_db = await get_graph_engine()
     subgraph = await graph_db.get_document_subgraph(document_id)

cognee/api/v1/update/update.py

@@ -17,7 +17,7 @@ async def update(
     graph_db_config: dict = None,
     preferred_loaders: dict[str, dict[str, Any]] = None,
     incremental_loading: bool = True,
-):
+) -> Any:
     """
     Update existing data in Cognee.
 

cognee/api/v1/visualize/start_visualization_server.py

@@ -1,7 +1,7 @@
 from cognee.shared.utils import start_visualization_server
 
 
-def visualization_server(port):
+def visualization_server(port) -> callable:
     """
     Start a visualization server on the specified port.
 

cognee/api/v1/visualize/visualize.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 from cognee.modules.visualization.cognee_network_visualization import (
     cognee_network_visualization,
 )
@@ -11,7 +13,7 @@
 logger = get_logger()
 
 
-async def visualize_graph(destination_file_path: str = None):
+async def visualize_graph(destination_file_path: str = None) -> Any:
     graph_engine = await get_graph_engine()
     graph_data = await graph_engine.get_graph_data()
 

cognee/modules/ingestion/classify.py

@@ -7,7 +7,9 @@
 from .data_types import TextData, BinaryData, S3BinaryData
 
 
-def classify(data: Union[str, BinaryIO], filename: str = None):
+def classify(
+    data: Union[str, BinaryIO], filename: str = None
+) -> Union[TextData, BinaryData, S3BinaryData]:
     if isinstance(data, str):
         return TextData(data)