feat(langfuse.py): Allow for individual call message/response redaction

This commit is contained in:
Alex Epstein 2024-05-12 22:38:29 -04:00
parent c5ca2619f9
commit f6e46a38d0
3 changed files with 54 additions and 6 deletions

View file

@ -213,8 +213,20 @@ chat(messages)
## Redacting Messages, Response Content from Langfuse Logging
### Redact Messages and Responses from all Langfuse Logging
Set `litellm.turn_off_message_logging=True` This will prevent the messages and responses from being logged to langfuse, but request metadata will still be logged.
### Redact Messages and Responses from specific Langfuse Logging
In the metadata typically passed for text completion or embedding calls you can set specific keys to mask the messages and responses for this call.
Setting `mask_input` to `True` will mask the input from being logged for this call
Setting `mask_output` to `True` will make the output from being logged for this call.
Be aware that if you are continuing an existing trace, and you set `update_trace_keys` to include either `input` or `output` and you set the corresponding `mask_input` or `mask_output`, then that trace will have its existing input and/or output replaced with a redacted message.
## Troubleshooting & Errors
### Data not getting logged to Langfuse ?
- Ensure you're on the latest version of langfuse `pip install langfuse -U`. The latest version allows litellm to log JSON input/outputs to langfuse

View file

@ -322,6 +322,8 @@ class LangFuseLogger:
existing_trace_id = clean_metadata.pop("existing_trace_id", None)
update_trace_keys = clean_metadata.pop("update_trace_keys", [])
debug = clean_metadata.pop("debug_langfuse", None)
mask_input = clean_metadata.pop("mask_input", False)
mask_output = clean_metadata.pop("mask_output", False)
if trace_name is None and existing_trace_id is None:
# just log `litellm-{call_type}` as the trace name
@ -349,15 +351,15 @@ class LangFuseLogger:
# Special keys that are found in the function arguments and not the metadata
if "input" in update_trace_keys:
trace_params["input"] = input
trace_params["input"] = input if not mask_input else "redacted-by-litellm"
if "output" in update_trace_keys:
trace_params["output"] = output
trace_params["output"] = output if not mask_output else "redacted-by-litellm"
else: # don't overwrite an existing trace
trace_params = {
"id": trace_id,
"name": trace_name,
"session_id": session_id,
"input": input,
"input": input if not mask_input else "redacted-by-litellm",
"version": clean_metadata.pop(
"trace_version", clean_metadata.get("version", None)
), # If provided just version, it will applied to the trace as well, if applied a trace version it will take precedence
@ -373,7 +375,7 @@ class LangFuseLogger:
if level == "ERROR":
trace_params["status_message"] = output
else:
trace_params["output"] = output
trace_params["output"] = output if not mask_output else "redacted-by-litellm"
if debug == True or (isinstance(debug, str) and debug.lower() == "true"):
if "metadata" in trace_params:
@ -463,8 +465,8 @@ class LangFuseLogger:
"end_time": end_time,
"model": kwargs["model"],
"model_parameters": optional_params,
"input": input,
"output": output,
"input": input if not mask_input else "redacted-by-litellm",
"output": output if not mask_output else "redacted-by-litellm",
"usage": usage,
"metadata": clean_metadata,
"level": level,

View file

@ -228,6 +228,40 @@ async def test_langfuse_logging_without_request_response(stream, langfuse_client
pytest.fail(f"An exception occurred - {e}")
@pytest.mark.asyncio
async def test_langfuse_masked_input_output(langfuse_client):
"""
Test that creates a trace with masked input and output
"""
import uuid
for mask_value in [True, False]:
_unique_trace_name = f"litellm-test-{str(uuid.uuid4())}"
litellm.set_verbose = True
litellm.success_callback = ["langfuse"]
response = await create_async_task(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "This is a test"}],
metadata={"trace_id": _unique_trace_name, "mask_input": mask_value, "mask_output": mask_value},
mock_response="This is a test response"
)
print(response)
expected_input = "redacted-by-litellm" if mask_value else {'messages': [{'content': 'This is a test', 'role': 'user'}]}
expected_output = "redacted-by-litellm" if mask_value else {'content': 'This is a test response', 'role': 'assistant'}
langfuse_client.flush()
await asyncio.sleep(2)
# get trace with _unique_trace_name
trace = langfuse_client.get_trace(id=_unique_trace_name)
generations = list(
reversed(langfuse_client.get_generations(trace_id=_unique_trace_name).data)
)
assert trace.input == expected_input
assert trace.output == expected_output
assert generations[0].input == expected_input
assert generations[0].output == expected_output
@pytest.mark.asyncio
async def test_langfuse_logging_metadata(langfuse_client):
"""