Coverage for haystack/tracing/logging_tracer.py: 100%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import contextlib 

6import dataclasses 

7import threading 

8from collections.abc import Iterator 

9from typing import Any 

10 

11from haystack import logging 

12from haystack.tracing import Span, Tracer 

13from haystack.tracing.utils import coerce_tag_value 

14 

15logger = logging.getLogger(__name__) 

16 

17RESET_COLOR = "\033[0m" 

18 

19 

20@dataclasses.dataclass 

21class LoggingSpan(Span): 

22 operation_name: str 

23 tags: dict[str, Any] = dataclasses.field(default_factory=dict) 

24 

25 def set_tag(self, key: str, value: Any) -> None: 

26 """ 

27 Set a single tag on the span. 

28 

29 :param key: the name of the tag. 

30 :param value: the value of the tag. 

31 """ 

32 self.tags[key] = value 

33 

34 

35class LoggingTracer(Tracer): 

36 """ 

37 A simple tracer that logs the operation name and tags of a span. 

38 """ 

39 

40 def __init__(self, tags_color_strings: dict[str, str] | None = None) -> None: 

41 """ 

42 Initialize the LoggingTracer. 

43 

44 :param tags_color_strings: 

45 A dictionary that maps tag names to color strings that should be used when logging the tags. 

46 The color strings should be in the format of 

47 [ANSI escape codes](https://en.wikipedia.org/wiki/ANSI_escape_code#Colors). 

48 For example, to color the tag "haystack.component.input" in red, you would pass 

49 `tags_color_strings={"haystack.component.input": "\x1b[1;31m"}`. 

50 """ 

51 

52 self.tags_color_strings = tags_color_strings or {} 

53 # Spans can be created and closed from parallel worker threads (e.g. one span per tool call in the Agent). 

54 # A span's operation name and its tags are emitted as separate log records, so we serialize the emission to 

55 # keep each span's records contiguous instead of interleaved with those of concurrently-closing spans. 

56 self._emit_lock = threading.Lock() 

57 

58 @contextlib.contextmanager 

59 def trace( 

60 self, 

61 operation_name: str, 

62 tags: dict[str, Any] | None = None, 

63 parent_span: Span | None = None, # noqa: ARG002 

64 ) -> Iterator[Span]: 

65 """ 

66 Trace the execution of a block of code. 

67 

68 :param operation_name: the name of the operation being traced. 

69 :param tags: tags to apply to the newly created span. 

70 :param parent_span: the parent span to use for the newly created span. Not used in this simple tracer. 

71 :returns: the newly created span. 

72 """ 

73 

74 custom_span = LoggingSpan(operation_name, tags=tags or {}) 

75 

76 try: 

77 yield custom_span 

78 except Exception as e: # noqa: TRY203 

79 raise e 

80 # we make sure to log the operation name and tags of the span when the context manager exits 

81 # both in case of success and error 

82 finally: 

83 operation_name = custom_span.operation_name 

84 tags = custom_span.tags or {} 

85 with self._emit_lock: 

86 logger.debug("Operation: {operation_name}", operation_name=operation_name) 

87 for tag_name, tag_value in tags.items(): 

88 color_string = self.tags_color_strings.get(tag_name, "") 

89 coerced_value = coerce_tag_value(tag_value) 

90 logger.debug( 

91 color_string + "{tag_name}={tag_value}" + RESET_COLOR, 

92 tag_name=tag_name, 

93 tag_value=coerced_value, 

94 ) 

95 

96 def current_span(self) -> Span | None: 

97 """Return the current active span, if any.""" 

98 # we don't store spans in this simple tracer 

99 return None