Protect trace fields from logged input collisions

This commit is contained in:
Andras Schmelczer 2026-10-27 16:51:54 +00:00
commit 1d088669d1
2 changed files with 29 additions and 0 deletions

View file

@ -76,6 +76,7 @@ class Trace(HashableBaseModel, Generic[T]):
return ",\n".join(self.tags) return ",\n".join(self.tags)
def to_flat_dict(self, include_original: bool = True) -> Dict[str, Any]: def to_flat_dict(self, include_original: bool = True) -> Dict[str, Any]:
reserved = set(type(self).model_fields) | {"_id"}
return { return {
**( **(
self.model_dump() self.model_dump()
@ -93,6 +94,7 @@ class Trace(HashableBaseModel, Generic[T]):
else pformat(v, indent=2, compact=True) else pformat(v, indent=2, compact=True)
) )
for k, v in self.logged_values.items() for k, v in self.logged_values.items()
if k not in reserved
}, },
"models_flat": self.models_flat, "models_flat": self.models_flat,
"exception_flat": self.exception_flat, "exception_flat": self.exception_flat,

View file

@ -1,3 +1,4 @@
from datetime import datetime
from typing import Any, Dict from typing import Any, Dict
from unittest.mock import patch from unittest.mock import patch
@ -7,6 +8,32 @@ from pymongo import MongoClient
from pymongo.synchronous.collection import Collection 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: def test_mongo_update_uses_valid_replacement() -> None:
trace = Trace( trace = Trace(
trace_id="id", trace_id="id",