agentic-composite-memory 0.1.0__py3-none-any.whl

agentic_memory_sdk/adapters/neptune_cypher.py

@@ -0,0 +1,309 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from datetime import datetime, timezone
+from typing import Any, Sequence
+
+from ..auth import AwsIamDatabaseAuthProvider
+from ..models import (
+    GraphEdgeMutation,
+    GraphFactReadResult,
+    GraphVertexUpsert,
+    PlatformTenantContext,
+    ReasoningTraceStep,
+)
+from ..ports import GraphStoreAdapter
+
+
+class NeptuneOpenCypherAdapter(GraphStoreAdapter):
+    """
+    Production Amazon Neptune adapter using parameterized openCypher
+    with mandatory tenant/namespace guards, idempotent batch MERGE,
+    bi-temporal functional edge invalidation, and reasoning trace recording.
+    Supports AWS IAM SigV4 and cross-account STS AssumeRole via `aws_auth_provider`.
+    """
+
+    def __init__(
+        self,
+        endpoint_by_tenant: dict[str, str],
+        region: str = "us-east-1",
+        aws_auth_provider: AwsIamDatabaseAuthProvider | None = None,
+    ):
+        self.endpoints = endpoint_by_tenant
+        self.region = region
+        self.aws_auth_provider = aws_auth_provider
+        self._clients: dict[str, Any] = {}
+
+    def _client(self, tenant_id: str):
+        import boto3
+
+        ep = self.endpoints.get(tenant_id) or self.endpoints["default"]
+        if ep not in self._clients:
+            endpoint_url = ep if ep.startswith("http") else f"https://{ep}:8182"
+            session = (
+                self.aws_auth_provider.get_boto3_session()
+                if self.aws_auth_provider is not None
+                else boto3.Session(region_name=self.region)
+            )
+            self._clients[ep] = session.client(
+                "neptunedata",
+                region_name=self.region,
+                endpoint_url=endpoint_url,
+            )
+        return self._clients[ep]
+
+    async def _run_cypher(
+        self, tenant_id: str, query: str, params: dict[str, Any]
+    ) -> list[dict[str, Any]]:
+        client = self._client(tenant_id)
+        resp = await asyncio.to_thread(
+            client.execute_open_cypher_query,
+            openCypherQuery=query,
+            parameters=json.dumps(params),
+        )
+        return resp.get("results", [])
+
+    async def upsert_vertices_and_bitemporal_edges(
+        self,
+        tenant_id: str,
+        vertices: Sequence[GraphVertexUpsert],
+        edges: Sequence[GraphEdgeMutation],
+    ) -> None:
+        if vertices:
+            vertex_dicts = [v.model_dump(exclude={"embedding"}) for v in vertices]
+            await self._run_cypher(
+                tenant_id,
+                """
+                UNWIND $vertices AS v
+                MERGE (n:Entity {`~id`: v.vertex_id, tenant_id: v.tenant_id})
+                ON CREATE SET n.canonical_name = v.canonical_name,
+                              n.entity_class = v.entity_class,
+                              n.upper_class = v.upper_class,
+                              n.summary = v.summary,
+                              n.namespace = v.namespace,
+                              n.updated_at = v.updated_at
+                ON MATCH SET  n.summary = v.summary,
+                              n.updated_at = v.updated_at
+                """,
+                {"vertices": vertex_dicts},
+            )
+
+        for e in edges:
+            if e.is_functional:
+                # Close t_invalid on any prior active functional edge with a different target
+                await self._run_cypher(
+                    tenant_id,
+                    """
+                    MATCH (s:Entity {`~id`: $src_id, tenant_id: $tenant_id})
+                          -[r:RELATES {predicate: $pred, namespace: $ns, t_invalid: 'ACTIVE'}]->
+                          (old_t:Entity)
+                    WHERE old_t.`~id` <> $tgt_id
+                    SET r.t_invalid = $t_valid
+                    """,
+                    {
+                        "src_id": e.source_vertex_id,
+                        "tenant_id": tenant_id,
+                        "pred": e.predicate,
+                        "ns": e.namespace,
+                        "tgt_id": e.target_vertex_id,
+                        "t_valid": e.t_valid,
+                    },
+                )
+
+            await self._run_cypher(
+                tenant_id,
+                """
+                MATCH (s:Entity {`~id`: $src_id, tenant_id: $tenant_id})
+                MATCH (t:Entity {`~id`: $tgt_id, tenant_id: $tenant_id})
+                MERGE (s)-[r:RELATES {
+                    predicate: $pred,
+                    namespace: $ns,
+                    tenant_id: $tenant_id,
+                    t_invalid: 'ACTIVE'
+                }]->(t)
+                ON CREATE SET r.edge_id = $edge_id,
+                              r.t_valid = $t_valid,
+                              r.t_created = $t_created,
+                              r.fact_summary = $summary,
+                              r.confidence = $conf,
+                              r.reinforcement_count = 1,
+                              r.source_episode_id = $ep_id,
+                              r.source_app = $app,
+                              r.user_id = $uid
+                ON MATCH SET  r.fact_summary = $summary,
+                              r.reinforcement_count = coalesce(r.reinforcement_count, 1) + 1,
+                              r.confidence = CASE
+                                  WHEN coalesce(r.confidence, 0.9) + 0.03 > 1.0 THEN 1.0
+                                  ELSE coalesce(r.confidence, 0.9) + 0.03
+                              END,
+                              r.updated_at = $t_valid
+                """,
+                {
+                    "src_id": e.source_vertex_id,
+                    "tgt_id": e.target_vertex_id,
+                    "tenant_id": tenant_id,
+                    "pred": e.predicate,
+                    "ns": e.namespace,
+                    "edge_id": e.edge_id,
+                    "t_valid": e.t_valid,
+                    "t_created": e.t_created,
+                    "summary": e.fact_summary,
+                    "conf": e.confidence,
+                    "ep_id": e.source_episode_id,
+                    "app": e.source_app,
+                    "uid": e.user_id,
+                },
+            )
+
+    async def expand_active_subgraph(
+        self,
+        ctx: PlatformTenantContext,
+        seed_vertex_ids: Sequence[str],
+        as_of_iso: str | None = None,
+        max_hops: int = 2,
+        limit: int = 25,
+    ) -> list[GraphFactReadResult]:
+        if not seed_vertex_ids:
+            return []
+
+        if as_of_iso is None:
+            temporal_filter = "r.t_invalid = 'ACTIVE'"
+        else:
+            temporal_filter = (
+                "r.t_valid <= $as_of AND (r.t_invalid = 'ACTIVE' OR r.t_invalid > $as_of)"
+            )
+
+        frontier = list(dict.fromkeys(seed_vertex_ids))
+        visited_vertices = set(frontier)
+        seen_edges: set[tuple[str, str, str, str, str]] = set()
+        out: list[GraphFactReadResult] = []
+
+        for _hop in range(max(1, max_hops)):
+            if not frontier or len(out) >= limit:
+                break
+
+            results = await self._run_cypher(
+                ctx.tenant_id,
+                f"""
+                MATCH (s:Entity)-[r:RELATES]->(t:Entity)
+                WHERE (s.`~id` IN $seed_ids OR t.`~id` IN $seed_ids)
+                  AND s.tenant_id = $tenant_id
+                  AND t.tenant_id = $tenant_id
+                  AND s.namespace IN $read_ns
+                  AND t.namespace IN $read_ns
+                  AND r.namespace IN $read_ns
+                  AND {temporal_filter}
+                RETURN s.`~id` AS src_id,
+                       s.canonical_name AS src,
+                       coalesce(s.entity_class, 'p:Concept') AS src_cls,
+                       r.predicate AS pred,
+                       t.`~id` AS tgt_id,
+                       t.canonical_name AS tgt,
+                       coalesce(t.entity_class, 'p:Concept') AS tgt_cls,
+                       r.fact_summary AS summary,
+                       r.t_valid AS t_valid,
+                       r.t_invalid AS t_invalid,
+                       r.namespace AS ns,
+                       coalesce(r.source_episode_id, '') AS ep_id,
+                       coalesce(r.confidence, 1.0) AS conf
+                LIMIT $limit
+                """,
+                {
+                    "seed_ids": frontier,
+                    "tenant_id": ctx.tenant_id,
+                    "read_ns": ctx.read_namespaces,
+                    "as_of": as_of_iso or "",
+                    "limit": limit,
+                },
+            )
+            next_frontier: list[str] = []
+            for r in results:
+                edge_key = (
+                    r["src_id"],
+                    r["pred"],
+                    r["tgt_id"],
+                    r["ns"],
+                    r["t_valid"],
+                )
+                if edge_key in seen_edges:
+                    continue
+                seen_edges.add(edge_key)
+                out.append(
+                    GraphFactReadResult(
+                        source_name=r["src"],
+                        source_class=r.get("src_cls", "p:Concept"),
+                        predicate=r["pred"],
+                        target_name=r["tgt"],
+                        target_class=r.get("tgt_cls", "p:Concept"),
+                        fact_summary=r["summary"],
+                        t_valid=r["t_valid"],
+                        t_invalid=r["t_invalid"],
+                        namespace=r["ns"],
+                        source_episode_id=r.get("ep_id", ""),
+                        confidence=float(r["conf"]),
+                    )
+                )
+                for vid in (r["src_id"], r["tgt_id"]):
+                    if vid not in visited_vertices:
+                        visited_vertices.add(vid)
+                        next_frontier.append(vid)
+                if len(out) >= limit:
+                    break
+            frontier = next_frontier
+
+        return out[:limit]
+
+    async def record_reasoning_trace(
+        self,
+        ctx: PlatformTenantContext,
+        run_id: str,
+        session_id: str,
+        steps: Sequence[ReasoningTraceStep],
+    ) -> None:
+        step_dicts = [s.model_dump() for s in steps]
+        await self._run_cypher(
+            ctx.tenant_id,
+            """
+            MERGE (run:AgentRun {`~id`: $run_id, tenant_id: $tenant_id})
+            ON CREATE SET run.app_name = $app,
+                          run.session_id = $session_id,
+                          run.namespace = $ns,
+                          run.created_at = $now
+            WITH run
+            UNWIND $steps AS st
+            MERGE (tc:ToolCall {`~id`: $run_id + '_' + toString(st.step_index), tenant_id: $tenant_id})
+            SET tc.tool_name = st.tool_name,
+                tc.input_summary = st.tool_input_summary,
+                tc.output_summary = st.tool_output_summary,
+                tc.status = st.status,
+                tc.namespace = $ns
+            MERGE (run)-[:EXECUTED_STEP {step_index: st.step_index}]->(tc)
+            WITH tc, st
+            UNWIND st.touched_vertex_ids AS vid
+            MATCH (e:Entity {`~id`: vid, tenant_id: $tenant_id})
+            MERGE (tc)-[:TOUCHED_ENTITY]->(e)
+            """,
+            {
+                "run_id": run_id,
+                "tenant_id": ctx.tenant_id,
+                "app": ctx.app_name,
+                "session_id": session_id,
+                "ns": ctx.app_private_namespace,
+                "now": datetime.now(timezone.utc).isoformat(),
+                "steps": step_dicts,
+            },
+        )
+
+    async def purge_episode_lineage(
+        self, tenant_id: str, episode_ids: Sequence[str], user_id: str
+    ) -> None:
+        await self._run_cypher(
+            tenant_id,
+            """
+            MATCH ()-[r:RELATES {tenant_id: $tenant_id}]->()
+            WHERE r.source_episode_id IN $ep_ids OR r.user_id = $uid
+            DELETE r
+            """,
+            {"tenant_id": tenant_id, "ep_ids": list(episode_ids), "uid": user_id},
+        )

agentic_memory_sdk/adapters/opensearch_seed.py

@@ -0,0 +1,322 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Any, Sequence
+
+from ..governance import EntityResolutionGuard
+from ..models import (
+    EpisodicRecord,
+    GraphVertexUpsert,
+    OutboxOperation,
+    PlatformTenantContext,
+)
+from ..ports import EpisodicAndSeedStoreAdapter
+
+
+class OpenSearchSeedAdapter(EpisodicAndSeedStoreAdapter):
+    """
+    Swappable Amazon OpenSearch adapter for Tier 1 Episodic Search and Tier 2 Hybrid
+    (BM25 + k-NN HNSW) Entry-Point Seeding & Entity Resolution Blocking.
+    Enforces tenant_id and namespace isolation filters on every query and write.
+    Can be paired with PostgreSQL (RegistryStoreAdapter) and Amazon Neptune (GraphStoreAdapter).
+    """
+
+    def __init__(
+        self,
+        hosts: list[dict[str, Any]] | list[str],
+        http_auth: tuple[str, str] | Any = None,
+        seed_index_name: str = "platform_graph_seeds",
+        episodic_index_name: str = "platform_episodic_memories",
+        outbox_index_name: str = "platform_graph_outbox",
+        use_ssl: bool = True,
+    ):
+        self.hosts = hosts
+        self.http_auth = http_auth
+        self.seed_index = seed_index_name
+        self.episodic_index = episodic_index_name
+        self.outbox_index = outbox_index_name
+        self.use_ssl = use_ssl
+        self._client: Any = None
+
+    def _get_client(self):
+        if self._client is None:
+            from opensearchpy import AsyncOpenSearch
+
+            self._client = AsyncOpenSearch(
+                hosts=self.hosts,
+                http_auth=self.http_auth,
+                use_ssl=self.use_ssl,
+                verify_certs=self.use_ssl,
+            )
+        return self._client
+
+    async def close(self) -> None:
+        if self._client is not None:
+            await self._client.close()
+            self._client = None
+
+    async def resolve_existing_vertex_id(
+        self,
+        ctx: PlatformTenantContext,
+        entity_class: str,
+        canonical_name: str,
+        embedding: list[float],
+        similarity_threshold: float = 0.91,
+    ) -> str | None:
+        client = self._get_client()
+        query_body = {
+            "size": 1,
+            "query": {
+                "bool": {
+                    "filter": [
+                        {"term": {"tenant_id": ctx.tenant_id}},
+                        {"terms": {"namespace": ctx.read_namespaces}},
+                        {"term": {"entity_class": entity_class}},
+                    ],
+                    "should": [
+                        {"term": {"canonical_name.keyword": canonical_name}},
+                        {
+                            "knn": {
+                                "embedding": {
+                                    "vector": embedding,
+                                    "k": 3,
+                                }
+                            }
+                        },
+                    ],
+                    "minimum_should_match": 1,
+                }
+            },
+        }
+        resp = await client.search(index=self.seed_index, body=query_body)
+        hits = resp.get("hits", {}).get("hits", [])
+        for hit in hits:
+            src = hit["_source"]
+            cand_name = src.get("canonical_name", "")
+            score = float(hit.get("_score", 0.0))
+            if EntityResolutionGuard.can_merge(
+                cand_name, canonical_name, score, similarity_threshold
+            ):
+                return src["vertex_id"]
+        return None
+
+    async def commit_episodes_and_outbox(
+        self,
+        ctx: PlatformTenantContext,
+        episodes: Sequence[EpisodicRecord],
+        vertex_upserts: Sequence[GraphVertexUpsert],
+        outbox_op: OutboxOperation | None,
+    ) -> None:
+        client = self._get_client()
+        allowed_write_ns = {ctx.write_namespace, ctx.app_private_namespace}
+
+        bulk_ops: list[dict[str, Any]] = []
+        for ep in episodes:
+            if ep.tenant_id != ctx.tenant_id or ep.app_name != ctx.app_name:
+                raise PermissionError("OpenSearch episodic write tenant/app isolation error")
+            bulk_ops.append({"index": {"_index": self.episodic_index, "_id": ep.episode_id}})
+            bulk_ops.append(
+                {
+                    "episode_id": ep.episode_id,
+                    "tenant_id": ep.tenant_id,
+                    "app_name": ep.app_name,
+                    "user_id": ep.user_id,
+                    "session_id": ep.session_id,
+                    "namespace": ep.namespace,
+                    "author": ep.author,
+                    "content_text": ep.content_text,
+                    "embedding": ep.embedding,
+                    "created_at": ep.created_at.isoformat(),
+                    "metadata": ep.metadata,
+                }
+            )
+
+        for v in vertex_upserts:
+            if v.tenant_id != ctx.tenant_id or v.namespace not in allowed_write_ns:
+                raise PermissionError(
+                    f"OpenSearch seed write namespace error: {v.namespace}"
+                )
+            doc_id = f"{v.tenant_id}:{v.vertex_id}"
+            bulk_ops.append({"index": {"_index": self.seed_index, "_id": doc_id}})
+            bulk_ops.append(v.model_dump())
+
+        if outbox_op is not None:
+            bulk_ops.append(
+                {"index": {"_index": self.outbox_index, "_id": outbox_op.outbox_id}}
+            )
+            bulk_ops.append(
+                {
+                    **outbox_op.model_dump(),
+                    "status": "PENDING",
+                    "created_at": datetime.now(timezone.utc).isoformat(),
+                }
+            )
+
+        if bulk_ops:
+            await client.bulk(body=bulk_ops, refresh="wait_for")
+
+    async def hybrid_seed_search(
+        self,
+        ctx: PlatformTenantContext,
+        query_text: str,
+        query_embedding: list[float],
+        limit: int = 8,
+    ) -> list[GraphVertexUpsert]:
+        client = self._get_client()
+        security_filter = [
+            {"term": {"tenant_id": ctx.tenant_id}},
+            {"terms": {"namespace": ctx.read_namespaces}},
+        ]
+        # Execute BM25 lexical search + k-NN vector search and fuse with RRF
+        body = {
+            "size": limit,
+            "query": {
+                "bool": {
+                    "filter": security_filter,
+                    "should": [
+                        {
+                            "multi_match": {
+                                "query": query_text,
+                                "fields": ["canonical_name^2.0", "summary"],
+                            }
+                        },
+                        {
+                            "knn": {
+                                "embedding": {
+                                    "vector": query_embedding,
+                                    "k": limit * 2,
+                                }
+                            }
+                        },
+                    ],
+                    "minimum_should_match": 1,
+                }
+            },
+            "_source": {"excludes": ["embedding"]},
+        }
+        resp = await client.search(index=self.seed_index, body=body)
+        hits = resp.get("hits", {}).get("hits", [])
+        return [
+            GraphVertexUpsert(**{**h["_source"], "embedding": []})
+            for h in hits
+        ]
+
+    async def search_episodes(
+        self,
+        ctx: PlatformTenantContext,
+        query_embedding: list[float],
+        limit: int = 4,
+    ) -> list[EpisodicRecord]:
+        client = self._get_client()
+        body = {
+            "size": limit,
+            "query": {
+                "bool": {
+                    "filter": [
+                        {"term": {"tenant_id": ctx.tenant_id}},
+                        {"term": {"app_name": ctx.app_name}},
+                        {"term": {"user_id": ctx.user_id}},
+                    ],
+                    "must": [
+                        {
+                            "knn": {
+                                "embedding": {
+                                    "vector": query_embedding,
+                                    "k": limit * 2,
+                                }
+                            }
+                        }
+                    ],
+                }
+            },
+            "_source": {"excludes": ["embedding"]},
+        }
+        resp = await client.search(index=self.episodic_index, body=body)
+        hits = resp.get("hits", {}).get("hits", [])
+        return [
+            EpisodicRecord(**{**h["_source"], "embedding": []})
+            for h in hits
+        ]
+
+    async def fetch_pending_outbox(self, batch_size: int = 25) -> list[OutboxOperation]:
+        client = self._get_client()
+        body = {
+            "size": batch_size,
+            "sort": [{"created_at": {"order": "asc"}}],
+            "query": {
+                "bool": {
+                    "filter": [
+                        {"terms": {"status": ["PENDING", "RETRY"]}},
+                        {"range": {"attempts": {"lt": 5}}},
+                    ]
+                }
+            },
+        }
+        resp = await client.search(index=self.outbox_index, body=body)
+        hits = resp.get("hits", {}).get("hits", [])
+        return [
+            OutboxOperation(
+                outbox_id=h["_source"]["outbox_id"],
+                tenant_id=h["_source"]["tenant_id"],
+                operation_type=h["_source"]["operation_type"],
+                payload=h["_source"]["payload"],
+                attempts=h["_source"].get("attempts", 0),
+            )
+            for h in hits
+        ]
+
+    async def mark_outbox_processed(self, outbox_id: str) -> None:
+        client = self._get_client()
+        await client.update(
+            index=self.outbox_index,
+            id=outbox_id,
+            body={
+                "doc": {
+                    "status": "PROCESSED",
+                    "processed_at": datetime.now(timezone.utc).isoformat(),
+                }
+            },
+        )
+
+    async def mark_outbox_failed(self, outbox_id: str, error_msg: str) -> None:
+        client = self._get_client()
+        await client.update(
+            index=self.outbox_index,
+            id=outbox_id,
+            body={
+                "script": {
+                    "source": """
+                        ctx._source.attempts += 1;
+                        ctx._source.last_error = params.err;
+                        ctx._source.status = (ctx._source.attempts >= 5) ? 'DEAD_LETTER' : 'RETRY';
+                    """,
+                    "lang": "painless",
+                    "params": {"err": error_msg[:500]},
+                }
+            },
+        )
+
+    async def purge_user_data(self, ctx: PlatformTenantContext) -> list[str]:
+        client = self._get_client()
+        query = {
+            "query": {
+                "bool": {
+                    "filter": [
+                        {"term": {"tenant_id": ctx.tenant_id}},
+                        {"term": {"user_id": ctx.user_id}},
+                    ]
+                }
+            },
+            "_source": ["episode_id"],
+            "size": 1000,
+        }
+        resp = await client.search(index=self.episodic_index, body=query)
+        hits = resp.get("hits", {}).get("hits", [])
+        episode_ids = [h["_source"]["episode_id"] for h in hits]
+        if episode_ids:
+            await client.delete_by_query(
+                index=self.episodic_index,
+                body={"query": {"terms": {"episode_id": episode_ids}}},
+                refresh=True,
+            )
+        return episode_ids