Protect trace fields from logged input collisions
This commit is contained in:
parent
1bfdf05081
commit
1d088669d1
2 changed files with 29 additions and 0 deletions
|
|
@ -76,6 +76,7 @@ class Trace(HashableBaseModel, Generic[T]):
|
|||
return ",\n".join(self.tags)
|
||||
|
||||
def to_flat_dict(self, include_original: bool = True) -> Dict[str, Any]:
|
||||
reserved = set(type(self).model_fields) | {"_id"}
|
||||
return {
|
||||
**(
|
||||
self.model_dump()
|
||||
|
|
@ -93,6 +94,7 @@ class Trace(HashableBaseModel, Generic[T]):
|
|||
else pformat(v, indent=2, compact=True)
|
||||
)
|
||||
for k, v in self.logged_values.items()
|
||||
if k not in reserved
|
||||
},
|
||||
"models_flat": self.models_flat,
|
||||
"exception_flat": self.exception_flat,
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from datetime import datetime
|
||||
from typing import Any, Dict
|
||||
from unittest.mock import patch
|
||||
|
||||
|
|
@ -7,6 +8,32 @@ from pymongo import MongoClient
|
|||
from pymongo.synchronous.collection import Collection
|
||||
|
||||
|
||||
def test_flattening_preserves_reserved_trace_fields() -> None:
|
||||
trace = Trace(
|
||||
trace_id="actual-id",
|
||||
created=datetime(2026, 1, 1).isoformat(),
|
||||
original_execution_time_ms=0,
|
||||
models=[],
|
||||
tags=["online"],
|
||||
output="prediction",
|
||||
logged_values={
|
||||
"trace_id": "input-id",
|
||||
"created": "input-created",
|
||||
"output": 123,
|
||||
"tags": ["input"],
|
||||
"models": "input",
|
||||
"_id": "input-id",
|
||||
"feature": 5,
|
||||
},
|
||||
)
|
||||
flat = trace.to_flat_dict()
|
||||
for name, value in trace.model_dump().items():
|
||||
assert flat[name] == value
|
||||
assert "_id" not in flat
|
||||
assert flat["feature"] == 5
|
||||
assert Trace.model_validate(flat) == trace
|
||||
|
||||
|
||||
def test_mongo_update_uses_valid_replacement() -> None:
|
||||
trace = Trace(
|
||||
trace_id="id",
|
||||
|
|
|
|||
Loading…
Reference in a new issue