cognee 0.3.2__py3-none-any.whl → 0.3.4__py3-none-any.whl

cognee/tasks/graph/extract_graph_from_data.py

@@ -3,8 +3,14 @@ from typing import Type, List, Optional
 from pydantic import BaseModel
 
 from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.modules.ontology.ontology_env_config import get_ontology_env_config
 from cognee.tasks.storage.add_data_points import add_data_points
-from cognee.modules.ontology.rdf_xml.OntologyResolver import OntologyResolver
+from cognee.modules.ontology.ontology_config import Config
+from cognee.modules.ontology.get_default_ontology_resolver import (
+    get_default_ontology_resolver,
+    get_ontology_resolver_from_env,
+)
+from cognee.modules.ontology.base_ontology_resolver import BaseOntologyResolver
 from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
 from cognee.modules.graph.utils import (
     expand_with_nodes_and_edges,
@@ -24,9 +30,28 @@ async def integrate_chunk_graphs(
     data_chunks: list[DocumentChunk],
     chunk_graphs: list,
     graph_model: Type[BaseModel],
-    ontology_adapter: OntologyResolver,
+    ontology_resolver: BaseOntologyResolver,
 ) -> List[DocumentChunk]:
-    """Updates DocumentChunk objects, integrates data points and edges into databases."""
+    """Integrate chunk graphs with ontology validation and store in databases.
+
+    This function processes document chunks and their associated knowledge graphs,
+    validates entities against an ontology resolver, and stores the integrated
+    data points and edges in the configured databases.
+
+    Args:
+        data_chunks: List of document chunks containing source data
+        chunk_graphs: List of knowledge graphs corresponding to each chunk
+        graph_model: Pydantic model class for graph data validation
+        ontology_resolver: Resolver for validating entities against ontology
+
+    Returns:
+        List of updated DocumentChunk objects with integrated data
+
+    Raises:
+        InvalidChunkGraphInputError: If input validation fails
+        InvalidGraphModelError: If graph model validation fails
+        InvalidOntologyAdapterError: If ontology resolver validation fails
+    """
 
     if not isinstance(data_chunks, list) or not isinstance(chunk_graphs, list):
         raise InvalidChunkGraphInputError("data_chunks and chunk_graphs must be lists.")
@@ -36,9 +61,9 @@ async def integrate_chunk_graphs(
         )
     if not isinstance(graph_model, type) or not issubclass(graph_model, BaseModel):
         raise InvalidGraphModelError(graph_model)
-    if ontology_adapter is None or not hasattr(ontology_adapter, "get_subgraph"):
+    if ontology_resolver is None or not hasattr(ontology_resolver, "get_subgraph"):
         raise InvalidOntologyAdapterError(
-            type(ontology_adapter).__name__ if ontology_adapter else "None"
+            type(ontology_resolver).__name__ if ontology_resolver else "None"
         )
 
     graph_engine = await get_graph_engine()
@@ -55,7 +80,7 @@ async def integrate_chunk_graphs(
     )
 
     graph_nodes, graph_edges = expand_with_nodes_and_edges(
-        data_chunks, chunk_graphs, ontology_adapter, existing_edges_map
+        data_chunks, chunk_graphs, ontology_resolver, existing_edges_map
     )
 
     if len(graph_nodes) > 0:
@@ -70,7 +95,7 @@ async def integrate_chunk_graphs(
 async def extract_graph_from_data(
     data_chunks: List[DocumentChunk],
     graph_model: Type[BaseModel],
-    ontology_adapter: OntologyResolver = None,
+    config: Config = None,
     custom_prompt: Optional[str] = None,
 ) -> List[DocumentChunk]:
     """
@@ -101,6 +126,24 @@ async def extract_graph_from_data(
                 if edge.source_node_id in valid_node_ids and edge.target_node_id in valid_node_ids
             ]
 
-    return await integrate_chunk_graphs(
-        data_chunks, chunk_graphs, graph_model, ontology_adapter or OntologyResolver()
-    )
+    # Extract resolver from config if provided, otherwise get default
+    if config is None:
+        ontology_config = get_ontology_env_config()
+        if (
+            ontology_config.ontology_file_path
+            and ontology_config.ontology_resolver
+            and ontology_config.matching_strategy
+        ):
+            config: Config = {
+                "ontology_config": {
+                    "ontology_resolver": get_ontology_resolver_from_env(**ontology_config.to_dict())
+                }
+            }
+        else:
+            config: Config = {
+                "ontology_config": {"ontology_resolver": get_default_ontology_resolver()}
+            }
+
+    ontology_resolver = config["ontology_config"]["ontology_resolver"]
+
+    return await integrate_chunk_graphs(data_chunks, chunk_graphs, graph_model, ontology_resolver)

cognee/tasks/graph/extract_graph_from_data_v2.py

@@ -3,7 +3,7 @@ from typing import List
 
 from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
 from cognee.shared.data_models import KnowledgeGraph
-from cognee.modules.ontology.rdf_xml.OntologyResolver import OntologyResolver
+from cognee.modules.ontology.base_ontology_resolver import BaseOntologyResolver
 from cognee.tasks.graph.cascade_extract.utils.extract_nodes import extract_nodes
 from cognee.tasks.graph.cascade_extract.utils.extract_content_nodes_and_relationship_names import (
     extract_content_nodes_and_relationship_names,
@@ -17,9 +17,21 @@ from cognee.tasks.graph.extract_graph_from_data import integrate_chunk_graphs
 async def extract_graph_from_data(
     data_chunks: List[DocumentChunk],
     n_rounds: int = 2,
-    ontology_adapter: OntologyResolver = None,
+    ontology_adapter: BaseOntologyResolver = None,
 ) -> List[DocumentChunk]:
-    """Extract and update graph data from document chunks in multiple steps."""
+    """Extract and update graph data from document chunks using cascade extraction.
+
+    This function performs multi-step graph extraction from document chunks,
+    using cascade extraction techniques to build comprehensive knowledge graphs.
+
+    Args:
+        data_chunks: List of document chunks to process
+        n_rounds: Number of extraction rounds to perform (default: 2)
+        ontology_adapter: Resolver for validating entities against ontology
+
+    Returns:
+        List of updated DocumentChunk objects with extracted graph data
+    """
     chunk_nodes = await asyncio.gather(
         *[extract_nodes(chunk.text, n_rounds) for chunk in data_chunks]
     )
@@ -44,5 +56,5 @@ async def extract_graph_from_data(
         data_chunks=data_chunks,
         chunk_graphs=chunk_graphs,
         graph_model=KnowledgeGraph,
-        ontology_adapter=ontology_adapter or OntologyResolver(),
+        ontology_adapter=ontology_adapter,
     )

cognee/tasks/ingestion/save_data_item_to_storage.py

@@ -41,6 +41,7 @@ async def save_data_item_to_storage(data_item: Union[BinaryIO, str, Any]) -> str
             abs_path.is_file()
         except (OSError, ValueError):
             # In case file path is too long it's most likely not a relative path
+            abs_path = data_item
             logger.debug(f"Data item was too long to be a possible file path: {abs_path}")
             abs_path = Path("")
 

cognee/tasks/temporal_graph/models.py

@@ -3,12 +3,17 @@ from pydantic import BaseModel, Field
 
 
 class Timestamp(BaseModel):
-    year: int = Field(..., ge=1, le=9999)
-    month: int = Field(..., ge=1, le=12)
-    day: int = Field(..., ge=1, le=31)
-    hour: int = Field(..., ge=0, le=23)
-    minute: int = Field(..., ge=0, le=59)
-    second: int = Field(..., ge=0, le=59)
+    year: int = Field(
+        ...,
+        ge=1,
+        le=9999,
+        description="Always required. If only a year is known, use it.",
+    )
+    month: int = Field(1, ge=1, le=12, description="If unknown, default to 1")
+    day: int = Field(1, ge=1, le=31, description="If unknown, default to 1")
+    hour: int = Field(0, ge=0, le=23, description="If unknown, default to 0")
+    minute: int = Field(0, ge=0, le=59, description="If unknown, default to 0")
+    second: int = Field(0, ge=0, le=59, description="If unknown, default to 0")
 
 
 class Interval(BaseModel):

cognee/tests/cli_tests/cli_unit_tests/test_cli_main.py

@@ -49,7 +49,7 @@ class TestCliMain:
     def test_main_no_command(self, mock_create_parser):
         """Test main function when no command is provided"""
         mock_parser = MagicMock()
-        mock_parser.parse_args.return_value = MagicMock(command=None)
+        mock_parser.parse_args.return_value = MagicMock(command=None, spec={})
         mock_create_parser.return_value = (mock_parser, {})
 
         result = main()
@@ -64,7 +64,7 @@ class TestCliMain:
         mock_command.execute.return_value = None
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})
@@ -84,7 +84,7 @@ class TestCliMain:
         mock_command.execute.side_effect = CliCommandException("Test error", error_code=2)
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})
@@ -103,7 +103,7 @@ class TestCliMain:
         mock_command.execute.side_effect = Exception("Generic error")
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})
@@ -126,7 +126,7 @@ class TestCliMain:
         mock_command.execute.side_effect = test_exception
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})

cognee/tests/test_cognee_server_start.py

@@ -48,7 +48,7 @@ class TestCogneeServerStart(unittest.TestCase):
         """Test that the server is running and can accept connections."""
         # Test health endpoint
         health_response = requests.get("http://localhost:8000/health", timeout=15)
-        self.assertIn(health_response.status_code, [200, 503])
+        self.assertIn(health_response.status_code, [200])
 
         # Test root endpoint
         root_response = requests.get("http://localhost:8000/", timeout=15)
@@ -88,7 +88,7 @@ class TestCogneeServerStart(unittest.TestCase):
         payload = {"datasets": [dataset_name]}
 
         add_response = requests.post(url, headers=headers, data=form_data, files=file, timeout=50)
-        if add_response.status_code not in [200, 201, 409]:
+        if add_response.status_code not in [200, 201]:
             add_response.raise_for_status()
 
         # Cognify request
@@ -99,7 +99,7 @@ class TestCogneeServerStart(unittest.TestCase):
         }
 
         cognify_response = requests.post(url, headers=headers, json=payload, timeout=150)
-        if cognify_response.status_code not in [200, 201, 409]:
+        if cognify_response.status_code not in [200, 201]:
             cognify_response.raise_for_status()
 
         # TODO: Add test to verify cognify pipeline is complete before testing search
@@ -115,7 +115,7 @@ class TestCogneeServerStart(unittest.TestCase):
         payload = {"searchType": "GRAPH_COMPLETION", "query": "What's in the document?"}
 
         search_response = requests.post(url, headers=headers, json=payload, timeout=50)
-        if search_response.status_code not in [200, 201, 409]:
+        if search_response.status_code not in [200, 201]:
             search_response.raise_for_status()
 
 

cognee/tests/test_temporal_graph.py

@@ -97,7 +97,7 @@ async def main():
         f"Expected exactly one DocumentChunk, but found {type_counts.get('DocumentChunk', 0)}"
     )
 
-    assert type_counts.get("Entity", 0) >= 20, (
+    assert type_counts.get("Entity", 0) >= 10, (
         f"Expected multiple entities (assert is set to 20), but found {type_counts.get('Entity', 0)}"
     )
 
@@ -105,52 +105,24 @@ async def main():
         f"Expected multiple entity types, but found {type_counts.get('EntityType', 0)}"
     )
 
-    assert type_counts.get("Event", 0) >= 20, (
+    assert type_counts.get("Event", 0) >= 10, (
         f"Expected multiple events (assert is set to 20), but found {type_counts.get('Event', 0)}"
     )
 
-    assert type_counts.get("Timestamp", 0) >= 20, (
-        f"Expected multiple timestamps (assert is set to 20), but found {type_counts.get('Timestamp', 0)}"
+    assert type_counts.get("Timestamp", 0) >= 10, (
+        f"Expected multiple timestamps (assert is set to 10), but found {type_counts.get('Timestamp', 0)}"
     )
 
-    assert type_counts.get("Interval", 0) >= 2, (
-        f"Expected multiple intervals, but found {type_counts.get('Interval', 0)}"
-    )
-
-    assert edge_type_counts.get("contains", 0) >= 20, (
+    assert edge_type_counts.get("contains", 0) >= 10, (
         f"Expected multiple 'contains' edge, but found {edge_type_counts.get('contains', 0)}"
     )
 
-    assert edge_type_counts.get("is_a", 0) >= 20, (
+    assert edge_type_counts.get("is_a", 0) >= 10, (
         f"Expected multiple 'is_a' edge, but found {edge_type_counts.get('is_a', 0)}"
     )
 
-    assert edge_type_counts.get("during", 0) == type_counts.get("Interval", 0), (
-        "Expected the same amount of during and interval objects in the graph"
-    )
-
-    assert edge_type_counts.get("during", 0) == type_counts.get("Interval", 0), (
-        "Expected the same amount of during and interval objects in the graph"
-    )
-
-    assert edge_type_counts.get("time_from", 0) == type_counts.get("Interval", 0), (
-        "Expected the same amount of time_from and interval objects in the graph"
-    )
-
-    assert edge_type_counts.get("time_to", 0) == type_counts.get("Interval", 0), (
-        "Expected the same amount of time_to and interval objects in the graph"
-    )
-
     retriever = TemporalRetriever()
 
-    result_before = await retriever.extract_time_from_query("What happened before 1890?")
-
-    assert result_before[0] is None
-
-    result_after = await retriever.extract_time_from_query("What happened after 1891?")
-
-    assert result_after[1] is None
-
     result_between = await retriever.extract_time_from_query("What happened between 1890 and 1900?")
 
     assert result_between[1]