/ tests / examples / test_utils / get_trace_data.py
get_trace_data.py
 1  import os
 2  import re
 3  import json
 4  import subprocess
 5  import logging
 6  from typing import Dict, Optional, List
 7  from dotenv import load_dotenv
 8  # Load environment variables from .env file
 9  load_dotenv()
10  
11  logging.basicConfig(level=logging.INFO)
12  logger = logging.getLogger(__name__)
13  
14  def run_command(command, cwd: Optional[str] = None):
15      cwd = cwd or os.getcwd()
16      logger.info(f"Running command: {command} in cwd: {cwd}")
17      try:
18          result = subprocess.run(
19              command,
20              shell=True,
21              cwd=cwd,
22              check=True,
23              capture_output=True,
24              text=True
25          )
26          logger.info(f"Command run successfully")
27          output = result.stdout + '\n' + result.stderr
28          return output
29      except Exception as e:
30          logger.error(f"Command failed: {e}")
31          raise
32  
33  
34  def extract_information(logs: str) -> str:
35      print("Extracting information from logs")
36      
37      # Define the patterns
38      patterns = [
39          re.compile(r"Trace saved to (.*)$"), 
40          # re.compile(r"Uploading trace metrics for (.*)$"),
41          # re.compile(r"Uploading agentic traces for (.*)$"),
42          re.compile(r"Submitting new upload task for file: (.*)$")
43      ]
44      
45      # Split the text into lines to process them individually
46      lines = logs.splitlines()
47      locations = []
48  
49      # Search each line for the patterns
50      for pattern in patterns: 
51          for line in lines:
52              match = pattern.search(line)
53              if match:
54                  # The captured group (.*) will contain the file path
55                  locations.append(match.group(1).strip())
56          if len(locations) > 0:
57              break
58      
59      return locations
60  
61  def load_trace_data(locations: List[str]) -> Dict:
62      final_data = {}
63      for location in locations:
64          try:
65              with open(location, 'r') as f:
66                  data = json.load(f)
67                  if len(str(data)) > len(str(final_data)):
68                      final_data = data
69          except Exception as e:
70              continue
71  
72      if final_data == {}:
73          raise ValueError("No trace data found")
74      return final_data
75  
76