From b89e85850094ef989ce2df0fb9f8eb5a679c023a Mon Sep 17 00:00:00 2001
From: Allen Joseph Antony <allen.antony@stud.th-deg.de>
Date: Fri, 30 Jan 2026 20:52:54 +0200
Subject: [PATCH] feat: add return type hints to API endpoint functions

---
 cognee/api/v1/delete/delete.py                        | 7 ++++---
 cognee/api/v1/update/update.py                        | 2 +-
 cognee/api/v1/visualize/start_visualization_server.py | 2 +-
 cognee/api/v1/visualize/visualize.py                  | 4 +++-
 cognee/modules/ingestion/classify.py                  | 4 +++-
 5 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/cognee/api/v1/delete/delete.py b/cognee/api/v1/delete/delete.py
index 73f264670d..46657915ad 100644
--- a/cognee/api/v1/delete/delete.py
+++ b/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)
diff --git a/cognee/api/v1/update/update.py b/cognee/api/v1/update/update.py
index 83b92c50f0..89958a1697 100644
--- a/cognee/api/v1/update/update.py
+++ b/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.
 
diff --git a/cognee/api/v1/visualize/start_visualization_server.py b/cognee/api/v1/visualize/start_visualization_server.py
index 00b2543932..efa25bd2da 100644
--- a/cognee/api/v1/visualize/start_visualization_server.py
+++ b/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.
 
diff --git a/cognee/api/v1/visualize/visualize.py b/cognee/api/v1/visualize/visualize.py
index 583530f923..4d541abea3 100644
--- a/cognee/api/v1/visualize/visualize.py
+++ b/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()
 
diff --git a/cognee/modules/ingestion/classify.py b/cognee/modules/ingestion/classify.py
index 5fc8ed83c5..182c02101a 100644
--- a/cognee/modules/ingestion/classify.py
+++ b/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)