xmemory-temporal 1.0.0__py3-none-any.whl

xmemory_temporal/__init__.py

@@ -0,0 +1,61 @@
+"""Temporal plugin for xmemory — durable agent memory as Temporal Activities.
+
+Register ``XmemoryPlugin`` on your Temporal ``Client`` (the Worker inherits its
+client's plugins) — or on the ``Worker`` — but **never both**: registering twice
+adds the activities twice and the worker crashes at boot with "More than one
+activity named xmemory_read". Then call ``xmemory_for_workflow`` inside a
+workflow to read and write memory through replay-safe activities.
+"""
+
+from xmemory_temporal.activities import (
+    ACTIVITY_READ,
+    ACTIVITY_WRITE,
+    ACTIVITY_WRITE_START,
+    ACTIVITY_WRITE_STATUS,
+    XmemoryActivities,
+)
+from xmemory_temporal.config import XmemoryConfig, XmemoryTimeouts
+from xmemory_temporal.dto import (
+    ReadInput,
+    ReadOutput,
+    ReadScope,
+    ScopeObject,
+    SubAnswer,
+    WriteInput,
+    WriteOutput,
+    WriteStartOutput,
+    WriteStatusInput,
+    WriteStatusOutput,
+)
+from xmemory_temporal.errors import NON_RETRYABLE_TYPES, to_application_error
+from xmemory_temporal.interceptor import AutoCaptureConfig
+from xmemory_temporal.plugin import XmemoryPlugin
+from xmemory_temporal.protocol import XmemoryInstanceProtocol
+from xmemory_temporal.workflow_api import WorkflowXmemory, xmemory_for_workflow
+
+__all__ = [
+    "XmemoryPlugin",
+    "XmemoryConfig",
+    "XmemoryTimeouts",
+    "AutoCaptureConfig",
+    "WorkflowXmemory",
+    "xmemory_for_workflow",
+    "XmemoryActivities",
+    "XmemoryInstanceProtocol",
+    "ReadInput",
+    "ReadOutput",
+    "ReadScope",
+    "ScopeObject",
+    "SubAnswer",
+    "WriteInput",
+    "WriteOutput",
+    "WriteStartOutput",
+    "WriteStatusInput",
+    "WriteStatusOutput",
+    "NON_RETRYABLE_TYPES",
+    "to_application_error",
+    "ACTIVITY_READ",
+    "ACTIVITY_WRITE",
+    "ACTIVITY_WRITE_START",
+    "ACTIVITY_WRITE_STATUS",
+]

xmemory_temporal/activities.py

@@ -0,0 +1,186 @@
+"""The xmemory activities: the only place this package does I/O.
+
+The client is injected per Worker; each call's client timeout is derived from
+the deadline Temporal assigned the activity. ``xmemory-ai`` is natively async,
+so these are plain ``async def`` with no thread pool.
+"""
+
+import contextvars
+from typing import Any, Callable
+
+from temporalio import activity
+from temporalio.exceptions import ApplicationError
+
+from xmemory_temporal.config import XmemoryConfig, client_timeout_seconds
+from xmemory_temporal.dto import (
+    ReadInput,
+    ReadOutput,
+    WriteInput,
+    WriteOutput,
+    WriteStartOutput,
+    WriteStatusInput,
+    WriteStatusOutput,
+    project_read,
+    project_write,
+    project_write_start,
+    project_write_status,
+)
+from xmemory_temporal.errors import TYPE_NO_DEADLINE, TYPE_NOT_BOUND, to_application_error
+from xmemory_temporal.protocol import XmemoryInstanceProtocol
+
+# Pinned so renaming a method cannot break replay of in-flight workflows.
+ACTIVITY_READ = "xmemory_read"
+ACTIVITY_WRITE = "xmemory_write"
+ACTIVITY_WRITE_START = "xmemory_write_start"
+ACTIVITY_WRITE_STATUS = "xmemory_write_status"
+
+# A ContextVar, not an attribute: one plugin object is shared across every Worker
+# built from a Client, so an attribute would be last-bind-wins and a Worker could
+# reach another's closed client. Each Worker's run_context binds its own.
+_bound_instance: contextvars.ContextVar[XmemoryInstanceProtocol | None] = contextvars.ContextVar(
+    "xmemory_bound_instance", default=None
+)
+
+
+class XmemoryActivities:
+    """The xmemory activity functions; the client is bound per-Worker (contextvar)."""
+
+    def __init__(self, config: XmemoryConfig) -> None:
+        self._config = config
+
+    @staticmethod
+    def bind(instance: XmemoryInstanceProtocol) -> contextvars.Token:
+        """Bind the live client for the current context. Returns a reset token."""
+        return _bound_instance.set(instance)
+
+    @staticmethod
+    def unbind(token: contextvars.Token) -> None:
+        _bound_instance.reset(token)
+
+    @property
+    def instance(self) -> XmemoryInstanceProtocol:
+        inst = _bound_instance.get()
+        if inst is None:
+            # Activities registered without the plugin. Non-retryable: no retry
+            # can bind a client.
+            raise ApplicationError(
+                "xmemory activities are not bound to a client — register XmemoryPlugin on the "
+                "Client (the Worker inherits it) rather than registering the activity functions "
+                "directly.",
+                type=TYPE_NOT_BOUND,
+                non_retryable=True,
+            )
+        return inst
+
+    def _client_timeout(self) -> float:
+        """Client budget for this call, derived from the activity's own deadline.
+
+        Deriving it, rather than keeping a second worker-side copy, is what
+        makes "the client gives up first" hold by construction when a workflow
+        lowers its timeout.
+
+        Temporal requires one of the two close timeouts on every activity, so
+        the ``None`` branch is unreachable in practice; both fields are typed
+        optional, and a silent default there is exactly the second copy this
+        design exists to avoid.
+        """
+        info = activity.info()
+        budget = info.start_to_close_timeout or info.schedule_to_close_timeout
+        if budget is None:
+            raise ApplicationError(
+                f"activity {info.activity_type} was scheduled without a deadline: set "
+                "start_to_close_timeout or schedule_to_close_timeout on it.",
+                type=TYPE_NO_DEADLINE,
+                non_retryable=True,
+            )
+        return client_timeout_seconds(budget.total_seconds(), self._config.client_margin_seconds)
+
+    @activity.defn(name=ACTIVITY_READ)
+    async def read(self, request: ReadInput) -> ReadOutput:
+        # Outside the try: an unbound-client or missing-deadline error must keep
+        # its non-retryable ApplicationError rather than being re-mapped.
+        instance = self.instance
+        timeout = self._client_timeout()
+        kwargs: dict[str, Any] = {}
+        if request.read_mode is not None:
+            kwargs["read_mode"] = request.read_mode
+        if request.scope is not None:
+            # The client validates this into its own ReadScope model; hand it a
+            # plain mapping so we never import a vendor type into the DTO layer.
+            kwargs["scope"] = {
+                "objects": [{"type": o.type, "key": o.key} for o in request.scope.objects],
+                "relations_scope": request.scope.relations_scope,
+            }
+        if request.read_id is not None:
+            kwargs["read_id"] = request.read_id
+        try:
+            result = await instance.read(
+                request.query,
+                timeout=timeout,
+                **kwargs,
+            )
+        except Exception as exc:
+            # `from None`, not `from exc`: Temporal serializes the cause chain
+            # into cleartext history and the client's message is unsanitized.
+            # Only `from None` suppresses the implicit __context__ too.
+            raise to_application_error(exc) from None
+        return project_read(result)
+
+    @activity.defn(name=ACTIVITY_WRITE)
+    async def write(self, request: WriteInput) -> WriteOutput:
+        instance = self.instance
+        timeout = self._client_timeout()
+        try:
+            result = await instance.write(
+                request.text,
+                timeout=timeout,
+                **self._write_kwargs(request),
+            )
+        except Exception as exc:
+            raise to_application_error(exc) from None
+        return project_write(result)
+
+    @activity.defn(name=ACTIVITY_WRITE_START)
+    async def write_start(self, request: WriteInput) -> WriteStartOutput:
+        instance = self.instance
+        timeout = self._client_timeout()
+        try:
+            result = await instance.write_async(
+                request.text,
+                timeout=timeout,
+                **self._write_kwargs(request),
+            )
+        except Exception as exc:
+            raise to_application_error(exc) from None
+        return project_write_start(result)
+
+    @activity.defn(name=ACTIVITY_WRITE_STATUS)
+    async def write_status(self, request: WriteStatusInput) -> WriteStatusOutput:
+        instance = self.instance
+        timeout = self._client_timeout()
+        try:
+            result = await instance.write_status(
+                request.write_id,
+                timeout=timeout,
+            )
+        except Exception as exc:
+            raise to_application_error(exc) from None
+        return project_write_status(result)
+
+    def _write_kwargs(self, request: WriteInput) -> dict[str, Any]:
+        kwargs: dict[str, Any] = {}
+        if request.structured_mutations is not None:
+            # A structured write carries its own keys, so the server applies it
+            # without running the extractor; text and extraction_logic are moot.
+            kwargs["structured_mutations"] = request.structured_mutations
+            return kwargs
+        logic = request.extraction_logic or self._config.default_extraction_logic
+        if logic is not None:
+            kwargs["extraction_logic"] = logic
+        if request.diff_engine is not None:
+            kwargs["diff_engine"] = request.diff_engine
+        return kwargs
+
+    def as_sequence(self) -> list[Callable[..., Any]]:
+        """The bound methods to hand to ``SimplePlugin(activities=...)``."""
+        return [self.read, self.write, self.write_start, self.write_status]

xmemory_temporal/client_factory.py

@@ -0,0 +1,43 @@
+"""One xmemory client per worker, shared by all concurrent activities.
+
+Per-invocation clients would mean a TCP+TLS handshake per memory op.
+"""
+
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import Any
+
+import httpx
+
+from xmemory_temporal.config import XmemoryConfig, XmemoryTimeouts, client_timeout_seconds
+from xmemory_temporal.protocol import XmemoryInstanceProtocol
+
+
+@asynccontextmanager
+async def open_instance(
+    config: XmemoryConfig,
+    *,
+    api_key: str | None = None,
+    http_client: httpx.AsyncClient | None = None,
+) -> AsyncGenerator[XmemoryInstanceProtocol, None]:
+    """Yield a bound instance handle, closing the client on exit.
+
+    ``api_key`` overrides the environment lookup; ``http_client`` supplies a
+    caller-owned transport (which the xmemory client will not close).
+    """
+    # Lazy import: injecting a fake must not require xmemory-ai at all.
+    from xmemory import AsyncXmemoryClient
+
+    key = api_key or config.resolve_api_key()
+    # Fallback only; every activity overrides this per call (activities.py).
+    default_timeout = client_timeout_seconds(XmemoryTimeouts().read_seconds, config.client_margin_seconds)
+    kwargs: dict[str, Any] = {"api_key": key, "timeout": default_timeout}
+    if http_client is not None:
+        # `url` + `http_client` together is rejected; the caller sets base_url.
+        kwargs["http_client"] = http_client
+    elif config.url is not None:
+        kwargs["url"] = config.url
+
+    # Closes only a client-owned transport, so a caller's stays open.
+    async with AsyncXmemoryClient(**kwargs) as client:
+        yield client.instance(config.instance_id)

xmemory_temporal/config.py

@@ -0,0 +1,96 @@
+"""Worker-side configuration.
+
+Carries the *name* of the env var holding the API key, never the key, so the
+config is safe to log, serialize, and persist into Temporal history.
+"""
+
+import os
+from datetime import timedelta
+
+from pydantic import BaseModel, ConfigDict
+
+DEFAULT_API_KEY_ENV = "XMEM_API_KEY"
+DEFAULT_CLIENT_MARGIN_SECONDS = 5
+
+
+def client_timeout_seconds(
+    activity_seconds: float,
+    margin_seconds: int = DEFAULT_CLIENT_MARGIN_SECONDS,
+) -> float:
+    """Client budget for an activity whose Temporal deadline is ``activity_seconds``.
+
+    Always strictly below that deadline, so the client fails first with an
+    attributable xmemory error. Budgets at or under the margin get a
+    proportional one, so the ordering holds for every positive budget.
+    """
+    if activity_seconds <= margin_seconds:
+        return max(0.1, activity_seconds * 0.8)
+    return float(activity_seconds - margin_seconds)
+
+
+class XmemoryTimeouts(BaseModel):
+    """Default ``start_to_close`` budgets applied by ``xmemory_for_workflow()``.
+
+    The workflow owns the real budget; activities derive their client timeout
+    from whatever Temporal assigned.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    read_seconds: int = 120
+    write_seconds: int = 180
+    write_start_seconds: int = 30
+    write_status_seconds: int = 30
+
+    @property
+    def read(self) -> timedelta:
+        return timedelta(seconds=self.read_seconds)
+
+    @property
+    def write(self) -> timedelta:
+        return timedelta(seconds=self.write_seconds)
+
+    @property
+    def write_start(self) -> timedelta:
+        return timedelta(seconds=self.write_start_seconds)
+
+    @property
+    def write_status(self) -> timedelta:
+        return timedelta(seconds=self.write_status_seconds)
+
+
+class XmemoryConfig(BaseModel):
+    """Worker-side configuration for the xmemory plugin.
+
+    No credential: the key is read from ``os.environ[api_key_env]``, or passed
+    in-process via ``XmemoryPlugin(config, api_key=...)``. Activity budgets are
+    not here either; they belong to ``xmemory_for_workflow``.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    instance_id: str
+    url: str | None = None
+    api_key_env: str = DEFAULT_API_KEY_ENV
+    # Gap between a call's Temporal deadline and its client timeout. If inverted,
+    # Temporal could abandon a write_start whose POST still enqueues server-side,
+    # and a later durable-write retry would double-enqueue.
+    client_margin_seconds: int = DEFAULT_CLIENT_MARGIN_SECONDS
+    default_extraction_logic: str = "fast"
+    # Summaries are visible to anyone with namespace access, and memory text is
+    # often personal, so content is redacted unless a caller opts in.
+    include_content_in_summary: bool = False
+
+    def resolve_api_key(self) -> str:
+        """Read the API key from the environment.
+
+        Raises at worker start rather than on the first activity, so a
+        misconfigured worker fails visibly.
+        """
+        key = os.environ.get(self.api_key_env)
+        if not key:
+            raise ValueError(
+                f"xmemory API key not found: environment variable {self.api_key_env!r} is unset or empty. "
+                f"Set it on the worker process, or pass XmemoryPlugin(config, api_key=...)."
+            )
+        return key

xmemory_temporal/dto.py

@@ -0,0 +1,160 @@
+"""Activity input/output types: ours, not the client's.
+
+Activity payloads are persisted verbatim into workflow history, so the type that
+crosses that boundary becomes a compatibility contract for every workflow that
+has ever run. Owning the wire format lets ``xmemory-ai`` evolve underneath us
+without breaking replay of completed workflows.
+
+Dataclasses, not pydantic: Temporal's default converter reconstructs those with
+no configuration, whereas pydantic models need a namespace-wide converter this
+plugin refuses to impose. For the same reason this module uses real annotations
+(no ``from __future__ import annotations``), since the converter resolves them via
+``typing.get_type_hints``, which fails on stringized ones.
+
+Outputs are flattened projections — only the fields a workflow can act on.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass(frozen=True)
+class ScopeObject:
+    """One record a scoped read may touch, addressed by its primary key."""
+
+    type: str
+    key: dict[str, str | int | float | bool] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class ReadScope:
+    """Restrict a read to specific records.
+
+    Mirrors the client's ``ReadScope`` so a malformed scope is a type error at
+    author time rather than a non-retryable ``XmemoryBadRequest`` at runtime.
+    ``relations_scope`` is ``no_relations`` (objects only) by default;
+    ``all_relations`` also exposes relations among the in-scope objects.
+    """
+
+    objects: list[ScopeObject] = field(default_factory=list)
+    relations_scope: str = "no_relations"
+
+
+@dataclass(frozen=True)
+class ReadInput:
+    query: str
+    read_mode: str | None = None
+    scope: ReadScope | None = None
+    read_id: str | None = None
+
+
+@dataclass(frozen=True)
+class SubAnswer:
+    """One decomposed sub-query and its own answer."""
+
+    sub_query: str
+    reader_result: Any = None
+    error: str | None = None
+
+
+@dataclass(frozen=True)
+class ReadOutput:
+    reader_result: Any = None
+    sub_answers: list[SubAnswer] = field(default_factory=list)
+    trace_id: str | None = None
+
+
+@dataclass(frozen=True)
+class WriteInput:
+    """Either free ``text`` for the extractor, or explicit ``structured_mutations``.
+
+    Mutations are plain JSON dicts in the client's ``WriteMutation`` shape. They
+    are kept as dicts rather than the client's pydantic models because activity
+    payloads must round-trip through Temporal's default converter.
+    """
+
+    text: str = ""
+    extraction_logic: str | None = None
+    diff_engine: bool | None = None
+    structured_mutations: list[dict[str, Any]] | None = None
+
+
+@dataclass(frozen=True)
+class WriteOutput:
+    write_id: str
+    trace_id: str | None = None
+    changes: Any = None
+
+
+@dataclass(frozen=True)
+class WriteStartOutput:
+    write_id: str
+
+
+@dataclass(frozen=True)
+class WriteStatusInput:
+    write_id: str
+
+
+@dataclass(frozen=True)
+class WriteStatusOutput:
+    write_id: str
+    write_status: str
+    error_detail: str | None = None
+    completed_at: str | None = None
+    # What the write applied. None until xmemory-ai surfaces it on write_status
+    # (see project_write_status); kept for symmetry with WriteOutput.changes.
+    changes: Any = None
+
+
+# --- Projections from the client's models ----------------------------------
+# `getattr` with defaults rather than attribute access: an older or newer client
+# release may not carry every field, and a missing one should degrade to `None`
+# rather than raise inside an activity.
+
+
+def project_read(result: Any) -> ReadOutput:
+    raw_sub = getattr(result, "reader_results", None) or []
+    return ReadOutput(
+        reader_result=getattr(result, "reader_result", None),
+        sub_answers=[
+            SubAnswer(
+                sub_query=getattr(item, "sub_query", ""),
+                reader_result=getattr(item, "reader_result", None),
+                error=getattr(item, "error", None),
+            )
+            for item in raw_sub
+        ],
+        trace_id=getattr(result, "trace_id", None),
+    )
+
+
+def project_write(result: Any) -> WriteOutput:
+    return WriteOutput(
+        write_id=getattr(result, "write_id", ""),
+        trace_id=getattr(result, "trace_id", None),
+        changes=getattr(result, "changes", None),
+    )
+
+
+def project_write_start(result: Any) -> WriteStartOutput:
+    return WriteStartOutput(write_id=getattr(result, "write_id", ""))
+
+
+def project_write_status(result: Any) -> WriteStatusOutput:
+    status = getattr(result, "write_status", None)
+    completed_at = getattr(result, "completed_at", None)
+    return WriteStatusOutput(
+        write_id=getattr(result, "write_id", ""),
+        # `WriteQueueStatus` is a `str` enum; normalize to its plain value so
+        # history never embeds an enum class the workflow side must import.
+        write_status=getattr(status, "value", status) or "",
+        error_detail=getattr(result, "error_detail", None),
+        completed_at=completed_at.isoformat() if completed_at is not None else None,
+        # The server returns what the write applied, but xmemory-ai's
+        # WriteStatusResult does not surface it yet — so `changes` is None here
+        # (unlike sync `write`, which carries WriteResult.changes). Picked up via
+        # getattr so it auto-populates if a future client exposes it. Surfacing it
+        # is an upstream client follow-up.
+        changes=getattr(result, "changes", None) or getattr(result, "result", None),
+    )