xpark.dataset.extensions.MemoryBuildAction#

class xpark.dataset.extensions.MemoryBuildAction(*, hint: str | list[str] | None = None, ensure_ascii: bool = False, base_url: str, model: str, api_key: str = 'NOT_SET', max_qps: int | None = None, max_concurrency: int | None = None, max_retries: int = 0, fallback_action: str = 'ADD', **kwargs: dict[str, Any])#

Decide what action to take ADD / UPDATE / DELETE / NONE for each newly extracted fact given existing memories.

Output schema (one action per input fact, in the same order):

{"actions": [
    {
        "event": "ADD" | "UPDATE" | "DELETE" | "NONE",
        "target_id": <str | None>,
        "text":      <str>,
        "old_memory": <str | None>,
        "rationale": <str | None>,
    },
    ...
]}

Event semantics:

  • ADD: record the new fact as a brand-new memory.

  • UPDATE: replace an existing memory’s text with an enriched / corrected version. target_id points to that memory. old_memory carries its prior text verbatim; text is the merged memory.

  • DELETE: invalidate an existing memory because the new fact contradicts or retracts it. target_id points to the contradicted memory; old_memory carries its prior text.

  • NONE: drop the new fact (e.g. already covered by an existing memory, or too trivial / generic to retain).

  • facts (string, required): each row is one of three accepted shapes, discriminated by attempting a JSON parse:

    1. Canonical MemoryExtract wrapper — {"memories": [{"text": <str>, "context"?: <str>}, ...]}. Each item is forwarded verbatim (only text / context survive; unknown keys are dropped). text is required on every item; missing it raises.

    2. JSON array of strings["fact a", "fact b", ...]. Each string becomes one fact’s text (no context). Empty array yields zero facts. Arrays containing any non-string element fall through to shape 3.

    3. Bare string fallback — anything else (parse failure, parsed scalar / object-without-memories / mixed-type array). The cell verbatim becomes a single fact’s text. The empty string is allowed and yields [] (semantically equivalent to an empty wrapper).

  • existing (string, optional per row): a JSON array of memory dicts retrieved from the upstream store:

    [{"id": <str>, "text": <str>, "context"?: <str>, ...}, ...]
    

    id and text are required on every item; missing either raises. Additional fields are tolerated on input but only id / text / context are forwarded into the prompt. Bare strings, wrapper objects, and non-array JSON are rejected. None / empty string mean “no prior memories for this row” and render as [].

Parameters:
  • hint – Optional override for the system instructions. If None, the default BUILD_ACTION_HINT is used

  • ensure_ascii – JSON output escaping flag.

  • base_url – LLM server base URL.

  • model – LLM model name.

  • api_key – LLM API key.

  • max_qps – Max queries per second.

  • max_concurrency – Max concurrent LLM requests.

  • max_retries – Max retries per request.

  • fallback_action – Per-fact event used when the LLM call fails. Must be "ADD" (default) or "NONE".

  • **kwargs – Extra arguments forwarded to the OpenAI chat completions API.

Examples

import os
from xpark.dataset import from_items
from xpark.dataset.expressions import col
from xpark.dataset.extensions import (
    MemoryBuildAction,
    MemoryExtract,
)

ds = from_items([
    {
        "messages": "[user]: I bought a Yamaha P-125 digital piano.",
        "context": "User: u1\nDate: 2025-01-01",
        "existing_json": '[{"id": "m_001", '
                         '"text": "Owns a digital keyboard"}]',
    }
])
ds = ds.with_column(
    "facts_json",
    MemoryExtract(
        model="deepseek-v3-0324",
        base_url=os.getenv("LLM_ENDPOINT"),
        api_key=os.getenv("LLM_API_KEY"),
    ).options(num_workers={"IO": 1}, batch_size=1)
     .with_column(col("messages"), col("context")),
)
ds = ds.with_column(
    "actions_json",
    MemoryBuildAction(
        model="deepseek-v3-0324",
        base_url=os.getenv("LLM_ENDPOINT"),
        api_key=os.getenv("LLM_API_KEY"),
    ).options(num_workers={"IO": 1}, batch_size=1)
     .with_column(col("facts_json"), col("existing_json")),
)

Methods

__call__(facts, existing)

Call self as a function.

options(**kwargs)

with_column(facts, existing)

__call__(facts: pa.ChunkedArray, existing: pa.ChunkedArray) pa.Array#

Call self as a function.

options(**kwargs: Unpack[ExprUDFOptions]) Self#
with_column(facts: pa.ChunkedArray, existing: pa.ChunkedArray) pa.Array#