mirror of
https://github.com/run-llama/workflows-py.git
synced 2026-08-24 10:23:00 -04:00
Share excluded pydantic fields by reference when copying state for an edit (#743)
State copies currently deep-copy pydantic fields marked `exclude=True`. Agent memory uses excluded fields for its tokenizer and stores, so each `ctx.store` write can copy live runtime dependencies. With tiktoken, the first copy detaches the encoding from its registry and later copies rebuild the BPE. State copies now share excluded fields and deep-copy everything else. Nested pydantic models use the same rule. The copy still handles cycles, shared references, custom `__deepcopy__`, and values that cannot be copied. Added focused state-copy tests and a guarded tiktoken regression test.
This commit is contained in:
Adrian Lyjak
@@ -0,0 +1,5 @@
|
||||
---
|
||||
"llama-index-workflows": patch
|
||||
---
|
||||
|
||||
Share pydantic fields marked exclude=True by reference when copying state for an edit
|
||||
@@ -45,18 +45,7 @@ def _mock_agent(
|
||||
|
||||
|
||||
async def test_multi_agent_handoff_streams_with_memory() -> None:
|
||||
"""End-to-end regression for issue 709.
|
||||
|
||||
The reported failure was a router-to-specialist handoff in a streamed
|
||||
multi-agent chat that passed a ``Memory`` object to ``run()``. The handoff
|
||||
tool and agent setup write to ``ctx.store``, which deep-copies state for
|
||||
edit isolation, and the live ``Memory`` (sqlalchemy/aiosqlite/tiktoken
|
||||
internals) used to crash that copy with ``cannot pickle 'module' object``.
|
||||
|
||||
This drives the whole pattern at once: a router hands off to a specialist,
|
||||
the specialist streams the final answer, and a ``Memory`` rides through
|
||||
``run()`` the way the original repro had it.
|
||||
"""
|
||||
"""A streamed handoff preserves the supplied memory."""
|
||||
router = _mock_agent(
|
||||
"router",
|
||||
"Routes the chat to a specialist.",
|
||||
@@ -99,16 +88,21 @@ async def test_multi_agent_handoff_streams_with_memory() -> None:
|
||||
assert specialist_stream == "specialist answer"
|
||||
assert result.response.content == "specialist answer"
|
||||
|
||||
stored_memory = await handler.ctx.store.get("memory")
|
||||
assert isinstance(stored_memory, Memory)
|
||||
assert stored_memory.tokenizer_fn is memory.tokenizer_fn
|
||||
assert stored_memory.sql_store is memory.sql_store
|
||||
messages = stored_memory.get_all()
|
||||
assert messages == memory.get_all()
|
||||
assert messages[0].content == "earlier turn"
|
||||
assert any(message.content == "help me" for message in messages)
|
||||
assert messages[-1].content == "specialist answer"
|
||||
|
||||
|
||||
async def test_store_set_accepts_non_serializable_object(
|
||||
create_workflow: WorkflowFactory,
|
||||
) -> None:
|
||||
"""Regression for issue 710: ctx.store.set with an unpicklable live object.
|
||||
|
||||
Storing an object that wraps a thread lock (e.g. an LLM client) used to
|
||||
raise ``TypeError: cannot pickle '_thread.lock' object`` from the edit-time
|
||||
whole-state deep copy. The object is kept by reference instead.
|
||||
"""
|
||||
"""The store keeps an uncopyable live object by reference."""
|
||||
lock = threading.Lock()
|
||||
captured = None
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ import functools
|
||||
import json
|
||||
import uuid
|
||||
import warnings
|
||||
from collections import deque
|
||||
from contextlib import asynccontextmanager
|
||||
from copy import copy, deepcopy
|
||||
from typing import (
|
||||
@@ -35,6 +36,8 @@ if TYPE_CHECKING:
|
||||
|
||||
MAX_DEPTH = 1000
|
||||
|
||||
_MODEL_DEEPCOPY = BaseModel.__deepcopy__
|
||||
|
||||
# Keys set by pre-built workflows that are known to be unserializable in some cases.
|
||||
KNOWN_UNSERIALIZABLE_KEYS: tuple[str, ...] = ("memory",)
|
||||
|
||||
@@ -627,38 +630,120 @@ class DictState(DictLikeModel):
|
||||
MODEL_T = TypeVar("MODEL_T", bound=BaseModel, default=DictState) # type: ignore[reportGeneralTypeIssues]
|
||||
|
||||
|
||||
def _copy_value_for_edit(value: Any) -> Any:
|
||||
"""Deep-copy a single state value, or keep the live reference if it can't be.
|
||||
def _deepcopy_or_share(value: Any, memo: dict[int, Any]) -> Any:
|
||||
"""Deep-copy a value, or share the live reference when it cannot be copied.
|
||||
|
||||
State can hold live workflow objects (memory, LLM clients) that wrap thread
|
||||
locks, modules, or sockets and raise on ``deepcopy``. Those are shared live
|
||||
handles, so there is nothing to isolate by copying — preserve the reference
|
||||
instead of failing the edit.
|
||||
State can hold live handles that raise on ``deepcopy``. Sharing those keeps
|
||||
the edit from crashing.
|
||||
"""
|
||||
try:
|
||||
return deepcopy(value)
|
||||
return deepcopy(value, memo)
|
||||
except Exception:
|
||||
# A copier that raised part-way may have left a broken entry behind.
|
||||
memo[id(value)] = value
|
||||
return value
|
||||
|
||||
|
||||
def _copy_value_for_edit(value: Any, memo: dict[int, Any]) -> Any:
|
||||
"""Copy one state value for an ``edit_state`` block.
|
||||
|
||||
``memo`` is the standard ``deepcopy`` memo, shared across the whole walk so
|
||||
cycles terminate and an object referenced twice stays one object.
|
||||
|
||||
Built-in containers are walked here rather than handed to ``deepcopy``, so
|
||||
a model nested in one still gets the exclude rule. Everything else,
|
||||
including container subclasses, goes to ``deepcopy``.
|
||||
"""
|
||||
if id(value) in memo:
|
||||
return memo[id(value)]
|
||||
|
||||
if isinstance(value, BaseModel):
|
||||
# A model with its own __deepcopy__ has already declared how it copies.
|
||||
if type(value).__deepcopy__ is not _MODEL_DEEPCOPY:
|
||||
return _deepcopy_or_share(value, memo)
|
||||
try:
|
||||
return _copy_model_for_edit(value, memo)
|
||||
except Exception:
|
||||
memo[id(value)] = value
|
||||
return value
|
||||
|
||||
kind = type(value)
|
||||
if kind is list:
|
||||
list_items: list[Any] = []
|
||||
memo[id(value)] = list_items
|
||||
list_items.extend(_copy_value_for_edit(item, memo) for item in value)
|
||||
return list_items
|
||||
if kind is dict:
|
||||
entries: dict[Any, Any] = {}
|
||||
memo[id(value)] = entries
|
||||
for key, item in value.items():
|
||||
entries[_copy_value_for_edit(key, memo)] = _copy_value_for_edit(item, memo)
|
||||
return entries
|
||||
if kind is tuple:
|
||||
# A tuple cannot be filled in after the fact, so it is memoized last —
|
||||
# a cycle through its elements may have already copied it.
|
||||
tuple_items = [_copy_value_for_edit(item, memo) for item in value]
|
||||
return memo.setdefault(id(value), tuple(tuple_items))
|
||||
if kind is set:
|
||||
set_items: set[Any] = set()
|
||||
memo[id(value)] = set_items
|
||||
set_items.update(_copy_value_for_edit(item, memo) for item in value)
|
||||
return set_items
|
||||
if kind is frozenset:
|
||||
frozen_items = frozenset(_copy_value_for_edit(item, memo) for item in value)
|
||||
return memo.setdefault(id(value), frozen_items)
|
||||
if kind is deque:
|
||||
deque_items: deque[Any] = deque(maxlen=value.maxlen)
|
||||
memo[id(value)] = deque_items
|
||||
deque_items.extend(_copy_value_for_edit(item, memo) for item in value)
|
||||
return deque_items
|
||||
|
||||
return _deepcopy_or_share(value, memo)
|
||||
|
||||
|
||||
def _copy_model_for_edit(value: BaseModel, memo: dict[int, Any]) -> BaseModel:
|
||||
"""Copy a model, sharing ``Field(exclude=True)`` values by reference.
|
||||
|
||||
``exclude=True`` marks a field as not part of the model's data, which is how
|
||||
a live handle hung off a model is already declared — a tokenizer, a chat
|
||||
store, a client. Copying one is wasted at best and ruinous at worst:
|
||||
deep-copying a tiktoken tokenizer detaches it from tiktoken's registry, and
|
||||
every copy after that rebuilds a 100k-entry BPE. Sharing them keeps the edit
|
||||
isolated on the data a reader could otherwise catch mid-edit.
|
||||
|
||||
``model_copy()`` lays out the new instance (fields, extras, private attrs,
|
||||
fields-set); this fills in the copies field by field.
|
||||
"""
|
||||
copied = value.model_copy()
|
||||
# Registered before recursing so a cycle back to this model resolves here.
|
||||
memo[id(value)] = copied
|
||||
|
||||
fields = type(value).model_fields
|
||||
contents = vars(copied)
|
||||
for name, child in value.__dict__.items():
|
||||
field = fields.get(name)
|
||||
if field is None or field.exclude is not True:
|
||||
contents[name] = _copy_value_for_edit(child, memo)
|
||||
|
||||
# Extras and private attrs carry data too: ``DictLikeModel._data`` is where
|
||||
# every dynamic entry of a ``DictState`` lives.
|
||||
for store in (copied.__pydantic_extra__, copied.__pydantic_private__):
|
||||
if not store:
|
||||
continue
|
||||
for name, child in list(store.items()):
|
||||
store[name] = _copy_value_for_edit(child, memo)
|
||||
return copied
|
||||
|
||||
|
||||
def copy_state_for_edit(state: MODEL_T) -> MODEL_T:
|
||||
"""Return an isolated copy of state for an ``edit_state`` block.
|
||||
|
||||
``edit_state`` mutates a copy so lockless readers keep seeing committed
|
||||
state until the block commits. Ordinary data entries are deep-copied for
|
||||
that isolation; entries holding non-deepcopyable live objects are kept by
|
||||
reference (see ``_copy_value_for_edit``) so the edit cannot crash on them.
|
||||
|
||||
Typed (non-``DictState``) state copies whole-model; if that model holds a
|
||||
non-deepcopyable field, fall back to a shallow copy rather than crash.
|
||||
state until the block commits. Data is copied for that isolation; live
|
||||
handles are shared instead — declared ones via ``Field(exclude=True)``,
|
||||
undeclared ones via the ``deepcopy`` failure fallback.
|
||||
"""
|
||||
if isinstance(state, DictState):
|
||||
copied = {key: _copy_value_for_edit(value) for key, value in state.items()}
|
||||
return cast(MODEL_T, DictState(**copied))
|
||||
try:
|
||||
return state.model_copy(deep=True)
|
||||
except Exception:
|
||||
return state.model_copy()
|
||||
return cast(MODEL_T, _copy_value_for_edit(state, {}))
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
|
||||
@@ -0,0 +1,386 @@
|
||||
# SPDX-License-Identifier: MIT
|
||||
# Copyright (c) 2026 LlamaIndex Inc.
|
||||
|
||||
"""Copy semantics of the state handed to an ``edit_state`` block."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from collections import deque
|
||||
from collections.abc import Callable, Iterator
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
|
||||
from workflows.context.state_store import DictState, copy_state_for_edit
|
||||
from workflows.events import Event
|
||||
|
||||
|
||||
class Handle:
|
||||
"""Stand-in for a live handle: a client, an engine, a tokenizer."""
|
||||
|
||||
|
||||
class Undeepcopyable:
|
||||
"""A live handle that raises on ``deepcopy``, like a lock or a module."""
|
||||
|
||||
def __deepcopy__(self, memo: dict[int, Any]) -> Undeepcopyable:
|
||||
raise TypeError("cannot pickle this object")
|
||||
|
||||
|
||||
class Memoryish(BaseModel):
|
||||
"""Shape of an agent memory: data next to declared live handles."""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
messages: list[str] = Field(default_factory=list)
|
||||
tokenizer: Any = Field(default=None, exclude=True)
|
||||
blocks: list[Blockish] = Field(default_factory=list)
|
||||
|
||||
|
||||
class Blockish(BaseModel):
|
||||
"""A nested, non-excluded model that owns a handle of its own."""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
facts: list[str] = Field(default_factory=list)
|
||||
llm: Any = Field(default=None, exclude=True)
|
||||
|
||||
|
||||
class FrozenBlockish(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)
|
||||
|
||||
name: str
|
||||
llm: Any = Field(default=None, exclude=True)
|
||||
|
||||
|
||||
class Node(BaseModel):
|
||||
"""Self-referencing model, for cycles."""
|
||||
|
||||
name: str
|
||||
peer: Node | None = None
|
||||
|
||||
|
||||
class Pair(BaseModel):
|
||||
left: list[int] = Field(default_factory=list)
|
||||
right: list[int] = Field(default_factory=list)
|
||||
|
||||
|
||||
class Privateer(BaseModel):
|
||||
"""Data in a field, a live handle in a private attribute."""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
nums: list[int] = Field(default_factory=list)
|
||||
_engine: Any = PrivateAttr(default=None)
|
||||
|
||||
|
||||
class TypedState(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
nums: list[int] = Field(default_factory=list)
|
||||
client: Any = Field(default=None, exclude=True)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Record:
|
||||
nums: list[int] = field(default_factory=list)
|
||||
|
||||
|
||||
def test_excluded_fields_are_shared_and_the_rest_is_isolated() -> None:
|
||||
tokenizer = Handle()
|
||||
memory = Memoryish(messages=["hi"], tokenizer=tokenizer)
|
||||
copied = copy_state_for_edit(DictState(memory=memory))["memory"]
|
||||
|
||||
assert copied.tokenizer is tokenizer
|
||||
assert copied.messages == ["hi"]
|
||||
assert copied.messages is not memory.messages
|
||||
copied.messages.append("bye")
|
||||
assert memory.messages == ["hi"]
|
||||
|
||||
|
||||
def test_nested_models_recurse_with_the_same_rule() -> None:
|
||||
llm = Handle()
|
||||
memory = Memoryish(blocks=[Blockish(facts=["a"], llm=llm)])
|
||||
copied = copy_state_for_edit(DictState(memory=memory))["memory"]
|
||||
|
||||
block = copied.blocks[0]
|
||||
assert block is not memory.blocks[0]
|
||||
assert block.llm is llm
|
||||
block.facts.append("b")
|
||||
assert memory.blocks[0].facts == ["a"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"container",
|
||||
[set, frozenset, deque],
|
||||
)
|
||||
def test_models_inside_builtin_containers_use_the_same_rule(
|
||||
container: Callable[[list[FrozenBlockish]], Any],
|
||||
) -> None:
|
||||
llm = Handle()
|
||||
block = FrozenBlockish(name="facts", llm=llm)
|
||||
|
||||
copied = copy_state_for_edit(DictState(items=container([block])))["items"]
|
||||
copied_block = next(iter(copied))
|
||||
|
||||
assert copied_block is not block
|
||||
assert copied_block.llm is llm
|
||||
|
||||
|
||||
def test_model_keeps_its_identity_and_validation_state() -> None:
|
||||
memory = Memoryish(messages=["hi"], tokenizer=Handle())
|
||||
copied = copy_state_for_edit(DictState(memory=memory))["memory"]
|
||||
|
||||
assert type(copied) is Memoryish
|
||||
assert copied.__pydantic_fields_set__ == memory.__pydantic_fields_set__
|
||||
assert copied.model_dump() == memory.model_dump()
|
||||
|
||||
|
||||
def test_private_attributes_holding_data_stay_isolated() -> None:
|
||||
"""``DictLikeModel`` keeps dynamic entries in a private attr — real data."""
|
||||
event = Event(payload={"a": 1})
|
||||
state = DictState(event=event)
|
||||
|
||||
copied = copy_state_for_edit(state)["event"]
|
||||
copied["payload"]["a"] = 2
|
||||
|
||||
assert event["payload"] == {"a": 1}
|
||||
|
||||
|
||||
def test_private_attribute_handles_are_shared_without_losing_siblings() -> None:
|
||||
"""A lazily-built engine cannot be copied, and must not be dropped either."""
|
||||
engine = Undeepcopyable()
|
||||
value = Privateer(nums=[1])
|
||||
value._engine = engine
|
||||
|
||||
copied = copy_state_for_edit(DictState(value=value))["value"]
|
||||
|
||||
assert copied._engine is engine
|
||||
copied.nums.append(2)
|
||||
assert value.nums == [1]
|
||||
|
||||
|
||||
def test_cycles_terminate_and_mirror_the_original_shape() -> None:
|
||||
a = Node(name="a")
|
||||
b = Node(name="b", peer=a)
|
||||
a.peer = b
|
||||
|
||||
copied = copy_state_for_edit(DictState(node=a))["node"]
|
||||
|
||||
assert copied is not a
|
||||
assert copied.peer.peer is copied
|
||||
assert copied.name == "a"
|
||||
|
||||
|
||||
def test_a_tuple_reachable_from_its_own_elements_is_copied_once() -> None:
|
||||
"""A tuple is built after its elements, so a cycle can reach it first."""
|
||||
items: list[Any] = []
|
||||
root = (items,)
|
||||
items.append(root)
|
||||
|
||||
copied = copy_state_for_edit(DictState(root=root))["root"]
|
||||
|
||||
assert copied is not root
|
||||
assert copied[0][0] is copied
|
||||
|
||||
|
||||
def test_a_failed_copy_does_not_leave_a_broken_entry_behind() -> None:
|
||||
"""A copier that raises part-way can strand a partial value in the memo."""
|
||||
|
||||
class Poison:
|
||||
def __deepcopy__(self, memo: dict[int, Any]) -> Poison:
|
||||
memo[id(self)] = "partial"
|
||||
raise TypeError("cannot pickle this object")
|
||||
|
||||
poison = Poison()
|
||||
copied = copy_state_for_edit(DictState(a=poison, b=poison))
|
||||
|
||||
assert copied["a"] is poison
|
||||
assert copied["b"] is poison
|
||||
|
||||
|
||||
def test_objects_referenced_twice_stay_one_object() -> None:
|
||||
shared = [1]
|
||||
pair = Pair()
|
||||
pair.left = shared
|
||||
pair.right = shared
|
||||
|
||||
copied = copy_state_for_edit(DictState(pair=pair))["pair"]
|
||||
|
||||
assert copied.left is copied.right
|
||||
assert copied.left is not shared
|
||||
|
||||
|
||||
def test_the_same_model_under_two_state_keys_is_copied_once() -> None:
|
||||
memory = Memoryish(messages=["hi"])
|
||||
copied = copy_state_for_edit(DictState(a=memory, b=memory))
|
||||
|
||||
assert copied["a"] is copied["b"]
|
||||
assert copied["a"] is not memory
|
||||
|
||||
|
||||
def test_non_model_values_are_deep_copied() -> None:
|
||||
record = Record(nums=[1])
|
||||
copied = copy_state_for_edit(DictState(record=record, plain={"k": [1]}))
|
||||
|
||||
assert copied["record"] is not record
|
||||
copied["record"].nums.append(2)
|
||||
copied["plain"]["k"].append(2)
|
||||
assert record.nums == [1]
|
||||
|
||||
|
||||
def test_non_deepcopyable_values_are_kept_by_reference() -> None:
|
||||
"""An edit must not crash on a live handle."""
|
||||
client = Undeepcopyable()
|
||||
copied = copy_state_for_edit(DictState(client=client, nums=[1]))
|
||||
|
||||
assert copied["client"] is client
|
||||
copied["nums"].append(2)
|
||||
|
||||
|
||||
def test_a_non_deepcopyable_field_no_longer_costs_its_siblings() -> None:
|
||||
"""Only the offending field is shared; declared data is still isolated."""
|
||||
|
||||
class Client(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
lock: Any = None
|
||||
calls: list[str] = Field(default_factory=list)
|
||||
|
||||
lock = Undeepcopyable()
|
||||
client = Client(lock=lock, calls=["a"])
|
||||
|
||||
copied = copy_state_for_edit(DictState(client=client))["client"]
|
||||
|
||||
assert copied.lock is lock
|
||||
copied.calls.append("b")
|
||||
assert client.calls == ["a"]
|
||||
|
||||
|
||||
def test_a_handle_inside_a_container_does_not_cost_its_neighbors() -> None:
|
||||
"""Containers are walked per element, so one bad entry is shared alone."""
|
||||
client = Undeepcopyable()
|
||||
copied = copy_state_for_edit(DictState(items=[client, [1]], by_key={"c": client}))
|
||||
|
||||
assert copied["items"][0] is client
|
||||
assert copied["by_key"]["c"] is client
|
||||
copied["items"][1].append(2)
|
||||
assert copied["items"][1] == [1, 2]
|
||||
|
||||
|
||||
class SelfSharing(BaseModel):
|
||||
"""`__deepcopy__` is the standard opt-out from being copied."""
|
||||
|
||||
nums: list[int] = Field(default_factory=list)
|
||||
|
||||
def __deepcopy__(self, memo: dict[int, Any] | None = None) -> SelfSharing:
|
||||
return self
|
||||
|
||||
|
||||
def test_a_model_with_its_own_deepcopy_keeps_it() -> None:
|
||||
value = SelfSharing(nums=[1])
|
||||
assert copy_state_for_edit(DictState(value=value))["value"] is value
|
||||
|
||||
|
||||
def test_the_deepcopy_opt_out_holds_for_typed_root_state() -> None:
|
||||
"""The root is copied by the same rule as any value nested under it."""
|
||||
state = SelfSharing(nums=[1])
|
||||
assert copy_state_for_edit(state) is state
|
||||
|
||||
|
||||
def test_typed_state_follows_the_same_rule() -> None:
|
||||
client = Handle()
|
||||
state = TypedState(nums=[1], client=client)
|
||||
|
||||
copied = copy_state_for_edit(state)
|
||||
|
||||
assert copied.client is client
|
||||
copied.nums.append(2)
|
||||
assert state.nums == [1]
|
||||
|
||||
|
||||
REGISTRY: dict[str, RegistryEncoding] = {}
|
||||
|
||||
|
||||
class RegistryEncoding:
|
||||
"""Copy-cost model of a tiktoken ``Encoding``.
|
||||
|
||||
Pickles by reference while it is the registry instance, and rebuilds from
|
||||
scratch once a copy has detached it — which is what turns one deep copy of
|
||||
an agent memory into an 80 ms BPE rebuild on every following copy.
|
||||
"""
|
||||
|
||||
rebuilds = 0
|
||||
|
||||
def __init__(self, name: str) -> None:
|
||||
self.name = name
|
||||
|
||||
def __deepcopy__(self, memo: dict[int, Any]) -> RegistryEncoding:
|
||||
if REGISTRY.get(self.name) is not self:
|
||||
RegistryEncoding.rebuilds += 1
|
||||
clone = RegistryEncoding.__new__(RegistryEncoding)
|
||||
clone.name = self.name
|
||||
return clone
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def encoding() -> Iterator[RegistryEncoding]:
|
||||
enc = RegistryEncoding("test-encoding")
|
||||
REGISTRY[enc.name] = enc
|
||||
RegistryEncoding.rebuilds = 0
|
||||
yield enc
|
||||
REGISTRY.pop(enc.name, None)
|
||||
|
||||
|
||||
def test_repeated_copies_of_a_declared_tokenizer_never_rebuild_it(
|
||||
encoding: RegistryEncoding,
|
||||
) -> None:
|
||||
"""An excluded tokenizer stays attached to its registry instance."""
|
||||
state = DictState(memory=Memoryish(messages=["hi"], tokenizer=encoding))
|
||||
|
||||
for _ in range(5):
|
||||
state = copy_state_for_edit(state)
|
||||
assert state["memory"].tokenizer is encoding
|
||||
|
||||
assert RegistryEncoding.rebuilds == 0
|
||||
|
||||
|
||||
def test_an_undeclared_tokenizer_is_still_copied(
|
||||
encoding: RegistryEncoding,
|
||||
) -> None:
|
||||
"""Control: without the marker there is nothing to go on, and the rebuilds
|
||||
are back. This is what makes the test above meaningful."""
|
||||
|
||||
class Undeclared(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
tokenizer: Any = None
|
||||
|
||||
state = DictState(memory=Undeclared(tokenizer=encoding))
|
||||
for _ in range(5):
|
||||
state = copy_state_for_edit(state)
|
||||
|
||||
assert RegistryEncoding.rebuilds > 0
|
||||
|
||||
|
||||
def test_tiktoken_backed_tokenizer_copies_stay_cheap() -> None:
|
||||
"""A real tiktoken tokenizer stays shared when tiktoken is installed."""
|
||||
tiktoken = pytest.importorskip("tiktoken")
|
||||
try:
|
||||
enc = tiktoken.get_encoding("cl100k_base")
|
||||
except Exception as exc: # no network, no cached BPE
|
||||
pytest.skip(f"tiktoken encoding unavailable: {exc}")
|
||||
|
||||
tokenizer = enc.encode
|
||||
state = DictState(memory=Memoryish(messages=["hi"], tokenizer=tokenizer))
|
||||
|
||||
started = time.perf_counter()
|
||||
for _ in range(8):
|
||||
state = copy_state_for_edit(state)
|
||||
elapsed = time.perf_counter() - started
|
||||
|
||||
assert state["memory"].tokenizer is tokenizer
|
||||
# A detached Encoding rebuilds in ~80 ms, so the unfixed path needs >500 ms
|
||||
# for these copies. The bound is loose enough for a loaded CI machine.
|
||||
assert elapsed < 0.25, f"copies degraded: {elapsed:.3f}s"
|
||||
Reference in New Issue
Block a user