feat: Add custom exceptions to more cognee-lib modules

Added custom exceptions to more modules

Feature COG-502
This commit is contained in:
Igor Ilic 2024-11-27 14:53:09 +01:00
parent d4236bf385
commit 6eecc39db0
5 changed files with 17 additions and 11 deletions

View file

@ -1,7 +1,7 @@
""" This module is used to set the configuration of the system.""" """ This module is used to set the configuration of the system."""
import os import os
from cognee.base_config import get_base_config 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.modules.cognify.config import get_cognify_config
from cognee.infrastructure.data.chunking.config import get_chunk_config from cognee.infrastructure.data.chunking.config import get_chunk_config
from cognee.infrastructure.databases.vector import get_vectordb_config from cognee.infrastructure.databases.vector import get_vectordb_config
@ -86,7 +86,7 @@ class config():
if hasattr(llm_config, key): if hasattr(llm_config, key):
object.__setattr__(llm_config, key, value) object.__setattr__(llm_config, key, value)
else: 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 @staticmethod
def set_chunk_strategy(chunk_strategy: object): def set_chunk_strategy(chunk_strategy: object):
@ -124,7 +124,7 @@ class config():
if hasattr(relational_db_config, key): if hasattr(relational_db_config, key):
object.__setattr__(relational_db_config, key, value) object.__setattr__(relational_db_config, key, value)
else: 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 @staticmethod
def set_vector_db_config(config_dict: dict): def set_vector_db_config(config_dict: dict):
@ -136,7 +136,7 @@ class config():
if hasattr(vector_db_config, key): if hasattr(vector_db_config, key):
object.__setattr__(vector_db_config, key, value) object.__setattr__(vector_db_config, key, value)
else: 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 @staticmethod
def set_vector_db_key(db_key: str): def set_vector_db_key(db_key: str):

View file

@ -4,6 +4,7 @@ import csv
import json import json
import logging import logging
from datetime import datetime, timezone from datetime import datetime, timezone
from fastapi import status
from typing import Any, Dict, List, Optional, Union, Type from typing import Any, Dict, List, Optional, Union, Type
import aiofiles import aiofiles
@ -78,7 +79,8 @@ class OntologyEngine:
else: else:
raise IngestionError(message="Unsupported file format") raise IngestionError(message="Unsupported file format")
except Exception as e: 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): 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.""" """Add graph ontology from a JSON or CSV file or infer from documents content."""

View file

@ -1,4 +1,6 @@
from typing import Union, BinaryIO, Any from typing import Union, BinaryIO, Any
from cognee.exceptions import IngestionError
from cognee.modules.ingestion import save_data_to_file 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: 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: else:
file_path = save_data_to_file(data_item, dataset_name) file_path = save_data_to_file(data_item, dataset_name)
else: 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 return file_path

View file

@ -1,6 +1,7 @@
import numpy as np import numpy as np
import pytest import pytest
from cognee.exceptions import InvalidValueError
from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge, Node 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.""" """Test that initializing an Edge with a non-positive dimension raises an error."""
node1 = Node("node1") node1 = Node("node1")
node2 = Node("node2") 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) Edge(node1, node2, dimension=0)
@ -124,7 +125,7 @@ def test_edge_alive_invalid_dimension():
node1 = Node("node1") node1 = Node("node1")
node2 = Node("node2") node2 = Node("node2")
edge = Edge(node1, node2, dimension=1) 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) edge.is_edge_alive_in_dimension(1)

View file

@ -1,5 +1,6 @@
import pytest import pytest
from cognee.exceptions import InvalidValueError
from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph
from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge, Node from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge, Node
@ -23,7 +24,7 @@ def test_add_duplicate_node(setup_graph):
graph = setup_graph graph = setup_graph
node = Node("node1") node = Node("node1")
graph.add_node(node) 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) graph.add_node(node)
@ -50,7 +51,7 @@ def test_add_duplicate_edge(setup_graph):
graph.add_node(node2) graph.add_node(node2)
edge = Edge(node1, node2) edge = Edge(node1, node2)
graph.add_edge(edge) 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) graph.add_edge(edge)
@ -83,5 +84,5 @@ def test_get_edges_success(setup_graph):
def test_get_edges_nonexistent_node(setup_graph): def test_get_edges_nonexistent_node(setup_graph):
"""Test retrieving edges for a nonexistent node raises an exception.""" """Test retrieving edges for a nonexistent node raises an exception."""
graph = setup_graph 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") graph.get_edges("nonexistent")