From a3fd079fd6429cee7edea09b025c03f66626f12b Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Tue, 26 Nov 2024 19:26:38 +0100 Subject: [PATCH 01/13] feat: Add Exceptions and exception handlers Add classes for exceptions and add exception handling Feature COG-502 --- cognee/api/client.py | 19 ++++++++++ cognee/exceptions/__init__.py | 0 cognee/exceptions/exceptions.py | 67 +++++++++++++++++++++++++++++++++ 3 files changed, 86 insertions(+) create mode 100644 cognee/exceptions/__init__.py create mode 100644 cognee/exceptions/exceptions.py diff --git a/cognee/api/client.py b/cognee/api/client.py index 53c9f9762..39dd72506 100644 --- a/cognee/api/client.py +++ b/cognee/api/client.py @@ -7,6 +7,9 @@ from fastapi import FastAPI from fastapi.responses import JSONResponse, Response from fastapi.middleware.cors import CORSMiddleware +from cognee.exceptions.exceptions import CogneeApiError +from traceback import format_exc + # Set up logging logging.basicConfig( level=logging.INFO, # Set the logging level (e.g., DEBUG, INFO, WARNING, ERROR, CRITICAL) @@ -76,6 +79,22 @@ async def request_validation_exception_handler(request: Request, exc: RequestVal content = jsonable_encoder({"detail": exc.errors(), "body": exc.body}), ) +@app.exception_handler(CogneeApiError) +async def exception_handler(_: Request, exc: CogneeApiError) -> JSONResponse: + #TODO: Add checking if all values exist for exception + detail = {} + if exc.message: + detail["message"] = exc.message + + if exc.name: + detail["message"] = f"{detail['message']} [{exc.name}]" + + # log the stack trace for easier serverside debugging + logger.error(format_exc()) + return JSONResponse( + status_code=exc.status_code, content={"detail": detail["message"]} + ) + app.include_router( get_auth_router(), prefix = "/api/v1/auth", diff --git a/cognee/exceptions/__init__.py b/cognee/exceptions/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cognee/exceptions/exceptions.py b/cognee/exceptions/exceptions.py new file mode 100644 index 000000000..ed90db735 --- /dev/null +++ b/cognee/exceptions/exceptions.py @@ -0,0 +1,67 @@ +from fastapi import status + + +class CogneeApiError(Exception): + """Base exception class""" + + def __init__( + self, + message: str = "Service is unavailable", + name: str = "Cognee", + status_code=status.HTTP_418_IM_A_TEAPOT, + ): + self.message = message + self.name = name + self.status_code = status_code + super().__init__(self.message, self.name) + + +class ServiceError(CogneeApiError): + """Failures in external services or APIs, like a database or a third-party service""" + + def __init__( + self, + message: str = "Service is unavailable", + name: str = "ServiceError", + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + ): + self.message = message + self.name = name + self.status_code = status_code + super().__init__(self.message, self.name) + + +class EntityDoesNotExistError(CogneeApiError): + """Database returns nothing""" + + pass + + +class GroupNotFound(CogneeApiError): + """User group not found""" + + pass + + +class EntityAlreadyExistsError(CogneeApiError): + """Conflict detected, like trying to create a resource that already exists""" + + pass + + +class InvalidOperationError(CogneeApiError): + """Invalid operations like trying to delete a non-existing entity, etc.""" + + pass + + +class AuthenticationFailed(CogneeApiError): + """Invalid authentication credentials""" + + pass + + +class InvalidTokenError(CogneeApiError): + """Invalid token""" + + pass From ae568409a7e7a6f7c11b3f99b05cba1606025862 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Wed, 27 Nov 2024 14:29:33 +0100 Subject: [PATCH 02/13] feat: Add custom exceptions to cognee lib Added use of custom exceptions to cognee lib --- cognee/api/client.py | 20 ++-- cognee/api/v1/cognify/code_graph_pipeline.py | 5 - cognee/api/v1/cognify/cognify_v2.py | 5 - cognee/api/v1/config/config.py | 3 +- .../datasets/routers/get_datasets_router.py | 16 ++- .../routers/get_permissions_router.py | 10 +- cognee/api/v1/search/search.legacy.py | 4 +- cognee/api/v1/search/search_v2.py | 6 +- cognee/exceptions/__init__.py | 20 ++++ cognee/exceptions/exceptions.py | 104 ++++++++++++++---- .../data/utils/extract_keywords.py | 4 +- .../hybrid/falkordb/FalkorDBAdapter.py | 3 +- .../sqlalchemy/SqlAlchemyAdapter.py | 5 +- .../vector/lancedb/LanceDBAdapter.py | 6 +- .../vector/pgvector/PGVectorAdapter.py | 5 +- .../databases/vector/qdrant/QDrantAdapter.py | 3 +- .../vector/weaviate_db/WeaviateAdapter.py | 3 +- .../infrastructure/llm/anthropic/adapter.py | 4 +- .../llm/generic_llm_api/adapter.py | 4 +- cognee/infrastructure/llm/get_llm_client.py | 10 +- cognee/infrastructure/llm/openai/adapter.py | 3 +- cognee/modules/data/methods/delete_data.py | 3 +- .../modules/data/operations/translate_text.py | 6 +- .../modules/graph/cognee_graph/CogneeGraph.py | 15 +-- .../graph/cognee_graph/CogneeGraphElements.py | 11 +- cognee/modules/ingestion/classify.py | 6 +- cognee/modules/ingestion/exceptions.py | 6 - cognee/modules/users/get_user_manager.py | 5 +- .../methods/check_permission_on_documents.py | 9 +- cognee/tasks/graph/infer_data_ontology.py | 5 +- .../ingestion/save_data_item_to_storage.py | 4 +- 31 files changed, 206 insertions(+), 107 deletions(-) delete mode 100644 cognee/modules/ingestion/exceptions.py diff --git a/cognee/api/client.py b/cognee/api/client.py index 39dd72506..95a273351 100644 --- a/cognee/api/client.py +++ b/cognee/api/client.py @@ -3,11 +3,11 @@ import os import uvicorn import logging import sentry_sdk -from fastapi import FastAPI +from fastapi import FastAPI, status from fastapi.responses import JSONResponse, Response from fastapi.middleware.cors import CORSMiddleware -from cognee.exceptions.exceptions import CogneeApiError +from cognee.exceptions import CogneeApiError from traceback import format_exc # Set up logging @@ -81,18 +81,22 @@ async def request_validation_exception_handler(request: Request, exc: RequestVal @app.exception_handler(CogneeApiError) async def exception_handler(_: Request, exc: CogneeApiError) -> JSONResponse: - #TODO: Add checking if all values exist for exception detail = {} - if exc.message: - detail["message"] = exc.message - if exc.name: - detail["message"] = f"{detail['message']} [{exc.name}]" + if exc.name and exc.message and exc.status_code: + status_code = exc.status_code + detail["message"] = f"{exc.message} [{exc.name}]" + else: + # Log an error indicating the exception is improperly defined + logger.error("Improperly defined exception: %s", exc) + # Provide a default error response + detail["message"] = "An unexpected error occurred." + status_code = status.HTTP_418_IM_A_TEAPOT # log the stack trace for easier serverside debugging logger.error(format_exc()) return JSONResponse( - status_code=exc.status_code, content={"detail": detail["message"]} + status_code=status_code, content={"detail": detail["message"]} ) app.include_router( diff --git a/cognee/api/v1/cognify/code_graph_pipeline.py b/cognee/api/v1/cognify/code_graph_pipeline.py index 2cbb606c1..59c658300 100644 --- a/cognee/api/v1/cognify/code_graph_pipeline.py +++ b/cognee/api/v1/cognify/code_graph_pipeline.py @@ -22,11 +22,6 @@ logger = logging.getLogger("code_graph_pipeline") update_status_lock = asyncio.Lock() -class PermissionDeniedException(Exception): - def __init__(self, message: str): - self.message = message - super().__init__(self.message) - async def code_graph_pipeline(datasets: Union[str, list[str]] = None, user: User = None): if user is None: user = await get_default_user() diff --git a/cognee/api/v1/cognify/cognify_v2.py b/cognee/api/v1/cognify/cognify_v2.py index be9ecd1ce..791ab516a 100644 --- a/cognee/api/v1/cognify/cognify_v2.py +++ b/cognee/api/v1/cognify/cognify_v2.py @@ -24,11 +24,6 @@ logger = logging.getLogger("cognify.v2") update_status_lock = asyncio.Lock() -class PermissionDeniedException(Exception): - def __init__(self, message: str): - self.message = message - super().__init__(self.message) - async def cognify(datasets: Union[str, list[str]] = None, user: User = None): if user is None: user = await get_default_user() diff --git a/cognee/api/v1/config/config.py b/cognee/api/v1/config/config.py index 1fbed9bdc..af78fe6ba 100644 --- a/cognee/api/v1/config/config.py +++ b/cognee/api/v1/config/config.py @@ -1,6 +1,7 @@ """ This module is used to set the configuration of the system.""" import os from cognee.base_config import get_base_config +from cognee.exceptions import InvalidValueError from cognee.modules.cognify.config import get_cognify_config from cognee.infrastructure.data.chunking.config import get_chunk_config from cognee.infrastructure.databases.vector import get_vectordb_config @@ -153,7 +154,7 @@ class config(): base_config = get_base_config() if "username" not in graphistry_config or "password" not in graphistry_config: - raise ValueError("graphistry_config dictionary must contain 'username' and 'password' keys.") + raise InvalidValueError(message="graphistry_config dictionary must contain 'username' and 'password' keys.") base_config.graphistry_username = graphistry_config.get("username") base_config.graphistry_password = graphistry_config.get("password") diff --git a/cognee/api/v1/datasets/routers/get_datasets_router.py b/cognee/api/v1/datasets/routers/get_datasets_router.py index f27c6c2ad..a3a30278b 100644 --- a/cognee/api/v1/datasets/routers/get_datasets_router.py +++ b/cognee/api/v1/datasets/routers/get_datasets_router.py @@ -9,6 +9,7 @@ from fastapi.responses import JSONResponse, FileResponse from pydantic import BaseModel from cognee.api.DTO import OutDTO +from cognee.exceptions import EntityNotFoundError from cognee.modules.users.models import User from cognee.modules.users.methods import get_authenticated_user from cognee.modules.pipelines.models import PipelineRunStatus @@ -55,9 +56,8 @@ def get_datasets_router() -> APIRouter: dataset = await get_dataset(user.id, dataset_id) if dataset is None: - raise HTTPException( - status_code=404, - detail=f"Dataset ({dataset_id}) not found." + raise EntityNotFoundError( + message=f"Dataset ({dataset_id}) not found." ) await delete_dataset(dataset) @@ -72,17 +72,15 @@ def get_datasets_router() -> APIRouter: #TODO: Handle situation differently if user doesn't have permission to access data? if dataset is None: - raise HTTPException( - status_code=404, - detail=f"Dataset ({dataset_id}) not found." + raise EntityNotFoundError( + message=f"Dataset ({dataset_id}) not found." ) data = await get_data(data_id) if data is None: - raise HTTPException( - status_code=404, - detail=f"Dataset ({data_id}) not found." + raise EntityNotFoundError( + message=f"Data ({data_id}) not found." ) await delete_data(data) diff --git a/cognee/api/v1/permissions/routers/get_permissions_router.py b/cognee/api/v1/permissions/routers/get_permissions_router.py index ab20fb1a2..30491ec7d 100644 --- a/cognee/api/v1/permissions/routers/get_permissions_router.py +++ b/cognee/api/v1/permissions/routers/get_permissions_router.py @@ -1,6 +1,8 @@ from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import JSONResponse from sqlalchemy.orm import Session + +from cognee.exceptions import GroupNotFoundError, UserNotFoundError from cognee.modules.users import get_user_db from cognee.modules.users.models import User, Group, Permission @@ -12,7 +14,7 @@ def get_permissions_router() -> APIRouter: group = db.query(Group).filter(Group.id == group_id).first() if not group: - raise HTTPException(status_code = 404, detail = "Group not found") + raise GroupNotFoundError permission = db.query(Permission).filter(Permission.name == permission).first() @@ -31,8 +33,10 @@ def get_permissions_router() -> APIRouter: user = db.query(User).filter(User.id == user_id).first() group = db.query(Group).filter(Group.id == group_id).first() - if not user or not group: - raise HTTPException(status_code = 404, detail = "User or group not found") + if not user: + raise UserNotFoundError + elif not group: + raise GroupNotFoundError user.groups.append(group) diff --git a/cognee/api/v1/search/search.legacy.py b/cognee/api/v1/search/search.legacy.py index aaa22fd62..cea3b3874 100644 --- a/cognee/api/v1/search/search.legacy.py +++ b/cognee/api/v1/search/search.legacy.py @@ -9,6 +9,8 @@ from cognee.modules.search.graph.search_adjacent import search_adjacent from cognee.modules.search.vector.search_traverse import search_traverse from cognee.modules.search.graph.search_summary import search_summary from cognee.modules.search.graph.search_similarity import search_similarity + +from cognee.exceptions import UserNotFoundError from cognee.shared.utils import send_telemetry from cognee.modules.users.permissions.methods import get_document_ids_for_user from cognee.modules.users.methods import get_default_user @@ -47,7 +49,7 @@ async def search(search_type: str, params: Dict[str, Any], user: User = None) -> user = await get_default_user() if user is None: - raise PermissionError("No user found in the system. Please create a user.") + raise UserNotFoundError own_document_ids = await get_document_ids_for_user(user.id) search_params = SearchParameters(search_type = search_type, params = params) diff --git a/cognee/api/v1/search/search_v2.py b/cognee/api/v1/search/search_v2.py index c1bc0ee4d..d18ba3b6e 100644 --- a/cognee/api/v1/search/search_v2.py +++ b/cognee/api/v1/search/search_v2.py @@ -2,6 +2,8 @@ import json from uuid import UUID from enum import Enum from typing import Callable, Dict + +from cognee.exceptions import UserNotFoundError, InvalidValueError from cognee.modules.search.operations import log_query, log_result from cognee.modules.storage.utils import JSONEncoder from cognee.shared.utils import send_telemetry @@ -22,7 +24,7 @@ async def search(query_type: SearchType, query_text: str, user: User = None) -> user = await get_default_user() if user is None: - raise PermissionError("No user found in the system. Please create a user.") + raise UserNotFoundError query = await log_query(query_text, str(query_type), user.id) @@ -52,7 +54,7 @@ async def specific_search(query_type: SearchType, query: str, user) -> list: search_task = search_tasks.get(query_type) if search_task is None: - raise ValueError(f"Unsupported search type: {query_type}") + raise InvalidValueError(message=f"Unsupported search type: {query_type}") send_telemetry("cognee.search EXECUTION STARTED", user.id) diff --git a/cognee/exceptions/__init__.py b/cognee/exceptions/__init__.py index e69de29bb..e07547508 100644 --- a/cognee/exceptions/__init__.py +++ b/cognee/exceptions/__init__.py @@ -0,0 +1,20 @@ +""" +Custom exceptions for the Cognee API. + +This module defines a set of exceptions for handling various application errors, +such as service failures, resource conflicts, and invalid operations. +""" + +from .exceptions import ( + CogneeApiError, + ServiceError, + EntityNotFoundError, + GroupNotFoundError, + UserNotFoundError, + EntityAlreadyExistsError, + InvalidOperationError, + PermissionDeniedError, + IngestionError, + InvalidValueError, + InvalidAttributeError, +) \ No newline at end of file diff --git a/cognee/exceptions/exceptions.py b/cognee/exceptions/exceptions.py index ed90db735..f9228873f 100644 --- a/cognee/exceptions/exceptions.py +++ b/cognee/exceptions/exceptions.py @@ -6,7 +6,7 @@ class CogneeApiError(Exception): def __init__( self, - message: str = "Service is unavailable", + message: str = "Service is unavailable.", name: str = "Cognee", status_code=status.HTTP_418_IM_A_TEAPOT, ): @@ -21,47 +21,105 @@ class ServiceError(CogneeApiError): def __init__( self, - message: str = "Service is unavailable", + message: str = "Service is unavailable.", name: str = "ServiceError", status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, ): - self.message = message - self.name = name - self.status_code = status_code - super().__init__(self.message, self.name) + super().__init__(message, name, status_code) -class EntityDoesNotExistError(CogneeApiError): - """Database returns nothing""" - - pass - - -class GroupNotFound(CogneeApiError): +class GroupNotFoundError(CogneeApiError): """User group not found""" - pass + def __init__( + self, + message: str = "User group not found.", + name: str = "GroupNotFoundError", + status_code=status.HTTP_404_NOT_FOUND, + ): + super().__init__(message, name, status_code) + + +class UserNotFoundError(CogneeApiError): + """User not found""" + + def __init__( + self, + message: str = "No user found in the system. Please create a user.", + name: str = "UserNotFoundError", + status_code=status.HTTP_404_NOT_FOUND, + ): + super().__init__(message, name, status_code) + + +class EntityNotFoundError(CogneeApiError): + """Database returns nothing""" + + def __init__( + self, + message: str = "The requested entity does not exist.", + name: str = "EntityNotFoundError", + status_code=status.HTTP_404_NOT_FOUND, + ): + super().__init__(message, name, status_code) class EntityAlreadyExistsError(CogneeApiError): """Conflict detected, like trying to create a resource that already exists""" - pass + def __init__( + self, + message: str = "The entity already exists.", + name: str = "EntityAlreadyExistsError", + status_code=status.HTTP_409_CONFLICT, + ): + super().__init__(message, name, status_code) class InvalidOperationError(CogneeApiError): """Invalid operations like trying to delete a non-existing entity, etc.""" - pass + def __init__( + self, + message: str = "Invalid operation attempted.", + name: str = "InvalidOperationError", + status_code=status.HTTP_400_BAD_REQUEST, + ): + super().__init__(message, name, status_code) -class AuthenticationFailed(CogneeApiError): - """Invalid authentication credentials""" +class PermissionDeniedError(CogneeApiError): + def __init__( + self, + message: str = "User does not have permission on documents.", + name: str = "PermissionDeniedError", + status_code=status.HTTP_403_FORBIDDEN, + ): + super().__init__(message, name, status_code) - pass +class IngestionError(CogneeApiError): + def __init__( + self, + message: str = "Type of data sent to classify not supported.", + name: str = "IngestionError", + status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, + ): + super().__init__(message, name, status_code) +class InvalidValueError(CogneeApiError): + def __init__( + self, + message: str = "Invalid Value.", + name: str = "InvalidValueError", + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + ): + super().__init__(message, name, status_code) -class InvalidTokenError(CogneeApiError): - """Invalid token""" - - pass +class InvalidAttributeError(CogneeApiError): + def __init__( + self, + message: str = "Invalid attribute.", + name: str = "InvalidAttributeError", + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + ): + super().__init__(message, name, status_code) \ No newline at end of file diff --git a/cognee/infrastructure/data/utils/extract_keywords.py b/cognee/infrastructure/data/utils/extract_keywords.py index ab32ddefb..11f061889 100644 --- a/cognee/infrastructure/data/utils/extract_keywords.py +++ b/cognee/infrastructure/data/utils/extract_keywords.py @@ -1,9 +1,11 @@ from sklearn.feature_extraction.text import TfidfVectorizer + +from cognee.exceptions import InvalidValueError from cognee.shared.utils import extract_pos_tags def extract_keywords(text: str) -> list[str]: if len(text) == 0: - raise ValueError("extract_keywords cannot extract keywords from empty text.") + raise InvalidValueError(message="extract_keywords cannot extract keywords from empty text.") tags = extract_pos_tags(text) nouns = [word for (word, tag) in tags if tag == "NN"] diff --git a/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py b/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py index ea5a75088..a28d827a1 100644 --- a/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py +++ b/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py @@ -4,6 +4,7 @@ from typing import Any from uuid import UUID from falkordb import FalkorDB +from cognee.exceptions import InvalidValueError from cognee.infrastructure.engine import DataPoint from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface from cognee.infrastructure.databases.vector.embeddings import EmbeddingEngine @@ -200,7 +201,7 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface): with_vector: bool = False, ): if query_text is None and query_vector is None: - raise ValueError("One of query_text or query_vector must be provided!") + raise InvalidValueError(message="One of query_text or query_vector must be provided!") if query_text and not query_vector: query_vector = (await self.embed_data([query_text]))[0] diff --git a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py index aa2a022d3..a5ae36988 100644 --- a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py +++ b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py @@ -7,6 +7,7 @@ from sqlalchemy import text, select, MetaData, Table from sqlalchemy.orm import joinedload from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker +from cognee.exceptions import EntityNotFoundError from ..ModelBase import Base class SQLAlchemyAdapter(): @@ -117,7 +118,7 @@ class SQLAlchemyAdapter(): if table_name in Base.metadata.tables: return Base.metadata.tables[table_name] else: - raise ValueError(f"Table '{table_name}' not found.") + raise EntityNotFoundError(message=f"Table '{table_name}' not found.") else: # Create a MetaData instance to load table information metadata = MetaData() @@ -128,7 +129,7 @@ class SQLAlchemyAdapter(): # Check if table is in list of tables for the given schema if full_table_name in metadata.tables: return metadata.tables[full_table_name] - raise ValueError(f"Table '{full_table_name}' not found.") + raise EntityNotFoundError(message=f"Table '{full_table_name}' not found.") async def get_table_names(self) -> List[str]: """ diff --git a/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py b/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py index 96f026b4f..507b51448 100644 --- a/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py +++ b/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py @@ -5,6 +5,8 @@ from uuid import UUID import lancedb from pydantic import BaseModel from lancedb.pydantic import Vector, LanceModel + +from cognee.exceptions import InvalidValueError, EntityNotFoundError from cognee.infrastructure.engine import DataPoint from cognee.infrastructure.files.storage import LocalStorage from cognee.modules.storage.utils import copy_model, get_own_properties @@ -122,7 +124,7 @@ class LanceDBAdapter(VectorDBInterface): new_size = await collection.count_rows() if new_size <= original_size: - raise ValueError( + raise InvalidValueError(message= "LanceDB create_datapoints error: data points did not get added.") @@ -150,7 +152,7 @@ class LanceDBAdapter(VectorDBInterface): with_vector: bool = False, ): if query_text is None and query_vector is None: - raise ValueError("One of query_text or query_vector must be provided!") + raise InvalidValueError(message="One of query_text or query_vector must be provided!") if query_text and not query_vector: query_vector = (await self.embedding_engine.embed_text([query_text]))[0] diff --git a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py index 01691714b..77dd03185 100644 --- a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py +++ b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py @@ -6,6 +6,7 @@ from sqlalchemy.orm import Mapped, mapped_column from sqlalchemy import JSON, Column, Table, select, delete from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker +from cognee.exceptions import EntityNotFoundError, InvalidValueError from cognee.infrastructure.engine import DataPoint from .serialize_data import serialize_data @@ -142,7 +143,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface): if collection_name in Base.metadata.tables: return Base.metadata.tables[collection_name] else: - raise ValueError(f"Table '{collection_name}' not found.") + raise EntityNotFoundError(message=f"Table '{collection_name}' not found.") async def retrieve(self, collection_name: str, data_point_ids: List[str]): # Get PGVectorDataPoint Table from database @@ -171,7 +172,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface): with_vector: bool = False, ) -> List[ScoredResult]: if query_text is None and query_vector is None: - raise ValueError("One of query_text or query_vector must be provided!") + raise InvalidValueError(message="One of query_text or query_vector must be provided!") if query_text and not query_vector: query_vector = (await self.embedding_engine.embed_text([query_text]))[0] diff --git a/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py b/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py index 1efcd47b3..74cd574eb 100644 --- a/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py +++ b/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py @@ -3,6 +3,7 @@ from uuid import UUID from typing import List, Dict, Optional from qdrant_client import AsyncQdrantClient, models +from cognee.exceptions import InvalidValueError from cognee.infrastructure.databases.vector.models.ScoredResult import ScoredResult from cognee.infrastructure.engine import DataPoint from ..vector_db_interface import VectorDBInterface @@ -151,7 +152,7 @@ class QDrantAdapter(VectorDBInterface): with_vector: bool = False ): if query_text is None and query_vector is None: - raise ValueError("One of query_text or query_vector must be provided!") + raise InvalidValueError(message="One of query_text or query_vector must be provided!") client = self.get_qdrant_client() diff --git a/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py b/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py index be356740f..4fa7fc646 100644 --- a/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py +++ b/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py @@ -3,6 +3,7 @@ import logging from typing import List, Optional from uuid import UUID +from cognee.exceptions import InvalidValueError from cognee.infrastructure.engine import DataPoint from ..vector_db_interface import VectorDBInterface from ..models.ScoredResult import ScoredResult @@ -164,7 +165,7 @@ class WeaviateAdapter(VectorDBInterface): import weaviate.classes as wvc if query_text is None and query_vector is None: - raise ValueError("One of query_text or query_vector must be provided!") + raise InvalidValueError(message="One of query_text or query_vector must be provided!") if query_vector is None: query_vector = (await self.embed_data([query_text]))[0] diff --git a/cognee/infrastructure/llm/anthropic/adapter.py b/cognee/infrastructure/llm/anthropic/adapter.py index 8df59e3e5..b74beaf6e 100644 --- a/cognee/infrastructure/llm/anthropic/adapter.py +++ b/cognee/infrastructure/llm/anthropic/adapter.py @@ -3,6 +3,8 @@ from pydantic import BaseModel import instructor from tenacity import retry, stop_after_attempt import anthropic + +from cognee.exceptions import InvalidValueError from cognee.infrastructure.llm.llm_interface import LLMInterface from cognee.infrastructure.llm.prompts import read_query_prompt @@ -45,7 +47,7 @@ class AnthropicAdapter(LLMInterface): if not text_input: text_input = "No user input provided." if not system_prompt: - raise ValueError("No system prompt path provided.") + raise InvalidValueError(message="No system prompt path provided.") system_prompt = read_query_prompt(system_prompt) diff --git a/cognee/infrastructure/llm/generic_llm_api/adapter.py b/cognee/infrastructure/llm/generic_llm_api/adapter.py index f65d559d5..7d3d97ebb 100644 --- a/cognee/infrastructure/llm/generic_llm_api/adapter.py +++ b/cognee/infrastructure/llm/generic_llm_api/adapter.py @@ -5,6 +5,8 @@ from pydantic import BaseModel import instructor from tenacity import retry, stop_after_attempt import openai + +from cognee.exceptions import InvalidValueError from cognee.infrastructure.llm.llm_interface import LLMInterface from cognee.infrastructure.llm.prompts import read_query_prompt from cognee.shared.data_models import MonitoringTool @@ -128,7 +130,7 @@ class GenericAPIAdapter(LLMInterface): if not text_input: text_input = "No user input provided." if not system_prompt: - raise ValueError("No system prompt path provided.") + raise InvalidValueError(message="No system prompt path provided.") system_prompt = read_query_prompt(system_prompt) formatted_prompt = f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n""" if system_prompt else None diff --git a/cognee/infrastructure/llm/get_llm_client.py b/cognee/infrastructure/llm/get_llm_client.py index 1449d33b3..9a23892f2 100644 --- a/cognee/infrastructure/llm/get_llm_client.py +++ b/cognee/infrastructure/llm/get_llm_client.py @@ -1,5 +1,7 @@ """Get the LLM client.""" from enum import Enum + +from cognee.exceptions import InvalidValueError from cognee.infrastructure.llm import get_llm_config # Define an Enum for LLM Providers @@ -17,7 +19,7 @@ def get_llm_client(): if provider == LLMProvider.OPENAI: if llm_config.llm_api_key is None: - raise ValueError("LLM API key is not set.") + raise InvalidValueError(message="LLM API key is not set.") from .openai.adapter import OpenAIAdapter @@ -32,7 +34,7 @@ def get_llm_client(): elif provider == LLMProvider.OLLAMA: if llm_config.llm_api_key is None: - raise ValueError("LLM API key is not set.") + raise InvalidValueError(message="LLM API key is not set.") from .generic_llm_api.adapter import GenericAPIAdapter return GenericAPIAdapter(llm_config.llm_endpoint, llm_config.llm_api_key, llm_config.llm_model, "Ollama") @@ -43,10 +45,10 @@ def get_llm_client(): elif provider == LLMProvider.CUSTOM: if llm_config.llm_api_key is None: - raise ValueError("LLM API key is not set.") + raise InvalidValueError(message="LLM API key is not set.") from .generic_llm_api.adapter import GenericAPIAdapter return GenericAPIAdapter(llm_config.llm_endpoint, llm_config.llm_api_key, llm_config.llm_model, "Custom") else: - raise ValueError(f"Unsupported LLM provider: {provider}") + raise InvalidValueError(message=f"Unsupported LLM provider: {provider}") diff --git a/cognee/infrastructure/llm/openai/adapter.py b/cognee/infrastructure/llm/openai/adapter.py index 28cdfff4e..4c4e0d933 100644 --- a/cognee/infrastructure/llm/openai/adapter.py +++ b/cognee/infrastructure/llm/openai/adapter.py @@ -7,6 +7,7 @@ import litellm import instructor from pydantic import BaseModel +from cognee.exceptions import InvalidValueError from cognee.infrastructure.llm.llm_interface import LLMInterface from cognee.infrastructure.llm.prompts import read_query_prompt @@ -121,7 +122,7 @@ class OpenAIAdapter(LLMInterface): if not text_input: text_input = "No user input provided." if not system_prompt: - raise ValueError("No system prompt path provided.") + raise InvalidValueError(message="No system prompt path provided.") system_prompt = read_query_prompt(system_prompt) formatted_prompt = f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n""" if system_prompt else None diff --git a/cognee/modules/data/methods/delete_data.py b/cognee/modules/data/methods/delete_data.py index 7560762e1..c0493a606 100644 --- a/cognee/modules/data/methods/delete_data.py +++ b/cognee/modules/data/methods/delete_data.py @@ -1,3 +1,4 @@ +from cognee.exceptions import InvalidAttributeError from cognee.modules.data.models import Data from cognee.infrastructure.databases.relational import get_relational_engine @@ -12,7 +13,7 @@ async def delete_data(data: Data): ValueError: If the data object is invalid. """ if not hasattr(data, '__tablename__'): - raise ValueError("The provided data object is missing the required '__tablename__' attribute.") + raise InvalidAttributeError(message="The provided data object is missing the required '__tablename__' attribute.") db_engine = get_relational_engine() diff --git a/cognee/modules/data/operations/translate_text.py b/cognee/modules/data/operations/translate_text.py index 411712648..d8c27e42a 100644 --- a/cognee/modules/data/operations/translate_text.py +++ b/cognee/modules/data/operations/translate_text.py @@ -1,5 +1,7 @@ import logging +from cognee.exceptions import InvalidValueError + logger = logging.getLogger(__name__) async def translate_text(text, source_language: str = "sr", target_language: str = "en", region_name = "eu-west-1"): @@ -18,10 +20,10 @@ async def translate_text(text, source_language: str = "sr", target_language: str from botocore.exceptions import BotoCoreError, ClientError if not text: - raise ValueError("No text to translate.") + raise InvalidValueError(message="No text to translate.") if not source_language or not target_language: - raise ValueError("Source and target language codes are required.") + raise InvalidValueError(message="Source and target language codes are required.") try: translate = boto3.client(service_name = "translate", region_name = region_name, use_ssl = True) diff --git a/cognee/modules/graph/cognee_graph/CogneeGraph.py b/cognee/modules/graph/cognee_graph/CogneeGraph.py index d15d93b73..255d18d4d 100644 --- a/cognee/modules/graph/cognee_graph/CogneeGraph.py +++ b/cognee/modules/graph/cognee_graph/CogneeGraph.py @@ -1,5 +1,6 @@ from typing import List, Dict, Union +from cognee.exceptions import EntityAlreadyExistsError, EntityNotFoundError, InvalidValueError from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface from cognee.modules.graph.cognee_graph.CogneeGraphElements import Node, Edge from cognee.modules.graph.cognee_graph.CogneeAbstractGraph import CogneeAbstractGraph @@ -26,7 +27,7 @@ class CogneeGraph(CogneeAbstractGraph): if node.id not in self.nodes: self.nodes[node.id] = node else: - raise ValueError(f"Node with id {node.id} already exists.") + raise EntityAlreadyExistsError(message=f"Node with id {node.id} already exists.") def add_edge(self, edge: Edge) -> None: if edge not in self.edges: @@ -34,7 +35,7 @@ class CogneeGraph(CogneeAbstractGraph): edge.node1.add_skeleton_edge(edge) edge.node2.add_skeleton_edge(edge) else: - raise ValueError(f"Edge {edge} already exists in the graph.") + raise EntityAlreadyExistsError(message=f"Edge {edge} already exists in the graph.") def get_node(self, node_id: str) -> Node: return self.nodes.get(node_id, None) @@ -44,7 +45,7 @@ class CogneeGraph(CogneeAbstractGraph): if node: return node.skeleton_edges else: - raise ValueError(f"Node with id {node_id} does not exist.") + raise EntityNotFoundError(message=f"Node with id {node_id} does not exist.") async def project_graph_from_db(self, adapter: Union[GraphDBInterface], @@ -55,15 +56,15 @@ class CogneeGraph(CogneeAbstractGraph): edge_dimension = 1) -> None: if node_dimension < 1 or edge_dimension < 1: - raise ValueError("Dimensions must be positive integers") + raise InvalidValueError(message="Dimensions must be positive integers") try: nodes_data, edges_data = await adapter.get_graph_data() if not nodes_data: - raise ValueError("No node data retrieved from the database.") + raise EntityNotFoundError(message="No node data retrieved from the database.") if not edges_data: - raise ValueError("No edge data retrieved from the database.") + raise EntityNotFoundError(message="No edge data retrieved from the database.") for node_id, properties in nodes_data: node_attributes = {key: properties.get(key) for key in node_properties_to_project} @@ -83,7 +84,7 @@ class CogneeGraph(CogneeAbstractGraph): target_node.add_skeleton_edge(edge) else: - raise ValueError(f"Edge references nonexistent nodes: {source_id} -> {target_id}") + raise EntityNotFoundError(message=f"Edge references nonexistent nodes: {source_id} -> {target_id}") except (ValueError, TypeError) as e: print(f"Error projecting graph: {e}") diff --git a/cognee/modules/graph/cognee_graph/CogneeGraphElements.py b/cognee/modules/graph/cognee_graph/CogneeGraphElements.py index 8235cb24d..e0c140630 100644 --- a/cognee/modules/graph/cognee_graph/CogneeGraphElements.py +++ b/cognee/modules/graph/cognee_graph/CogneeGraphElements.py @@ -1,6 +1,9 @@ import numpy as np from typing import List, Dict, Optional, Any +from cognee.exceptions import InvalidValueError + + class Node: """ Represents a node in a graph. @@ -18,7 +21,7 @@ class Node: def __init__(self, node_id: str, attributes: Optional[Dict[str, Any]] = None, dimension: int = 1): if dimension <= 0: - raise ValueError("Dimension must be a positive integer") + raise InvalidValueError(message="Dimension must be a positive integer") self.id = node_id self.attributes = attributes if attributes is not None else {} self.skeleton_neighbours = [] @@ -52,7 +55,7 @@ class Node: def is_node_alive_in_dimension(self, dimension: int) -> bool: if dimension < 0 or dimension >= len(self.status): - raise ValueError(f"Dimension {dimension} is out of range. Valid range is 0 to {len(self.status) - 1}.") + raise InvalidValueError(message=f"Dimension {dimension} is out of range. Valid range is 0 to {len(self.status) - 1}.") return self.status[dimension] == 1 def __repr__(self) -> str: @@ -83,7 +86,7 @@ class Edge: def __init__(self, node1: "Node", node2: "Node", attributes: Optional[Dict[str, Any]] = None, directed: bool = True, dimension: int = 1): if dimension <= 0: - raise ValueError("Dimensions must be a positive integer.") + raise InvalidValueError(message="Dimensions must be a positive integer.") self.node1 = node1 self.node2 = node2 self.attributes = attributes if attributes is not None else {} @@ -92,7 +95,7 @@ class Edge: def is_edge_alive_in_dimension(self, dimension: int) -> bool: if dimension < 0 or dimension >= len(self.status): - raise ValueError(f"Dimension {dimension} is out of range. Valid range is 0 to {len(self.status) - 1}.") + raise InvalidValueError(message=f"Dimension {dimension} is out of range. Valid range is 0 to {len(self.status) - 1}.") return self.status[dimension] == 1 def __repr__(self) -> str: diff --git a/cognee/modules/ingestion/classify.py b/cognee/modules/ingestion/classify.py index 8e8c9fb00..972082765 100644 --- a/cognee/modules/ingestion/classify.py +++ b/cognee/modules/ingestion/classify.py @@ -1,9 +1,11 @@ from io import BufferedReader from typing import Union, BinaryIO -from .exceptions import IngestionException from .data_types import TextData, BinaryData from tempfile import SpooledTemporaryFile +from cognee.exceptions import IngestionError + + def classify(data: Union[str, BinaryIO], filename: str = None): if isinstance(data, str): return TextData(data) @@ -11,4 +13,4 @@ def classify(data: Union[str, BinaryIO], filename: str = None): if isinstance(data, BufferedReader) or isinstance(data, SpooledTemporaryFile): return BinaryData(data, data.name.split("/")[-1] if data.name else filename) - raise IngestionException(f"Type of data sent to classify(data: Union[str, BinaryIO) not supported: {type(data)}") + raise IngestionError(message=f"Type of data sent to classify(data: Union[str, BinaryIO) not supported: {type(data)}") diff --git a/cognee/modules/ingestion/exceptions.py b/cognee/modules/ingestion/exceptions.py deleted file mode 100644 index 0a189fb81..000000000 --- a/cognee/modules/ingestion/exceptions.py +++ /dev/null @@ -1,6 +0,0 @@ - -class IngestionException(Exception): - message: str - - def __init__(self, message: str): - self.message = message diff --git a/cognee/modules/users/get_user_manager.py b/cognee/modules/users/get_user_manager.py index b538535ca..9f443110e 100644 --- a/cognee/modules/users/get_user_manager.py +++ b/cognee/modules/users/get_user_manager.py @@ -2,13 +2,14 @@ import os import uuid from typing import Optional from fastapi import Depends, Request -from fastapi_users.exceptions import UserNotExists from fastapi_users import BaseUserManager, UUIDIDMixin, models from fastapi_users.db import SQLAlchemyUserDatabase from .get_user_db import get_user_db from .models import User from .methods import get_user +from ...exceptions import UserNotFoundError + class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): reset_password_token_secret = os.getenv("FASTAPI_USERS_RESET_PASSWORD_TOKEN_SECRET", "super_secret") @@ -25,7 +26,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): user = await get_user(id) if user is None: - raise UserNotExists() + raise UserNotFoundError return user diff --git a/cognee/modules/users/permissions/methods/check_permission_on_documents.py b/cognee/modules/users/permissions/methods/check_permission_on_documents.py index c8c283e4a..e6cf2ef24 100644 --- a/cognee/modules/users/permissions/methods/check_permission_on_documents.py +++ b/cognee/modules/users/permissions/methods/check_permission_on_documents.py @@ -2,6 +2,8 @@ import logging from uuid import UUID from sqlalchemy import select from sqlalchemy.orm import joinedload + +from cognee.exceptions import PermissionDeniedError from cognee.infrastructure.databases.relational import get_relational_engine from ...models.User import User @@ -9,11 +11,6 @@ from ...models.ACL import ACL logger = logging.getLogger(__name__) -class PermissionDeniedException(Exception): - def __init__(self, message: str): - self.message = message - super().__init__(self.message) - async def check_permission_on_documents(user: User, permission_type: str, document_ids: list[UUID]): user_group_ids = [group.id for group in user.groups] @@ -33,4 +30,4 @@ async def check_permission_on_documents(user: User, permission_type: str, docume has_permissions = all(document_id in resource_ids for document_id in document_ids) if not has_permissions: - raise PermissionDeniedException(f"User {user.email} does not have {permission_type} permission on documents") + raise PermissionDeniedError(message=f"User {user.email} does not have {permission_type} permission on documents") diff --git a/cognee/tasks/graph/infer_data_ontology.py b/cognee/tasks/graph/infer_data_ontology.py index eea378eb1..45214800b 100644 --- a/cognee/tasks/graph/infer_data_ontology.py +++ b/cognee/tasks/graph/infer_data_ontology.py @@ -10,6 +10,7 @@ import aiofiles import pandas as pd from pydantic import BaseModel +from cognee.exceptions import IngestionError, EntityNotFoundError from cognee.infrastructure.llm.prompts import read_query_prompt from cognee.infrastructure.llm.get_llm_client import get_llm_client from cognee.infrastructure.data.chunking.config import get_chunk_config @@ -75,7 +76,7 @@ class OntologyEngine: reader = csv.DictReader(content.splitlines()) return list(reader) else: - raise ValueError("Unsupported file format") + raise IngestionError(message="Unsupported file format") except Exception as e: raise RuntimeError(f"Failed to load data from {file_path}: {e}") @@ -148,7 +149,7 @@ class OntologyEngine: if node_id in valid_ids: await graph_client.add_node(node_id, node_data) if node_id not in valid_ids: - raise ValueError(f"Node ID {node_id} not found in the dataset") + raise EntityNotFoundError(message=f"Node ID {node_id} not found in the dataset") if pd.notna(row.get("relationship_source")) and pd.notna(row.get("relationship_target")): await graph_client.add_edge( row["relationship_source"], diff --git a/cognee/tasks/ingestion/save_data_item_to_storage.py b/cognee/tasks/ingestion/save_data_item_to_storage.py index 4782f271f..450acb3df 100644 --- a/cognee/tasks/ingestion/save_data_item_to_storage.py +++ b/cognee/tasks/ingestion/save_data_item_to_storage.py @@ -1,4 +1,6 @@ from typing import Union, BinaryIO + +from cognee.exceptions import IngestionError from cognee.modules.ingestion import save_data_to_file def save_data_item_to_storage(data_item: Union[BinaryIO, str], dataset_name: str) -> str: @@ -15,6 +17,6 @@ def save_data_item_to_storage(data_item: Union[BinaryIO, str], dataset_name: str else: file_path = save_data_to_file(data_item, dataset_name) else: - raise ValueError(f"Data type not supported: {type(data_item)}") + raise IngestionError(message=f"Data type not supported: {type(data_item)}") return file_path \ No newline at end of file From d4236bf385e0309509a48d8c880e0e62c34447dc Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Wed, 27 Nov 2024 14:40:31 +0100 Subject: [PATCH 03/13] refactor: Change import of UserNotFoundError Changed import of UserNotFoundError exception Refactor COG-502 --- cognee/modules/users/get_user_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cognee/modules/users/get_user_manager.py b/cognee/modules/users/get_user_manager.py index 9f443110e..0e8adaebd 100644 --- a/cognee/modules/users/get_user_manager.py +++ b/cognee/modules/users/get_user_manager.py @@ -8,7 +8,7 @@ from fastapi_users.db import SQLAlchemyUserDatabase from .get_user_db import get_user_db from .models import User from .methods import get_user -from ...exceptions import UserNotFoundError +from cognee.exceptions import UserNotFoundError class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): From 6eecc39db079dd23b1f6f1ef57249f460a101e7f Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Wed, 27 Nov 2024 14:53:09 +0100 Subject: [PATCH 04/13] feat: Add custom exceptions to more cognee-lib modules Added custom exceptions to more modules Feature COG-502 --- cognee/api/v1/config/config.py | 8 ++++---- cognee/tasks/graph/infer_data_ontology.py | 4 +++- .../ingestion/save_data_item_with_metadata_to_storage.py | 4 +++- .../unit/modules/graph/cognee_graph_elements_test.py | 5 +++-- cognee/tests/unit/modules/graph/cognee_graph_test.py | 7 ++++--- 5 files changed, 17 insertions(+), 11 deletions(-) diff --git a/cognee/api/v1/config/config.py b/cognee/api/v1/config/config.py index af78fe6ba..1347fcba8 100644 --- a/cognee/api/v1/config/config.py +++ b/cognee/api/v1/config/config.py @@ -1,7 +1,7 @@ """ This module is used to set the configuration of the system.""" import os from cognee.base_config import get_base_config -from cognee.exceptions import InvalidValueError +from cognee.exceptions import InvalidValueError, InvalidAttributeError from cognee.modules.cognify.config import get_cognify_config from cognee.infrastructure.data.chunking.config import get_chunk_config from cognee.infrastructure.databases.vector import get_vectordb_config @@ -86,7 +86,7 @@ class config(): if hasattr(llm_config, key): object.__setattr__(llm_config, key, value) else: - raise AttributeError(f"'{key}' is not a valid attribute of the config.") + raise InvalidAttributeError(message=f"'{key}' is not a valid attribute of the config.") @staticmethod def set_chunk_strategy(chunk_strategy: object): @@ -124,7 +124,7 @@ class config(): if hasattr(relational_db_config, key): object.__setattr__(relational_db_config, key, value) else: - raise AttributeError(f"'{key}' is not a valid attribute of the config.") + raise InvalidAttributeError(message=f"'{key}' is not a valid attribute of the config.") @staticmethod def set_vector_db_config(config_dict: dict): @@ -136,7 +136,7 @@ class config(): if hasattr(vector_db_config, key): object.__setattr__(vector_db_config, key, value) else: - raise AttributeError(f"'{key}' is not a valid attribute of the config.") + raise InvalidAttributeError(message=f"'{key}' is not a valid attribute of the config.") @staticmethod def set_vector_db_key(db_key: str): diff --git a/cognee/tasks/graph/infer_data_ontology.py b/cognee/tasks/graph/infer_data_ontology.py index 45214800b..fd71aa05c 100644 --- a/cognee/tasks/graph/infer_data_ontology.py +++ b/cognee/tasks/graph/infer_data_ontology.py @@ -4,6 +4,7 @@ import csv import json import logging from datetime import datetime, timezone +from fastapi import status from typing import Any, Dict, List, Optional, Union, Type import aiofiles @@ -78,7 +79,8 @@ class OntologyEngine: else: raise IngestionError(message="Unsupported file format") except Exception as e: - raise RuntimeError(f"Failed to load data from {file_path}: {e}") + raise IngestionError(message=f"Failed to load data from {file_path}: {e}", + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY) async def add_graph_ontology(self, file_path: str = None, documents: list = None): """Add graph ontology from a JSON or CSV file or infer from documents content.""" diff --git a/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py b/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py index ec29edb89..e60716fac 100644 --- a/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py +++ b/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py @@ -1,4 +1,6 @@ from typing import Union, BinaryIO, Any + +from cognee.exceptions import IngestionError from cognee.modules.ingestion import save_data_to_file def save_data_item_with_metadata_to_storage(data_item: Union[BinaryIO, str, Any], dataset_name: str) -> str: @@ -23,6 +25,6 @@ def save_data_item_with_metadata_to_storage(data_item: Union[BinaryIO, str, Any] else: file_path = save_data_to_file(data_item, dataset_name) else: - raise ValueError(f"Data type not supported: {type(data_item)}") + raise IngestionError(message=f"Data type not supported: {type(data_item)}") return file_path \ No newline at end of file diff --git a/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py b/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py index d2a1b6c59..d7dfcec15 100644 --- a/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py +++ b/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py @@ -1,6 +1,7 @@ import numpy as np import pytest +from cognee.exceptions import InvalidValueError from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge, Node @@ -105,7 +106,7 @@ def test_edge_invalid_dimension(): """Test that initializing an Edge with a non-positive dimension raises an error.""" node1 = Node("node1") node2 = Node("node2") - with pytest.raises(ValueError, match="Dimensions must be a positive integer."): + with pytest.raises(InvalidValueError, match="Dimensions must be a positive integer."): Edge(node1, node2, dimension=0) @@ -124,7 +125,7 @@ def test_edge_alive_invalid_dimension(): node1 = Node("node1") node2 = Node("node2") edge = Edge(node1, node2, dimension=1) - with pytest.raises(ValueError, match="Dimension 1 is out of range"): + with pytest.raises(InvalidValueError, match="Dimension 1 is out of range"): edge.is_edge_alive_in_dimension(1) diff --git a/cognee/tests/unit/modules/graph/cognee_graph_test.py b/cognee/tests/unit/modules/graph/cognee_graph_test.py index d05292d75..d77bab8ad 100644 --- a/cognee/tests/unit/modules/graph/cognee_graph_test.py +++ b/cognee/tests/unit/modules/graph/cognee_graph_test.py @@ -1,5 +1,6 @@ import pytest +from cognee.exceptions import InvalidValueError from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge, Node @@ -23,7 +24,7 @@ def test_add_duplicate_node(setup_graph): graph = setup_graph node = Node("node1") graph.add_node(node) - with pytest.raises(ValueError, match="Node with id node1 already exists."): + with pytest.raises(InvalidValueError, match="Node with id node1 already exists."): graph.add_node(node) @@ -50,7 +51,7 @@ def test_add_duplicate_edge(setup_graph): graph.add_node(node2) edge = Edge(node1, node2) graph.add_edge(edge) - with pytest.raises(ValueError, match="Edge .* already exists in the graph."): + with pytest.raises(InvalidValueError, match="Edge .* already exists in the graph."): graph.add_edge(edge) @@ -83,5 +84,5 @@ def test_get_edges_success(setup_graph): def test_get_edges_nonexistent_node(setup_graph): """Test retrieving edges for a nonexistent node raises an exception.""" graph = setup_graph - with pytest.raises(ValueError, match="Node with id nonexistent does not exist."): + with pytest.raises(InvalidValueError, match="Node with id nonexistent does not exist."): graph.get_edges("nonexistent") From 90287e0dac72e9146698193ea9927fd2e005a2b3 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Wed, 27 Nov 2024 15:07:30 +0100 Subject: [PATCH 05/13] refactor: Use fastapi exception instead of custom exception Return the use of fastapi exception instead of custom exception Refactor COG-502 --- cognee/modules/users/get_user_manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cognee/modules/users/get_user_manager.py b/cognee/modules/users/get_user_manager.py index 0e8adaebd..30410a985 100644 --- a/cognee/modules/users/get_user_manager.py +++ b/cognee/modules/users/get_user_manager.py @@ -8,7 +8,7 @@ from fastapi_users.db import SQLAlchemyUserDatabase from .get_user_db import get_user_db from .models import User from .methods import get_user -from cognee.exceptions import UserNotFoundError +from fastapi_users.exceptions import UserNotExists class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): @@ -26,7 +26,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): user = await get_user(id) if user is None: - raise UserNotFoundError + raise UserNotExists() return user From 5d297c50f45bfd2757d056df58765e145b06d630 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Wed, 27 Nov 2024 15:29:32 +0100 Subject: [PATCH 06/13] test: Update unit tests regrading exceptions Updated unit tests to check for custom exceptions instead Test COG-502 --- .../unit/modules/graph/cognee_graph_elements_test.py | 4 ++-- cognee/tests/unit/modules/graph/cognee_graph_test.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py b/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py index d7dfcec15..7f7493ed1 100644 --- a/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py +++ b/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py @@ -16,7 +16,7 @@ def test_node_initialization(): def test_node_invalid_dimension(): """Test that initializing a Node with a non-positive dimension raises an error.""" - with pytest.raises(ValueError, match="Dimension must be a positive integer"): + with pytest.raises(InvalidValueError, match="Dimension must be a positive integer"): Node("node1", dimension=0) @@ -69,7 +69,7 @@ def test_is_node_alive_in_dimension(): def test_node_alive_invalid_dimension(): """Test that checking alive status with an invalid dimension raises an error.""" node = Node("node1", dimension=1) - with pytest.raises(ValueError, match="Dimension 1 is out of range"): + with pytest.raises(InvalidValueError, match="Dimension 1 is out of range"): node.is_node_alive_in_dimension(1) diff --git a/cognee/tests/unit/modules/graph/cognee_graph_test.py b/cognee/tests/unit/modules/graph/cognee_graph_test.py index d77bab8ad..bdcd9725c 100644 --- a/cognee/tests/unit/modules/graph/cognee_graph_test.py +++ b/cognee/tests/unit/modules/graph/cognee_graph_test.py @@ -1,6 +1,6 @@ import pytest -from cognee.exceptions import InvalidValueError +from cognee.exceptions import EntityNotFoundError, EntityAlreadyExistsError from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge, Node @@ -24,7 +24,7 @@ def test_add_duplicate_node(setup_graph): graph = setup_graph node = Node("node1") graph.add_node(node) - with pytest.raises(InvalidValueError, match="Node with id node1 already exists."): + with pytest.raises(EntityAlreadyExistsError, match="Node with id node1 already exists."): graph.add_node(node) @@ -51,7 +51,7 @@ def test_add_duplicate_edge(setup_graph): graph.add_node(node2) edge = Edge(node1, node2) graph.add_edge(edge) - with pytest.raises(InvalidValueError, match="Edge .* already exists in the graph."): + with pytest.raises(EntityAlreadyExistsError, match="Edge .* already exists in the graph."): graph.add_edge(edge) @@ -84,5 +84,5 @@ def test_get_edges_success(setup_graph): def test_get_edges_nonexistent_node(setup_graph): """Test retrieving edges for a nonexistent node raises an exception.""" graph = setup_graph - with pytest.raises(InvalidValueError, match="Node with id nonexistent does not exist."): + with pytest.raises(EntityNotFoundError, match="Node with id nonexistent does not exist."): graph.get_edges("nonexistent") From e67bd91bd21283dcbafd3e1255ec80006824c7be Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Wed, 27 Nov 2024 15:51:38 +0100 Subject: [PATCH 07/13] feat: Add brief logging on raising exception Added brief logging of exception raised on raising exception Feature COG-502 --- cognee/exceptions/exceptions.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/cognee/exceptions/exceptions.py b/cognee/exceptions/exceptions.py index f9228873f..752376f60 100644 --- a/cognee/exceptions/exceptions.py +++ b/cognee/exceptions/exceptions.py @@ -1,5 +1,7 @@ from fastapi import status +import logging +logger = logging.getLogger(__name__) class CogneeApiError(Exception): """Base exception class""" @@ -13,6 +15,10 @@ class CogneeApiError(Exception): self.message = message self.name = name self.status_code = status_code + + # Automatically log the exception details + logger.error(f"{self.name}: {self.message} (Status code: {self.status_code})") + super().__init__(self.message, self.name) From 7d1210c8892b9b321dd6daeeaa9d2b58c5345725 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Wed, 27 Nov 2024 15:57:37 +0100 Subject: [PATCH 08/13] feat: Add custom exception handling to dataset router Added custom exceptions for dataset router Feature COG-502 --- cognee/api/v1/datasets/routers/get_datasets_router.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/cognee/api/v1/datasets/routers/get_datasets_router.py b/cognee/api/v1/datasets/routers/get_datasets_router.py index a3a30278b..1e62f7f49 100644 --- a/cognee/api/v1/datasets/routers/get_datasets_router.py +++ b/cognee/api/v1/datasets/routers/get_datasets_router.py @@ -156,18 +156,13 @@ def get_datasets_router() -> APIRouter: dataset_data = await get_dataset_data(dataset.id) if dataset_data is None: - raise HTTPException(status_code=404, detail=f"No data found in dataset ({dataset_id}).") + raise EntityNotFoundError(message=f"No data found in dataset ({dataset_id}).") matching_data = [data for data in dataset_data if str(data.id) == data_id] # Check if matching_data contains an element if len(matching_data) == 0: - return JSONResponse( - status_code=404, - content={ - "detail": f"Data ({data_id}) not found in dataset ({dataset_id})." - } - ) + raise EntityNotFoundError(message= f"Data ({data_id}) not found in dataset ({dataset_id}).") data = matching_data[0] From df0b4b4820adfece1f32ab700fcedc500da40706 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Fri, 29 Nov 2024 17:06:34 +0100 Subject: [PATCH 09/13] refactor: Move user and group errors to users module Moved user and group errors to users module Refactor #COG-502 --- .../routers/get_permissions_router.py | 2 +- cognee/api/v1/search/search_v2.py | 3 ++- cognee/exceptions/exceptions.py | 24 ----------------- cognee/modules/users/exceptions/__init__.py | 10 +++++++ cognee/modules/users/exceptions/exceptions.py | 26 +++++++++++++++++++ 5 files changed, 39 insertions(+), 26 deletions(-) create mode 100644 cognee/modules/users/exceptions/__init__.py create mode 100644 cognee/modules/users/exceptions/exceptions.py diff --git a/cognee/api/v1/permissions/routers/get_permissions_router.py b/cognee/api/v1/permissions/routers/get_permissions_router.py index 30491ec7d..8d012d600 100644 --- a/cognee/api/v1/permissions/routers/get_permissions_router.py +++ b/cognee/api/v1/permissions/routers/get_permissions_router.py @@ -2,7 +2,7 @@ from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import JSONResponse from sqlalchemy.orm import Session -from cognee.exceptions import GroupNotFoundError, UserNotFoundError +from cognee.modules.users.exceptions import UserNotFoundError, GroupNotFoundError from cognee.modules.users import get_user_db from cognee.modules.users.models import User, Group, Permission diff --git a/cognee/api/v1/search/search_v2.py b/cognee/api/v1/search/search_v2.py index d18ba3b6e..d77aa5fa8 100644 --- a/cognee/api/v1/search/search_v2.py +++ b/cognee/api/v1/search/search_v2.py @@ -3,10 +3,11 @@ from uuid import UUID from enum import Enum from typing import Callable, Dict -from cognee.exceptions import UserNotFoundError, InvalidValueError +from cognee.exceptions import InvalidValueError from cognee.modules.search.operations import log_query, log_result from cognee.modules.storage.utils import JSONEncoder from cognee.shared.utils import send_telemetry +from cognee.modules.users.exceptions import UserNotFoundError from cognee.modules.users.models import User from cognee.modules.users.methods import get_default_user from cognee.modules.users.permissions.methods import get_document_ids_for_user diff --git a/cognee/exceptions/exceptions.py b/cognee/exceptions/exceptions.py index 752376f60..2ea9ded1e 100644 --- a/cognee/exceptions/exceptions.py +++ b/cognee/exceptions/exceptions.py @@ -34,30 +34,6 @@ class ServiceError(CogneeApiError): super().__init__(message, name, status_code) -class GroupNotFoundError(CogneeApiError): - """User group not found""" - - def __init__( - self, - message: str = "User group not found.", - name: str = "GroupNotFoundError", - status_code=status.HTTP_404_NOT_FOUND, - ): - super().__init__(message, name, status_code) - - -class UserNotFoundError(CogneeApiError): - """User not found""" - - def __init__( - self, - message: str = "No user found in the system. Please create a user.", - name: str = "UserNotFoundError", - status_code=status.HTTP_404_NOT_FOUND, - ): - super().__init__(message, name, status_code) - - class EntityNotFoundError(CogneeApiError): """Database returns nothing""" diff --git a/cognee/modules/users/exceptions/__init__.py b/cognee/modules/users/exceptions/__init__.py new file mode 100644 index 000000000..ee4f99eda --- /dev/null +++ b/cognee/modules/users/exceptions/__init__.py @@ -0,0 +1,10 @@ +""" +Custom exceptions for the Cognee API. + +This module defines a set of exceptions for handling various user errors +""" + +from .exceptions import ( + GroupNotFoundError, + UserNotFoundError, +) \ No newline at end of file diff --git a/cognee/modules/users/exceptions/exceptions.py b/cognee/modules/users/exceptions/exceptions.py new file mode 100644 index 000000000..d45531746 --- /dev/null +++ b/cognee/modules/users/exceptions/exceptions.py @@ -0,0 +1,26 @@ +from cognee.exceptions import CogneeApiError +from fastapi import status + + +class GroupNotFoundError(CogneeApiError): + """User group not found""" + + def __init__( + self, + message: str = "User group not found.", + name: str = "GroupNotFoundError", + status_code=status.HTTP_404_NOT_FOUND, + ): + super().__init__(message, name, status_code) + + +class UserNotFoundError(CogneeApiError): + """User not found""" + + def __init__( + self, + message: str = "No user found in the system. Please create a user.", + name: str = "UserNotFoundError", + status_code=status.HTTP_404_NOT_FOUND, + ): + super().__init__(message, name, status_code) From 1b2bdd9b83235f165d4f3bcb910b1b18bd426e14 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Fri, 29 Nov 2024 17:07:27 +0100 Subject: [PATCH 10/13] refactor: Update __init__.py of base exception module Updated init file of cognee base exception module Refactor COG-502 --- cognee/exceptions/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/cognee/exceptions/__init__.py b/cognee/exceptions/__init__.py index e07547508..bbc71c4d2 100644 --- a/cognee/exceptions/__init__.py +++ b/cognee/exceptions/__init__.py @@ -9,8 +9,6 @@ from .exceptions import ( CogneeApiError, ServiceError, EntityNotFoundError, - GroupNotFoundError, - UserNotFoundError, EntityAlreadyExistsError, InvalidOperationError, PermissionDeniedError, From eb09e5ad890deb9735e5d7e9e9560c894d8467a5 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Fri, 29 Nov 2024 17:15:54 +0100 Subject: [PATCH 11/13] refactor: Moved ingestion exceptions to ingestion module Moved custom ingestion exceptions to ingestion module Refactor COG-502 --- cognee/exceptions/__init__.py | 2 -- cognee/exceptions/exceptions.py | 18 ------------------ cognee/modules/ingestion/classify.py | 2 +- .../modules/ingestion/exceptions/__init__.py | 9 +++++++++ .../modules/ingestion/exceptions/exceptions.py | 11 +++++++++++ cognee/modules/users/exceptions/__init__.py | 1 + cognee/modules/users/exceptions/exceptions.py | 10 ++++++++++ .../methods/check_permission_on_documents.py | 2 +- cognee/tasks/graph/infer_data_ontology.py | 3 ++- .../ingestion/save_data_item_to_storage.py | 2 +- .../save_data_item_with_metadata_to_storage.py | 2 +- 11 files changed, 37 insertions(+), 25 deletions(-) create mode 100644 cognee/modules/ingestion/exceptions/__init__.py create mode 100644 cognee/modules/ingestion/exceptions/exceptions.py diff --git a/cognee/exceptions/__init__.py b/cognee/exceptions/__init__.py index bbc71c4d2..4d76a530c 100644 --- a/cognee/exceptions/__init__.py +++ b/cognee/exceptions/__init__.py @@ -11,8 +11,6 @@ from .exceptions import ( EntityNotFoundError, EntityAlreadyExistsError, InvalidOperationError, - PermissionDeniedError, - IngestionError, InvalidValueError, InvalidAttributeError, ) \ No newline at end of file diff --git a/cognee/exceptions/exceptions.py b/cognee/exceptions/exceptions.py index 2ea9ded1e..fd83ed6ab 100644 --- a/cognee/exceptions/exceptions.py +++ b/cognee/exceptions/exceptions.py @@ -70,24 +70,6 @@ class InvalidOperationError(CogneeApiError): super().__init__(message, name, status_code) -class PermissionDeniedError(CogneeApiError): - def __init__( - self, - message: str = "User does not have permission on documents.", - name: str = "PermissionDeniedError", - status_code=status.HTTP_403_FORBIDDEN, - ): - super().__init__(message, name, status_code) - -class IngestionError(CogneeApiError): - def __init__( - self, - message: str = "Type of data sent to classify not supported.", - name: str = "IngestionError", - status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, - ): - super().__init__(message, name, status_code) - class InvalidValueError(CogneeApiError): def __init__( self, diff --git a/cognee/modules/ingestion/classify.py b/cognee/modules/ingestion/classify.py index 972082765..dbb191cc3 100644 --- a/cognee/modules/ingestion/classify.py +++ b/cognee/modules/ingestion/classify.py @@ -3,7 +3,7 @@ from typing import Union, BinaryIO from .data_types import TextData, BinaryData from tempfile import SpooledTemporaryFile -from cognee.exceptions import IngestionError +from cognee.modules.ingestion.exceptions import IngestionError def classify(data: Union[str, BinaryIO], filename: str = None): diff --git a/cognee/modules/ingestion/exceptions/__init__.py b/cognee/modules/ingestion/exceptions/__init__.py new file mode 100644 index 000000000..33d59e113 --- /dev/null +++ b/cognee/modules/ingestion/exceptions/__init__.py @@ -0,0 +1,9 @@ +""" +Custom exceptions for the Cognee API. + +This module defines a set of exceptions for handling various ingestion errors +""" + +from .exceptions import ( + IngestionError, +) \ No newline at end of file diff --git a/cognee/modules/ingestion/exceptions/exceptions.py b/cognee/modules/ingestion/exceptions/exceptions.py new file mode 100644 index 000000000..4901be110 --- /dev/null +++ b/cognee/modules/ingestion/exceptions/exceptions.py @@ -0,0 +1,11 @@ +from cognee.exceptions import CogneeApiError +from fastapi import status + +class IngestionError(CogneeApiError): + def __init__( + self, + message: str = "Type of data sent to classify not supported.", + name: str = "IngestionError", + status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, + ): + super().__init__(message, name, status_code) \ No newline at end of file diff --git a/cognee/modules/users/exceptions/__init__.py b/cognee/modules/users/exceptions/__init__.py index ee4f99eda..70e6e9d2a 100644 --- a/cognee/modules/users/exceptions/__init__.py +++ b/cognee/modules/users/exceptions/__init__.py @@ -7,4 +7,5 @@ This module defines a set of exceptions for handling various user errors from .exceptions import ( GroupNotFoundError, UserNotFoundError, + PermissionDeniedError, ) \ No newline at end of file diff --git a/cognee/modules/users/exceptions/exceptions.py b/cognee/modules/users/exceptions/exceptions.py index d45531746..7dda702db 100644 --- a/cognee/modules/users/exceptions/exceptions.py +++ b/cognee/modules/users/exceptions/exceptions.py @@ -24,3 +24,13 @@ class UserNotFoundError(CogneeApiError): status_code=status.HTTP_404_NOT_FOUND, ): super().__init__(message, name, status_code) + + +class PermissionDeniedError(CogneeApiError): + def __init__( + self, + message: str = "User does not have permission on documents.", + name: str = "PermissionDeniedError", + status_code=status.HTTP_403_FORBIDDEN, + ): + super().__init__(message, name, status_code) diff --git a/cognee/modules/users/permissions/methods/check_permission_on_documents.py b/cognee/modules/users/permissions/methods/check_permission_on_documents.py index e6cf2ef24..f9c5a2258 100644 --- a/cognee/modules/users/permissions/methods/check_permission_on_documents.py +++ b/cognee/modules/users/permissions/methods/check_permission_on_documents.py @@ -3,7 +3,7 @@ from uuid import UUID from sqlalchemy import select from sqlalchemy.orm import joinedload -from cognee.exceptions import PermissionDeniedError +from cognee.modules.users.exceptions import PermissionDeniedError from cognee.infrastructure.databases.relational import get_relational_engine from ...models.User import User diff --git a/cognee/tasks/graph/infer_data_ontology.py b/cognee/tasks/graph/infer_data_ontology.py index fd71aa05c..ee267a83b 100644 --- a/cognee/tasks/graph/infer_data_ontology.py +++ b/cognee/tasks/graph/infer_data_ontology.py @@ -11,7 +11,8 @@ import aiofiles import pandas as pd from pydantic import BaseModel -from cognee.exceptions import IngestionError, EntityNotFoundError +from cognee.exceptions import EntityNotFoundError +from cognee.modules.ingestion.exceptions import IngestionError from cognee.infrastructure.llm.prompts import read_query_prompt from cognee.infrastructure.llm.get_llm_client import get_llm_client from cognee.infrastructure.data.chunking.config import get_chunk_config diff --git a/cognee/tasks/ingestion/save_data_item_to_storage.py b/cognee/tasks/ingestion/save_data_item_to_storage.py index 450acb3df..e2a7c8ee7 100644 --- a/cognee/tasks/ingestion/save_data_item_to_storage.py +++ b/cognee/tasks/ingestion/save_data_item_to_storage.py @@ -1,6 +1,6 @@ from typing import Union, BinaryIO -from cognee.exceptions import IngestionError +from cognee.modules.ingestion.exceptions import IngestionError from cognee.modules.ingestion import save_data_to_file def save_data_item_to_storage(data_item: Union[BinaryIO, str], dataset_name: str) -> str: diff --git a/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py b/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py index e60716fac..bf5a1f093 100644 --- a/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py +++ b/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py @@ -1,6 +1,6 @@ from typing import Union, BinaryIO, Any -from cognee.exceptions import IngestionError +from cognee.modules.ingestion.exceptions import IngestionError from cognee.modules.ingestion import save_data_to_file def save_data_item_with_metadata_to_storage(data_item: Union[BinaryIO, str, Any], dataset_name: str) -> str: From 6b97e95e1489f0b4b436775e2cb1b5c85495ad9b Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Fri, 29 Nov 2024 17:40:48 +0100 Subject: [PATCH 12/13] refactor: Split entity related exceptions into graph and database exceptions Move and split database entity related exceptions into graph and database exceptions Refactor COG-502 --- .../datasets/routers/get_datasets_router.py | 2 +- cognee/exceptions/__init__.py | 3 -- cognee/exceptions/exceptions.py | 37 +------------------ .../databases/exceptions/__init__.py | 10 +++++ .../databases/exceptions/exceptions.py | 25 +++++++++++++ .../sqlalchemy/SqlAlchemyAdapter.py | 2 +- .../vector/pgvector/PGVectorAdapter.py | 3 +- .../modules/graph/cognee_graph/CogneeGraph.py | 5 ++- cognee/modules/graph/exceptions/__init__.py | 10 +++++ cognee/modules/graph/exceptions/exceptions.py | 25 +++++++++++++ cognee/tasks/graph/infer_data_ontology.py | 2 +- 11 files changed, 79 insertions(+), 45 deletions(-) create mode 100644 cognee/infrastructure/databases/exceptions/__init__.py create mode 100644 cognee/infrastructure/databases/exceptions/exceptions.py create mode 100644 cognee/modules/graph/exceptions/__init__.py create mode 100644 cognee/modules/graph/exceptions/exceptions.py diff --git a/cognee/api/v1/datasets/routers/get_datasets_router.py b/cognee/api/v1/datasets/routers/get_datasets_router.py index 1e62f7f49..31e3fa67d 100644 --- a/cognee/api/v1/datasets/routers/get_datasets_router.py +++ b/cognee/api/v1/datasets/routers/get_datasets_router.py @@ -9,7 +9,7 @@ from fastapi.responses import JSONResponse, FileResponse from pydantic import BaseModel from cognee.api.DTO import OutDTO -from cognee.exceptions import EntityNotFoundError +from cognee.infrastructure.databases.exceptions import EntityNotFoundError from cognee.modules.users.models import User from cognee.modules.users.methods import get_authenticated_user from cognee.modules.pipelines.models import PipelineRunStatus diff --git a/cognee/exceptions/__init__.py b/cognee/exceptions/__init__.py index 4d76a530c..40120e0e1 100644 --- a/cognee/exceptions/__init__.py +++ b/cognee/exceptions/__init__.py @@ -8,9 +8,6 @@ such as service failures, resource conflicts, and invalid operations. from .exceptions import ( CogneeApiError, ServiceError, - EntityNotFoundError, - EntityAlreadyExistsError, - InvalidOperationError, InvalidValueError, InvalidAttributeError, ) \ No newline at end of file diff --git a/cognee/exceptions/exceptions.py b/cognee/exceptions/exceptions.py index fd83ed6ab..f94daf8c9 100644 --- a/cognee/exceptions/exceptions.py +++ b/cognee/exceptions/exceptions.py @@ -34,42 +34,6 @@ class ServiceError(CogneeApiError): super().__init__(message, name, status_code) -class EntityNotFoundError(CogneeApiError): - """Database returns nothing""" - - def __init__( - self, - message: str = "The requested entity does not exist.", - name: str = "EntityNotFoundError", - status_code=status.HTTP_404_NOT_FOUND, - ): - super().__init__(message, name, status_code) - - -class EntityAlreadyExistsError(CogneeApiError): - """Conflict detected, like trying to create a resource that already exists""" - - def __init__( - self, - message: str = "The entity already exists.", - name: str = "EntityAlreadyExistsError", - status_code=status.HTTP_409_CONFLICT, - ): - super().__init__(message, name, status_code) - - -class InvalidOperationError(CogneeApiError): - """Invalid operations like trying to delete a non-existing entity, etc.""" - - def __init__( - self, - message: str = "Invalid operation attempted.", - name: str = "InvalidOperationError", - status_code=status.HTTP_400_BAD_REQUEST, - ): - super().__init__(message, name, status_code) - - class InvalidValueError(CogneeApiError): def __init__( self, @@ -79,6 +43,7 @@ class InvalidValueError(CogneeApiError): ): super().__init__(message, name, status_code) + class InvalidAttributeError(CogneeApiError): def __init__( self, diff --git a/cognee/infrastructure/databases/exceptions/__init__.py b/cognee/infrastructure/databases/exceptions/__init__.py new file mode 100644 index 000000000..5836e7d11 --- /dev/null +++ b/cognee/infrastructure/databases/exceptions/__init__.py @@ -0,0 +1,10 @@ +""" +Custom exceptions for the Cognee API. + +This module defines a set of exceptions for handling various database errors +""" + +from .exceptions import ( + EntityNotFoundError, + EntityAlreadyExistsError, +) \ No newline at end of file diff --git a/cognee/infrastructure/databases/exceptions/exceptions.py b/cognee/infrastructure/databases/exceptions/exceptions.py new file mode 100644 index 000000000..af15bb616 --- /dev/null +++ b/cognee/infrastructure/databases/exceptions/exceptions.py @@ -0,0 +1,25 @@ +from cognee.exceptions import CogneeApiError +from fastapi import status + +class EntityNotFoundError(CogneeApiError): + """Database returns nothing""" + + def __init__( + self, + message: str = "The requested entity does not exist.", + name: str = "EntityNotFoundError", + status_code=status.HTTP_404_NOT_FOUND, + ): + super().__init__(message, name, status_code) + + +class EntityAlreadyExistsError(CogneeApiError): + """Conflict detected, like trying to create a resource that already exists""" + + def __init__( + self, + message: str = "The entity already exists.", + name: str = "EntityAlreadyExistsError", + status_code=status.HTTP_409_CONFLICT, + ): + super().__init__(message, name, status_code) \ No newline at end of file diff --git a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py index a5ae36988..4150eed7b 100644 --- a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py +++ b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py @@ -7,7 +7,7 @@ from sqlalchemy import text, select, MetaData, Table from sqlalchemy.orm import joinedload from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker -from cognee.exceptions import EntityNotFoundError +from cognee.infrastructure.databases.exceptions import EntityNotFoundError from ..ModelBase import Base class SQLAlchemyAdapter(): diff --git a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py index 77dd03185..3a041f6fe 100644 --- a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py +++ b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py @@ -6,7 +6,8 @@ from sqlalchemy.orm import Mapped, mapped_column from sqlalchemy import JSON, Column, Table, select, delete from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker -from cognee.exceptions import EntityNotFoundError, InvalidValueError +from cognee.exceptions import InvalidValueError +from cognee.infrastructure.databases.exceptions import EntityNotFoundError from cognee.infrastructure.engine import DataPoint from .serialize_data import serialize_data diff --git a/cognee/modules/graph/cognee_graph/CogneeGraph.py b/cognee/modules/graph/cognee_graph/CogneeGraph.py index 255d18d4d..76bb602f8 100644 --- a/cognee/modules/graph/cognee_graph/CogneeGraph.py +++ b/cognee/modules/graph/cognee_graph/CogneeGraph.py @@ -1,10 +1,11 @@ from typing import List, Dict, Union -from cognee.exceptions import EntityAlreadyExistsError, EntityNotFoundError, InvalidValueError +from cognee.exceptions import InvalidValueError +from cognee.modules.graph.exceptions import EntityNotFoundError, EntityAlreadyExistsError from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface from cognee.modules.graph.cognee_graph.CogneeGraphElements import Node, Edge from cognee.modules.graph.cognee_graph.CogneeAbstractGraph import CogneeAbstractGraph -from cognee.infrastructure.databases.graph import get_graph_engine + class CogneeGraph(CogneeAbstractGraph): """ diff --git a/cognee/modules/graph/exceptions/__init__.py b/cognee/modules/graph/exceptions/__init__.py new file mode 100644 index 000000000..e8330caf3 --- /dev/null +++ b/cognee/modules/graph/exceptions/__init__.py @@ -0,0 +1,10 @@ +""" +Custom exceptions for the Cognee API. + +This module defines a set of exceptions for handling various graph errors +""" + +from .exceptions import ( + EntityNotFoundError, + EntityAlreadyExistsError, +) \ No newline at end of file diff --git a/cognee/modules/graph/exceptions/exceptions.py b/cognee/modules/graph/exceptions/exceptions.py new file mode 100644 index 000000000..af15bb616 --- /dev/null +++ b/cognee/modules/graph/exceptions/exceptions.py @@ -0,0 +1,25 @@ +from cognee.exceptions import CogneeApiError +from fastapi import status + +class EntityNotFoundError(CogneeApiError): + """Database returns nothing""" + + def __init__( + self, + message: str = "The requested entity does not exist.", + name: str = "EntityNotFoundError", + status_code=status.HTTP_404_NOT_FOUND, + ): + super().__init__(message, name, status_code) + + +class EntityAlreadyExistsError(CogneeApiError): + """Conflict detected, like trying to create a resource that already exists""" + + def __init__( + self, + message: str = "The entity already exists.", + name: str = "EntityAlreadyExistsError", + status_code=status.HTTP_409_CONFLICT, + ): + super().__init__(message, name, status_code) \ No newline at end of file diff --git a/cognee/tasks/graph/infer_data_ontology.py b/cognee/tasks/graph/infer_data_ontology.py index ee267a83b..4e11cd9af 100644 --- a/cognee/tasks/graph/infer_data_ontology.py +++ b/cognee/tasks/graph/infer_data_ontology.py @@ -11,7 +11,7 @@ import aiofiles import pandas as pd from pydantic import BaseModel -from cognee.exceptions import EntityNotFoundError +from cognee.modules.graph.exceptions import EntityNotFoundError, EntityAlreadyExistsError from cognee.modules.ingestion.exceptions import IngestionError from cognee.infrastructure.llm.prompts import read_query_prompt from cognee.infrastructure.llm.get_llm_client import get_llm_client From 343ac47fd487bfbb68f72cf4fbf3c5f7eeed5c90 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Mon, 2 Dec 2024 13:19:55 +0100 Subject: [PATCH 13/13] fix: Update import location for LanceDB Updated import path for LanceDB exceptions Fix COG-502 --- .../infrastructure/databases/vector/lancedb/LanceDBAdapter.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py b/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py index 22190a8bb..3cea3bc27 100644 --- a/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py +++ b/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py @@ -6,7 +6,7 @@ import lancedb from pydantic import BaseModel from lancedb.pydantic import Vector, LanceModel -from cognee.exceptions import InvalidValueError, EntityNotFoundError +from cognee.exceptions import InvalidValueError from cognee.infrastructure.engine import DataPoint from cognee.infrastructure.files.storage import LocalStorage from cognee.modules.storage.utils import copy_model, get_own_properties @@ -151,7 +151,7 @@ class LanceDBAdapter(VectorDBInterface): query_vector: List[float] = None ): if query_text is None and query_vector is None: - raise ValueError("One of query_text or query_vector must be provided!") + raise InvalidValueError(message="One of query_text or query_vector must be provided!") if query_text and not query_vector: query_vector = (await self.embedding_engine.embed_text([query_text]))[0]