xpark.dataset.extensions.MemoryExtract#

class xpark.dataset.extensions.MemoryExtract(*, hint: str | list[str] | None = None, extraction_target: str = 'user', 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_response: str | None = '{"memories": []}', **kwargs: dict[str, Any])#

Extract memory facts from conversation messages.

This operator takes per-row conversation messages (and an optional context block carrying identity / reference-date metadata) and produces a JSON list of newly extracted memory facts.

The output schema:

{“memories”: [

{“text”: “<self-contained fact>”}, …

]}

Time information lives inside text (resolved against the reference date carried by context, if available).

Parameters:

  • hint – Optional extra instructions or constraints to guide the model. When None, the default fine-tuned hint is selected based on extraction_target. Accepts either a single string or a list of strings, where each item is one hint written in plain text.

  • extraction_target – Selects the role to extract conversation for (OpenAI standard). "user" (default) or "agent". This only affects prompt selection, not message filtering.

  • 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_response – Fallback value on LLM failure. MUST be a valid JSON string conforming to FLAT_EXTRACTION_SCHEMA, or None to disable the fallback.

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

Inputs (passed positionally via with_column, in the following order):

  • messages (string, required): conversation text ready for extraction (any role-based filtering should be done upstream).

  • context (string, optional): per-conversation metadata (identity, reference date). When omitted (or empty), the <context> block is dropped from the prompt and all relative-time expressions in messages will be left unresolved unless the model can infer dates from the text itself.

Examples

import os
from xpark.dataset import from_items
from xpark.dataset.expressions import col
from xpark.dataset.extensions import MemoryExtract

ds = from_items([
    {
        "messages": "[user]: I went to the park yesterday.",
        "context": "User: u1\nDate: 2025-01-01",
    }
])
ds = ds.with_column(
    "memories",
    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=32)
    .with_column(col("messages"), col("context")),
)

print(ds.take_all())

Methods

__call__(messages[, context])

Call self as a function.

options(**kwargs)

with_column(messages[, context])

__call__(messages: pa.ChunkedArray, context: pa.ChunkedArray | None = None) → pa.Array#

Call self as a function.

options(**kwargs: Unpack[ExprUDFOptions]) → Self#
with_column(messages: pa.ChunkedArray, context: pa.ChunkedArray | None = None) → pa.Array#