mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
588 lines
27 KiB
Python
588 lines
27 KiB
Python
import asyncio
|
|
import copy
|
|
import gzip
|
|
import json
|
|
import threading
|
|
from collections import OrderedDict
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from functools import wraps
|
|
from typing import Literal
|
|
from unittest import TestCase
|
|
|
|
import httpx
|
|
|
|
import litellm
|
|
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.integrations.datadog.datadog import DataDogLogger
|
|
from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
|
|
from litellm.integrations.literal_ai import LiteralAILogger
|
|
from litellm.integrations.rubrik import RubrikLogger
|
|
from litellm.litellm_core_utils.litellm_logging import Logging, create_dummy_standard_logging_payload
|
|
from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr.crowdstrike_aidr import CrowdStrikeAIDRHandler
|
|
from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview.purview_dlp import MicrosoftPurviewDLPGuardrail
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
from litellm.types.utils import ModelResponse
|
|
|
|
|
|
async def integration_invoke(owners, callback, *args, **kwargs):
|
|
owner = owners.prepare(callback, args, kwargs, awaited=True)
|
|
pending = owner.invoke()
|
|
owner.close()
|
|
return await pending
|
|
|
|
|
|
def integration_response(url, body, status=200, headers=None):
|
|
return httpx.Response(status, json=body, headers=headers, request=httpx.Request("POST", url))
|
|
|
|
|
|
def integration_callback_scope(scenario):
|
|
@wraps(scenario)
|
|
async def run(owners):
|
|
callbacks = tuple(litellm.callbacks)
|
|
try:
|
|
return await scenario(owners)
|
|
finally:
|
|
litellm.callbacks[:] = callbacks
|
|
|
|
return run
|
|
|
|
|
|
@integration_callback_scope
|
|
async def real_logging_queue_chain(owners):
|
|
return await integration_logging_queue_case(owners)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class QueueObservation:
|
|
gcs_model_parameters: str
|
|
datadog_snapshot: str
|
|
literal_prepared_settings: str
|
|
|
|
|
|
@integration_callback_scope
|
|
async def real_logging_queue_copy_control(owners):
|
|
baseline = await integration_logging_queue_case(owners)
|
|
copied = await integration_logging_queue_case(owners, literal_copy="payload")
|
|
envelope = await integration_logging_queue_case(owners, literal_copy="envelope")
|
|
assert envelope == baseline
|
|
assert json.loads(baseline.gcs_model_parameters) == {"stream": True, "temperature": 0.25}
|
|
assert copied.datadog_snapshot == baseline.datadog_snapshot
|
|
assert copied.literal_prepared_settings == baseline.literal_prepared_settings
|
|
assert json.loads(copied.literal_prepared_settings) == {"stream": True}
|
|
assert json.loads(copied.gcs_model_parameters) == {
|
|
**json.loads(baseline.gcs_model_parameters),
|
|
"tools": [{"type": "function", "function": {"name": "lookup"}}],
|
|
}
|
|
return copied
|
|
|
|
|
|
def queue_loggers(entered, release, uploads):
|
|
class VertexTransport:
|
|
async def _ensure_access_token_async(self, **kwargs):
|
|
entered.set()
|
|
await release.wait()
|
|
return "fixture-token", "fixture-project"
|
|
|
|
def _get_token_and_url(self, **kwargs):
|
|
return kwargs["auth_header"], None
|
|
|
|
class Transport:
|
|
async def post(self, url, **kwargs):
|
|
wire = {**kwargs, "json": json.loads(json.dumps(kwargs["json"]))} if "json" in kwargs else kwargs
|
|
uploads.append((url, wire))
|
|
return integration_response(url, {}, 202 if "datadog" in url else 200)
|
|
|
|
datadog = DataDogLogger.__new__(DataDogLogger)
|
|
CustomBatchLogger.__init__(datadog, batch_size=100, flush_lock=asyncio.Lock())
|
|
datadog.intake_url, datadog.DD_API_KEY, datadog.is_mock_mode = "https://datadog.invalid/logs", "test", False
|
|
datadog.async_client = Transport()
|
|
gcs = GCSBucketLogger.__new__(GCSBucketLogger)
|
|
CustomBatchLogger.__init__(gcs, batch_size=100)
|
|
gcs.log_queue = asyncio.Queue()
|
|
gcs.BUCKET_NAME, gcs.path_service_account_json = "fixture-bucket", None
|
|
gcs.vertex_instances = {"IAM_AUTH": VertexTransport()}
|
|
gcs.use_batched_logging = True
|
|
gcs.async_httpx_client = Transport()
|
|
literal = LiteralAILogger.__new__(LiteralAILogger)
|
|
CustomBatchLogger.__init__(literal, batch_size=100, flush_lock=asyncio.Lock())
|
|
literal.literalai_api_url, literal.headers = "https://literal.invalid", {}
|
|
literal.async_httpx_client = Transport()
|
|
return datadog, gcs, literal
|
|
|
|
|
|
async def integration_logging_queue_case(owners, *, literal_copy: Literal["direct", "envelope", "payload"] = "direct"):
|
|
entered, release = asyncio.Event(), asyncio.Event()
|
|
uploads = []
|
|
datadog, gcs, literal = queue_loggers(entered, release, uploads)
|
|
copy_payload = literal_copy == "payload"
|
|
|
|
class LiteralCallback(CustomLogger):
|
|
def __init__(self, delegate):
|
|
super().__init__()
|
|
self.delegate = delegate
|
|
self.calls = self.completed = 0
|
|
|
|
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
self.calls += 1
|
|
self.received = kwargs
|
|
self.forwarded = {
|
|
**kwargs,
|
|
"standard_logging_object": (
|
|
copy.deepcopy(kwargs["standard_logging_object"])
|
|
if copy_payload
|
|
else kwargs["standard_logging_object"]
|
|
),
|
|
}
|
|
await self.delegate.async_log_failure_event(self.forwarded, response_obj, start_time, end_time)
|
|
self.completed += 1
|
|
|
|
literal_callback = literal if literal_copy == "direct" else LiteralCallback(literal)
|
|
|
|
payload = create_dummy_standard_logging_payload()
|
|
payload.update(status="failure", error_str="x" * 10001)
|
|
messages, settings, metadata = payload["messages"], payload["model_parameters"], payload["metadata"]
|
|
completion = payload["response"]["choices"][0]["message"]
|
|
tools = [{"type": "function", "function": {"name": "lookup"}}]
|
|
settings["tools"] = tools
|
|
now = datetime.now()
|
|
logging = Logging(
|
|
model="fixture-model",
|
|
messages=messages,
|
|
stream=False,
|
|
call_type="acompletion",
|
|
start_time=now,
|
|
litellm_call_id="fixture-queue",
|
|
function_id="fixture",
|
|
dynamic_async_failure_callbacks=[datadog, gcs, literal_callback],
|
|
)
|
|
error = RuntimeError("fixture failure")
|
|
kwargs = logging.model_call_details
|
|
kwargs.update(standard_logging_object=payload, model="fixture-model", exception=error, end_time=now)
|
|
await integration_invoke(owners, logging.async_failure_handler, error, "fixture traceback", now, now)
|
|
if literal_copy != "direct":
|
|
assert literal_callback.calls == literal_callback.completed == 1
|
|
assert literal_callback.received is kwargs and literal_callback.forwarded is not kwargs
|
|
assert (literal_callback.forwarded["standard_logging_object"] is payload) is (not copy_payload)
|
|
assert len(datadog.log_queue) == gcs.log_queue.qsize() == len(literal.log_queue) == 1
|
|
assert kwargs["standard_logging_object"] is payload
|
|
assert payload["messages"] is messages and payload["model_parameters"] is settings
|
|
assert payload["error_str"].endswith("truncated by litellm, this logger does not support large content")
|
|
assert ("tools" in settings) is copy_payload
|
|
dd_snapshot = json.loads(datadog.log_queue[0]["message"])
|
|
assert dd_snapshot["model_parameters"]["tools"] == tools
|
|
queued = gcs.log_queue.get_nowait()
|
|
assert queued["payload"] is payload and queued["kwargs"] is kwargs and queued["response_obj"] is None
|
|
gcs.log_queue.put_nowait(queued)
|
|
generation = literal.log_queue[0]["generation"]
|
|
prepared_settings = json.dumps(generation["settings"], sort_keys=True)
|
|
assert "tools" not in generation["settings"] and generation["tools"] == tools
|
|
if copy_payload:
|
|
assert generation["settings"] is not settings and generation["tools"] is not tools
|
|
assert generation["messages"] is not messages and generation["messageCompletion"] is not completion
|
|
assert literal.log_queue[0]["metadata"] is not metadata
|
|
else:
|
|
assert generation["settings"] is settings and generation["tools"] is tools
|
|
assert generation["messages"] is messages and generation["messageCompletion"] is completion
|
|
assert literal.log_queue[0]["metadata"] is metadata
|
|
|
|
flush = asyncio.create_task(integration_invoke(owners, gcs.flush_queue))
|
|
try:
|
|
await entered.wait()
|
|
assert not uploads and not flush.done() and gcs.log_queue.empty()
|
|
messages[0]["content"] = "mutated before serialization"
|
|
settings["temperature"] = 0.25
|
|
completion["content"] = "late completion"
|
|
payload["messages"] = [{"role": "user", "content": "replacement field"}]
|
|
kwargs["standard_logging_object"] = {"replacement": True}
|
|
release.set()
|
|
await flush
|
|
finally:
|
|
release.set()
|
|
if not flush.done():
|
|
flush.cancel()
|
|
await asyncio.gather(flush, return_exceptions=True)
|
|
assert len(uploads) == 1
|
|
gcs_snapshot = json.loads(uploads[0][1]["data"])
|
|
assert gcs_snapshot["messages"] == payload["messages"]
|
|
assert gcs_snapshot["model_parameters"] == settings
|
|
assert "replacement" not in gcs_snapshot
|
|
await integration_invoke(owners, datadog.flush_queue)
|
|
await integration_invoke(owners, literal.flush_queue)
|
|
assert len(uploads) == 3 and not datadog.log_queue and not literal.log_queue
|
|
sent_dd = json.loads(gzip.decompress(uploads[1][1]["data"]))
|
|
assert json.loads(sent_dd[0]["message"]) == dd_snapshot
|
|
literal_wire = uploads[2][1]["json"]
|
|
sent_generation = literal_wire["variables"]["generation_0"]
|
|
assert sent_generation["messages"] != gcs_snapshot["messages"]
|
|
if copy_payload:
|
|
assert sent_generation["messages"] == dd_snapshot["messages"]
|
|
assert json.dumps(sent_generation["settings"], sort_keys=True) == prepared_settings
|
|
assert sent_generation["messageCompletion"] == dd_snapshot["response"]["choices"][0]["message"]
|
|
else:
|
|
assert sent_generation["messages"] == messages
|
|
assert sent_generation["settings"]["temperature"] == 0.25
|
|
assert sent_generation["messageCompletion"]["content"] == "late completion"
|
|
messages[0]["content"] = "after serialization"
|
|
assert sent_generation["messages"][0]["content"] == (
|
|
"Hello, world!" if copy_payload else "mutated before serialization"
|
|
)
|
|
assert dd_snapshot["messages"][0]["content"] == "Hello, world!"
|
|
return QueueObservation(
|
|
gcs_model_parameters=json.dumps(gcs_snapshot["model_parameters"], sort_keys=True),
|
|
datadog_snapshot=json.dumps(dd_snapshot, sort_keys=True),
|
|
literal_prepared_settings=prepared_settings,
|
|
)
|
|
|
|
|
|
@integration_callback_scope
|
|
async def real_crowdstrike_translator_identity(owners):
|
|
entered, release = asyncio.Event(), asyncio.Event()
|
|
calls = []
|
|
|
|
class Transport:
|
|
async def post(self, url, json, **kwargs):
|
|
calls.append(json)
|
|
entered.set()
|
|
await release.wait()
|
|
return integration_response(
|
|
url,
|
|
{
|
|
"result": {
|
|
"blocked": False,
|
|
"transformed": True,
|
|
"guard_output": {"messages": [{"role": "user", "content": "redacted"}]},
|
|
}
|
|
},
|
|
)
|
|
|
|
guardrail = CrowdStrikeAIDRHandler.__new__(CrowdStrikeAIDRHandler)
|
|
CustomGuardrail.__init__(guardrail, guardrail_name="fixture-crowdstrike", event_hook=GuardrailEventHooks.pre_call)
|
|
guardrail.api_base, guardrail.api_key, guardrail.fail_on_error = "https://crowdstrike.invalid", "test", True
|
|
guardrail.skip_system_message_in_guardrail = True
|
|
guardrail.async_handler = Transport()
|
|
system = {"role": "system", "content": "internal policy"}
|
|
user = {"role": "user", "content": "private text", "extra": {"retained": True}}
|
|
messages = [system, user]
|
|
data = {"model": "fixture-model", "messages": messages}
|
|
task = asyncio.create_task(
|
|
integration_invoke(
|
|
owners,
|
|
OpenAIChatCompletionsHandler().process_input_messages,
|
|
data,
|
|
guardrail,
|
|
)
|
|
)
|
|
try:
|
|
await entered.wait()
|
|
assert data["messages"] is messages and not task.done()
|
|
assert calls[0]["guard_input"]["messages"] == [{"role": "user", "content": "private text"}]
|
|
user["extra"]["during_http"] = True
|
|
release.set()
|
|
assert await task is data
|
|
finally:
|
|
release.set()
|
|
if not task.done():
|
|
task.cancel()
|
|
await asyncio.gather(task, return_exceptions=True)
|
|
assert data["messages"] is not messages
|
|
assert data["messages"][0] is system
|
|
assert data["messages"][1] is not user
|
|
assert data["messages"][1]["extra"] is user["extra"]
|
|
assert data["messages"][1]["content"] == "redacted"
|
|
assert user["content"] == "private text" and messages[1] is user
|
|
detached = copy.deepcopy(messages[1:])
|
|
inputs = {"texts": ["private text"], "structured_messages": detached}
|
|
control = await integration_invoke(owners, guardrail.apply_guardrail, inputs, {"messages": messages}, "request")
|
|
assert len(calls) == 2
|
|
assert control["structured_messages"] is detached and detached[0] is not user
|
|
assert control["texts"] == ["redacted"]
|
|
assert detached[0]["content"] == user["content"] == "private text"
|
|
|
|
|
|
@integration_callback_scope
|
|
async def real_rubrik_block_lifecycle(owners):
|
|
for input_type, populated in (("request", False), ("response", True)):
|
|
entered, release = asyncio.Event(), asyncio.Event()
|
|
moderation, uploads = [], []
|
|
|
|
class Transport:
|
|
def __init__(self, moderation, uploads, entered, release):
|
|
self.moderation, self.uploads = moderation, uploads
|
|
self.entered, self.release = entered, release
|
|
|
|
async def post(self, url, json, **kwargs):
|
|
if url.endswith("/batch"):
|
|
self.uploads.append(json)
|
|
return integration_response(url, {})
|
|
self.moderation.append(json)
|
|
self.entered.set()
|
|
await self.release.wait()
|
|
return integration_response(url, {"choices": [{"message": {"content": "blocked by policy"}}]})
|
|
|
|
rubrik = RubrikLogger.__new__(RubrikLogger)
|
|
CustomGuardrail.__init__(
|
|
rubrik,
|
|
guardrail_name="fixture-rubrik",
|
|
event_hook=GuardrailEventHooks.post_call,
|
|
flush_lock=asyncio.Lock(),
|
|
batch_size=100,
|
|
)
|
|
rubrik._periodic_flush_task = None
|
|
rubrik.sampling_rate, rubrik._headers = 1.0, {}
|
|
rubrik._dropped_since_warning, rubrik._last_drop_warning_time = 0, 0.0
|
|
rubrik.prompt_moderation_endpoint = "https://rubrik.invalid/before"
|
|
rubrik.response_moderation_endpoint = "https://rubrik.invalid/after"
|
|
rubrik.logging_endpoint = "https://rubrik.invalid/batch"
|
|
rubrik.moderation_client = rubrik.async_httpx_client = Transport(moderation, uploads, entered, release)
|
|
other = RubrikLogger.__new__(RubrikLogger)
|
|
CustomGuardrail.__init__(other, guardrail_name="fixture-other-rubrik", event_hook=GuardrailEventHooks.post_call)
|
|
logging = Logging(
|
|
model="fixture-model",
|
|
messages=[{"role": "user", "content": "original prompt"}],
|
|
stream=False,
|
|
call_type="acompletion",
|
|
start_time=datetime.now(),
|
|
litellm_call_id="fixture-correlation",
|
|
function_id="fixture",
|
|
)
|
|
details = logging.model_call_details
|
|
details.update(messages=logging.messages, model="fixture-model", litellm_call_id="fixture-correlation")
|
|
details["system"] = "system scaffold"
|
|
if populated:
|
|
details["standard_logging_object"] = create_dummy_standard_logging_payload()
|
|
messages = details["messages"]
|
|
request = {"model": "fixture-model", "litellm_call_id": "fixture-correlation", "messages": messages}
|
|
inputs = {"texts": ["original response"], "structured_messages": messages}
|
|
success = owners.prepare(rubrik.async_log_success_event, (details, None, None, None), awaited=True)
|
|
task = asyncio.create_task(
|
|
integration_invoke(owners, rubrik.apply_guardrail, inputs, request, input_type, logging)
|
|
)
|
|
try:
|
|
await entered.wait()
|
|
assert not task.done() and "_rubrik_logging_obj" not in request
|
|
assert "_rubrik_blocked" not in details
|
|
if input_type == "request":
|
|
assert moderation[0]["correlation_key"] == "fixture-correlation"
|
|
assert moderation[0]["messages"][0]["content"] == "original prompt"
|
|
else:
|
|
assert moderation[0]["request"]["messages"] is messages
|
|
assert moderation[0]["response"]["id"] == "fixture-correlation"
|
|
release.set()
|
|
with TestCase().assertRaises(ModifyResponseException) as caught:
|
|
await task
|
|
error = caught.exception
|
|
assert error.request_data is request and error.message == "blocked by policy"
|
|
assert request["_rubrik_logging_obj"] is logging and details["_rubrik_blocked"] is True
|
|
await integration_invoke(
|
|
owners, other.async_post_call_failure_hook, request, error, UserAPIKeyAuth(user_id="fixture-user")
|
|
)
|
|
assert request["_rubrik_logging_obj"] is logging and details["_rubrik_blocked"] is True
|
|
assert not other.log_queue and not rubrik.log_queue and not uploads
|
|
await integration_invoke(
|
|
owners, rubrik.async_post_call_failure_hook, request, error, UserAPIKeyAuth(user_id="fixture-user")
|
|
)
|
|
assert "_rubrik_logging_obj" not in request and details["_rubrik_blocked"] is True
|
|
assert len(rubrik.log_queue) == 1
|
|
queued = rubrik.log_queue[0]
|
|
assert queued["id"] == "fixture-correlation"
|
|
assert queued["response"] == "ModifyResponseException: blocked by policy"
|
|
assert queued["messages"][0] == {"role": "system", "content": "system scaffold"}
|
|
assert messages[0] == {"role": "user", "content": "original prompt"}
|
|
if populated:
|
|
base = details["standard_logging_object"]
|
|
assert queued["metadata"] is not base["metadata"]
|
|
assert queued["messages"][1] is not base["messages"][0]
|
|
assert isinstance(base["response"], dict)
|
|
else:
|
|
assert queued["messages"][1] is messages[0]
|
|
assert queued["status"] == "failure"
|
|
assert queued["metadata"]["user_api_key_user_id"] == "fixture-user"
|
|
pending = success.invoke()
|
|
success.close()
|
|
assert await pending is None
|
|
assert len(rubrik.log_queue) == 1 and rubrik.log_queue[0] is queued
|
|
await integration_invoke(owners, rubrik.flush_queue)
|
|
assert not rubrik.log_queue and len(uploads) == 1 and uploads[0][0] is queued
|
|
finally:
|
|
success.close()
|
|
release.set()
|
|
if not task.done():
|
|
task.cancel()
|
|
await asyncio.gather(task, return_exceptions=True)
|
|
await rubrik.aclose()
|
|
if rubrik._periodic_flush_task is not None:
|
|
await asyncio.gather(rubrik._periodic_flush_task, return_exceptions=True)
|
|
|
|
|
|
@integration_callback_scope
|
|
async def real_parallel_guardrail_snapshots(owners):
|
|
original_mode = litellm.safe_memory_mode
|
|
try:
|
|
for safe_memory_mode in (False, True):
|
|
litellm.safe_memory_mode = safe_memory_mode
|
|
await integration_parallel_snapshot_case(owners)
|
|
finally:
|
|
litellm.safe_memory_mode = original_mode
|
|
|
|
|
|
async def integration_parallel_snapshot_case(owners):
|
|
from litellm.caching.dual_cache import DualCache
|
|
from litellm.litellm_core_utils.core_helpers import independent_snapshot
|
|
from litellm.proxy.utils import ProxyLogging
|
|
|
|
class Uncopyable:
|
|
def __init__(self):
|
|
self.attempts = 0
|
|
self.observed = []
|
|
|
|
def __deepcopy__(self, memo):
|
|
self.attempts += 1
|
|
raise TypeError("fixture cannot be copied")
|
|
|
|
arrived, release, mutated = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
|
observations = {}
|
|
sentinel = Uncopyable()
|
|
live = {"messages": [{"role": "user", "content": "original"}], "uncopyable": sentinel}
|
|
raw = independent_snapshot(live)
|
|
assert raw is not live and raw["messages"][0] is not live["messages"][0]
|
|
assert raw["uncopyable"] is sentinel and sentinel.attempts == 1
|
|
live["messages"][0]["content"] = "masked"
|
|
|
|
class Inspect(CustomGuardrail):
|
|
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
|
observations[self.guardrail_name] = data
|
|
if len(observations) == 4:
|
|
arrived.set()
|
|
await release.wait()
|
|
if self.guardrail_name == "fixture-live-writer":
|
|
data["messages"][0]["content"] = "shared mutation"
|
|
mutated.set()
|
|
await mutated.wait()
|
|
if self.scan_raw_request:
|
|
assert data["messages"][0]["content"] == "original"
|
|
data["messages"][0]["content"] = self.guardrail_name
|
|
assert data["uncopyable"] is sentinel
|
|
data["uncopyable"].observed.append(self.guardrail_name)
|
|
else:
|
|
assert data is live and data["messages"][0]["content"] == "shared mutation"
|
|
return {"discarded": self.guardrail_name}
|
|
|
|
guardrails = tuple(
|
|
Inspect(
|
|
guardrail_name=name,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
default_on=True,
|
|
run_in_parallel=True,
|
|
scan_raw_request="raw" in name,
|
|
)
|
|
for name in ("fixture-live-writer", "fixture-live-reader", "fixture-raw-a", "fixture-raw-b")
|
|
)
|
|
proxy = ProxyLogging.__new__(ProxyLogging)
|
|
proxy.call_details = {"user_api_key_cache": DualCache()}
|
|
task = asyncio.create_task(
|
|
integration_invoke(
|
|
owners, proxy._run_parallel_pre_call_guardrails, guardrails, live, raw, UserAPIKeyAuth(), "acompletion"
|
|
)
|
|
)
|
|
try:
|
|
await arrived.wait()
|
|
assert not task.done()
|
|
assert observations["fixture-live-writer"] is observations["fixture-live-reader"] is live
|
|
first, second = observations["fixture-raw-a"], observations["fixture-raw-b"]
|
|
assert first is not second and first is not raw and second is not raw
|
|
assert first["messages"][0] is not second["messages"][0]
|
|
assert first["messages"][0] is not raw["messages"][0]
|
|
assert first["uncopyable"] is second["uncopyable"] is raw["uncopyable"] is sentinel
|
|
assert sentinel.attempts == 3 and not sentinel.observed
|
|
release.set()
|
|
assert await task is None
|
|
assert raw["messages"] == [{"role": "user", "content": "original"}]
|
|
assert live["messages"][0]["content"] == "shared mutation"
|
|
assert set(sentinel.observed) == {"fixture-raw-a", "fixture-raw-b"} and len(sentinel.observed) == 2
|
|
assert "discarded" not in live
|
|
assert first["messages"][0]["content"] == "fixture-raw-a"
|
|
assert second["messages"][0]["content"] == "fixture-raw-b"
|
|
assert all(guardrail._pre_call_hook_already_ran(live) for guardrail in guardrails if guardrail.scan_raw_request)
|
|
finally:
|
|
release.set()
|
|
mutated.set()
|
|
if not task.done():
|
|
task.cancel()
|
|
await asyncio.gather(task, return_exceptions=True)
|
|
|
|
|
|
@integration_callback_scope
|
|
async def real_purview_sync_background(owners):
|
|
entered, release = threading.Event(), threading.Event()
|
|
calls, workers = [], []
|
|
main_thread = threading.get_ident()
|
|
|
|
class Transport:
|
|
async def post(self, url, **kwargs):
|
|
workers.append(threading.current_thread())
|
|
calls.append((url, kwargs))
|
|
if url.endswith("/token"):
|
|
entered.set()
|
|
assert release.wait(5), "background audit was not released"
|
|
return integration_response(url, {"access_token": "fixture-token", "expires_in": 3600})
|
|
assert kwargs["headers"]["Authorization"] == "Bearer fixture-token"
|
|
if url.endswith("/compute"):
|
|
return integration_response(url, {}, headers={"etag": "fixture-etag"})
|
|
assert url.endswith("/processContent")
|
|
assert kwargs["headers"]["If-None-Match"] == "fixture-etag"
|
|
return integration_response(
|
|
url, {"policyActions": [{"action": "restrictAccess", "restrictionAction": "block"}]}
|
|
)
|
|
|
|
purview = MicrosoftPurviewDLPGuardrail.__new__(MicrosoftPurviewDLPGuardrail)
|
|
CustomGuardrail.__init__(purview, guardrail_name="fixture-purview", event_hook=GuardrailEventHooks.logging_only)
|
|
purview.async_handler = Transport()
|
|
purview.tenant_id, purview.client_id, purview.client_secret = "fixture-tenant", "fixture-client", "test"
|
|
purview.purview_app_name, purview.user_id_field, purview.guardrail_provider = (
|
|
"fixture",
|
|
"user_id",
|
|
"microsoft_purview",
|
|
)
|
|
purview._token_cache, purview._scope_cache = None, OrderedDict()
|
|
purview._scope_cache_maxsize, purview._cache_lock = 1000, threading.Lock()
|
|
metadata = {"user_api_key_user_id": "fixture-user"}
|
|
kwargs = {
|
|
"messages": [{"role": "user", "content": "prompt at dispatch"}],
|
|
"litellm_params": {"metadata": metadata},
|
|
"litellm_call_id": "before-http",
|
|
}
|
|
result = ModelResponse(model="fixture-model", choices=[{"message": {"role": "assistant", "content": "before"}}])
|
|
owner = owners.prepare(purview.logging_hook, (kwargs, result, "completion"), awaited=False)
|
|
try:
|
|
returned = await asyncio.to_thread(owner.invoke)
|
|
owner.close()
|
|
assert returned[0] is kwargs and returned[1] is result
|
|
assert await asyncio.to_thread(entered.wait, 5)
|
|
assert len(calls) == 1 and workers[0].ident != main_thread and workers[0].daemon
|
|
assert workers[0].is_alive()
|
|
kwargs["messages"][0]["content"] = "too late for prompt extraction"
|
|
kwargs["litellm_call_id"] = "after-http"
|
|
result.choices[0].message.content = "response mutated while audit waits"
|
|
release.set()
|
|
await asyncio.to_thread(workers[0].join, 5)
|
|
assert not workers[0].is_alive() and all(worker is workers[0] for worker in workers)
|
|
assert len(calls) == 4
|
|
entries = [call[1]["json"]["contentToProcess"] for call in calls[2:]]
|
|
assert [entry["activityMetadata"]["activity"] for entry in entries] == ["uploadText", "downloadText"]
|
|
assert entries[0]["contentEntries"][0]["content"]["data"] == "prompt at dispatch"
|
|
assert entries[1]["contentEntries"][0]["content"]["data"] == "response mutated while audit waits"
|
|
assert all(entry["contentEntries"][0]["correlationId"] == "after-http" for entry in entries)
|
|
assert kwargs["litellm_params"]["metadata"] is metadata
|
|
info = kwargs["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(info) == 2 and all(item["guardrail_status"] == "guardrail_intervened" for item in info)
|
|
assert all(item["end_time"] >= item["start_time"] and item["duration"] >= 0 for item in info)
|
|
finally:
|
|
owner.close()
|
|
release.set()
|
|
if workers:
|
|
await asyncio.to_thread(workers[0].join, 5)
|