feat: fix token usage calls and tracing config lint

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")

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

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}")