mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
feat(storage): tag the native lane with its producer ({producer, blocks})
Persist provider_data as a {producer, blocks} envelope (producer = the generating
provider's name) so the lowering layer can later replay the native lane verbatim only
to its producer and rebuild from neutral fields for any other.
The envelope is storage-only: prepare_provider_data_for_save runs the P1 mirror on the
bare block list and then wraps; reconstruct_messages dual-reads (new envelope OR legacy
bare list), unwraps to a bare _provider_content list (every consumer's contract), and
surfaces the producer on the stripped-before-wire _producer side channel. The producer
threads the same four save layers as is_error (facade -> protocol -> SQLite + Postgres);
the live assistant save tags it from self._provider.provider_name and the fork carries
_producer. Legacy rows need no migration to keep working (dual-read); the one-shot
backfill that tags them is a follow-up.
Sub-commit 2a of the canonical-trajectory storage cut.
This commit is contained in:
@@ -0,0 +1,94 @@
|
||||
"""provider_data ``{producer, blocks}`` storage envelope (#5 sub-commit 2a).
|
||||
|
||||
Native blocks persist wrapped with the generating provider's name so the lowering layer
|
||||
can later replay them verbatim only to their producer. The envelope is storage-only:
|
||||
``reconstruct_messages`` unwraps it back to a bare block list (every ``_provider_content``
|
||||
consumer requires a plain list) and surfaces the producer on the ``_producer`` side
|
||||
channel. Legacy bare-list rows are dual-read unchanged.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from turnstone.core.storage._utils import prepare_provider_data_for_save, wrap_provider_data
|
||||
|
||||
_BLOCKS = [{"type": "thinking", "thinking": "reasoning", "signature": "s"}]
|
||||
_BLOCKS_JSON = json.dumps(_BLOCKS)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# helpers
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_wrap_adds_envelope() -> None:
|
||||
out = wrap_provider_data(_BLOCKS_JSON, "anthropic")
|
||||
assert out is not None
|
||||
assert json.loads(out) == {"producer": "anthropic", "blocks": _BLOCKS}
|
||||
|
||||
|
||||
def test_wrap_no_producer_keeps_bare_list() -> None:
|
||||
assert wrap_provider_data(_BLOCKS_JSON, None) == _BLOCKS_JSON
|
||||
|
||||
|
||||
def test_wrap_none_passthrough() -> None:
|
||||
assert wrap_provider_data(None, "anthropic") is None
|
||||
|
||||
|
||||
def test_wrap_already_wrapped_is_unchanged() -> None:
|
||||
wrapped = json.dumps({"producer": "x", "blocks": _BLOCKS})
|
||||
assert wrap_provider_data(wrapped, "anthropic") == wrapped
|
||||
|
||||
|
||||
def test_prepare_strips_orphan_then_wraps() -> None:
|
||||
# A tool_use with no matching tool_calls is an orphan (P1 mirror) → stripped,
|
||||
# and what survives is wrapped with the producer.
|
||||
pd = json.dumps(
|
||||
[
|
||||
{"type": "thinking", "thinking": "t"},
|
||||
{"type": "tool_use", "id": "c1", "name": "x", "input": {}},
|
||||
]
|
||||
)
|
||||
out = prepare_provider_data_for_save("assistant", pd, None, "anthropic")
|
||||
assert out is not None
|
||||
env = json.loads(out)
|
||||
assert env["producer"] == "anthropic"
|
||||
assert [b["type"] for b in env["blocks"]] == ["thinking"]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# round-trip via a real (ephemeral) backend: save → reconstruct dual-read
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_round_trip_surfaces_bare_blocks_and_producer(backend: Any) -> None:
|
||||
ws = "ws-env-1"
|
||||
backend.save_message(ws, "assistant", "hi", provider_data=_BLOCKS_JSON, producer="anthropic")
|
||||
a = next(m for m in backend.load_messages(ws, repair=False) if m["role"] == "assistant")
|
||||
assert a["_provider_content"] == _BLOCKS # bare list — the consumer contract
|
||||
assert a["_producer"] == "anthropic"
|
||||
|
||||
|
||||
def test_legacy_bare_list_dual_read(backend: Any) -> None:
|
||||
ws = "ws-env-2"
|
||||
# producer omitted → stored as a bare list (legacy shape), read back unchanged.
|
||||
backend.save_message(ws, "assistant", "hi", provider_data=_BLOCKS_JSON)
|
||||
a = next(m for m in backend.load_messages(ws, repair=False) if m["role"] == "assistant")
|
||||
assert a["_provider_content"] == _BLOCKS
|
||||
assert "_producer" not in a
|
||||
|
||||
|
||||
def test_bulk_round_trip_with_producer(backend: Any) -> None:
|
||||
ws = "ws-env-3"
|
||||
backend.save_messages_bulk(
|
||||
[
|
||||
{
|
||||
"ws_id": ws,
|
||||
"role": "assistant",
|
||||
"content": "hi",
|
||||
"provider_data": _BLOCKS_JSON,
|
||||
"producer": "google",
|
||||
}
|
||||
]
|
||||
)
|
||||
a = next(m for m in backend.load_messages(ws, repair=False) if m["role"] == "assistant")
|
||||
assert a["_provider_content"] == _BLOCKS
|
||||
assert a["_producer"] == "google"
|
||||
@@ -44,6 +44,7 @@ def save_message(
|
||||
source: str | None = None,
|
||||
event_id: int | None = None,
|
||||
is_error: bool = False,
|
||||
producer: str | None = None,
|
||||
) -> int:
|
||||
"""Log a message to the conversations table.
|
||||
|
||||
@@ -71,6 +72,7 @@ def save_message(
|
||||
source=source,
|
||||
event_id=event_id,
|
||||
is_error=is_error,
|
||||
producer=producer,
|
||||
)
|
||||
except Exception:
|
||||
log.warning("Failed to save message for ws=%s role=%s", ws_id, role, exc_info=True)
|
||||
|
||||
@@ -2591,6 +2591,7 @@ class ChatSession:
|
||||
"provider_data": pd_str,
|
||||
"source": src if isinstance(src, str) and src else None,
|
||||
"is_error": bool(msg.get("is_error", False)),
|
||||
"producer": msg.get("_producer"),
|
||||
}
|
||||
)
|
||||
save_messages_bulk(bulk_rows)
|
||||
@@ -3935,6 +3936,7 @@ class ChatSession:
|
||||
provider_data=provider_data,
|
||||
tool_calls=tool_calls_json,
|
||||
event_id=self._ui_event_id(),
|
||||
producer=self._provider.provider_name if self._provider else None,
|
||||
)
|
||||
|
||||
tool_calls = assistant_msg.get("tool_calls")
|
||||
|
||||
@@ -109,10 +109,10 @@ from turnstone.core.storage._utils import (
|
||||
from turnstone.core.storage._utils import (
|
||||
escape_like as _escape_like,
|
||||
)
|
||||
from turnstone.core.storage._utils import normalize_native_for_save, sanitize_text
|
||||
from turnstone.core.storage._utils import (
|
||||
normalize_search_terms as _normalize_search_terms,
|
||||
)
|
||||
from turnstone.core.storage._utils import prepare_provider_data_for_save, sanitize_text
|
||||
from turnstone.core.storage._utils import (
|
||||
reconstruct_messages as _reconstruct_messages,
|
||||
)
|
||||
@@ -292,10 +292,13 @@ class PostgreSQLBackend:
|
||||
source: str | None = None,
|
||||
event_id: int | None = None,
|
||||
is_error: bool = False,
|
||||
producer: str | None = None,
|
||||
) -> int:
|
||||
now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S")
|
||||
content = sanitize_text(content)
|
||||
provider_data = normalize_native_for_save(role, sanitize_text(provider_data), tool_calls)
|
||||
provider_data = prepare_provider_data_for_save(
|
||||
role, sanitize_text(provider_data), tool_calls, producer
|
||||
)
|
||||
source = sanitize_text(source)
|
||||
with self._conn() as conn:
|
||||
result = conn.execute(
|
||||
@@ -339,8 +342,11 @@ class PostgreSQLBackend:
|
||||
"content": sanitize_text(row["content"]),
|
||||
"tool_name": row.get("tool_name"),
|
||||
"tool_call_id": row.get("tool_call_id"),
|
||||
"provider_data": normalize_native_for_save(
|
||||
row["role"], sanitize_text(row.get("provider_data")), row.get("tool_calls")
|
||||
"provider_data": prepare_provider_data_for_save(
|
||||
row["role"],
|
||||
sanitize_text(row.get("provider_data")),
|
||||
row.get("tool_calls"),
|
||||
row.get("producer"),
|
||||
),
|
||||
"tool_calls": row.get("tool_calls"),
|
||||
"_source": sanitize_text(row.get("source")),
|
||||
|
||||
@@ -160,6 +160,7 @@ class StorageBackend(Protocol):
|
||||
source: str | None = None,
|
||||
event_id: int | None = None,
|
||||
is_error: bool = False,
|
||||
producer: str | None = None,
|
||||
) -> int:
|
||||
"""Log a message to the conversations table.
|
||||
|
||||
|
||||
@@ -109,10 +109,10 @@ from turnstone.core.storage._utils import (
|
||||
from turnstone.core.storage._utils import (
|
||||
escape_like as _escape_like,
|
||||
)
|
||||
from turnstone.core.storage._utils import normalize_native_for_save, sanitize_text
|
||||
from turnstone.core.storage._utils import (
|
||||
normalize_search_terms as _normalize_search_terms,
|
||||
)
|
||||
from turnstone.core.storage._utils import prepare_provider_data_for_save, sanitize_text
|
||||
from turnstone.core.storage._utils import (
|
||||
reconstruct_messages as _reconstruct_messages,
|
||||
)
|
||||
@@ -326,10 +326,13 @@ class SQLiteBackend:
|
||||
source: str | None = None,
|
||||
event_id: int | None = None,
|
||||
is_error: bool = False,
|
||||
producer: str | None = None,
|
||||
) -> int:
|
||||
now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S")
|
||||
content = sanitize_text(content)
|
||||
provider_data = normalize_native_for_save(role, sanitize_text(provider_data), tool_calls)
|
||||
provider_data = prepare_provider_data_for_save(
|
||||
role, sanitize_text(provider_data), tool_calls, producer
|
||||
)
|
||||
source = sanitize_text(source)
|
||||
with self._conn() as conn:
|
||||
result = conn.execute(
|
||||
@@ -387,8 +390,11 @@ class SQLiteBackend:
|
||||
"content": sanitize_text(row["content"]),
|
||||
"tool_name": row.get("tool_name"),
|
||||
"tool_call_id": row.get("tool_call_id"),
|
||||
"provider_data": normalize_native_for_save(
|
||||
row["role"], sanitize_text(row.get("provider_data")), row.get("tool_calls")
|
||||
"provider_data": prepare_provider_data_for_save(
|
||||
row["role"],
|
||||
sanitize_text(row.get("provider_data")),
|
||||
row.get("tool_calls"),
|
||||
row.get("producer"),
|
||||
),
|
||||
"tool_calls": row.get("tool_calls"),
|
||||
"_source": sanitize_text(row.get("source")),
|
||||
|
||||
@@ -83,6 +83,44 @@ def normalize_native_for_save(
|
||||
return json.dumps(kept) if kept else None
|
||||
|
||||
|
||||
def wrap_provider_data(provider_data: str | None, producer: str | None) -> str | None:
|
||||
"""Wrap a bare native-block list in the storage envelope ``{producer, blocks}``.
|
||||
|
||||
``producer`` (the provider that generated the turn) lets the lowering layer replay the
|
||||
native lane verbatim only to its producing provider. Storage-only: the envelope is
|
||||
unwrapped back to a bare block list by :func:`reconstruct_messages`, so every
|
||||
``_provider_content`` consumer still sees a plain list. Input that is empty, already
|
||||
wrapped, unparseable, or has no ``producer`` is returned unchanged (the last keeps the
|
||||
legacy bare-list shape, which reconstruct dual-reads).
|
||||
"""
|
||||
if not provider_data or not producer:
|
||||
return provider_data
|
||||
try:
|
||||
blocks = json.loads(provider_data)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return provider_data
|
||||
if isinstance(blocks, dict) and "blocks" in blocks:
|
||||
return provider_data
|
||||
return json.dumps({"producer": producer, "blocks": blocks})
|
||||
|
||||
|
||||
def prepare_provider_data_for_save(
|
||||
role: str | None,
|
||||
provider_data: str | None,
|
||||
tool_calls_json: str | None,
|
||||
producer: str | None,
|
||||
) -> str | None:
|
||||
"""Save-boundary preparation of the native lane: enforce the mirror, then wrap.
|
||||
|
||||
The single entry point both backends' save paths call:
|
||||
:func:`normalize_native_for_save` (the native↔tool_calls mirror, on the bare block
|
||||
list) followed by :func:`wrap_provider_data` (the ``{producer, blocks}`` envelope).
|
||||
"""
|
||||
return wrap_provider_data(
|
||||
normalize_native_for_save(role, provider_data, tool_calls_json), producer
|
||||
)
|
||||
|
||||
|
||||
def _attachment_to_content_part(att: dict[str, Any]) -> dict[str, Any] | None:
|
||||
"""Convert a stored attachment row into an OpenAI-style content part.
|
||||
|
||||
@@ -491,7 +529,16 @@ def reconstruct_messages(
|
||||
msg: dict[str, Any] = {"role": "assistant", "content": content or ""}
|
||||
if provider_data:
|
||||
with contextlib.suppress(json.JSONDecodeError, TypeError):
|
||||
msg["_provider_content"] = json.loads(provider_data)
|
||||
parsed = json.loads(provider_data)
|
||||
# Storage envelope ``{producer, blocks}`` (new) vs bare list (legacy):
|
||||
# surface bare blocks as ``_provider_content`` (every consumer expects a
|
||||
# plain list) and carry the producer on a stripped-before-wire side channel.
|
||||
if isinstance(parsed, dict) and "blocks" in parsed:
|
||||
msg["_provider_content"] = parsed["blocks"]
|
||||
if parsed.get("producer"):
|
||||
msg["_producer"] = parsed["producer"]
|
||||
else:
|
||||
msg["_provider_content"] = parsed
|
||||
if tool_calls_json:
|
||||
with contextlib.suppress(json.JSONDecodeError, TypeError):
|
||||
msg["tool_calls"] = json.loads(tool_calls_json)
|
||||
|
||||
Reference in New Issue
Block a user