diff --git a/letta/services/telemetry_manager.py b/letta/services/telemetry_manager.py index 213bdb09..bf8efd80 100644 --- a/letta/services/telemetry_manager.py +++ b/letta/services/telemetry_manager.py @@ -1,6 +1,7 @@ from letta.helpers.json_helpers import json_dumps, json_loads from letta.helpers.singleton import singleton from letta.orm.provider_trace import ProviderTrace as ProviderTraceModel +from letta.otel.tracing import trace_method from letta.schemas.provider_trace import ProviderTrace as PydanticProviderTrace from letta.schemas.provider_trace import ProviderTraceCreate from letta.schemas.step import Step as PydanticStep @@ -10,7 +11,9 @@ from letta.utils import enforce_types class TelemetryManager: + @enforce_types + @trace_method async def get_provider_trace_by_step_id_async( self, step_id: str, @@ -21,6 +24,7 @@ class TelemetryManager: return provider_trace.to_pydantic() @enforce_types + @trace_method async def create_provider_trace_async(self, actor: PydanticUser, provider_trace_create: ProviderTraceCreate) -> PydanticProviderTrace: async with db_registry.async_session() as session: provider_trace = ProviderTraceModel(**provider_trace_create.model_dump()) @@ -31,10 +35,13 @@ class TelemetryManager: if provider_trace_create.response_json: response_json_str = json_dumps(provider_trace_create.response_json) provider_trace.response_json = json_loads(response_json_str) - await provider_trace.create_async(session, actor=actor) - return provider_trace.to_pydantic() + await provider_trace.create_async(session, actor=actor, no_commit=True, no_refresh=True) + pydantic_provider_trace = provider_trace.to_pydantic() + await session.commit() + return pydantic_provider_trace @enforce_types + @trace_method def create_provider_trace(self, actor: PydanticUser, provider_trace_create: ProviderTraceCreate) -> PydanticProviderTrace: with db_registry.session() as session: provider_trace = ProviderTraceModel(**provider_trace_create.model_dump())