fix key issues

Signed-off-by: kiranscaria <kiranscaria@outlook.com>
This commit is contained in:
kiranscaria
2025-03-03 10:48:40 +05:30
parent 402a06eceb
commit 5046f7d956
3 changed files with 27 additions and 24 deletions
View File
@@ -54,15 +54,7 @@ class RAGATraceExporter(SpanExporter):
self.trace_spans.clear()
def process_complete_trace(self, spans, trace_id):
# # Save trace to temp file with trace_id.jsonl name
# trace_file_path = os.path.join(self.tmp_dir, f"{trace_id}.jsonl")
# with open(trace_file_path, "w") as file:
# for record in spans:
# file.write(json.dumps(record) + "\n")
# Convert the trace to ragaai trace format
ragaai_trace = None
try:
ragaai_trace_details = self.prepare_trace(spans, trace_id)
except Exception as e:
@@ -70,10 +62,8 @@ class RAGATraceExporter(SpanExporter):
# Upload the trace if upload_trace function is provided
try:
# Upload either the cleaned trace or the original file path
self.upload_trace(ragaai_trace_details, trace_id)
except Exception as e:
# Handle or log the error
print(f"Error uploading trace {trace_id}: {e}")
def prepare_trace(self, spans, trace_id):
@@ -85,9 +75,15 @@ class RAGATraceExporter(SpanExporter):
hash_id, zip_path = zip_list_of_unique_files(
self.files_to_zip, output_dir=self.tmp_dir
)
ragaai_trace["metadata"]["system_info"] = asdict(self.system_monitor.get_system_info())
ragaai_trace["metadata"]["resources"] = asdict(self.system_monitor.get_resources())
ragaai_trace["metadata"]["system_info"]["source_code"] = hash_id
ragaai_trace["data"][0]["start_time"] = ragaai_trace["start_time"]
ragaai_trace["data"][0]["end_time"] = ragaai_trace["end_time"]
ragaai_trace["project_name"] = self.project_name
# Save the trace_json
trace_file_path = os.path.join(self.tmp_dir, f"{trace_id}.json")
@@ -64,6 +64,7 @@ def get_spans(input_trace):
final_span["data"]={}
final_span["info"]={}
final_span["metrics"] =[]
final_span["extra_info"]={}
if span_type=="agent":
if "input.value" in span["attributes"]:
try:
@@ -121,7 +122,7 @@ def get_spans(input_trace):
try:
output_data[key] = json.loads(span['attributes'][key])
except json.JSONDecodeError as e:
input_data[key] = span['attributes'].get(key, None)
output_data[key] = span['attributes'].get(key, None)
final_span["data"]["output"] = output_data
if "llm.model_name" in span["attributes"]:
@@ -129,9 +130,14 @@ def get_spans(input_trace):
else:
final_span["info"]["model_name"] = None
if "llm.invocation_parameters" in span["attributes"]:
final_span["info"]["llm_parameters"] = span["attributes"]["llm.invocation_parameters"]
try:
final_span["info"].update(**json.loads(span["attributes"]["llm.invocation_parameters"]))
except json.JSONDecodeError as e:
print(f"Error in parsing: {e}")
final_span["extra_info"]["llm_parameters"] = span["attributes"]["llm.invocation_parameters"]
else:
final_span["info"]["llm_parameters"] = None
final_span["extra_info"]["llm_parameters"] = None
else:
if "input.value" in span["attributes"]:
@@ -166,11 +172,11 @@ def convert_json_format(input_trace):
final_trace: The converted JSON, or None if an error occurs.
"""
final_trace = {
"id": input_trace[0]["context"]["trace_id"],
"trace_name": "",
"project_name": "",
"start_time": convert_time_format(min(item["start_time"] for item in input_trace)), # Find the minimum start_time of all spans
"end_time": convert_time_format(max(item["end_time"] for item in input_trace)) # Find the maximum end_time of all spans
"id": input_trace[0]["context"]["trace_id"],
"trace_name": "",
"project_name": "",
"start_time": convert_time_format(min(item["start_time"] for item in input_trace)), # Find the minimum start_time of all spans
"end_time": convert_time_format(max(item["end_time"] for item in input_trace)) # Find the maximum end_time of all spans
}
final_trace["metadata"] ={"tokens": {
"prompt_tokens": 0,
@@ -183,12 +189,13 @@ def convert_json_format(input_trace):
final_trace["network_calls"] =[]
final_trace["interactions"] = []
for itr in final_trace["data"][0]["spans"]:
if "prompt_tokens" in itr["info"]:
final_trace["metadata"]["prompt_tokens"]+=itr["info"]['prompt_tokens']
if "completion_tokens" in itr["info"]:
final_trace["metadata"]["completion_tokens"]+=itr["info"]['completion_tokens']
if "total_tokens" in itr["info"]:
final_trace["metadata"]["total_tokens"]+=itr["info"]['total_tokens']
if itr["type"]=="llm":
if "prompt_tokens" in itr["info"]:
final_trace["metadata"]["tokens"]["prompt_tokens"]+=itr["info"]['prompt_tokens']
if "completion_tokens" in itr["info"]:
final_trace["metadata"]["tokens"]["completion_tokens"]+=itr["info"]['completion_tokens']
if "total_tokens" in itr["info"]:
final_trace["metadata"]["tokens"]["total_tokens"]+=itr["info"]['total_tokens']
return final_trace
if __name__ == "__main__":