great-ai/great_ai/views/trace.py
Andras Schmelczer c55eba2077
All checks were successful
Publish documentation / publish (push) Successful in 51s
Check / Test on Python 3.10 (push) Successful in 1m1s
Check / Lint, format & type checks (push) Successful in 1m9s
Check / Test on Python 3.11 (push) Successful in 51s
Check / Test on Python 3.12 (push) Successful in 58s
Check / Test on Python 3.13 (push) Successful in 55s
Add max deps and bump old ones
2026-06-06 22:29:16 +01:00

108 lines
3.5 KiB
Python

from pprint import pformat
from typing import Any, Dict, Generic, List, Optional, TypeVar
from pydantic import ConfigDict
from .hashable_base_model import HashableBaseModel
from .model import Model
T = TypeVar("T")
class Trace(HashableBaseModel, Generic[T]):
"""Universal structure for storing prediction traces and training data.
Attributes:
trace_id: UUID4 identifier for uniquely referring to a trace.
created: Timestamp of its (original) construction.
original_execution_time_ms: Wall-time elapsed while its generating
TracingContext was alive.
logged_values: Values persisted through using `@parameter` or `log_metric()`.
models: Marks left by each encountered `@use_model` decorated function.
exception: Exception description if any was encountered.
output: Return value of the function wrapped by GreatAI.
feedback: Feedback obtained using the REST API of `add_ground_truth`.
tags: Tags used for filtering traces. Contains the name of the original
function, value of `ENVIRONMENT`, its split if has any, and either
`ground_truth` or `online` depending on the origin of the Trace.
"""
trace_id: str
created: str
original_execution_time_ms: float
logged_values: Dict[str, Any]
models: List[Model]
exception: Optional[str] = None
output: Optional[T] = None
feedback: Any = None
tags: List[str]
model_config = ConfigDict(extra="ignore")
@property
def input(self) -> Any:
return (
self.logged_values["input"]
if list(self.logged_values.keys()) == ["input"]
else self.logged_values
)
@property
def models_flat(self) -> str:
return ", ".join(f"{m.key}:{m.version}" for m in self.models)
@property
def output_flat(self) -> str:
return pformat(self.output, indent=2, compact=True)
@property
def exception_flat(self) -> str:
return (
"null"
if self.exception is None
else pformat(self.exception, indent=2, compact=True)
)
@property
def feedback_flat(self) -> str:
return (
"null"
if self.feedback is None
else pformat(self.feedback, indent=2, compact=True)
)
@property
def tags_flat(self) -> str:
return ",\n".join(self.tags)
def to_flat_dict(self, include_original: bool = True) -> Dict[str, Any]:
return {
**(
self.model_dump()
if include_original
else {
"trace_id": self.trace_id,
"created": self.created,
"original_execution_time_ms": self.original_execution_time_ms,
}
),
**{
k: (
v
if (isinstance(v, float) or isinstance(v, int))
else pformat(v, indent=2, compact=True)
)
for k, v in self.logged_values.items()
},
"models_flat": self.models_flat,
"exception_flat": self.exception_flat,
"output_flat": self.output_flat,
"feedback_flat": self.feedback_flat,
"tags_flat": self.tags_flat,
}
def __repr__(self) -> str:
formatted = pformat(self.model_dump(), indent=2, compact=True).replace(
"{ ", "{", 1
)
return f"Trace[{type(self.output).__name__}]({formatted})"