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