cognite-neat 0.96.6__py3-none-any.whl → 0.97.0__py3-none-any.whl

cognite/neat/_session/engine/_load.py

@@ -0,0 +1,129 @@
+import os
+import re
+import shutil
+import sys
+import tempfile
+import warnings
+from collections.abc import Callable
+from pathlib import Path
+from typing import Literal, cast
+
+from cognite.client import CogniteClient
+from packaging.version import Version
+from packaging.version import parse as parse_version
+
+from cognite.neat._issues.errors import NeatValueError
+from cognite.neat._version import __engine__
+
+ENVIRONMENT_VARIABLE = "NEATENGINE"
+PACKAGE_NAME = "neatengine"
+PYVERSION = f"{sys.version_info.major}{sys.version_info.minor}"
+
+
+def load_neat_engine(client: CogniteClient | None, location: Literal["newest", "cache"]) -> str | None:
+    if location not in ["newest", "cache"]:
+        raise NeatValueError(f"Cannot load engine from location: {location}")
+
+    if not __engine__.startswith("^"):
+        # Using value error as this is a developer error
+        raise ValueError(f"Invalid engine version: {__engine__}")
+
+    lower_bound = parse_version(__engine__[1:])
+    upper_bound = Version(f"{lower_bound.major + 1}.0.0")
+
+    cache_dir = Path(tempfile.gettempdir()) / PACKAGE_NAME
+    cache_dir.mkdir(exist_ok=True)
+    pattern = re.compile(rf"{PACKAGE_NAME}-(\d+\.\d+\.\d+)-{PYVERSION}.zip")
+
+    candidates: dict[Version, Callable[[], Path]] = {}
+    if location == "cache" and cache_dir.exists():
+        candidates = _load_from_path(cache_dir, pattern, lower_bound, upper_bound)
+
+    if location == "newest" or not candidates:
+        # Loading in reverse order of priority
+        # 3. Downloads folder
+        candidates = _load_from_path(Path.home() / "Downloads", pattern, lower_bound, upper_bound)
+        # 2. CDF
+        if client:
+            candidates.update(_load_from_cdf(client, pattern, lower_bound, upper_bound, cache_dir))
+        # 1. Environment variable
+        if ENVIRONMENT_VARIABLE in os.environ:
+            environ_path = Path(os.environ[ENVIRONMENT_VARIABLE])
+            if environ_path.exists():
+                candidates.update(_load_from_path(environ_path, pattern, lower_bound, upper_bound))
+            else:
+                warnings.warn(
+                    f"Environment variable {ENVIRONMENT_VARIABLE} points to non-existing path: {environ_path}",
+                    UserWarning,
+                    stacklevel=2,
+                )
+
+    if not candidates:
+        return None
+
+    selected_version = max(candidates.keys(), default=None)
+    if not selected_version:
+        return None
+    source_path = candidates[selected_version]()
+    destination_path = cache_dir / source_path.name
+    if not destination_path.exists():
+        shutil.copy(source_path, destination_path)
+    sys.path.append(str(destination_path))
+    try:
+        from neatengine._version import __version__ as engine_version  # type: ignore[import-not-found]
+    except ImportError:
+        return None
+    return engine_version
+
+
+def _load_from_path(
+    path: Path, pattern: re.Pattern[str], lower_bound: Version, upper_bound: Version
+) -> dict[Version, Callable[[], Path]]:
+    if path.is_file() and (match := pattern.match(path.name)):
+        version = parse_version(match.group(1))
+        if lower_bound <= version < upper_bound:
+            return {parse_version(match.group(1)): lambda: path}
+        return {}
+    elif path.is_dir():
+        output: dict[Version, Callable[[], Path]] = {}
+        for candidate in path.iterdir():
+            if candidate.is_file() and (match := pattern.match(candidate.name)):
+                version = parse_version(match.group(1))
+                if lower_bound <= version < upper_bound:
+                    # Setting default value to ensure we use the candidate from the current iteration
+                    # If not set, the function will use the last candidate from the loop
+                    def return_path(the_path: Path = candidate) -> Path:
+                        return the_path
+
+                    output[parse_version(match.group(1))] = return_path
+
+        return output
+    return {}
+
+
+def _load_from_cdf(
+    client: CogniteClient, pattern: re.Pattern[str], lower_bound: Version, upper_bound: Version, cache_dir: Path
+) -> dict[Version, Callable[[], Path]]:
+    file_metadata = client.files.list(
+        limit=-1,
+        data_set_external_ids=PACKAGE_NAME,
+        external_id_prefix=PACKAGE_NAME,
+        metadata={"python_version": PYVERSION},
+    )
+    output: dict[Version, Callable[[], Path]] = {}
+    for file in file_metadata:
+        name = cast(str, file.name)
+
+        # Use function to lazily download file
+        # Setting default value to ensure we use the file_id from the current iteration
+        # If not set, the function will use the last file_id from the loop
+        def download_file(file_id: int = file.id, filename: str = name) -> Path:
+            client.files.download(cache_dir, file_id)
+            return cache_dir / filename
+
+        if match := pattern.match(name):
+            version = parse_version(match.group(1))
+            if lower_bound <= version < upper_bound:
+                output[version] = download_file
+
+    return output

cognite/neat/_session/exceptions.py

@@ -22,10 +22,20 @@ def _intercept_session_exceptions(func: Callable):
         try:
             return func(*args, **kwargs)
         except NeatSessionError as e:
-            action = func.__name__
-            if action == "__call__":
-                action = func.__qualname__.removesuffix(".__call__").removesuffix("API")
+            action = _get_action()
             print(f"{_PREFIX} Cannot {action}: {e}")
+        except ModuleNotFoundError as e:
+            if e.name == "neatengine":
+                action = _get_action()
+                print(f"{_PREFIX} The functionality {action} requires the NeatEngine.")
+            else:
+                raise e
+
+    def _get_action():
+        action = func.__name__
+        if action == "__call__":
+            action = func.__qualname__.removesuffix(".__call__").removesuffix("API")
+        return action
 
     return wrapper
 

cognite/neat/_shared.py

@@ -1,9 +1,10 @@
 from abc import abstractmethod
 from collections.abc import Hashable, Sequence
 from dataclasses import dataclass
-from typing import Any, TypeVar
+from typing import Any, TypeAlias, TypeVar
 
 import pandas as pd
+from rdflib import Literal, URIRef
 
 T_ID = TypeVar("T_ID", bound=Hashable)
 
@@ -49,3 +50,7 @@ class NeatList(list, Sequence[T_NeatObject]):
 
     def _repr_html_(self) -> str:
         return self.to_pandas()._repr_html_()  # type: ignore[operator]
+
+
+Triple: TypeAlias = tuple[URIRef, URIRef, Literal | URIRef]
+InstanceType: TypeAlias = URIRef

cognite/neat/_store/_base.py

@@ -12,13 +12,14 @@ from rdflib.plugins.stores.sparqlstore import SPARQLUpdateStore
 
 from cognite.neat._constants import DEFAULT_NAMESPACE
 from cognite.neat._graph.extractors import RdfFileExtractor, TripleExtractors
-from cognite.neat._graph.models import InstanceType, Triple
 from cognite.neat._graph.queries import Queries
 from cognite.neat._graph.transformers import Transformers
 from cognite.neat._rules.analysis import InformationAnalysis
 from cognite.neat._rules.models import InformationRules
 from cognite.neat._rules.models.entities import ClassEntity
+from cognite.neat._shared import InstanceType, Triple
 from cognite.neat._utils.auxiliary import local_import
+from cognite.neat._utils.rdf_ import add_triples_in_batch
 
 from ._provenance import Change, Provenance
 
@@ -288,29 +289,7 @@ class NeatGraphStore:
             batch_size: Batch size of triples per commit, by default 10_000
             verbose: Verbose mode, by default False
         """
-
-        commit_counter = 0
-        number_of_written_triples = 0
-
-        def check_commit(force_commit: bool = False):
-            """Commit nodes to the graph if batch counter is reached or if force_commit is True"""
-            nonlocal commit_counter
-            nonlocal number_of_written_triples
-            if force_commit:
-                number_of_written_triples += commit_counter
-                self.graph.commit()
-                return
-            commit_counter += 1
-            if commit_counter >= batch_size:
-                number_of_written_triples += commit_counter
-                self.graph.commit()
-                commit_counter = 0
-
-        for triple in triples:
-            self.graph.add(triple)
-            check_commit()
-
-        check_commit(force_commit=True)
+        add_triples_in_batch(self.graph, triples, batch_size)
 
     def transform(self, transformer: Transformers) -> None:
         """Transforms the graph store using a transformer."""

cognite/neat/_store/_provenance.py

@@ -1,22 +1,33 @@
-# we will use prov-o to represent the provenance of the neat graph store
-# basically tracking changes that occur in the graph store
-# prov-o use concepts of Agent, Activity and Entity to represent provenance
-# where in case of neat we have:
-# Agent: triples extractors, graph enhancers, contextualizers, etc.
-# Activity: write triple, add connection, etc.
-# Entity: neat graph store
+"""
+We use prov-o to represent the provenance of instances and data models
+basically tracking changes that occur.
+prov-o use concepts of Agent, Activity and Entity to represent provenance
+where in case of neat when dealing with instances we have:
 
+ * Agent: triples extractors, graph enhancers, contextualizers, etc.
+ * Activity: write/remove triples such as add connection, etc.
+ * Entity: neat graph store
+
+ and in case of data models we have:
+
+ * Agent: Rules importers, exporters, transformers, etc.
+ * Activity: convert, verify, etc.
+ * Entity: data model (aka Rules)
+
+"""
 
 import uuid
-from collections.abc import Sequence
-from dataclasses import dataclass
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass, field
 from datetime import datetime
-from typing import TypeVar
+from typing import Optional, TypeVar
 
+from cognite.client.data_classes.data_modeling import DataModelId, DataModelIdentifier
 from rdflib import PROV, RDF, Literal, URIRef
 
-from cognite.neat._constants import DEFAULT_NAMESPACE
-from cognite.neat._shared import FrozenNeatObject, NeatList
+from cognite.neat._constants import CDF_NAMESPACE, DEFAULT_NAMESPACE
+from cognite.neat._rules._shared import JustRules, ReadRules, VerifiedRules
+from cognite.neat._shared import FrozenNeatObject, NeatList, Triple
 
 
 @dataclass(frozen=True)
@@ -24,61 +35,137 @@ class Agent:
     id_: URIRef = DEFAULT_NAMESPACE.agent
     acted_on_behalf_of: str = "NEAT"
 
-    def as_triples(self):
+    def as_triples(self) -> list[Triple]:
         return [
             (self.id_, RDF.type, PROV[type(self).__name__]),
-            (self.id_, PROV.actedOnBehalfOf, self.acted_on_behalf_of),
+            (self.id_, PROV.actedOnBehalfOf, Literal(self.acted_on_behalf_of)),
+        ]
+
+
+CDF_AGENT = Agent(acted_on_behalf_of="UNKNOWN", id_=CDF_NAMESPACE["agent"])
+NEAT_AGENT = Agent(acted_on_behalf_of="UNKNOWN", id_=DEFAULT_NAMESPACE["agent"])
+UNKNOWN_AGENT = Agent(acted_on_behalf_of="UNKNOWN", id_=DEFAULT_NAMESPACE["unknown-agent"])
+
+
+@dataclass(frozen=True)
+class Entity:
+    was_attributed_to: Agent
+    was_generated_by: Optional["Activity"] = field(default=None, repr=False)
+    id_: URIRef = DEFAULT_NAMESPACE["graph-store"]
+
+    def as_triples(self) -> list[Triple]:
+        output: list[tuple[URIRef, URIRef, Literal | URIRef]] = [
+            (self.id_, RDF.type, PROV[type(self).__name__]),
+            (self.id_, PROV.wasAttributedTo, self.was_attributed_to.id_),
         ]
 
+        if self.was_generated_by:
+            output.append(
+                (
+                    self.id_,
+                    PROV.wasGeneratedBy,
+                    self.was_generated_by.id_,
+                )
+            )
+
+        return output
+
+    @classmethod
+    def from_data_model_id(cls, data_model_id: DataModelIdentifier) -> "Entity":
+        data_model_id = DataModelId.load(data_model_id)
+
+        return cls(
+            was_attributed_to=CDF_AGENT,
+            id_=CDF_NAMESPACE[
+                f"dms/data-model/{data_model_id.space}/{data_model_id.external_id}/{data_model_id.version}"
+            ],
+        )
+
+    @classmethod
+    def from_rules(
+        cls,
+        rules: ReadRules | JustRules | VerifiedRules,
+        agent: Agent | None = None,
+        activity: "Activity | None" = None,
+    ) -> "Entity":
+        agent = agent or UNKNOWN_AGENT
+        if isinstance(rules, VerifiedRules):
+            return cls(
+                was_attributed_to=agent,
+                was_generated_by=activity,
+                id_=rules.id_,
+            )
+
+        elif isinstance(rules, ReadRules | JustRules) and rules.rules is not None:
+            return cls(
+                was_attributed_to=agent,
+                was_generated_by=activity,
+                id_=rules.rules.id_,
+            )
+        else:
+            return cls(
+                was_attributed_to=agent,
+                was_generated_by=activity,
+                id_=DEFAULT_NAMESPACE["unknown-entity"],
+            )
+
+    @classmethod
+    def new_unknown_entity(cls) -> "Entity":
+        return cls(
+            was_attributed_to=UNKNOWN_AGENT,
+            id_=DEFAULT_NAMESPACE[f"unknown-entity/{uuid.uuid4()}"],
+        )
+
+
+INSTANCES_ENTITY = Entity(was_attributed_to=NEAT_AGENT, id_=CDF_NAMESPACE["instances"])
+
 
 @dataclass(frozen=True)
 class Activity:
     was_associated_with: Agent
     ended_at_time: datetime
     started_at_time: datetime
-    used: str  # this would be set to for example Extractor, Enhancer, Contextualizer, etc.
-    id_: URIRef = DEFAULT_NAMESPACE[f"activity-{uuid.uuid4()}"]
+    used: str | Entity | None = None
+    id_: URIRef = field(default_factory=lambda: DEFAULT_NAMESPACE[f"activity-{uuid.uuid4()}"])
 
-    def as_triples(self):
-        return [
+    def as_triples(self) -> list[Triple]:
+        output: list[tuple[URIRef, URIRef, Literal | URIRef]] = [
             (self.id_, RDF.type, PROV[type(self).__name__]),
             (self.id_, PROV.wasAssociatedWith, self.was_associated_with.id_),
             (self.id_, PROV.startedAtTime, Literal(self.started_at_time)),
             (self.id_, PROV.endedAtTime, Literal(self.ended_at_time)),
-            (self.id_, PROV.used, self.used),
         ]
 
+        if self.used:
+            output.append(
+                (
+                    self.id_,
+                    PROV.used,
+                    (self.used.id_ if isinstance(self.used, Entity) else Literal(self.used)),
+                )
+            )
 
-@dataclass(frozen=True)
-class Entity:
-    was_generated_by: Activity
-    was_attributed_to: Agent
-    id_: URIRef = DEFAULT_NAMESPACE["graph-store"]
-
-    def as_triples(self):
-        return [
-            (self.id_, RDF.type, PROV[type(self).__name__]),
-            (self.id_, PROV.wasGeneratedBy, self.was_generated_by.id_),
-            (self.id_, PROV.wasAttributedTo, self.was_attributed_to.id_),
-        ]
+        return output
 
 
 @dataclass(frozen=True)
 class Change(FrozenNeatObject):
     agent: Agent
     activity: Activity
-    entity: Entity
+    target_entity: Entity
     description: str
-    # triples that were added to the graph store
-    addition: list[tuple[URIRef, URIRef, URIRef | Literal]] | None = None
-    # triples that were removed from the graph store
-    subtraction: list[tuple[URIRef, URIRef, URIRef | Literal]] | None = None
+    source_entity: Entity = field(default_factory=Entity.new_unknown_entity)
 
-    def as_triples(self):
-        return self.agent.as_triples() + self.activity.as_triples() + self.entity.as_triples()
+    def as_triples(self) -> list[Triple]:
+        return (
+            self.source_entity.as_triples()
+            + self.agent.as_triples()
+            + self.activity.as_triples()
+            + self.target_entity.as_triples()  # type: ignore[operator]
+        )
 
     @classmethod
-    def record(cls, activity: str, start: datetime, end: datetime, description: str):
+    def record(cls, activity: str, start: datetime, end: datetime, description: str) -> "Change":
         """User friendly method to record a change that occurred in the graph store."""
         agent = Agent()
         activity = Activity(
@@ -87,14 +174,48 @@ class Change(FrozenNeatObject):
             started_at_time=start,
             ended_at_time=end,
         )
-        entity = Entity(was_generated_by=activity, was_attributed_to=agent)
-        return cls(agent, activity, entity, description)
+        target_entity = Entity(was_generated_by=activity, was_attributed_to=agent)
+        return cls(
+            agent=agent,
+            activity=activity,
+            target_entity=target_entity,
+            description=description,
+        )
+
+    @classmethod
+    def from_rules_activity(
+        cls,
+        rules: ReadRules | JustRules | VerifiedRules,
+        agent: Agent,
+        start: datetime,
+        end: datetime,
+        description: str,
+        source_entity: Entity | None = None,
+    ) -> "Change":
+        source_entity = source_entity or Entity.new_unknown_entity()
+        activity = Activity(
+            started_at_time=start,
+            ended_at_time=end,
+            was_associated_with=agent,
+            used=source_entity,
+        )
+
+        target_entity = Entity.from_rules(rules, agent, activity)
+
+        return cls(
+            agent=agent,
+            activity=activity,
+            target_entity=target_entity,
+            description=description,
+            source_entity=source_entity,
+        )
 
     def dump(self, aggregate: bool = True) -> dict[str, str]:
         return {
+            "Source Entity": self.source_entity.id_,
             "Agent": self.agent.id_,
             "Activity": self.activity.id_,
-            "Entity": self.entity.id_,
+            "Target Entity": self.target_entity.id_,
             "Description": self.description,
         }
 
@@ -124,3 +245,25 @@ class Provenance(NeatList[Change]):
         text += "</ul>"
 
         return text
+
+    def activity(self, id_: URIRef) -> Activity | None:
+        return next((change.activity for change in self if change.activity.id_ == id_), None)
+
+    def agent(self, id_: URIRef) -> Agent | None:
+        return next((change.agent for change in self if change.agent.id_ == id_), None)
+
+    def target_entity(self, id_: URIRef) -> Entity | None:
+        return next(
+            (change.target_entity for change in self if change.target_entity.id_ == id_),
+            None,
+        )
+
+    def source_entity(self, id_: URIRef) -> Entity | None:
+        return next(
+            (change.source_entity for change in self if change.source_entity.id_ == id_),
+            None,
+        )
+
+    def as_triples(self) -> Iterable[Triple]:
+        for change in self:
+            yield from change.as_triples()

cognite/neat/_utils/rdf_.py

@@ -4,8 +4,8 @@ from typing import Any, Literal, TypeAlias, overload
 
 from cognite.client.utils.useful_types import SequenceNotStr
 from pydantic import HttpUrl, TypeAdapter, ValidationError
+from rdflib import Graph, Namespace, URIRef
 from rdflib import Literal as RdfLiteral
-from rdflib import Namespace, URIRef
 
 Triple: TypeAlias = tuple[URIRef, URIRef, RdfLiteral | URIRef]
 
@@ -171,3 +171,36 @@ def get_inheritance_path(child: Any, child_parent: dict[Any, list[Any]]) -> list
         for parent in child_parent[child]:
             path.extend(get_inheritance_path(parent, child_parent))
     return path
+
+
+def add_triples_in_batch(graph: Graph, triples: Iterable[Triple], batch_size: int = 10_000):
+    """Adds triples to the graph store in batches.
+
+    Args:
+        triples: list of triples to be added to the graph store
+        batch_size: Batch size of triples per commit, by default 10_000
+        verbose: Verbose mode, by default False
+    """
+
+    commit_counter = 0
+    number_of_written_triples = 0
+
+    def check_commit(force_commit: bool = False):
+        """Commit nodes to the graph if batch counter is reached or if force_commit is True"""
+        nonlocal commit_counter
+        nonlocal number_of_written_triples
+        if force_commit:
+            number_of_written_triples += commit_counter
+            graph.commit()
+            return
+        commit_counter += 1
+        if commit_counter >= batch_size:
+            number_of_written_triples += commit_counter
+            graph.commit()
+            commit_counter = 0
+
+    for triple in triples:
+        graph.add(triple)
+        check_commit()
+
+    check_commit(force_commit=True)

cognite/neat/_utils/reader/__init__.py

@@ -0,0 +1,3 @@
+from ._base import GitHubReader, NeatReader, PathReader
+
+__all__ = ["NeatReader", "PathReader", "GitHubReader"]