64 lines
3.0 KiB
Python
64 lines
3.0 KiB
Python
import json
|
|
from types import SimpleNamespace
|
|
|
|
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
|
|
|
from artemis.llm.prefix_cache import PromptPrefixTracker
|
|
from artemis.llm.router import ModelEndpoint
|
|
from artemis.services.llm import RobustChatModelWrapper
|
|
|
|
|
|
def test_prefix_diagnostics_distinguish_append_rewrite_and_tool_changes_without_payloads():
|
|
tracker = PromptPrefixTracker()
|
|
messages = [SystemMessage(content="private rules"), HumanMessage(content="private screenshot")]
|
|
options = {"tools": [{"name": "click"}, {"name": "observe"}]}
|
|
assert tracker.observe("task:flash", messages, options)["cold_start"]
|
|
messages.extend(
|
|
[
|
|
AIMessage(content="", tool_calls=[{"name": "click", "args": {"x": 5}, "id": "c1"}]),
|
|
ToolMessage(content="clicked", tool_call_id="c1"),
|
|
]
|
|
)
|
|
extended = tracker.observe("task:flash", messages, options)
|
|
assert extended["append_only"] and extended["shared_prefix_messages"] == 2
|
|
messages[1] = HumanMessage(content="compressed summary")
|
|
rewritten = tracker.observe("task:flash", messages, options)
|
|
assert rewritten["first_changed_message"] == 2 and not rewritten["append_only"]
|
|
reordered = tracker.observe("task:flash", messages, {"tools": list(reversed(options["tools"]))})
|
|
assert reordered["tools_or_format_changed"] and not reordered["append_only"]
|
|
assert tracker.observe("task:checker", messages, options)["cold_start"]
|
|
assert "private" not in json.dumps(tracker._previous)
|
|
|
|
|
|
def test_bound_wrappers_reuse_task_prefix_tracker_after_message_adaptation():
|
|
class Model:
|
|
kwargs = {"tools": [{"type": "function", "name": "observe"}]}
|
|
|
|
def bind_tools(self, *_args, **_kwargs):
|
|
return self
|
|
|
|
ctx = SimpleNamespace(data_engine=SimpleNamespace(current_session_id="session"))
|
|
wrapper = RobustChatModelWrapper(Model(), ctx=ctx, endpoint=ModelEndpoint(), role="flash")
|
|
messages = [SystemMessage(content="rules"), HumanMessage(content="step 1")]
|
|
|
|
def diagnostics(items):
|
|
bound = wrapper.bind_tools([]) # Flash binds again for each call
|
|
args, kwargs = bound._prepare_model_call((items,), {})
|
|
return bound._diagnostics(args, kwargs)
|
|
|
|
first = diagnostics(messages)
|
|
messages.append(HumanMessage(content="step 2"))
|
|
second = diagnostics(messages)
|
|
assert first["prefix_cache"]["cold_start"]
|
|
assert second["prefix_cache"]["append_only"]
|
|
assert first["prompt_cache_key"] == second["prompt_cache_key"]
|
|
assert not any("rules" in str(value) for value in ctx.prompt_prefix_tracker._previous.values())
|
|
|
|
ctx.transcript_ledger = SimpleNamespace(
|
|
context_phase=1, last_compaction={"reason": "context_pressure"}
|
|
)
|
|
third = diagnostics(messages[:1] + [HumanMessage(content="history summary")])
|
|
assert third["prefix_cache"]["context_phase"] == 1
|
|
assert third["prefix_cache"]["last_compaction"]["reason"] == "context_pressure"
|
|
assert not third["prefix_cache"]["append_only"]
|