@pentatonic-ai/ai-agent-sdk 0.8.0 → 0.8.2

package/packages/memory-engine/tests/test_channel_stat_rollups.py

@@ -0,0 +1,308 @@
+"""Integration tests for ChannelStat rollup writes + reader fast-path.
+
+These exercise the actual Cypher running against a live Neo4j: the
+writer block in /index-internal-batch maintains (:ChannelStat) nodes
+on every store with contact_email metadata, and /aggregate-internal
+reads from those nodes when present (falling back to the edge walk
+when stats haven't been populated yet).
+
+Gated on NEO4J_TEST_URI + NEO4J_TEST_PASSWORD; skip cleanly when
+those env vars are absent so unit-only test runs stay fast.
+
+Run:
+
+    cd packages/memory-engine
+    NEO4J_TEST_URI=bolt://localhost:17687 \\
+    NEO4J_TEST_PASSWORD=testpassword \\
+    .venv/bin/python -m pytest tests/test_channel_stat_rollups.py -v
+"""
+from __future__ import annotations
+
+import os
+import uuid
+
+import pytest
+
+
+_NEO4J_URI = os.environ.get("NEO4J_TEST_URI")
+_NEO4J_USER = os.environ.get("NEO4J_TEST_USER", "neo4j")
+_NEO4J_PASSWORD = os.environ.get("NEO4J_TEST_PASSWORD")
+
+_skip_no_neo4j = pytest.mark.skipif(
+    not (_NEO4J_URI and _NEO4J_PASSWORD),
+    reason="set NEO4J_TEST_URI + NEO4J_TEST_PASSWORD to run integration tests",
+)
+
+
+@pytest.fixture
+def neo4j_driver():
+    """Per-test driver + cleanup. Two arenas so isolation tests can
+    run side by side without trampling each other."""
+    from neo4j import GraphDatabase
+
+    driver = GraphDatabase.driver(_NEO4J_URI, auth=(_NEO4J_USER, _NEO4J_PASSWORD))
+    arenas = [f"stat_a_{uuid.uuid4().hex[:8]}", f"stat_b_{uuid.uuid4().hex[:8]}"]
+    yield driver, arenas
+    with driver.session() as session:
+        for arena in arenas:
+            # ChannelStat lives alongside Person/Chunk in the arena;
+            # the per-arena DETACH DELETE catches all of them.
+            session.run(
+                "MATCH (n) WHERE n.arena = $arena DETACH DELETE n",
+                arena=arena,
+            )
+    driver.close()
+
+
+def _store_chunk_with_email(
+    session,
+    arena: str,
+    cid: str,
+    email: str,
+    channel: str,
+    direction: str,
+    occurred_at: str,
+) -> None:
+    """Mirror exactly the Cypher /index-internal-batch runs for the
+    email-keyed Person path, including the ChannelStat rollup. The
+    real call sites this is testing are inside that endpoint; this
+    helper lets the integration test exercise the same Cypher
+    without spinning up the FastAPI app."""
+    now_iso = "2026-05-10T00:00:00Z"
+    # Ensure the chunk exists (matches what /index-internal-batch
+    # does immediately before writing the Person edge).
+    session.run(
+        """
+        MERGE (c:Chunk {id: $cid})
+        SET c.text = 't', c.path = 'p', c.arena = $arena,
+            c.created_at = $now
+        """,
+        cid=cid, arena=arena, now=now_iso,
+    )
+    session.run(
+        """
+        MERGE (p:Entity:Person {arena: $arena, email: $email})
+        ON CREATE SET p.created_at = $now,
+                      p.first_seen = $occurred_at,
+                      p.last_seen = $occurred_at
+        ON MATCH SET p.last_seen = CASE
+              WHEN $occurred_at > coalesce(p.last_seen, '')
+                THEN $occurred_at
+              ELSE p.last_seen END
+        WITH p
+        MATCH (c:Chunk {arena: $arena, id: $cid})
+        MERGE (p)-[r:COMMUNICATED]->(c)
+        ON CREATE SET r.channel = $channel,
+                      r.direction = $direction,
+                      r.occurred_at = $occurred_at,
+                      r.weight = 1.0,
+                      r._counted = false
+        WITH p, r
+        FOREACH (_ IN CASE WHEN r._counted = false THEN [1] ELSE [] END |
+          MERGE (s:ChannelStat {arena: $arena, person_email: $email, channel: $channel})
+          ON CREATE SET s.count = 0,
+                        s.inbound = 0,
+                        s.outbound = 0,
+                        s.first_seen = $occurred_at,
+                        s.last_seen = $occurred_at,
+                        s.created_at = $now
+          SET s.count = s.count + 1,
+              s.inbound = s.inbound + (CASE WHEN $direction = 'inbound' THEN 1 ELSE 0 END),
+              s.outbound = s.outbound + (CASE WHEN $direction = 'outbound' THEN 1 ELSE 0 END),
+              s.first_seen = CASE
+                WHEN $occurred_at < coalesce(s.first_seen, $occurred_at)
+                  THEN $occurred_at
+                ELSE s.first_seen END,
+              s.last_seen = CASE
+                WHEN $occurred_at > coalesce(s.last_seen, '')
+                  THEN $occurred_at
+                ELSE s.last_seen END,
+              s.updated_at = $now
+          MERGE (p)-[:HAS_STAT]->(s)
+          SET r._counted = true
+        )
+        """,
+        arena=arena, email=email, cid=cid,
+        channel=channel, direction=direction,
+        occurred_at=occurred_at, now=now_iso,
+    )
+
+
+def _read_stats(session, arena: str, email: str) -> list[dict]:
+    return [
+        dict(rec)
+        for rec in session.run(
+            "MATCH (s:ChannelStat {arena: $arena, person_email: $email})\n"
+            "RETURN s.channel AS channel, s.count AS count,\n"
+            "       s.inbound AS inbound, s.outbound AS outbound,\n"
+            "       s.first_seen AS first_seen, s.last_seen AS last_seen\n"
+            "ORDER BY s.channel\n",
+            arena=arena, email=email,
+        )
+    ]
+
+
+# ---------------------------------------------------------------------------
+# Writer behaviour.
+# ---------------------------------------------------------------------------
+
+
+@_skip_no_neo4j
+def test_first_store_creates_channelstat_with_count_one(neo4j_driver) -> None:
+    driver, (arena, _) = neo4j_driver
+    with driver.session() as session:
+        _store_chunk_with_email(
+            session, arena, cid="c1", email="alex@x.io",
+            channel="email", direction="inbound",
+            occurred_at="2026-05-09T10:00:00Z",
+        )
+        stats = _read_stats(session, arena, "alex@x.io")
+        assert len(stats) == 1
+        assert stats[0]["channel"] == "email"
+        assert stats[0]["count"] == 1
+        assert stats[0]["inbound"] == 1
+        assert stats[0]["outbound"] == 0
+        assert stats[0]["first_seen"] == "2026-05-09T10:00:00Z"
+        assert stats[0]["last_seen"] == "2026-05-09T10:00:00Z"
+
+
+@_skip_no_neo4j
+def test_repeated_stores_accumulate_counts_and_widen_time_bounds(neo4j_driver) -> None:
+    driver, (arena, _) = neo4j_driver
+    with driver.session() as session:
+        _store_chunk_with_email(
+            session, arena, cid="c1", email="alex@x.io",
+            channel="email", direction="inbound",
+            occurred_at="2026-05-08T09:00:00Z",
+        )
+        _store_chunk_with_email(
+            session, arena, cid="c2", email="alex@x.io",
+            channel="email", direction="outbound",
+            occurred_at="2026-05-09T10:00:00Z",
+        )
+        _store_chunk_with_email(
+            session, arena, cid="c3", email="alex@x.io",
+            channel="email", direction="inbound",
+            occurred_at="2026-05-07T08:00:00Z",  # earliest — should drive first_seen
+        )
+        stats = _read_stats(session, arena, "alex@x.io")
+        assert len(stats) == 1
+        assert stats[0]["count"] == 3
+        assert stats[0]["inbound"] == 2
+        assert stats[0]["outbound"] == 1
+        assert stats[0]["first_seen"] == "2026-05-07T08:00:00Z"
+        assert stats[0]["last_seen"] == "2026-05-09T10:00:00Z"
+
+
+@_skip_no_neo4j
+def test_per_channel_buckets_kept_separate(neo4j_driver) -> None:
+    driver, (arena, _) = neo4j_driver
+    with driver.session() as session:
+        _store_chunk_with_email(
+            session, arena, "c1", "alex@x.io",
+            "email", "inbound", "2026-05-09T10:00:00Z",
+        )
+        _store_chunk_with_email(
+            session, arena, "c2", "alex@x.io",
+            "slack", "outbound", "2026-05-09T11:00:00Z",
+        )
+        _store_chunk_with_email(
+            session, arena, "c3", "alex@x.io",
+            "slack", "inbound", "2026-05-09T12:00:00Z",
+        )
+        stats = sorted(_read_stats(session, arena, "alex@x.io"), key=lambda s: s["channel"])
+        assert [s["channel"] for s in stats] == ["email", "slack"]
+        email = next(s for s in stats if s["channel"] == "email")
+        slack = next(s for s in stats if s["channel"] == "slack")
+        assert email["count"] == 1
+        assert slack["count"] == 2
+        assert slack["inbound"] == 1
+        assert slack["outbound"] == 1
+
+
+@_skip_no_neo4j
+def test_replay_of_same_chunk_does_not_double_count(neo4j_driver) -> None:
+    """Idempotency under retry: storing the same (chunk_id, email)
+    twice should not double the count. Guarded by the r._counted
+    flag set on first edge creation."""
+    driver, (arena, _) = neo4j_driver
+    with driver.session() as session:
+        for _ in range(3):
+            _store_chunk_with_email(
+                session, arena, "c1", "alex@x.io",
+                "email", "inbound", "2026-05-09T10:00:00Z",
+            )
+        stats = _read_stats(session, arena, "alex@x.io")
+        assert len(stats) == 1
+        # Three identical stores → still count == 1.
+        assert stats[0]["count"] == 1
+        assert stats[0]["inbound"] == 1
+
+
+@_skip_no_neo4j
+def test_arena_isolation_on_channelstat(neo4j_driver) -> None:
+    """A's stats never reflect B's writes, even when both arenas
+    have the same email — the multi-tenancy invariant that #28
+    landed for Person/Chunk extends to ChannelStat."""
+    driver, (arena_a, arena_b) = neo4j_driver
+    with driver.session() as session:
+        _store_chunk_with_email(
+            session, arena_a, "c1", "shared@example.com",
+            "email", "inbound", "2026-05-09T10:00:00Z",
+        )
+        _store_chunk_with_email(
+            session, arena_b, "c2", "shared@example.com",
+            "email", "inbound", "2026-05-09T11:00:00Z",
+        )
+        _store_chunk_with_email(
+            session, arena_b, "c3", "shared@example.com",
+            "slack", "outbound", "2026-05-09T12:00:00Z",
+        )
+        a_stats = _read_stats(session, arena_a, "shared@example.com")
+        b_stats = _read_stats(session, arena_b, "shared@example.com")
+        # A only saw the email message.
+        assert len(a_stats) == 1
+        assert a_stats[0]["channel"] == "email"
+        assert a_stats[0]["count"] == 1
+        # B saw both.
+        assert len(b_stats) == 2
+
+
+@_skip_no_neo4j
+def test_communicated_edge_has_counted_flag_after_write(neo4j_driver) -> None:
+    """The idempotency flag must be true after the first store so
+    re-runs of the writer (which use the same MERGE pattern) skip
+    the increment subquery."""
+    driver, (arena, _) = neo4j_driver
+    with driver.session() as session:
+        _store_chunk_with_email(
+            session, arena, "c1", "alex@x.io",
+            "email", "inbound", "2026-05-09T10:00:00Z",
+        )
+        rec = session.run(
+            "MATCH (:Person {arena: $arena, email: $email})-[r:COMMUNICATED]->(:Chunk {arena: $arena, id: $cid})\n"
+            "RETURN r._counted AS counted",
+            arena=arena, email="alex@x.io", cid="c1",
+        ).single()
+        assert rec["counted"] is True
+
+
+@_skip_no_neo4j
+def test_has_stat_edge_links_person_to_channelstat(neo4j_driver) -> None:
+    """The Person→ChannelStat traversal exists for ops queries
+    ('show me all stat rollups for this person') and for cascade
+    deletes when a tenant is offboarded."""
+    driver, (arena, _) = neo4j_driver
+    with driver.session() as session:
+        _store_chunk_with_email(
+            session, arena, "c1", "alex@x.io",
+            "email", "inbound", "2026-05-09T10:00:00Z",
+        )
+        rows = list(session.run(
+            "MATCH (p:Person {arena: $arena, email: $email})-[:HAS_STAT]->(s:ChannelStat {arena: $arena})\n"
+            "RETURN s.channel AS channel, s.count AS count\n",
+            arena=arena, email="alex@x.io",
+        ))
+        assert len(rows) == 1
+        assert rows[0]["channel"] == "email"
+        assert rows[0]["count"] == 1

package/packages/memory-engine/tests/test_compat_nv_embed_probe.py

@@ -0,0 +1,48 @@
+"""Regression test for the nv_embed /health probe URL construction.
+
+Before v0.8.1 the probe was built via str.replace("/v1/embeddings", "/health"),
+which silently produced the wrong URL when NV_EMBED_URL was anything other
+than ".../v1/embeddings" (e.g. /v1/embed for the Pentatonic AI Gateway, or
+a bare host). The probe would end up GET-ing the embed endpoint itself and
+the gateway would return 401 — making /health falsely report a degraded
+nv_embed layer even when embeddings worked fine.
+
+The fix replaces the substring replace with a urlparse-based path swap.
+"""
+
+from urllib.parse import urlparse, urlunparse
+
+
+def probe_url(nv_embed_url: str) -> str:
+    """Re-implementation of the probe URL builder from compat/server.py.
+    Kept tiny + isolated so it can be unit-tested without booting FastAPI."""
+    u = urlparse(nv_embed_url)
+    return urlunparse((u.scheme, u.netloc, "/health", "", "", ""))
+
+
+def test_probe_with_v1_embeddings_path():
+    assert probe_url("https://nv-embed:8041/v1/embeddings") == "https://nv-embed:8041/health"
+
+
+def test_probe_with_v1_embed_path():
+    """The bug case — gateway moved to /v1/embed and old replace() left URL untouched."""
+    assert probe_url("https://lambda-gateway.pentatonic.com/v1/embed") == (
+        "https://lambda-gateway.pentatonic.com/health"
+    )
+
+
+def test_probe_with_bare_host():
+    """No path at all — provider supplies path_default, but probe just wants /health."""
+    assert probe_url("https://lambda-gateway.pentatonic.com") == (
+        "https://lambda-gateway.pentatonic.com/health"
+    )
+
+
+def test_probe_strips_query_and_fragment():
+    assert probe_url("https://gw.example.com/v1/embeddings?token=x#frag") == (
+        "https://gw.example.com/health"
+    )
+
+
+def test_probe_preserves_port():
+    assert probe_url("http://nv-embed:8041/v1/embeddings") == "http://nv-embed:8041/health"