90 lines
3.9 KiB
Python
90 lines
3.9 KiB
Python
from concurrent.futures import ThreadPoolExecutor
|
|
from types import SimpleNamespace
|
|
from uuid import uuid4
|
|
import re
|
|
|
|
import pytest
|
|
|
|
from artemis.llm.cache_affinity import cache_affinity_options
|
|
|
|
|
|
def context(session, exploration=None):
|
|
return SimpleNamespace(
|
|
data_engine=SimpleNamespace(current_session_id=session), exploration_id=exploration
|
|
)
|
|
|
|
|
|
def test_routing_identity_survives_resume_but_isolates_tasks_roles_and_models():
|
|
ctx = context(uuid4(), "exploration")
|
|
first = cache_affinity_options(ctx, "flash", "gpt-6-luna")
|
|
ctx.data_engine.current_session_id = uuid4()
|
|
assert first == cache_affinity_options(ctx, "flash", "gpt-6-luna")
|
|
key = first["prompt_cache_key"]
|
|
assert re.fullmatch(r"[-0-9A-Za-z._:]{1,128}", key)
|
|
assert "exploration" not in key
|
|
assert first["extra_headers"] == {"Thread-Id": key, "Session-Id": key}
|
|
for role, model in [("checker", "gpt-6-luna"), ("flash", "gpt-6.1-sol")]:
|
|
assert cache_affinity_options(ctx, role, model)["prompt_cache_key"] != key
|
|
assert (
|
|
cache_affinity_options(context("another"), "flash", "gpt-6-luna")["prompt_cache_key"] != key
|
|
)
|
|
|
|
|
|
def test_explicit_options_are_preserved_without_mutating_headers_or_adding_aliases():
|
|
headers = {"thread_ID": "explicit-thread", "session-id": "explicit-session", "X-Test": "ok"}
|
|
options = {"prompt_cache_key": "explicit.key", "extra_headers": headers, "temperature": 0}
|
|
actual = cache_affinity_options(context("task"), "flash", "test", options)
|
|
assert actual == options and actual is not options
|
|
assert actual["extra_headers"] is not headers
|
|
assert options["extra_headers"] == headers
|
|
assert cache_affinity_options(None, "flash", "test") == {}
|
|
for key in (None, "", "a" * 129, "unsafe\r\nheader"):
|
|
assert cache_affinity_options(
|
|
context("task"), "flash", "test", {"prompt_cache_key": key}
|
|
) == {"prompt_cache_key": key}
|
|
|
|
|
|
def test_concurrent_calls_do_not_share_or_mutate_task_headers():
|
|
shared = {"extra_headers": {"X-Test": "shared"}}
|
|
with ThreadPoolExecutor(max_workers=2) as pool:
|
|
results = list(
|
|
pool.map(
|
|
lambda task: cache_affinity_options(context(task), "flash", "test", shared),
|
|
["task-a", "task-b"],
|
|
)
|
|
)
|
|
assert results[0]["extra_headers"]["Thread-Id"] != results[1]["extra_headers"]["Thread-Id"]
|
|
assert shared == {"extra_headers": {"X-Test": "shared"}}
|
|
|
|
|
|
@pytest.mark.parametrize("name", ["Thread-Id", "THREAD_ID", "session-id", "Session_id"])
|
|
def test_header_only_override_aligns_body_and_missing_headers(name):
|
|
headers = {name: "caller-conversation", "X-Test": "untouched"}
|
|
actual = cache_affinity_options(context("task"), "flash", "test", {"extra_headers": headers})
|
|
assert actual["prompt_cache_key"] == "caller-conversation"
|
|
routing = {
|
|
key.lower().replace("_", "-"): value for key, value in actual["extra_headers"].items()
|
|
}
|
|
assert routing["thread-id"] == routing["session-id"] == actual["prompt_cache_key"]
|
|
assert headers == {name: "caller-conversation", "X-Test": "untouched"}
|
|
assert cache_affinity_options(None, None, None, {"extra_headers": headers}) == actual
|
|
|
|
|
|
def test_explicit_body_and_header_options_are_not_overwritten():
|
|
options = {
|
|
"prompt_cache_key": "body-override",
|
|
"extra_headers": {"thread_id": "header-override"},
|
|
}
|
|
actual = cache_affinity_options(context("task"), "flash", "test", options)
|
|
assert actual["prompt_cache_key"] == "body-override"
|
|
assert actual["extra_headers"]["thread_id"] == "header-override"
|
|
assert "Thread-Id" not in actual["extra_headers"]
|
|
|
|
|
|
@pytest.mark.parametrize("value", [None, "", "unsafe\r\nheader", "a" * 129])
|
|
def test_invalid_header_is_not_promoted_to_body(value):
|
|
actual = cache_affinity_options(
|
|
context("task"), "flash", "test", {"extra_headers": {"Thread-Id": value}}
|
|
)
|
|
assert actual["prompt_cache_key"].startswith("artemis:")
|