From d155757efa3b0dc8856561c946481a73b00bb8c4 Mon Sep 17 00:00:00 2001 From: Bharat Ramanathan Date: Thu, 27 Feb 2025 13:06:17 +0530 Subject: [PATCH] feat: fix token usage calls and tracing config lint --- api/core/ops/weave_trace/entities/weave_trace_entity.py | 8 ++++---- api/core/ops/weave_trace/weave_trace.py | 7 +++++-- api/tasks/ops_trace_task.py | 5 ++++- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/api/core/ops/weave_trace/entities/weave_trace_entity.py b/api/core/ops/weave_trace/entities/weave_trace_entity.py index e8a36c5fb8..1f6b907e74 100644 --- a/api/core/ops/weave_trace/entities/weave_trace_entity.py +++ b/api/core/ops/weave_trace/entities/weave_trace_entity.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List, Optional, Union +from typing import Any, Optional, Union from pydantic import BaseModel, Field, field_validator from pydantic_core.core_schema import ValidationInfo @@ -19,9 +19,9 @@ class WeaveMultiModel(BaseModel): class WeaveTraceModel(WeaveTokenUsage, WeaveMultiModel): id: str = Field(..., description="ID of the trace") op: str = Field(..., description="Name of the operation") - inputs: Optional[Union[str, Dict[str, Any], List, None]] = Field(None, description="Inputs of the trace") - outputs: Optional[Union[str, Dict[str, Any], List, None]] = Field(None, description="Outputs of the trace") - attributes: Optional[Union[str, Dict[str, Any], List, None]] = Field( + inputs: Optional[Union[str, dict[str, Any], list, None]] = Field(None, description="Inputs of the trace") + outputs: Optional[Union[str, dict[str, Any], list, None]] = Field(None, description="Outputs of the trace") + attributes: Optional[Union[str, dict[str, Any], list, None]] = Field( None, description="Metadata and attributes associated with trace" ) exception: Optional[str] = Field(None, description="Exception message of the trace") diff --git a/api/core/ops/weave_trace/weave_trace.py b/api/core/ops/weave_trace/weave_trace.py index 6dd9a126cc..33547f2092 100644 --- a/api/core/ops/weave_trace/weave_trace.py +++ b/api/core/ops/weave_trace/weave_trace.py @@ -92,6 +92,7 @@ class WeaveDataTrace(BaseTraceInstance): op=str(TraceTaskName.MESSAGE_TRACE.value), inputs=dict(trace_info.workflow_run_inputs), outputs=dict(trace_info.workflow_run_outputs), + total_tokens=trace_info.total_tokens, attributes=message_attributes, exception=trace_info.error, file_list=[], @@ -355,8 +356,10 @@ class WeaveDataTrace(BaseTraceInstance): attributes=attributes, exception=trace_info.error, ) - - self.start_call(tool_run, parent_run_id=trace_info.message_id) + message_id = trace_info.message_id or trace_info.conversation_id + message_id = message_id or None + self.start_call(tool_run, parent_run_id=message_id) + self.finish_call(tool_run) def generate_name_trace(self, trace_info: GenerateNameTraceInfo): attributes = trace_info.metadata diff --git a/api/tasks/ops_trace_task.py b/api/tasks/ops_trace_task.py index bb3b9e17ea..bbff288ec5 100644 --- a/api/tasks/ops_trace_task.py +++ b/api/tasks/ops_trace_task.py @@ -46,7 +46,10 @@ def process_trace_tasks(file_info): trace_info = trace_type(**trace_info) trace_instance.trace(trace_info) logging.info(f"Processing trace tasks success, app_id: {app_id}") - except Exception: + except Exception as e: + logging.info( + f"error:\n\n\n{e}\n\n\n\n", + ) failed_key = f"{OPS_TRACE_FAILED_KEY}_{app_id}" redis_client.incr(failed_key) logging.info(f"Processing trace tasks failed, app_id: {app_id}")