mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-25 08:04:57 +08:00
83 lines
2.3 KiB
Python
83 lines
2.3 KiB
Python
|
"""Directive checkpoint bands and append-only inject."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from plyngent.agent.directive_checkpoint import (
|
||
|
|
bands_to_fire,
|
||
|
|
inject_directive_checkpoints,
|
||
|
|
parse_checkpoint_bands,
|
||
|
|
token_band,
|
||
|
|
)
|
||
|
|
from plyngent.lmproto.openai_compatible.model import (
|
||
|
|
AnyChatMessage,
|
||
|
|
DeveloperChatMessage,
|
||
|
|
UserChatMessage,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_token_band() -> None:
|
||
|
|
assert token_band(0, 100_000) == 0
|
||
|
|
assert token_band(99_999, 100_000) == 0
|
||
|
|
assert token_band(100_000, 100_000) == 1
|
||
|
|
assert token_band(250_000, 100_000) == 2
|
||
|
|
assert token_band(100_000, 0) == 0
|
||
|
|
|
||
|
|
|
||
|
|
def test_bands_to_fire_fill_gaps() -> None:
|
||
|
|
assert bands_to_fire(last_fired_band=0, current_band=0) == []
|
||
|
|
assert bands_to_fire(last_fired_band=0, current_band=1) == [1]
|
||
|
|
assert bands_to_fire(last_fired_band=0, current_band=2) == [1, 2]
|
||
|
|
assert bands_to_fire(last_fired_band=2, current_band=2) == []
|
||
|
|
assert bands_to_fire(last_fired_band=1, current_band=3) == [2, 3]
|
||
|
|
|
||
|
|
|
||
|
|
def test_inject_append_only_and_parse() -> None:
|
||
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="hi")]
|
||
|
|
band, appended = inject_directive_checkpoints(
|
||
|
|
messages,
|
||
|
|
prompt_tokens=100_000,
|
||
|
|
source="api",
|
||
|
|
interval=100_000,
|
||
|
|
last_fired_band=0,
|
||
|
|
)
|
||
|
|
assert band == 1
|
||
|
|
assert len(appended) == 1
|
||
|
|
assert isinstance(messages[-1], DeveloperChatMessage)
|
||
|
|
assert "band=1" in messages[-1].content
|
||
|
|
assert parse_checkpoint_bands(messages) == 1
|
||
|
|
|
||
|
|
band2, appended2 = inject_directive_checkpoints(
|
||
|
|
messages,
|
||
|
|
prompt_tokens=250_000,
|
||
|
|
source="api",
|
||
|
|
interval=100_000,
|
||
|
|
last_fired_band=band,
|
||
|
|
)
|
||
|
|
assert band2 == 2
|
||
|
|
assert len(appended2) == 1
|
||
|
|
assert parse_checkpoint_bands(messages) == 2
|
||
|
|
|
||
|
|
# Same band: no re-fire
|
||
|
|
band3, appended3 = inject_directive_checkpoints(
|
||
|
|
messages,
|
||
|
|
prompt_tokens=250_000,
|
||
|
|
source="api",
|
||
|
|
interval=100_000,
|
||
|
|
last_fired_band=band2,
|
||
|
|
)
|
||
|
|
assert band3 == 2
|
||
|
|
assert appended3 == []
|
||
|
|
|
||
|
|
|
||
|
|
def test_inject_disabled() -> None:
|
||
|
|
messages: list[AnyChatMessage] = []
|
||
|
|
band, appended = inject_directive_checkpoints(
|
||
|
|
messages,
|
||
|
|
prompt_tokens=500_000,
|
||
|
|
source="api",
|
||
|
|
interval=0,
|
||
|
|
last_fired_band=0,
|
||
|
|
)
|
||
|
|
assert band == 0
|
||
|
|
assert appended == []
|