qixinbo
66 lines
1.7 KiB
Python
66 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, Dict, Mapping, Optional
|
|
|
|
|
|
def sanitize_attributes(attributes: Optional[Mapping[str, Any]]) -> Dict[str, Any]:
|
|
if not attributes:
|
|
return {}
|
|
normalized: Dict[str, Any] = {}
|
|
for key, value in attributes.items():
|
|
if value is None:
|
|
continue
|
|
name = str(key).strip()
|
|
if not name:
|
|
continue
|
|
if isinstance(value, (str, int, float, bool)):
|
|
normalized[name] = value
|
|
continue
|
|
normalized[name] = str(value)
|
|
return normalized
|
|
|
|
|
|
def build_chat_trace_attributes(
|
|
*,
|
|
session_id: str,
|
|
project_id: Optional[int],
|
|
model_id: Optional[str],
|
|
route_mode: str,
|
|
source: str,
|
|
knowledge_base_id: Optional[str],
|
|
) -> Dict[str, Any]:
|
|
return sanitize_attributes(
|
|
{
|
|
"session_id": session_id,
|
|
"project_id": project_id,
|
|
"model_id": model_id,
|
|
"route_mode": route_mode,
|
|
"source": source,
|
|
"knowledge_base_id": knowledge_base_id,
|
|
"component": "chat_stream",
|
|
}
|
|
)
|
|
|
|
|
|
def build_usage_attributes(usage: Optional[Mapping[str, Any]]) -> Dict[str, Any]:
|
|
if not usage:
|
|
return {}
|
|
return sanitize_attributes(
|
|
{
|
|
"usage.prompt_tokens": usage.get("prompt_tokens"),
|
|
"usage.completion_tokens": usage.get("completion_tokens"),
|
|
"usage.total_tokens": usage.get("total_tokens"),
|
|
}
|
|
)
|
|
|
|
|
|
def build_error_attributes(exc: Exception, *, stage: str) -> Dict[str, Any]:
|
|
return sanitize_attributes(
|
|
{
|
|
"error": True,
|
|
"error.stage": stage,
|
|
"error.type": exc.__class__.__name__,
|
|
"error.message": str(exc),
|
|
}
|
|
)
|