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

agentic_memory_sdk/adk_integration.py

@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any, Sequence
+
+import httpx
+
+from .models import AppMemoryManifest, ReasoningTraceStep
+from .service import BaseMemoryService, Event, MemoryEntry, SearchMemoryResponse, Session, types
+
+logger = logging.getLogger(__name__)
+
+
+class RemoteCompositeMemoryService(BaseMemoryService):
+    """
+    Lightweight HTTP Client implementing Google ADK's `BaseMemoryService` contract.
+    Designed for platforms hosting 100s of ADK agent apps:
+      - Agents do NOT need direct Postgres/OpenSearch/Neptune credentials or connection pools.
+      - Plugs directly into `google.adk.runners.Runner(memory_service=RemoteCompositeMemoryService(...))`.
+    """
+
+    def __init__(
+        self,
+        base_url: str = "http://localhost:8080",
+        api_key: str | None = None,
+        timeout_sec: float = 10.0,
+        client: httpx.AsyncClient | None = None,
+    ):
+        self.base_url = base_url.rstrip("/")
+        self.api_key = api_key
+        self.timeout_sec = timeout_sec
+        self._external_client = client
+
+    def _headers(self) -> dict[str, str]:
+        headers = {"Content-Type": "application/json"}
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+        return headers
+
+    async def _request(self, method: str, path: str, json_body: dict[str, Any] | None = None) -> dict[str, Any]:
+        if self._external_client is not None:
+            resp = await self._external_client.request(
+                method, f"{self.base_url}{path}", json=json_body, headers=self._headers()
+            )
+            resp.raise_for_status()
+            return resp.json()
+
+        async with httpx.AsyncClient(timeout=self.timeout_sec) as client:
+            resp = await client.request(
+                method, f"{self.base_url}{path}", json=json_body, headers=self._headers()
+            )
+            resp.raise_for_status()
+            return resp.json()
+
+    async def register_manifest(self, manifest: AppMemoryManifest) -> dict[str, Any]:
+        """Registers an ADK app's ontology manifest with the remote Memory Platform Service."""
+        return await self._request("POST", "/v1/apps/register", manifest.model_dump(mode="json"))
+
+    async def add_session_to_memory(self, session: Session) -> None:
+        await self.add_events_to_memory(
+            app_name=session.app_name,
+            user_id=session.user_id,
+            events=session.events,
+            session_id=session.id,
+        )
+
+    async def add_events_to_memory(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        events: Sequence[Event],
+        session_id: str | None = None,
+        custom_metadata: dict[str, Any] | None = None,
+    ) -> None:
+        serialized_events = []
+        for ev in events:
+            if not getattr(ev, "content", None) or not getattr(ev.content, "parts", None):
+                continue
+            text_parts = [p.text for p in ev.content.parts if getattr(p, "text", None)]
+            if text_parts:
+                serialized_events.append(
+                    {
+                        "author": getattr(ev, "author", None) or "agent",
+                        "text": "\n".join(text_parts),
+                    }
+                )
+        if not serialized_events:
+            return
+
+        await self._request(
+            "POST",
+            "/v1/memory/events",
+            {
+                "app_name": app_name,
+                "user_id": user_id,
+                "session_id": session_id or "default_session",
+                "events": serialized_events,
+                "custom_metadata": custom_metadata or {},
+            },
+        )
+
+    async def add_memory(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        memories: Sequence[MemoryEntry],
+        custom_metadata: dict[str, Any] | None = None,
+    ) -> None:
+        synthetic_events = [
+            Event(author=m.author or "user", content=m.content)
+            for m in memories
+            if getattr(m, "content", None)
+        ]
+        await self.add_events_to_memory(
+            app_name=app_name,
+            user_id=user_id,
+            events=synthetic_events,
+            custom_metadata=custom_metadata,
+        )
+
+    async def search_memory(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        query: str,
+        as_of_iso: str | None = None,
+        max_hops: int = 2,
+    ) -> SearchMemoryResponse:
+        data = await self._request(
+            "POST",
+            "/v1/memory/search",
+            {
+                "app_name": app_name,
+                "user_id": user_id,
+                "query": query,
+                "as_of_iso": as_of_iso,
+                "max_hops": max_hops,
+            },
+        )
+        entries: list[MemoryEntry] = []
+        for item in data.get("memories", []):
+            entries.append(
+                MemoryEntry(
+                    author=item.get("author", "memory"),
+                    timestamp=item.get("timestamp"),
+                    content=types.Content(
+                        role="user",
+                        parts=[types.Part.from_text(text=item.get("text", ""))],
+                    ),
+                )
+            )
+        return SearchMemoryResponse(memories=entries)
+
+    async def record_decision_trace(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        run_id: str,
+        session_id: str,
+        steps: Sequence[ReasoningTraceStep],
+    ) -> dict[str, Any]:
+        return await self._request(
+            "POST",
+            "/v1/traces",
+            {
+                "app_name": app_name,
+                "user_id": user_id,
+                "run_id": run_id,
+                "session_id": session_id,
+                "steps": [s.model_dump(mode="json") for s in steps],
+            },
+        )
+
+
+# ============================================================================
+# DROP-IN GOOGLE ADK CALLBACKS & EXPLICIT AGENT TOOLS
+# ============================================================================
+async def non_blocking_memory_commit_callback(callback_context: Any) -> None:
+    """
+    Google ADK `after_agent_callback`.
+    Commits the latest user + agent turn delta asynchronously off the critical user response path.
+    """
+    session = getattr(callback_context, "session", None)
+    if not session or not getattr(session, "events", None):
+        return None
+    recent_events = session.events[-2:]
+    if recent_events and hasattr(callback_context, "add_events_to_memory"):
+        asyncio.create_task(callback_context.add_events_to_memory(events=recent_events))
+    return None
+
+
+def get_adk_memory_tools(
+    memory_service: BaseMemoryService,
+    default_app_name: str = "adk_agent",
+    default_user_id: str = "default_user",
+) -> list[Any]:
+    """
+    Returns ready-to-attach Google ADK tool functions (`recall_memory_context`,
+    `commit_memory_episode`, `record_agent_decision_trace`) backed by either
+    `RemoteCompositeMemoryService` (no DB credentials needed) or in-process `CompositeMemoryService`.
+    Can be added alongside `PreloadMemoryTool()` in `LlmAgent(tools=[...])`.
+    """
+
+    async def recall_memory_context(
+        query: str,
+        as_of_timestamp: str | None = None,
+        app_name: str = default_app_name,
+        user_id: str = default_user_id,
+    ) -> dict[str, Any]:
+        """Search composite memory (Bi-Temporal Knowledge Graph + Episodic Memory), optionally at a historical point in time (`as_of_timestamp`)."""
+        resp = await memory_service.search_memory(  # type: ignore[attr-defined]
+            app_name=app_name,
+            user_id=user_id,
+            query=query,
+            as_of_iso=as_of_timestamp,
+        )
+        blocks = []
+        for m in resp.memories:
+            if getattr(m, "content", None) and getattr(m.content, "parts", None):
+                text = "\n".join(
+                    p.text for p in m.content.parts if getattr(p, "text", None)
+                )
+                blocks.append({"author": m.author, "timestamp": m.timestamp, "text": text})
+        return {"query": query, "as_of": as_of_timestamp, "memories": blocks}
+
+    async def commit_memory_episode(
+        fact_or_observation: str,
+        app_name: str = default_app_name,
+        user_id: str = default_user_id,
+    ) -> dict[str, Any]:
+        """Explicitly commit an important observation or domain fact into Episodic Memory and the Bi-Temporal Knowledge Graph."""
+        ev = Event(
+            author="agent",
+            content=types.Content(
+                role="user", parts=[types.Part.from_text(text=fact_or_observation)]
+            ),
+        )
+        await memory_service.add_events_to_memory(  # type: ignore[attr-defined]
+            app_name=app_name,
+            user_id=user_id,
+            events=[ev],
+        )
+        return {"status": "committed", "app_name": app_name}
+
+    return [recall_memory_context, commit_memory_episode]

agentic_memory_sdk/auth.py

@@ -0,0 +1,380 @@
+from __future__ import annotations
+
+import asyncio
+import ssl
+import time
+from dataclasses import dataclass
+from typing import Any, Awaitable, Callable, Protocol, runtime_checkable
+
+
+@dataclass(frozen=True)
+class EnterpriseTlsConfig:
+    """
+    Enterprise TLS / mTLS configuration for cloud-managed databases
+    (AWS RDS/Aurora, ElastiCache, MemoryDB, GCP Cloud SQL/AlloyDB, Azure Postgres, Neo4j Aura).
+    """
+
+    enabled: bool = True
+    verify_hostname: bool = True
+    ca_file: str | None = None
+    client_cert_file: str | None = None
+    client_key_file: str | None = None
+    min_tls_version: int = ssl.TLSVersion.TLSv1_2
+
+    def build_ssl_context(self) -> ssl.SSLContext | bool:
+        if not self.enabled:
+            return False
+        ctx = ssl.create_default_context(
+            purpose=ssl.Purpose.SERVER_AUTH,
+            cafile=self.ca_file,
+        )
+        ctx.minimum_version = ssl.TLSVersion(self.min_tls_version)
+        if self.verify_hostname:
+            ctx.check_hostname = True
+            ctx.verify_mode = ssl.CERT_REQUIRED
+        else:
+            ctx.check_hostname = False
+            ctx.verify_mode = ssl.CERT_NONE
+        if self.client_cert_file:
+            ctx.load_cert_chain(
+                certfile=self.client_cert_file,
+                keyfile=self.client_key_file,
+            )
+        return ctx
+
+
+@runtime_checkable
+class CloudAuthProvider(Protocol):
+    """
+    Protocol for enterprise cloud database authentication (IAM tokens, Entra ID,
+    GCP Workload Identity, Vault dynamic secrets, or rotating Secrets Manager credentials).
+    """
+
+    async def get_credentials(self) -> tuple[str, str]:
+        """Returns `(username, password_or_token)` with automatic caching & pre-expiry refresh."""
+        ...
+
+    async def get_username(self) -> str:
+        """Returns the current database principal/username."""
+        ...
+
+    async def get_password(self) -> str:
+        """
+        Returns the current database password or short-lived IAM/OAuth token.
+        Directly compatible with `asyncpg.create_pool(..., password=auth_provider.get_password)`.
+        """
+        ...
+
+    def invalidate(self) -> None:
+        """Forces immediate token/secret rotation on next connection attempt."""
+        ...
+
+
+class CachedTokenAuthProviderBase:
+    """
+    Base class providing concurrency-safe token caching and proactive pre-expiry rotation
+    for short-lived cloud IAM tokens (typically 15–60 minutes TTL).
+    """
+
+    def __init__(
+        self,
+        username: str,
+        token_ttl_sec: int = 900,
+        refresh_margin_sec: int = 120,
+    ):
+        self._username = username
+        self.token_ttl_sec = max(10, token_ttl_sec)
+        self.refresh_margin_sec = min(refresh_margin_sec, self.token_ttl_sec // 2)
+        self._cached_username: str | None = None
+        self._cached_token: str | None = None
+        self._expires_at: float = 0.0
+        self._refresh_count: int = 0
+        self._lock = asyncio.Lock()
+
+    @property
+    def refresh_count(self) -> int:
+        return self._refresh_count
+
+    def invalidate(self) -> None:
+        self._cached_token = None
+        self._expires_at = 0.0
+
+    async def _mint_credentials(self) -> tuple[str, str]:
+        raise NotImplementedError
+
+    async def get_credentials(self) -> tuple[str, str]:
+        now = time.monotonic()
+        if (
+            self._cached_token is not None
+            and self._cached_username is not None
+            and now < (self._expires_at - self.refresh_margin_sec)
+        ):
+            return self._cached_username, self._cached_token
+
+        async with self._lock:
+            now = time.monotonic()
+            if (
+                self._cached_token is not None
+                and self._cached_username is not None
+                and now < (self._expires_at - self.refresh_margin_sec)
+            ):
+                return self._cached_username, self._cached_token
+
+            user, token = await self._mint_credentials()
+            self._cached_username = user
+            self._cached_token = token
+            self._expires_at = time.monotonic() + self.token_ttl_sec
+            self._refresh_count += 1
+            return user, token
+
+    async def get_username(self) -> str:
+        user, _ = await self.get_credentials()
+        return user
+
+    async def get_password(self) -> str:
+        _, token = await self.get_credentials()
+        return token
+
+
+class StaticOrEnvAuthProvider(CachedTokenAuthProviderBase):
+    """
+    Static or environment-backed credentials provider implementing `CloudAuthProvider`.
+    """
+
+    def __init__(self, username: str, password: str):
+        super().__init__(username=username, token_ttl_sec=86400 * 365, refresh_margin_sec=0)
+        self._password = password
+
+    async def _mint_credentials(self) -> tuple[str, str]:
+        return self._username, self._password
+
+
+class AwsIamDatabaseAuthProvider(CachedTokenAuthProviderBase):
+    """
+    Enterprise AWS IAM Authentication Provider for:
+      - Amazon RDS / Aurora PostgreSQL (`service="rds"`, uses `generate_db_auth_token`)
+      - Amazon ElastiCache / MemoryDB for Redis (`service="elasticache"` or `"memorydb"`, SigV4 presigned token)
+    Supports optional cross-account IAM Role assumption (`assume_role_arn` + `external_id`)
+    and custom token generators for offline/unit testing.
+    """
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        username: str,
+        region: str = "us-east-1",
+        service: str = "rds",
+        cluster_name: str | None = None,
+        assume_role_arn: str | None = None,
+        external_id: str | None = None,
+        token_ttl_sec: int = 900,
+        refresh_margin_sec: int = 120,
+        custom_token_generator: Callable[[], str | Awaitable[str]] | None = None,
+    ):
+        super().__init__(
+            username=username,
+            token_ttl_sec=token_ttl_sec,
+            refresh_margin_sec=refresh_margin_sec,
+        )
+        self.host = host
+        self.port = port
+        self.region = region
+        self.service = service.lower()
+        self.cluster_name = cluster_name or host
+        self.assume_role_arn = assume_role_arn
+        self.external_id = external_id
+        self.custom_token_generator = custom_token_generator
+
+    def get_boto3_session(self) -> Any:
+        import boto3
+
+        if not self.assume_role_arn:
+            return boto3.Session(region_name=self.region)
+
+        sts = boto3.client("sts", region_name=self.region)
+        assume_kwargs: dict[str, Any] = {
+            "RoleArn": self.assume_role_arn,
+            "RoleSessionName": "adk-memory-platform-db-auth",
+            "DurationSeconds": max(900, self.token_ttl_sec),
+        }
+        if self.external_id:
+            assume_kwargs["ExternalId"] = self.external_id
+        creds = sts.assume_role(**assume_kwargs)["Credentials"]
+        return boto3.Session(
+            aws_access_key_id=creds["AccessKeyId"],
+            aws_secret_access_key=creds["SecretAccessKey"],
+            aws_session_token=creds["SessionToken"],
+            region_name=self.region,
+        )
+
+    def _generate_sync(self) -> str:
+        session = self.get_boto3_session()
+        if self.service == "rds":
+            rds = session.client("rds", region_name=self.region)
+            return rds.generate_db_auth_token(
+                DBHostname=self.host,
+                Port=self.port,
+                DBUsername=self. _username,
+                Region=self.region,
+            )
+        # ElastiCache / MemoryDB IAM auth uses a SigV4 presigned URL token
+        from botocore.signers import RequestSigner
+
+        credentials = session.get_credentials().get_frozen_credentials()
+        signer = RequestSigner(
+            service_name=self.service,
+            region_name=self.region,
+            signing_name=self.service,
+            signature_version="v4",
+            credentials=credentials,
+            event_emitter=session.events,
+        )
+        url = signer.generate_presigned_url(
+            request_dict={
+                "url_path": "/",
+                "query_string": {"Action": "connect", "User": self._username},
+                "headers": {"Host": self.cluster_name},
+                "method": "GET",
+                "body": b"",
+            },
+            region_name=self.region,
+            expires_in=self.token_ttl_sec,
+            operation_name="connect",
+        )
+        return url.replace("https://", "").replace("http://", "")
+
+    async def _mint_credentials(self) -> tuple[str, str]:
+        if self.custom_token_generator is not None:
+            res = self.custom_token_generator()
+            token = await res if asyncio.iscoroutine(res) else str(res)
+            return self._username, token
+        token = await asyncio.to_thread(self._generate_sync)
+        return self._username, token
+
+
+class GcpCloudSqlIamAuthProvider(CachedTokenAuthProviderBase):
+    """
+    Enterprise Google Cloud IAM Database Authentication Provider for:
+      - Cloud SQL for PostgreSQL / AlloyDB (`https://www.googleapis.com/auth/cloud-platform`)
+      - Cloud Memorystore for Redis IAM Auth
+    Uses Application Default Credentials (ADC) / Workload Identity Federation.
+    """
+
+    def __init__(
+        self,
+        username: str,
+        scopes: list[str] | None = None,
+        target_service_account: str | None = None,
+        token_ttl_sec: int = 3300,
+        refresh_margin_sec: int = 300,
+        custom_token_generator: Callable[[], str | Awaitable[str]] | None = None,
+    ):
+        super().__init__(
+            username=username,
+            token_ttl_sec=token_ttl_sec,
+            refresh_margin_sec=refresh_margin_sec,
+        )
+        self.scopes = scopes or ["https://www.googleapis.com/auth/cloud-platform"]
+        self.target_service_account = target_service_account
+        self.custom_token_generator = custom_token_generator
+
+    def _generate_sync(self) -> str:
+        import google.auth  # type: ignore
+        from google.auth.transport.requests import Request  # type: ignore
+
+        creds, _ = google.auth.default(scopes=self.scopes)
+        if self.target_service_account:
+            from google.auth import impersonated_credentials  # type: ignore
+
+            creds = impersonated_credentials.Credentials(
+                source_credentials=creds,
+                target_principal=self.target_service_account,
+                target_scopes=self.scopes,
+                lifetime=self.token_ttl_sec,
+            )
+        creds.refresh(Request())
+        return str(creds.token)
+
+    async def _mint_credentials(self) -> tuple[str, str]:
+        if self.custom_token_generator is not None:
+            res = self.custom_token_generator()
+            token = await res if asyncio.iscoroutine(res) else str(res)
+            return self._username, token
+        token = await asyncio.to_thread(self._generate_sync)
+        return self._username, token
+
+
+class AzureEntraDatabaseAuthProvider(CachedTokenAuthProviderBase):
+    """
+    Enterprise Microsoft Entra ID (Azure AD / Managed Identity) Authentication Provider for:
+      - Azure Database for PostgreSQL Flexible Server (`scope="https://ossrdbms-aad.database.windows.net/.default"`)
+      - Azure Managed Redis / Azure Cache for Redis (`scope="https://redis.azure.com/.default"`)
+    """
+
+    POSTGRES_SCOPE = "https://ossrdbms-aad.database.windows.net/.default"
+    REDIS_SCOPE = "https://redis.azure.com/.default"
+
+    def __init__(
+        self,
+        username: str,
+        scope: str = POSTGRES_SCOPE,
+        managed_identity_client_id: str | None = None,
+        token_ttl_sec: int = 3300,
+        refresh_margin_sec: int = 300,
+        custom_token_generator: Callable[[], str | Awaitable[str]] | None = None,
+    ):
+        super().__init__(
+            username=username,
+            token_ttl_sec=token_ttl_sec,
+            refresh_margin_sec=refresh_margin_sec,
+        )
+        self.scope = scope
+        self.managed_identity_client_id = managed_identity_client_id
+        self.custom_token_generator = custom_token_generator
+
+    def _generate_sync(self) -> str:
+        from azure.identity import DefaultAzureCredential  # type: ignore
+
+        cred = DefaultAzureCredential(
+            managed_identity_client_id=self.managed_identity_client_id
+        )
+        token_obj = cred.get_token(self.scope)
+        return str(token_obj.token)
+
+    async def _mint_credentials(self) -> tuple[str, str]:
+        if self.custom_token_generator is not None:
+            res = self.custom_token_generator()
+            token = await res if asyncio.iscoroutine(res) else str(res)
+            return self._username, token
+        token = await asyncio.to_thread(self._generate_sync)
+        return self._username, token
+
+
+class VaultOrSecretManagerAuthProvider(CachedTokenAuthProviderBase):
+    """
+    Dynamic Secret Rotation Provider for HashiCorp Vault Database Secrets Engine,
+    AWS Secrets Manager, or GCP Secret Manager.
+    Accepts an async or sync callable `secret_fetcher` returning `(username, password)`.
+    """
+
+    def __init__(
+        self,
+        secret_fetcher: Callable[[], tuple[str, str] | Awaitable[tuple[str, str]]],
+        token_ttl_sec: int = 1800,
+        refresh_margin_sec: int = 180,
+    ):
+        super().__init__(
+            username="dynamic",
+            token_ttl_sec=token_ttl_sec,
+            refresh_margin_sec=refresh_margin_sec,
+        )
+        self.secret_fetcher = secret_fetcher
+
+    async def _mint_credentials(self) -> tuple[str, str]:
+        res = self.secret_fetcher()
+        if asyncio.iscoroutine(res):
+            user, pwd = await res
+        else:
+            user, pwd = res  # type: ignore
+        return str(user), str(pwd)