cortexdb-airflow 0.1.0__py3-none-any.whl

cortexdb_airflow/__init__.py

@@ -0,0 +1,29 @@
+"""Apache Airflow provider for CortexDB.
+
+Provides a Hook for connection management and an Operator for executing
+CortexDB operations (experience, recall, forget, search) within Airflow DAGs.
+"""
+
+from cortexdb_airflow.hooks import CortexDBHook
+from cortexdb_airflow.operators import CortexDBOperator
+
+__all__ = [
+    "CortexDBHook",
+    "CortexDBOperator",
+]
+
+
+def get_provider_info() -> dict:
+    """Return Airflow provider metadata for auto-discovery."""
+    return {
+        "package-name": "cortexdb-airflow",
+        "name": "CortexDB",
+        "description": "Apache Airflow provider for CortexDB long-term memory.",
+        "connection-types": [
+            {
+                "connection-type": "cortexdb",
+                "hook-class-name": "cortexdb_airflow.hooks.CortexDBHook",
+            }
+        ],
+        "versions": ["0.1.0"],
+    }

cortexdb_airflow/hooks.py

@@ -0,0 +1,124 @@
+"""Airflow Hook for CortexDB connections.
+
+Manages CortexDB client lifecycle through Airflow's connection management
+system, enabling operators and tasks to share connection configuration.
+
+Example::
+
+    from cortexdb_airflow import CortexDBHook
+
+    hook = CortexDBHook(cortexdb_conn_id="cortexdb_default")
+    client = hook.get_client()
+    result = client.recall("user:default", query="recent events")
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from airflow.hooks.base import BaseHook
+
+from cortexdb import Cortex
+
+
+class CortexDBHook(BaseHook):
+    """Airflow Hook for managing CortexDB connections.
+
+    Reads connection parameters from Airflow's connection store and creates
+    a configured :class:`cortexdb.Cortex` client instance.
+
+    The Airflow connection should be configured as:
+    - **Host**: CortexDB server hostname (e.g., ``localhost``)
+    - **Port**: CortexDB server port (e.g., ``3141``)
+    - **Password**: PASETO bearer token (optional)
+    - **Login**: Actor id (optional, defaults to ``"user:default"``)
+    - **Schema**: Default scope path (optional, defaults to ``"user:default"``)
+    - **Extra**: JSON with ``timeout`` (optional)
+
+    Args:
+        cortexdb_conn_id: The Airflow connection ID. Defaults to
+            ``"cortexdb_default"``.
+    """
+
+    conn_name_attr = "cortexdb_conn_id"
+    default_conn_name = "cortexdb_default"
+    conn_type = "cortexdb"
+    hook_name = "CortexDB"
+
+    def __init__(
+        self,
+        cortexdb_conn_id: str = "cortexdb_default",
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.cortexdb_conn_id = cortexdb_conn_id
+        self._client: Cortex | None = None
+
+    def get_client(self) -> Cortex:
+        """Create or return a cached CortexDB client.
+
+        Returns:
+            An initialized :class:`cortexdb.Cortex` client using the Airflow
+            connection configuration.
+        """
+        if self._client is not None:
+            return self._client
+
+        conn = self.get_connection(self.cortexdb_conn_id)
+        extra = conn.extra_dejson if conn.extra else {}
+
+        scheme = extra.get("scheme", "http")
+        host = conn.host or "localhost"
+        port = conn.port or 3141
+        api_url = f"{scheme}://{host}:{port}"
+
+        self._client = Cortex(
+            api_url,
+            actor=conn.login or "user:default",
+            bearer=conn.password or None,
+            timeout=float(extra.get("timeout", 30.0)),
+        )
+        return self._client
+
+    def get_scope(self) -> str:
+        """Return the default scope path from the connection schema.
+
+        Returns:
+            The scope path, defaulting to ``"user:default"`` if not configured.
+        """
+        conn = self.get_connection(self.cortexdb_conn_id)
+        return conn.schema or "user:default"
+
+    def test_connection(self) -> tuple[bool, str]:
+        """Test the CortexDB connection.
+
+        Returns:
+            A tuple of (success, message).
+        """
+        try:
+            client = self.get_client()
+            who = client.whoami()
+            return True, f"Connection successful: {who}"
+        except Exception as exc:
+            return False, f"Connection failed: {exc}"
+
+    @classmethod
+    def get_ui_field_behaviour(cls) -> dict[str, Any]:
+        """Return custom UI field configuration for the Airflow connection form."""
+        return {
+            "relabeling": {
+                "login": "Actor",
+                "password": "Bearer Token",
+                "schema": "Default Scope",
+                "host": "CortexDB Host",
+                "port": "CortexDB Port",
+            },
+            "placeholders": {
+                "host": "localhost",
+                "port": "3141",
+                "login": "user:default",
+                "password": "v4.public... (optional)",
+                "schema": "user:default",
+                "extra": '{"timeout": 30, "scheme": "http"}',
+            },
+        }

cortexdb_airflow/operators.py

@@ -0,0 +1,142 @@
+"""Airflow Operator for CortexDB operations.
+
+Provides a single operator that can execute experience (write), recall,
+forget, and search operations against CortexDB within Airflow DAGs.
+
+Example::
+
+    from airflow.decorators import dag
+    from cortexdb_airflow import CortexDBOperator
+
+    @dag(schedule=None)
+    def memory_dag():
+        store = CortexDBOperator(
+            task_id="store_result",
+            operation="experience",
+            content="Pipeline completed successfully with 1,000 records.",
+            scope="org:data-team",
+        )
+
+        recall = CortexDBOperator(
+            task_id="recall_context",
+            operation="recall",
+            query="recent pipeline results",
+            scope="org:data-team",
+        )
+
+        store >> recall
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional, Sequence
+
+from airflow.models import BaseOperator
+from airflow.utils.context import Context
+
+from cortexdb_airflow.hooks import CortexDBHook
+
+
+class CortexDBOperator(BaseOperator):
+    """Airflow Operator for executing CortexDB operations.
+
+    Supports four operations: ``experience``, ``recall``, ``forget``, and
+    ``search``. The operation result is pushed to XCom for downstream
+    task consumption.
+
+    Args:
+        operation: The CortexDB operation to execute. One of
+            ``"experience"``, ``"recall"``, ``"forget"``, ``"search"``.
+        cortexdb_conn_id: Airflow connection ID. Defaults to
+            ``"cortexdb_default"``.
+        scope: Scope path. If ``None``, uses the connection's
+            default scope.
+        content: Content to store (for ``experience`` operation).
+        query: Query string (for ``recall``, ``search``).
+        labels: Optional labels for the stored experience.
+        reason: Reason for forgetting (for ``forget``).
+        view: Recall view (for ``recall``/``search``). Defaults to ``holistic``.
+    """
+
+    template_fields: Sequence[str] = (
+        "content",
+        "query",
+        "scope",
+        "reason",
+    )
+
+    def __init__(
+        self,
+        *,
+        operation: str,
+        cortexdb_conn_id: str = "cortexdb_default",
+        scope: Optional[str] = None,
+        content: Optional[str] = None,
+        query: Optional[str] = None,
+        labels: Optional[Sequence[str]] = None,
+        reason: str = "",
+        view: str = "holistic",
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+
+        valid_operations = ("experience", "recall", "forget", "search")
+        if operation not in valid_operations:
+            raise ValueError(
+                f"Invalid operation '{operation}'. "
+                f"Must be one of {valid_operations}."
+            )
+
+        self.operation = operation
+        self.cortexdb_conn_id = cortexdb_conn_id
+        self.scope = scope
+        self.content = content
+        self.query = query
+        self.labels = labels
+        self.reason = reason
+        self.view = view
+
+    def execute(self, context: Context) -> Any:
+        """Execute the configured CortexDB operation.
+
+        Args:
+            context: The Airflow task execution context.
+
+        Returns:
+            The result of the CortexDB operation, pushed to XCom.
+        """
+        hook = CortexDBHook(cortexdb_conn_id=self.cortexdb_conn_id)
+        client = hook.get_client()
+        scope = self.scope or hook.get_scope()
+
+        if self.operation == "experience":
+            if not self.content:
+                raise ValueError("'content' is required for the experience operation.")
+            result = client.experience(
+                scope,
+                text=self.content,
+                labels=self.labels,
+            )
+
+        elif self.operation in ("recall", "search"):
+            if not self.query:
+                raise ValueError(f"'query' is required for the {self.operation} operation.")
+            result = client.recall(
+                scope,
+                query=self.query,
+                view=self.view,
+            )
+
+        elif self.operation == "forget":
+            result = client.forget(
+                scope,
+                confirm_all=True,
+                cascade="redact_events",
+                reason=self.reason,
+            )
+
+        else:
+            raise ValueError(f"Unsupported operation: {self.operation}")
+
+        self.log.info("CortexDB %s operation completed.", self.operation)
+        return result

cortexdb_airflow-0.1.0.dist-info/METADATA

@@ -0,0 +1,13 @@
+Metadata-Version: 2.4
+Name: cortexdb-airflow
+Version: 0.1.0
+Summary: Apache Airflow provider for CortexDB — long-term memory for data orchestration
+License-Expression: Apache-2.0
+Requires-Python: >=3.10
+Requires-Dist: apache-airflow>=2.8
+Requires-Dist: cortexdbai>=0.1.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# cortexdb-airflow

cortexdb_airflow-0.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+cortexdb_airflow/__init__.py,sha256=rsp3oDKSx9ZIe5Mbh-xJb5dE5_zQDRpdwcqBFV_ppPA,880
+cortexdb_airflow/hooks.py,sha256=q3GJi0_o48Mmgr_PZGxon-9EtakVTMbpLFjJ_A0P79o,4101
+cortexdb_airflow/operators.py,sha256=m7zEZm8Nfg6l4c4p03eDfTUtw_d-i02xY5RDOssZGIo,4577
+cortexdb_airflow-0.1.0.dist-info/METADATA,sha256=qk9CC3JfgNEZAmbTQs1KqoTQoIknY1zWjBMht24cyIE,396
+cortexdb_airflow-0.1.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+cortexdb_airflow-0.1.0.dist-info/entry_points.txt,sha256=E5H0b6TawCb3ykL9hk-IbR1Cdiu6TfdxkmyDKY9H3qk,75
+cortexdb_airflow-0.1.0.dist-info/RECORD,,

cortexdb_airflow-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any

cortexdb_airflow-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[airflow.providers]
+cortexdb = cortexdb_airflow.__init__:get_provider_info