Coverage for src / agent / cli / utils.py: 74%

1# Copyright 2025-2026 Microsoft Corporation 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15"""Utility functions for CLI module.""" 

16 

17import logging 

18import os 

19import platform 

20import sys 

21from typing import Any 

22 

23from rich.console import Console 

24 

25from agent.config.schema import AgentSettings 

26 

27logger = logging.getLogger(__name__) 

28 

29 

30def get_console() -> Console: 

31 """Create Rich console with proper encoding for Windows. 

32 

33 On Windows in non-interactive mode (subprocess, pipe, etc.), the default 

34 encoding is often CP1252 which cannot handle Unicode characters. This 

35 function detects such cases and forces UTF-8 encoding when possible. 

36 

37 Returns: 

38 Console: Configured Rich console instance 

39 """ 

40 if platform.system() == "Windows" and not sys.stdout.isatty(): 

41 # Running in non-interactive mode (subprocess, pipe, etc) 

42 # Try to use UTF-8 if available 

43 try: 

44 import locale 

45 

46 encoding = locale.getpreferredencoding() or "" 

47 if "utf" not in encoding.lower(): 

48 # Force UTF-8 for better Unicode support 

49 os.environ["PYTHONIOENCODING"] = "utf-8" 

50 # Create console with legacy Windows mode disabled 

51 return Console(force_terminal=True, legacy_windows=False) 

52 else: 

53 return Console() 

54 except Exception: 

55 # Fallback to safe ASCII mode if encoding detection fails 

56 return Console(legacy_windows=True, safe_box=True) 

57 else: 

58 # Normal interactive mode or non-Windows 

59 return Console() 

60 

61 

62def hide_connection_string_if_otel_disabled(config: AgentSettings) -> str | None: 

63 """Conditionally hide Azure Application Insights connection string. 

64 

65 The agent_framework auto-enables OpenTelemetry when it sees 

66 APPLICATIONINSIGHTS_CONNECTION_STRING in the environment, which causes 

67 1-3s exit lag from daemon threads flushing metrics. 

68 

69 This helper hides the connection string ONLY when telemetry is disabled, 

70 allowing users who explicitly enable OTEL to still use it. 

71 

72 Args: 

73 config: Loaded AgentConfig (must be loaded first to check enable_otel) 

74 

75 Returns: 

76 The connection string if it was hidden, None otherwise 

77 

78 Example: 

79 >>> config = AgentConfig.from_combined() 

80 >>> saved = hide_connection_string_if_otel_disabled(config) 

81 >>> # ... create agent ... 

82 >>> if saved: 

83 ... os.environ["APPLICATIONINSIGHTS_CONNECTION_STRING"] = saved 

84 """ 

85 should_enable_otel = config.enable_otel and config.enable_otel_explicit 

86 

87 if not should_enable_otel and config.applicationinsights_connection_string: 

88 saved = os.environ.pop("APPLICATIONINSIGHTS_CONNECTION_STRING", None) 

89 if saved: 

90 logger.debug( 

91 "[PERF] Hiding Azure connection string to prevent OpenTelemetry " 

92 "auto-init (set ENABLE_OTEL=true to enable telemetry)" 

93 ) 

94 return saved 

95 

96 return None 

97 

98 

99def set_model_span_attributes(span: Any, config: AgentSettings) -> None: 

100 """Set OpenTelemetry span attributes based on provider and model configuration. 

101 

102 Args: 

103 span: OpenTelemetry span to set attributes on 

104 config: Agent configuration containing provider and model information 

105 """ 

106 span.set_attribute("gen_ai.system", config.llm_provider or "unknown") 

107 

108 if config.llm_provider == "openai" and config.openai_model: 

109 span.set_attribute("gen_ai.request.model", config.openai_model) 

110 elif config.llm_provider == "anthropic" and config.anthropic_model: 

111 span.set_attribute("gen_ai.request.model", config.anthropic_model) 

112 elif config.llm_provider == "azure" and config.azure_openai_deployment: 

113 span.set_attribute("gen_ai.request.model", config.azure_openai_deployment) 

114 elif config.llm_provider == "foundry" and config.azure_model_deployment: 

115 span.set_attribute("gen_ai.request.model", config.azure_model_deployment)