cognite-neat 0.123.2__py3-none-any.whl → 0.127.30__py3-none-any.whl

cognite/neat/_session/_usage_analytics/_collector.py

@@ -0,0 +1,131 @@
+import os
+import platform
+import threading
+import uuid
+from contextlib import suppress
+from typing import Any, Literal
+
+from mixpanel import Consumer, Mixpanel  # type: ignore[import-untyped]
+
+from cognite.neat._session._usage_analytics._constants import IN_NOTEBOOK, IN_PYODIDE
+from cognite.neat._session._usage_analytics._storage import ManualReadResult, ReadResult, get_storage
+
+_NEAT_MIXPANEL_TOKEN: str = "bd630ad149d19999df3989c3a3750c4f"
+
+
+class Collector:
+    """Collects usage data and sends it to Mixpanel."""
+
+    _instance: "Collector | None" = None
+    _initialized: bool = False
+
+    def __new__(cls, *args: Any, **kwargs: Any) -> "Collector":
+        # Implementing Singleton pattern
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self, skip_tracking: bool = False) -> None:
+        if self._initialized:
+            return
+        self.mp = Mixpanel(_NEAT_MIXPANEL_TOKEN, consumer=Consumer(api_host="api-eu.mixpanel.com"))
+        self._storage = get_storage()
+        self._opt_status_key = "neat-opt-status"
+        self._distinct_id_key = "neat-distinct-id"
+        self.skip_tracking = skip_tracking
+        self._initialized = True
+        # We trigger these reads immediately to have them ready when needed.
+        # This is necessary in a Pyodide environment were the read will not be available
+        self._opt_status_read: ReadResult | None = self._storage.read(self._opt_status_key)
+        self._distinct_id_read: ReadResult | None = self._storage.read(self._distinct_id_key)
+
+    @property
+    def _opt_status(self) -> str:
+        if self._opt_status_read is None:
+            self._opt_status_read = self._storage.read(self._opt_status_key)
+        if self._opt_status_read.is_ready:
+            return self._opt_status_read.get_data()
+        return ""  # We do not have an opt status yet.
+
+    @property
+    def _distinct_id(self) -> str | None:
+        if self._distinct_id_read is None:
+            self._distinct_id_read = self._storage.read(self._distinct_id_key)
+        if self._distinct_id_read.is_ready:
+            return self._distinct_id_read.get_data()
+        return None
+
+    def bust_opt_status(self) -> None:
+        self._storage.delete(self._opt_status_key)
+        self._opt_status_read = None
+
+    @property
+    def is_opted_out(self) -> bool:
+        return self._opt_status == "opted-out"
+
+    @property
+    def is_opted_in(self) -> bool:
+        return self._opt_status == "opted-in"
+
+    def enable(self) -> None:
+        self._storage.write(self._opt_status_key, "opted-in")
+        # Override cached property
+        self._opt_status_read = ManualReadResult("opted-in")
+
+    def disable(self) -> None:
+        self._storage.write(self._opt_status_key, "opted-out")
+        # Override cached property
+        self._opt_status_read = ManualReadResult("opted-out")
+
+    def get_distinct_id(self) -> str:
+        existing_id = self._distinct_id
+        if existing_id:
+            return existing_id
+
+        distinct_id = f"{platform.system()}-{platform.python_version()}-{uuid.uuid4()!s}"
+        self._storage.write(self._distinct_id_key, distinct_id)
+        with suppress(ConnectionError):
+            self.mp.people_set(
+                distinct_id,
+                {
+                    "$os": platform.system(),
+                    "$python_version": platform.python_version(),
+                    "$distinct_id": distinct_id,
+                    "environment": self._get_environment(),
+                },
+            )
+        self._distinct_id_read = ManualReadResult(distinct_id)
+        return distinct_id
+
+    @staticmethod
+    def _get_environment() -> Literal["pyodide", "notebook", "python"]:
+        """Get the current environment the user is running in."""
+        if IN_PYODIDE:
+            return "pyodide"
+        if IN_NOTEBOOK:
+            return "notebook"
+        return "python"
+
+    @property
+    def can_collect(self) -> bool:
+        """Check if tracking is possible."""
+        return not self.skip_tracking and self.is_opted_in and "PYTEST_CURRENT_TEST" not in os.environ
+
+    def collect(
+        self, event_name: Literal["action", "initSession", "deployment"], event_properties: dict[str, Any]
+    ) -> None:
+        distinct_id = self.get_distinct_id()
+
+        def track() -> None:
+            # If we are unable to connect to Mixpanel, we don't want to crash the program
+            with suppress(ConnectionError):
+                self.mp.track(distinct_id, event_name, event_properties)
+
+        if IN_PYODIDE:
+            # We cannot spawn threads in Pyodide
+            track()
+            return None
+
+        thread = threading.Thread(target=track, daemon=True)
+        thread.start()
+        return None

cognite/neat/_session/_usage_analytics/_constants.py

@@ -0,0 +1,23 @@
+def _is_in_notebook() -> bool:
+    try:
+        from IPython import get_ipython
+
+        if "IPKernelApp" not in get_ipython().config:  # pragma: no cover
+            return False
+    except ImportError:
+        return False
+    except AttributeError:
+        return False
+    return True
+
+
+def _is_in_browser() -> bool:
+    try:
+        from pyodide.ffi import IN_BROWSER  # type: ignore [import-not-found]
+    except ModuleNotFoundError:
+        return False
+    return IN_BROWSER
+
+
+IN_PYODIDE = _is_in_browser()
+IN_NOTEBOOK = _is_in_notebook() or IN_PYODIDE

cognite/neat/_session/_usage_analytics/_storage.py

@@ -0,0 +1,240 @@
+"""Storage abstraction for persisting data in both local filesystem and pyodide environments."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    try:
+        import asyncio
+
+        from pyodide.webloop import PyodideTask  # type: ignore[import-not-found]
+    except ImportError:
+        pass
+
+from cognite.neat._session._usage_analytics._constants import IN_PYODIDE
+
+
+@dataclass
+class ReadResult(ABC):
+    @property
+    @abstractmethod
+    def is_ready(self) -> bool: ...
+
+    @abstractmethod
+    def get_data(self) -> str: ...
+
+
+@dataclass
+class ManualReadResult(ReadResult):
+    data: str
+
+    @property
+    def is_ready(self) -> bool:
+        return True
+
+    def get_data(self) -> str:
+        return self.data
+
+
+class Storage(Protocol):
+    """Protocol for storage implementations."""
+
+    def read(self, key: str) -> ReadResult:
+        """Read a value from storage."""
+        ...
+
+    def write(self, key: str, value: str) -> None:
+        """Write a value to storage."""
+        ...
+
+    def delete(self, key: str) -> None:
+        """Delete a value from storage."""
+        ...
+
+
+@dataclass
+class FileReadResult(ReadResult):
+    data: str
+
+    @property
+    def is_ready(self) -> bool:
+        return True
+
+    def get_data(self) -> str:
+        return self.data
+
+
+class FileSystemStorage:
+    """Storage implementation using the filesystem."""
+
+    ENCODING = "utf-8"
+
+    def __init__(self, base_dir: Path) -> None:
+        self._base_dir = base_dir
+
+    def _get_path(self, key: str) -> Path:
+        return self._base_dir / f"{key}.bin"
+
+    def read(self, key: str) -> ReadResult:
+        path = self._get_path(key)
+        if path.exists():
+            return FileReadResult(path.read_text(encoding=self.ENCODING))
+        return FileReadResult("")
+
+    def write(self, key: str, value: str) -> None:
+        path = self._get_path(key)
+        path.write_text(value, encoding=self.ENCODING)
+
+    def delete(self, key: str) -> None:
+        path = self._get_path(key)
+        path.unlink(missing_ok=True)
+
+
+@dataclass
+class PyodideResult(ReadResult):
+    task: "PyodideTask"
+
+    @property
+    def is_ready(self) -> bool:
+        return self.task.done()
+
+    def get_data(self) -> str:
+        return self.task.result()
+
+
+class LocalStorageAdapter:
+    """
+    Storage implementation using browser IndexedDB for pyodide environments.
+
+    This adapter provides a synchronous interface to the asynchronous IndexedDB API
+    by managing an asyncio event loop.
+    """
+
+    _db_name = "neat-storage"
+    _store_name = "keyval"
+    _db_version = 1
+
+    def __init__(self) -> None:
+        self._loop = self._get_or_create_event_loop()
+
+    @staticmethod
+    def _get_or_create_event_loop() -> "asyncio.AbstractEventLoop":
+        """Gets the current asyncio event loop or creates a new one."""
+        import asyncio
+
+        try:
+            return asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            return loop
+
+    def _execute_db_operation(self, operation: str, key: str, value: str | None = None) -> "PyodideTask":
+        """Executes a database operation by creating and calling a JS function."""
+
+        async def db_operation() -> str:
+            import js  # type: ignore[import-not-found]
+            from pyodide.code import run_js  # type: ignore[import-not-found]
+            from pyodide.ffi import to_js  # type: ignore[import-not-found]
+
+            # Convert Python values to JS
+            js_key = to_js(key)
+            js_value = to_js(value if value is not None else "")
+
+            # Store values in js namespace temporarily to pass them safely
+            js._idb_key = js_key
+            js._idb_value = js_value
+
+            try:
+                # Create and execute the function directly in JavaScript context using run_js
+                js_code = f"""
+                (async () => {{
+                    const dbName = "{self._db_name}";
+                    const dbVersion = {self._db_version};
+                    const storeName = "{self._store_name}";
+                    const operation = "{operation}";
+
+                    return new Promise((resolve, reject) => {{
+                        const request = indexedDB.open(dbName, dbVersion);
+
+                        request.onerror = (event) => reject(event.target.error);
+
+                        request.onupgradeneeded = (event) => {{
+                            const db = event.target.result;
+                            if (!db.objectStoreNames.contains(storeName)) {{
+                                db.createObjectStore(storeName);
+                            }}
+                        }};
+
+                        request.onsuccess = (event) => {{
+                            const db = event.target.result;
+                            const mode = (operation === "read") ? "readonly" : "readwrite";
+                            try {{
+                                const transaction = db.transaction([storeName], mode);
+                                const store = transaction.objectStore(storeName);
+
+                                let storeRequest;
+                                if (operation === "read") {{
+                                    storeRequest = store.get(_idb_key);
+                                }} else if (operation === "write") {{
+                                    storeRequest = store.put(_idb_value, _idb_key);
+                                }} else if (operation === "delete") {{
+                                    storeRequest = store.delete(_idb_key);
+                                }} else {{
+                                    db.close();
+                                    return reject(new Error(`Unknown operation: ${{operation}}`));
+                                }}
+
+                                storeRequest.onsuccess = () => resolve(storeRequest.result ?? "");
+                                storeRequest.onerror = (event) => reject(event.target.error);
+
+                                transaction.oncomplete = () => db.close();
+                                transaction.onerror = (event) => reject(event.target.error);
+
+                            }} catch (error) {{
+                                db.close();
+                                reject(error);
+                            }}
+                        }};
+                    }});
+                }})()
+                """
+
+                # Use run_js which properly returns an awaitable promise
+                result = await run_js(js_code)
+                return str(result) if result is not None else ""
+            finally:
+                # Clean up js namespace
+                try:
+                    del js._idb_key
+                    del js._idb_value
+                except Exception:
+                    pass
+
+        try:
+            return self._loop.run_until_complete(db_operation())
+        except Exception:
+            return ""
+
+    def read(self, key: str) -> ReadResult:
+        return PyodideResult(self._execute_db_operation("read", key))
+
+    def write(self, key: str, value: str) -> None:
+        self._execute_db_operation("write", key, value)
+
+    def delete(self, key: str) -> None:
+        self._execute_db_operation("delete", key)
+
+
+def get_storage(base_dir: Path | None = None) -> Storage:
+    """Get the appropriate storage implementation for the current environment."""
+    if IN_PYODIDE:
+        return LocalStorageAdapter()
+    else:
+        if base_dir is None:
+            import tempfile
+
+            base_dir = Path(tempfile.gettempdir()).resolve()
+        return FileSystemStorage(base_dir)

cognite/neat/_session/_wrappers.py

@@ -0,0 +1,82 @@
+from collections.abc import Callable
+from functools import wraps
+from typing import Any, Protocol, TypeVar
+
+from cognite.neat._session._usage_analytics._collector import Collector
+from cognite.neat._store._store import NeatStore
+from cognite.neat._utils.text import NEWLINE, split_on_capitals
+
+
+class HasStore(Protocol):
+    _store: NeatStore
+
+
+T_Class = TypeVar("T_Class", bound=HasStore)
+
+_COLLECTOR = Collector()
+
+
+def session_wrapper(cls: type[T_Class]) -> type[T_Class]:
+    # 1. Define the method decorator inside
+    def _handle_method_call(func: Callable[..., Any]) -> Callable[..., Any]:
+        """Decorator to handle exceptions and print provenance length"""
+
+        @wraps(func)
+        def wrapper(self: HasStore, *args: Any, **kwargs: Any) -> Any:
+            display_name = f"{' '.join(split_on_capitals(cls.__name__))} - {func.__name__}"
+            identifier = f"{cls.__name__}.{func.__name__}"
+            try:
+                res = func(self, *args, **kwargs)
+                if not self._store.provenance or "DataModel" not in identifier:
+                    print(f"{display_name} ✅")
+                    if _COLLECTOR.can_collect:
+                        _COLLECTOR.collect("action", {"action": identifier, "success": True})
+                    return res
+                change = self._store.provenance[-1]
+                issues_count = len(change.issues) if change.issues else 0
+                errors_count = len(change.errors) if change.errors else 0
+                total_issues = issues_count + errors_count
+
+                data_model_not_read = not change.successful and "ReadPhysicalDataModel" in identifier
+
+                print(
+                    f"{display_name} "
+                    f"{'✅' if change.successful else '❌'}"
+                    f"{f' | Issues: {total_issues} (of which {errors_count} critical)' if total_issues > 0 else ''}"
+                    f"{NEWLINE + '⚠️ Data model not read into session' if data_model_not_read else ''}"
+                    f"{NEWLINE + '📋 For details on issues run neat.issues' if change.issues or change.errors else ''}"
+                    f"{NEWLINE + '📊 For details on result run neat.result' if change.result else ''}"
+                )
+                if _COLLECTOR.can_collect:
+                    event = change.as_mixpanel_event()
+                    event["action"] = identifier
+                    _COLLECTOR.collect("action", event)
+                    if change.result:
+                        event = change.result.as_mixpanel_event()
+                        event["action"] = identifier
+                        _COLLECTOR.collect("deployment", event)
+
+                return res
+
+            # if an error occurs, we catch it and print it out instead of
+            # getting a full traceback
+            except Exception as e:
+                print(f"{display_name} ❌")
+                print(f"{e!s}")
+                if _COLLECTOR.can_collect:
+                    _COLLECTOR.collect("action", {"action": identifier, "success": False, "error_message": str(e)})
+
+        return wrapper
+
+    # Iterate through all attributes of the class
+    for attr_name in dir(cls):
+        # Skip private/protected methods (starting with _)
+        if not attr_name.startswith("_"):
+            attr = getattr(cls, attr_name)
+            # Only wrap callable methods
+            if callable(attr):
+                # Replace the original method with wrapped version
+                setattr(cls, attr_name, _handle_method_call(attr))
+
+    # Return the modified class
+    return cls

cognite/neat/_state_machine/__init__.py

@@ -0,0 +1,10 @@
+from ._base import State
+from ._states import EmptyState, ForbiddenState, PhysicalState, Undo
+
+__all__ = [
+    "EmptyState",
+    "ForbiddenState",
+    "PhysicalState",
+    "State",
+    "Undo",
+]

cognite/neat/_state_machine/_base.py

@@ -0,0 +1,37 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class State(ABC):
+    def __init__(self) -> None:
+        # this will be reference to the actual store in the session
+        # used to store data models and instances, here only as a placeholder
+        self._store = None
+
+    @abstractmethod
+    def transition(self, event: Any) -> "State":
+        """
+        Handle events that are delegated to this State.
+        """
+        raise NotImplementedError("on_event() must be implemented by the subclass.")
+
+    def can_transition(self, event: Any) -> bool:
+        """
+        Check if the state can transition on the given event.
+        """
+        # avoiding circular import
+        from cognite.neat._state_machine._states import ForbiddenState
+
+        return not isinstance(self.transition(event), ForbiddenState)
+
+    def __repr__(self) -> str:
+        """
+        Leverages the __str__ method to describe the State.
+        """
+        return self.__str__()
+
+    def __str__(self) -> str:
+        """
+        Returns the name of the State.
+        """
+        return self.__class__.__name__

cognite/neat/_state_machine/_states.py

@@ -0,0 +1,52 @@
+from typing import Any
+
+from cognite.neat._data_model.exporters import DMSExporter
+from cognite.neat._data_model.importers import DMSImporter
+
+from ._base import State
+
+
+class Undo:
+    """
+    Event to trigger undoing the last action.
+    """
+
+    pass
+
+
+class ForbiddenState(State):
+    """
+    State representing forbidden transitions - returns to previous state.
+    """
+
+    def __init__(self, previous_state: State):
+        self.previous_state = previous_state
+
+    def transition(self, event: Any) -> State:
+        # only "undo" to trigger going back to previous state
+        if isinstance(event, Undo):
+            return self.previous_state
+        return self
+
+
+class EmptyState(State):
+    """
+    The initial state with empty NEAT store.
+    """
+
+    def transition(self, event: Any) -> State:
+        if isinstance(event, DMSImporter):
+            return PhysicalState()
+        return ForbiddenState(self)
+
+
+class PhysicalState(State):
+    """
+    State with physical model loaded.
+    """
+
+    def transition(self, event: Any) -> State:
+        if isinstance(event, DMSExporter):
+            return PhysicalState()
+
+        return ForbiddenState(self)

cognite/neat/_store/__init__.py

@@ -0,0 +1,3 @@
+from ._store import NeatStore
+
+__all__ = ["NeatStore"]

cognite/neat/_store/_provenance.py

@@ -0,0 +1,81 @@
+import itertools
+from collections import UserList
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+from cognite.neat._data_model.deployer.data_classes import DeploymentResult
+from cognite.neat._issues import ConsistencyError, IssueList, ModelSyntaxError
+from cognite.neat._state_machine import State
+
+
+@dataclass
+class Change:
+    agent: str
+    activity: str
+    source_state: State
+    start: datetime
+    end: datetime
+
+    target_state: State | None = field(default=None)
+    source_entity: str | None = field(default="ExternalEntity")
+    target_entity: str | None = field(default="FailedEntity")
+    issues: IssueList | None = field(default=None)
+    errors: IssueList | None = field(default=None)
+    # for time being setting to Any, can be refined later
+    result: DeploymentResult | None = field(default=None)
+    description: str | None = field(default=None)
+
+    @staticmethod
+    def standardize_activity_name(activity: str, start: datetime, end: datetime) -> str:
+        """Create standardized activity name"""
+        return f"{activity}_{start.timestamp()}-{end.timestamp()}"
+
+    @property
+    def successful(self) -> bool:
+        """Check if change was successful"""
+        return not self.errors
+
+    @property
+    def error_count(self) -> int:
+        """Get number of errors"""
+        return sum(
+            1
+            for issue in itertools.chain(self.issues or [], self.errors or [])
+            if isinstance(issue, ModelSyntaxError | ConsistencyError)
+        )
+
+    def as_mixpanel_event(self) -> dict[str, Any]:
+        """Convert change to mixpanel event format"""
+        return {
+            "agent": self.agent,
+            "activity": self.activity,
+            "sourceEntity": self.source_entity,
+            "targetEntity": self.target_entity,
+            "sourceState": type(self.source_state).__name__,
+            "targetState": type(self.target_state).__name__ if self.target_state else "None",
+            "duration_ms": int((self.end - self.start).total_seconds() * 1000),
+            "successful": self.successful,
+            "issues": [issue.code or "<no code>" for issue in self.issues] if self.issues else [],
+            "errors": [error.code or "<no code>" for error in self.errors] if self.errors else [],
+        }
+
+
+class Provenance(UserList[Change]):
+    def __delitem__(self, *args: Any, **kwargs: Any) -> None:
+        raise TypeError("Cannot delete change from provenance")
+
+    def __setitem__(self, *args: Any, **kwargs: Any) -> None:
+        raise TypeError("Cannot modify change from provenance")
+
+    @property
+    def last_change(self) -> Change | None:
+        return self[-1] if len(self) > 0 else None
+
+    def can_agent_do_activity(self, activity: Any) -> bool:
+        "Check if activity can be performed based on provenance"
+        raise NotImplementedError("Not implemented yet")
+
+    def provenance_without_failures(self) -> "Provenance":
+        """This method removes all the failed changes from the provenance list."""
+        raise NotImplementedError("Not implemented yet")