cortexdb-airflow 0.1.0__tar.gz

cortexdb_airflow-0.1.0/.gitignore

@@ -0,0 +1,79 @@
+# Rust
+/target
+**/*.rs.bk
+
+# Environment / secrets
+.env
+.env.local
+.env*.local
+*.pem
+*.key
+.npmrc
+
+# SQLite database
+*.sqlite
+*.sqlite-wal
+*.sqlite-shm
+
+# OS
+.DS_Store
+Thumbs.db
+desktop.ini
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# Data directories
+cortexdb_data*/
+/data/
+# Per-bench tenant stores (RocksDB + Tantivy + HNSW state; regeneratable per run)
+/data_*/
+# Experimental per-branch stores (not tracked on this branch but left gitignored
+# so checkout from other branches doesn't surface them in git status)
+/event_memory_store/
+/llm_cache/
+
+# Benchmark inputs and per-run outputs (kept local, regenerated each run)
+benchmarks/longmemeval/data/
+benchmarks/longmemeval/server_results/
+benchmarks/longmemeval/fast_results/
+benchmarks/longmemeval/micro_results/
+benchmarks/longmemeval/server_logs/
+benchmarks/longmemeval/*.log
+benchmarks/locomo/locomo_results*.json
+benchmarks/locomo/server_results/
+benchmarks/locomo/*.log
+/answer_out.json
+
+# Local Claude Code state
+.claude/
+.tmp/
+
+# Python
+__pycache__/
+*.pyc
+.venv/
+venv/
+
+# Node
+node_modules/
+dist/
+.next/
+
+# Egg info
+*.egg-info/
+
+# Scratch/debug text files at root
+/*.txt
+/*.log
+
+# Local debug / marketing / private content (not for repo)
+harness/.reports/
+harness_data_*/
+blog/
+sales/
+videos/
+local-instance/

cortexdb_airflow-0.1.0/PKG-INFO

@@ -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/README.md

@@ -0,0 +1 @@
+# cortexdb-airflow

cortexdb_airflow-0.1.0/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-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-0.1.0/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/pyproject.toml

@@ -0,0 +1,23 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "cortexdb-airflow"
+version = "0.1.0"
+description = "Apache Airflow provider for CortexDB — long-term memory for data orchestration"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "Apache-2.0"
+dependencies = [
+    "cortexdbai>=0.1.0",
+    "apache-airflow>=2.8",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0",
+]
+
+[project.entry-points."airflow.providers"]
+cortexdb = "cortexdb_airflow.__init__:get_provider_info"

cortexdb_airflow-0.1.0/tests/test_hooks.py

@@ -0,0 +1,148 @@
+"""Tests for CortexDB Airflow Hook."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from cortexdb_airflow.hooks import CortexDBHook
+
+
+class TestCortexDBHook:
+    """Tests for the CortexDBHook class."""
+
+    def test_default_conn_id(self) -> None:
+        """Hook should use the default connection ID."""
+        hook = CortexDBHook()
+        assert hook.cortexdb_conn_id == "cortexdb_default"
+
+    def test_custom_conn_id(self) -> None:
+        """Hook should accept a custom connection ID."""
+        hook = CortexDBHook(cortexdb_conn_id="my_cortex")
+        assert hook.cortexdb_conn_id == "my_cortex"
+
+    @patch("cortexdb_airflow.hooks.Cortex")
+    @patch.object(CortexDBHook, "get_connection")
+    def test_get_client(
+        self, mock_get_conn: MagicMock, mock_cortex_cls: MagicMock
+    ) -> None:
+        """get_client should create a Cortex client from Airflow connection."""
+        mock_conn = MagicMock()
+        mock_conn.host = "cortex.example.com"
+        mock_conn.port = 3141
+        mock_conn.login = "user:svc"
+        mock_conn.password = "v4.public.token"
+        mock_conn.extra = '{"timeout": 15}'
+        mock_conn.extra_dejson = {"timeout": 15}
+        mock_get_conn.return_value = mock_conn
+
+        hook = CortexDBHook()
+        client = hook.get_client()
+
+        mock_cortex_cls.assert_called_once_with(
+            "http://cortex.example.com:3141",
+            actor="user:svc",
+            bearer="v4.public.token",
+            timeout=15.0,
+        )
+
+    @patch("cortexdb_airflow.hooks.Cortex")
+    @patch.object(CortexDBHook, "get_connection")
+    def test_get_client_defaults(
+        self, mock_get_conn: MagicMock, mock_cortex_cls: MagicMock
+    ) -> None:
+        """get_client should use defaults when connection fields are empty."""
+        mock_conn = MagicMock()
+        mock_conn.host = None
+        mock_conn.port = None
+        mock_conn.login = None
+        mock_conn.password = None
+        mock_conn.extra = None
+        mock_conn.extra_dejson = {}
+        mock_get_conn.return_value = mock_conn
+
+        hook = CortexDBHook()
+        client = hook.get_client()
+
+        mock_cortex_cls.assert_called_once_with(
+            "http://localhost:3141",
+            actor="user:default",
+            bearer=None,
+            timeout=30.0,
+        )
+
+    @patch.object(CortexDBHook, "get_connection")
+    def test_get_scope(self, mock_get_conn: MagicMock) -> None:
+        """get_scope should return schema or default."""
+        mock_conn = MagicMock()
+        mock_conn.schema = "org:production"
+        mock_get_conn.return_value = mock_conn
+
+        hook = CortexDBHook()
+        assert hook.get_scope() == "org:production"
+
+    @patch.object(CortexDBHook, "get_connection")
+    def test_get_scope_default(self, mock_get_conn: MagicMock) -> None:
+        """get_scope should return 'user:default' when schema is empty."""
+        mock_conn = MagicMock()
+        mock_conn.schema = None
+        mock_get_conn.return_value = mock_conn
+
+        hook = CortexDBHook()
+        assert hook.get_scope() == "user:default"
+
+    @patch.object(CortexDBHook, "get_client")
+    def test_test_connection_success(
+        self, mock_get_client: MagicMock
+    ) -> None:
+        """test_connection should return success on healthy client."""
+        mock_client = MagicMock()
+        mock_client.whoami.return_value = {"actor": "user:svc"}
+        mock_get_client.return_value = mock_client
+
+        hook = CortexDBHook()
+        success, msg = hook.test_connection()
+
+        assert success is True
+        assert "successful" in msg
+
+    @patch.object(CortexDBHook, "get_client")
+    def test_test_connection_failure(
+        self, mock_get_client: MagicMock
+    ) -> None:
+        """test_connection should return failure on error."""
+        mock_get_client.side_effect = ConnectionError("refused")
+
+        hook = CortexDBHook()
+        success, msg = hook.test_connection()
+
+        assert success is False
+        assert "failed" in msg.lower()
+
+    def test_ui_field_behaviour(self) -> None:
+        """UI field behaviour should define expected fields."""
+        behaviour = CortexDBHook.get_ui_field_behaviour()
+        assert "relabeling" in behaviour
+        assert "placeholders" in behaviour
+
+    @patch("cortexdb_airflow.hooks.Cortex")
+    @patch.object(CortexDBHook, "get_connection")
+    def test_client_caching(
+        self, mock_get_conn: MagicMock, mock_cortex_cls: MagicMock
+    ) -> None:
+        """get_client should cache and reuse the client instance."""
+        mock_conn = MagicMock()
+        mock_conn.host = "localhost"
+        mock_conn.port = 3141
+        mock_conn.password = None
+        mock_conn.extra = None
+        mock_conn.extra_dejson = {}
+        mock_get_conn.return_value = mock_conn
+
+        hook = CortexDBHook()
+        client1 = hook.get_client()
+        client2 = hook.get_client()
+
+        assert client1 is client2
+        assert mock_cortex_cls.call_count == 1

cortexdb_airflow-0.1.0/tests/test_operators.py

@@ -0,0 +1,166 @@
+"""Tests for CortexDB Airflow Operator (v1 surface)."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from cortexdb_airflow.operators import CortexDBOperator
+
+
+@pytest.fixture
+def mock_hook():
+    """Create a mock CortexDBHook with a mock client."""
+    with patch("cortexdb_airflow.operators.CortexDBHook") as mock_hook_cls:
+        mock_client = MagicMock()
+        mock_client.experience.return_value = {"status": "ok"}
+        mock_client.recall.return_value = {"context_block": ""}
+        mock_client.forget.return_value = {"deleted": 0}
+
+        mock_hook_instance = MagicMock()
+        mock_hook_instance.get_client.return_value = mock_client
+        mock_hook_instance.get_scope.return_value = "user:default"
+        mock_hook_cls.return_value = mock_hook_instance
+
+        yield mock_client, mock_hook_instance
+
+
+class TestCortexDBOperator:
+    """Tests for the CortexDBOperator class."""
+
+    def test_invalid_operation(self) -> None:
+        """Operator should reject invalid operations."""
+        with pytest.raises(ValueError, match="Invalid operation"):
+            CortexDBOperator(
+                task_id="test",
+                operation="invalid",
+            )
+
+    def test_experience(self, mock_hook) -> None:
+        """Experience operation should call client.experience."""
+        mock_client, _ = mock_hook
+
+        op = CortexDBOperator(
+            task_id="test_experience",
+            operation="experience",
+            content="Test content",
+            scope="org:test",
+            labels=["a", "b"],
+        )
+        result = op.execute(context=MagicMock())
+
+        mock_client.experience.assert_called_once_with(
+            "org:test",
+            text="Test content",
+            labels=["a", "b"],
+        )
+        assert result == {"status": "ok"}
+
+    def test_experience_missing_content(self, mock_hook) -> None:
+        """Experience operation should raise if content is missing."""
+        op = CortexDBOperator(
+            task_id="test_experience",
+            operation="experience",
+        )
+        with pytest.raises(ValueError, match="content"):
+            op.execute(context=MagicMock())
+
+    def test_recall(self, mock_hook) -> None:
+        """Recall operation should call client.recall."""
+        mock_client, _ = mock_hook
+
+        op = CortexDBOperator(
+            task_id="test_recall",
+            operation="recall",
+            query="What happened?",
+            scope="org:test",
+        )
+        op.execute(context=MagicMock())
+
+        mock_client.recall.assert_called_once_with(
+            "org:test",
+            query="What happened?",
+            view="holistic",
+        )
+
+    def test_recall_missing_query(self, mock_hook) -> None:
+        """Recall operation should raise if query is missing."""
+        op = CortexDBOperator(
+            task_id="test_recall",
+            operation="recall",
+        )
+        with pytest.raises(ValueError, match="query"):
+            op.execute(context=MagicMock())
+
+    def test_forget(self, mock_hook) -> None:
+        """Forget operation should call client.forget."""
+        mock_client, _ = mock_hook
+
+        op = CortexDBOperator(
+            task_id="test_forget",
+            operation="forget",
+            reason="cleanup",
+            scope="org:test",
+        )
+        op.execute(context=MagicMock())
+
+        mock_client.forget.assert_called_once_with(
+            "org:test",
+            confirm_all=True,
+            cascade="redact_events",
+            reason="cleanup",
+        )
+
+    def test_search(self, mock_hook) -> None:
+        """Search operation should map to client.recall."""
+        mock_client, _ = mock_hook
+
+        op = CortexDBOperator(
+            task_id="test_search",
+            operation="search",
+            query="pipeline events",
+            scope="org:test",
+        )
+        op.execute(context=MagicMock())
+
+        mock_client.recall.assert_called_once_with(
+            "org:test",
+            query="pipeline events",
+            view="holistic",
+        )
+
+    def test_search_missing_query(self, mock_hook) -> None:
+        """Search operation should raise if query is missing."""
+        op = CortexDBOperator(
+            task_id="test_search",
+            operation="search",
+        )
+        with pytest.raises(ValueError, match="query"):
+            op.execute(context=MagicMock())
+
+    def test_scope_fallback_to_hook(self, mock_hook) -> None:
+        """Operator should use hook's scope when scope is not set."""
+        mock_client, mock_hook_instance = mock_hook
+        mock_hook_instance.get_scope.return_value = "user:hook"
+
+        op = CortexDBOperator(
+            task_id="test_fallback",
+            operation="experience",
+            content="test",
+        )
+        op.execute(context=MagicMock())
+
+        mock_client.experience.assert_called_once()
+        assert mock_client.experience.call_args.args[0] == "user:hook"
+
+    def test_template_fields(self) -> None:
+        """Operator should expose expected template fields."""
+        assert "content" in CortexDBOperator.template_fields
+        assert "query" in CortexDBOperator.template_fields
+        assert "scope" in CortexDBOperator.template_fields
+
+
+if __name__ == "__main__":
+    import sys
+    sys.exit(pytest.main([__file__]))