Skip to content

Commit

Permalink
feat (backend): create remove_trace_prefix utility function
Browse files Browse the repository at this point in the history
  • Loading branch information
aybruhm committed Nov 26, 2024
1 parent 3e93e11 commit 36d311f
Showing 1 changed file with 46 additions and 1 deletion.
47 changes: 46 additions & 1 deletion agenta-backend/agenta_backend/utils/traces.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,59 @@
import logging
import traceback
from copy import deepcopy
from typing import Any, Dict, Union
from collections import OrderedDict
from typing import Any, Dict, Union, Optional


logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)


def remove_trace_prefix(
mapping_dict: Optional[Dict] = None, settings_values: Optional[Dict] = None
) -> Dict:
"""
Modify the values of the mapping dictionary to remove 'trace.' prefix if it exists.
Args:
mapping_dict (Optional[Dict]): A dictionary containing the mapping values.
settings_values (Optional[Dict]): A dictionary with keys like "answer_key",
"contexts_key", "question_key" to override
specific mapping values.
Returns:
Dict: A dictionary with the 'trace.' prefix removed from any string values.
Raises:
ValueError: If neither `mapping_dict` nor `settings_values` is provided.
"""

if mapping_dict is None and settings_values is None:
raise ValueError("No mapping dictionary or settings values provided")

# Determine which dictionary to use
if settings_values:
mapping_values = {
"answer_key": settings_values.get("answer_key"),
"contexts_key": settings_values.get("contexts_key"),
"question_key": settings_values.get("question_key"),
}
elif mapping_dict:
mapping_values = mapping_dict
else:
mapping_values = {}

# Update the mapping by removing the 'trace.' prefix from string values
updated_mapping_dict = {
key: value.replace("trace.", "") if isinstance(value, str) else value
for key, value in mapping_values.items()
if value is not None
}

return updated_mapping_dict


def _make_spans_id_tree(trace):
"""
Creates spans tree (id only) from flat spans list
Expand Down

0 comments on commit 36d311f

Please sign in to comment.