diff --git a/README.md b/README.md index a87e3d3..d693166 100644 --- a/README.md +++ b/README.md @@ -197,6 +197,34 @@ ctx.update( ) # The enriched record now shows up in vector search and hybrid retrieve. +# Raw interaction-log ingestion: map raw rows into ContextRecord fields while +# preserving the original row under metadata["raw_record"]. Use mode="upsert" +# with stable external_id values for idempotent re-ingestion. +ctx.ingest_jsonl( + "chat-log.jsonl", + field_map={ + "external_id": "event_id", + "role": "speaker", + "content": "message", + "session_id": "conversation_id", + "run_id": "trace_id", + "created_at": "timestamp", + }, + defaults={"tenant": "example-org"}, + mode="upsert", + batch_size=500, +) + +# OpenAI-style messages can be ingested directly as raw records. +ctx.ingest_messages( + [ + {"role": "system", "content": "Be concise."}, + {"role": "user", "content": "What changed?"}, + ], + session_id="conversation-2026-03-01", + external_id_prefix="conversation-2026-03-01", +) + # Time-travel to prior state first_version = ctx.version() ctx.add("assistant", "Let me fetch suggestions…") diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index 24e2274..8bfd43d 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -6,7 +6,10 @@ from collections.abc import Iterable, Mapping from datetime import datetime from io import BytesIO -from typing import Any +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from os import PathLike from ._internal import Context as _Context # pyright: ignore[reportMissingImports] from ._internal import ( # pyright: ignore[reportMissingImports] @@ -30,6 +33,31 @@ __version__ = _version() _ARROW_STREAM_MIME = "application/vnd.apache.arrow.stream" +_DEFAULT_INGEST_BATCH_SIZE = 1000 +_MISSING = object() +_INGEST_DIRECT_FIELDS = { + "role", + "content", + "content_type", + "data_type", + "embedding", + "bot_id", + "session_id", + "tenant", + "source", + "external_id", + "run_id", + "created_at", + "state_metadata", + "relationships", + "expires_at", + "retention_policy", + "lifecycle_status", + "retired_at", + "retired_reason", + "supersedes_id", + "superseded_by_id", +} def _is_module(value: Any, prefix: str) -> bool: @@ -274,6 +302,107 @@ def _merge_storage_options( return options +def _normalize_field_map(field_map: Mapping[str, str] | None) -> dict[str, str]: + if field_map is None: + return {} + if not isinstance(field_map, Mapping): + raise TypeError("field_map must be a mapping") + + normalized: dict[str, str] = {} + for key, value in field_map.items(): + if not isinstance(key, str): + raise TypeError("field_map keys must be strings") + if not isinstance(value, str): + raise TypeError("field_map values must be strings") + normalized[key] = value + return normalized + + +def _mapped_value( + raw: Mapping[str, Any], field_map: Mapping[str, str], field: str +) -> Any: + source_field = field_map.get(field, field) + if source_field in raw: + return raw[source_field] + return _MISSING + + +def _jsonable_mapping(value: Mapping[str, Any], *, name: str) -> dict[str, Any]: + normalized = dict(value) + _json_dumps(normalized, name) + return normalized + + +def _normalize_ingest_record( + raw: Mapping[str, Any], + *, + field_map: Mapping[str, str], + defaults: Mapping[str, Any], + source: str, + preserve_raw: bool, + raw_metadata_key: str, + index: int, +) -> dict[str, Any]: + record: dict[str, Any] = {} + + for field in _INGEST_DIRECT_FIELDS: + value = _mapped_value(raw, field_map, field) + if value is not _MISSING: + record[field] = value + + for field, value in defaults.items(): + if field != "metadata": + record.setdefault(field, value) + + record.setdefault("source", source) + record.setdefault("role", "source") + + if "content" not in record: + raise ValueError(f"records[{index}] is missing required ingest field 'content'") + + metadata: dict[str, Any] = {} + mapped_metadata = _mapped_value(raw, field_map, "metadata") + if mapped_metadata is not _MISSING: + if not isinstance(mapped_metadata, Mapping): + raise TypeError(f"records[{index}].metadata must be a mapping") + metadata.update( + _jsonable_mapping(mapped_metadata, name=f"records[{index}].metadata") + ) + + default_metadata = defaults.get("metadata") + if default_metadata is not None: + if not isinstance(default_metadata, Mapping): + raise TypeError("defaults['metadata'] must be a mapping") + merged = _jsonable_mapping(default_metadata, name="defaults['metadata']") + merged.update(metadata) + metadata = merged + + if preserve_raw: + metadata.setdefault( + raw_metadata_key, + _jsonable_mapping(raw, name=f"records[{index}].raw"), + ) + + if metadata: + record["metadata"] = metadata + + return record + + +def _validate_ingest_defaults(defaults: Mapping[str, Any] | None) -> dict[str, Any]: + if defaults is None: + return {} + if not isinstance(defaults, Mapping): + raise TypeError("defaults must be a mapping") + + normalized = dict(defaults) + invalid = set(normalized) - (_INGEST_DIRECT_FIELDS | {"metadata"}) + if invalid: + joined = ", ".join(sorted(invalid)) + raise ValueError(f"defaults contains unsupported record field(s): {joined}") + return normalized + + class Context: """Multimodal, versioned context store backed by Lance. @@ -457,6 +586,8 @@ def add( tenant: str | None = None, source: str | None = None, external_id: str | None = None, + run_id: str | None = None, + created_at: datetime | str | None = None, state_metadata: Mapping[str, Any] | None = None, metadata: dict[str, Any] | None = None, relationships: list[dict[str, Any]] | None = None, @@ -495,6 +626,8 @@ def add( tenant, source, external_id, + run_id, + _coerce_timestamp(created_at, field_name="created_at"), _normalize_state_metadata(state_metadata), _json_dumps(metadata, "metadata"), _coerce_timestamp(expires_at, field_name="expires_at"), @@ -519,6 +652,9 @@ def upsert( tenant: str | None = None, source: str | None = None, external_id: str | None = None, + run_id: str | None = None, + created_at: datetime | str | None = None, + state_metadata: Mapping[str, Any] | None = None, metadata: dict[str, Any] | None = None, relationships: list[dict[str, Any]] | None = None, expires_at: datetime | str | None = None, @@ -562,6 +698,9 @@ def upsert( tenant, source, external_id, + run_id, + _coerce_timestamp(created_at, field_name="created_at"), + _normalize_state_metadata(state_metadata), _json_dumps(metadata, "metadata"), _coerce_timestamp(expires_at, field_name="expires_at"), retention_policy, @@ -650,7 +789,7 @@ def add_many(self, records: Iterable[Mapping[str, Any]]) -> None: Each record accepts the same fields as :meth:`add`: ``role``, ``content``, optional ``content_type``/``data_type``, ``embedding``, ``bot_id``, ``session_id``, ``tenant``, ``source``, ``external_id``, - ``state_metadata``, + ``run_id``, ``created_at``, ``state_metadata``, ``metadata``, ``relationships``, and lifecycle fields such as ``expires_at`` and ``lifecycle_status``. """ @@ -693,6 +832,11 @@ def add_many(self, records: Iterable[Mapping[str, Any]]) -> None: "source", getattr(self, "_default_fields", {}).get("source") ), "external_id": record.get("external_id"), + "run_id": record.get("run_id"), + "created_at": _coerce_timestamp( + record.get("created_at"), + field_name=f"records[{index}].created_at", + ), "state_metadata": _normalize_state_metadata( record.get("state_metadata") ), @@ -719,6 +863,203 @@ def add_many(self, records: Iterable[Mapping[str, Any]]) -> None: self._auto_embed_batch(normalized) self._inner.add_many(normalized) + def ingest_records( + self, + records: Iterable[Mapping[str, Any]], + *, + field_map: Mapping[str, str] | None = None, + defaults: Mapping[str, Any] | None = None, + source: str = "raw", + batch_size: int = _DEFAULT_INGEST_BATCH_SIZE, + mode: str = "append", + preserve_raw: bool = True, + raw_metadata_key: str = "raw_record", + ) -> dict[str, int]: + """Ingest raw row-like records into faithful ``ContextRecord`` entries. + + ``field_map`` maps context-record fields to raw input fields, for + example ``{"content": "message", "external_id": "event_id"}``. + ``mode="append"`` batches through :meth:`add_many`; ``mode="upsert"`` + uses each row's ``external_id`` for idempotent re-ingestion. + """ + if batch_size <= 0: + raise ValueError("batch_size must be positive") + if mode not in {"append", "upsert"}: + raise ValueError("mode must be 'append' or 'upsert'") + if not source: + raise ValueError("source must be non-empty") + if not raw_metadata_key: + raise ValueError("raw_metadata_key must be non-empty") + + normalized_field_map = _normalize_field_map(field_map) + normalized_defaults = _validate_ingest_defaults(defaults) + processed = 0 + inserted = 0 + updated = 0 + batches = 0 + batch: list[dict[str, Any]] = [] + + def flush() -> None: + nonlocal inserted, updated, batches, batch + if not batch: + return + if mode == "append": + self.add_many(batch) + inserted += len(batch) + else: + for record in batch: + external_id = record.get("external_id") + if not external_id: + raise ValueError( + "mode='upsert' requires external_id on every record" + ) + result = self.upsert( + record["role"], + record["content"], + content_type=record.get("content_type"), + data_type=record.get("data_type"), + embedding=record.get("embedding"), + bot_id=record.get("bot_id"), + session_id=record.get("session_id"), + tenant=record.get("tenant"), + source=record.get("source"), + external_id=external_id, + run_id=record.get("run_id"), + created_at=record.get("created_at"), + state_metadata=record.get("state_metadata"), + metadata=record.get("metadata"), + relationships=record.get("relationships"), + expires_at=record.get("expires_at"), + retention_policy=record.get("retention_policy"), + lifecycle_status=record.get("lifecycle_status"), + retired_at=record.get("retired_at"), + retired_reason=record.get("retired_reason"), + ) + if result["inserted"]: + inserted += 1 + else: + updated += 1 + batches += 1 + batch = [] + + for index, raw in enumerate(records): + if not isinstance(raw, Mapping): + raise TypeError(f"records[{index}] must be a mapping") + batch.append( + _normalize_ingest_record( + raw, + field_map=normalized_field_map, + defaults=normalized_defaults, + source=source, + preserve_raw=preserve_raw, + raw_metadata_key=raw_metadata_key, + index=index, + ) + ) + processed += 1 + if len(batch) >= batch_size: + flush() + + flush() + return { + "processed": processed, + "inserted": inserted, + "updated": updated, + "batches": batches, + } + + def ingest_jsonl( + self, + path: str | PathLike[str], + *, + field_map: Mapping[str, str] | None = None, + defaults: Mapping[str, Any] | None = None, + source: str = "raw", + batch_size: int = _DEFAULT_INGEST_BATCH_SIZE, + mode: str = "append", + preserve_raw: bool = True, + raw_metadata_key: str = "raw_record", + ) -> dict[str, int]: + """Stream-ingest newline-delimited JSON objects from ``path``.""" + + def iter_rows() -> Iterable[Mapping[str, Any]]: + with open(path, encoding="utf-8") as stream: + for line_number, line in enumerate(stream, start=1): + stripped = line.strip() + if not stripped: + continue + row = json.loads(stripped) + if not isinstance(row, Mapping): + raise ValueError(f"JSONL line {line_number} must be an object") + yield row + + return self.ingest_records( + iter_rows(), + field_map=field_map, + defaults=defaults, + source=source, + batch_size=batch_size, + mode=mode, + preserve_raw=preserve_raw, + raw_metadata_key=raw_metadata_key, + ) + + def ingest_messages( + self, + messages: Iterable[Mapping[str, Any]], + *, + session_id: str | None = None, + tenant: str | None = None, + bot_id: str | None = None, + external_id_prefix: str | None = None, + source: str = "raw", + batch_size: int = _DEFAULT_INGEST_BATCH_SIZE, + mode: str = "append", + metadata: Mapping[str, Any] | None = None, + ) -> dict[str, int]: + """Ingest OpenAI-style chat messages as raw context records.""" + default_metadata = dict(metadata or {}) + + def iter_rows() -> Iterable[Mapping[str, Any]]: + for index, message in enumerate(messages): + if not isinstance(message, Mapping): + raise TypeError(f"messages[{index}] must be a mapping") + content = message.get("content") + content_type = message.get("content_type") + if content is None: + content = {k: v for k, v in message.items() if k != "role"} + content_type = content_type or "application/json" + if not isinstance(content, (bytes, str)): + content = json.dumps(content, sort_keys=True, separators=(",", ":")) + content_type = content_type or "application/json" + + row_metadata = dict(default_metadata) + row_metadata["raw_message"] = dict(message) + row: dict[str, Any] = { + "role": message.get("role", "message"), + "content": content, + "content_type": content_type, + "source": source, + "metadata": row_metadata, + } + if session_id is not None: + row["session_id"] = session_id + if tenant is not None: + row["tenant"] = tenant + if bot_id is not None: + row["bot_id"] = bot_id + if external_id_prefix is not None: + row["external_id"] = f"{external_id_prefix}#message-{index + 1}" + yield row + + return self.ingest_records( + iter_rows(), + batch_size=batch_size, + mode=mode, + source=source, + preserve_raw=False, + ) + def _auto_embed_batch(self, records: list[dict[str, Any]]) -> None: """Embed text records without an embedding in one provider call.""" provider = getattr(self, "_embedding_provider", None) diff --git a/python/src/lib.rs b/python/src/lib.rs index e3615c5..c7acaae 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -45,6 +45,8 @@ struct RecordInput { tenant: Option, source: Option, external_id: Option, + run_id: Option, + created_at: Option>, state_metadata: Option, metadata_json: Option, relationships: Vec, @@ -290,7 +292,7 @@ impl Context { } #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (role, content, data_type = None, embedding = None, bot_id = None, session_id = None, tenant = None, source = None, external_id = None, state_metadata = None, metadata_json = None, expires_at = None, retention_policy = None, lifecycle_status = None, retired_at = None, retired_reason = None, supersedes_id = None, superseded_by_id = None, relationships_json = None))] + #[pyo3(signature = (role, content, data_type = None, embedding = None, bot_id = None, session_id = None, tenant = None, source = None, external_id = None, run_id = None, created_at = None, state_metadata = None, metadata_json = None, expires_at = None, retention_policy = None, lifecycle_status = None, retired_at = None, retired_reason = None, supersedes_id = None, superseded_by_id = None, relationships_json = None))] fn add( &mut self, py: Python<'_>, @@ -303,6 +305,8 @@ impl Context { tenant: Option, source: Option, external_id: Option, + run_id: Option, + created_at: Option, state_metadata: Option<&Bound<'_, PyDict>>, metadata_json: Option, expires_at: Option, @@ -334,6 +338,8 @@ impl Context { tenant, source, external_id, + run_id, + created_at: parse_optional_datetime(created_at, "created_at")?, state_metadata: state_metadata_from_dict(state_metadata)?, metadata_json, relationships: relationships_from_json(relationships_json)?, @@ -356,7 +362,7 @@ impl Context { } #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (role, content, data_type = None, embedding = None, bot_id = None, session_id = None, tenant = None, source = None, external_id = None, metadata_json = None, expires_at = None, retention_policy = None, lifecycle_status = None, retired_at = None, retired_reason = None, relationships_json = None, key = "external_id"))] + #[pyo3(signature = (role, content, data_type = None, embedding = None, bot_id = None, session_id = None, tenant = None, source = None, external_id = None, run_id = None, created_at = None, state_metadata = None, metadata_json = None, expires_at = None, retention_policy = None, lifecycle_status = None, retired_at = None, retired_reason = None, relationships_json = None, key = "external_id"))] fn upsert( &mut self, py: Python<'_>, @@ -369,6 +375,9 @@ impl Context { tenant: Option, source: Option, external_id: Option, + run_id: Option, + created_at: Option, + state_metadata: Option<&Bound<'_, PyDict>>, metadata_json: Option, expires_at: Option, retention_policy: Option, @@ -409,7 +418,9 @@ impl Context { tenant, source, external_id, - state_metadata: None, + run_id, + created_at: parse_optional_datetime(created_at, "created_at")?, + state_metadata: state_metadata_from_dict(state_metadata)?, metadata_json, relationships: relationships_from_json(relationships_json)?, lifecycle, @@ -889,6 +900,8 @@ impl Context { let source = optional_item(dict, "source")?.map(|value| value.extract::()); let external_id = optional_item(dict, "external_id")?.map(|value| value.extract::()); + let run_id = optional_item(dict, "run_id")?.map(|value| value.extract::()); + let created_at = optional_item(dict, "created_at")?.map(|value| value.extract::()); let state_metadata = match optional_item(dict, "state_metadata")? { Some(value) => { let metadata = value.downcast::().map_err(|_| { @@ -936,6 +949,8 @@ impl Context { tenant: tenant.transpose()?, source: source.transpose()?, external_id: external_id.transpose()?, + run_id: run_id.transpose()?, + created_at: parse_optional_datetime(created_at.transpose()?, "created_at")?, state_metadata, metadata_json: metadata_json.transpose()?, relationships: relationships_from_json(relationships_json.transpose()?)?, @@ -961,6 +976,8 @@ impl Context { tenant, source, external_id, + run_id, + created_at, state_metadata, metadata_json, relationships, @@ -990,18 +1007,19 @@ impl Context { } }; - let record_id = format!("{}-{}", self.run_id, self.inner.entries() + offset); + let record_run_id = run_id.unwrap_or_else(|| self.run_id.clone()); + let record_id = format!("{}-{}", record_run_id, self.inner.entries() + offset); let metadata = metadata_from_json(metadata_json)?; Ok(PreparedRecord { record: ContextRecord { id: record_id, external_id, - run_id: self.run_id.clone(), + run_id: record_run_id, bot_id, session_id, tenant, source, - created_at: Utc::now(), + created_at: created_at.unwrap_or_else(Utc::now), role: role.clone(), state_metadata, metadata, diff --git a/python/tests/test_ingestion.py b/python/tests/test_ingestion.py new file mode 100644 index 0000000..73150cf --- /dev/null +++ b/python/tests/test_ingestion.py @@ -0,0 +1,122 @@ +"""Tests for raw interaction-log ingestion helpers.""" + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pathlib import Path + +from lance_context.api import Context + + +def test_ingest_records_maps_raw_rows_and_preserves_provenance(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "context.lance")) + created_at = "2026-06-23T20:15:00Z" + + result = ctx.ingest_records( + [ + { + "event_id": "chat-1#turn-1", + "speaker": "user", + "message": "hello", + "conversation": "chat-1", + "trace_id": "run-1", + "timestamp": created_at, + } + ], + field_map={ + "external_id": "event_id", + "role": "speaker", + "content": "message", + "session_id": "conversation", + "run_id": "trace_id", + "created_at": "timestamp", + }, + defaults={"tenant": "acme"}, + ) + + assert result == {"processed": 1, "inserted": 1, "updated": 0, "batches": 1} + [record] = ctx.list() + assert record["external_id"] == "chat-1#turn-1" + assert record["role"] == "user" + assert record["text"] == "hello" + assert record["source"] == "raw" + assert record["tenant"] == "acme" + assert record["session_id"] == "chat-1" + assert record["run_id"] == "run-1" + assert record["created_at"] == datetime(2026, 6, 23, 20, 15, tzinfo=timezone.utc) + assert record["metadata"]["raw_record"]["event_id"] == "chat-1#turn-1" + + +def test_ingest_jsonl_streams_batches_and_upserts_idempotently(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "context.lance")) + path = tmp_path / "chat.jsonl" + rows = [ + {"id": "turn-1", "role": "user", "text": "first"}, + {"id": "turn-2", "role": "assistant", "text": "second"}, + ] + path.write_text("\n".join(json.dumps(row) for row in rows), encoding="utf-8") + + first = ctx.ingest_jsonl( + path, + field_map={"external_id": "id", "content": "text"}, + mode="upsert", + batch_size=1, + ) + second = ctx.ingest_jsonl( + path, + field_map={"external_id": "id", "content": "text"}, + mode="upsert", + batch_size=1, + ) + + assert first == {"processed": 2, "inserted": 2, "updated": 0, "batches": 2} + assert second == {"processed": 2, "inserted": 0, "updated": 2, "batches": 2} + records = ctx.list() + assert len(records) == 2 + assert {record["external_id"] for record in records} == {"turn-1", "turn-2"} + first_turn = ctx.get(external_id="turn-1") + assert first_turn is not None + assert first_turn["text"] == "first" + + +def test_ingest_messages_loads_openai_style_chat(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "context.lance")) + + result = ctx.ingest_messages( + [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "status?"}, + {"role": "assistant", "tool_calls": [{"name": "lookup", "arguments": {}}]}, + ], + session_id="session-1", + tenant="acme", + external_id_prefix="chat-99", + ) + + assert result == {"processed": 3, "inserted": 3, "updated": 0, "batches": 1} + records = ctx.list() + assert [record["role"] for record in records] == ["system", "user", "assistant"] + assert [record["external_id"] for record in records] == [ + "chat-99#message-1", + "chat-99#message-2", + "chat-99#message-3", + ] + assert {record["source"] for record in records} == {"raw"} + assert {record["session_id"] for record in records} == {"session-1"} + assert records[2]["content_type"] == "application/json" + assert records[2]["metadata"]["raw_message"]["tool_calls"][0]["name"] == "lookup" + + +def test_upsert_ingestion_requires_external_id(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "context.lance")) + + try: + ctx.ingest_records([{"content": "missing id"}], mode="upsert") + except ValueError as exc: + assert "requires external_id" in str(exc) + else: # pragma: no cover + raise AssertionError("expected upsert ingestion to require external_id")