From a200846a2af3bf89ab36e2979fd384591bf0fe4e Mon Sep 17 00:00:00 2001 From: devjpt23 Date: Wed, 18 Sep 2024 20:24:05 +0300 Subject: [PATCH] Added response metadata capture to trace directory Signed-off-by: devjpt23 --- kai/kai_trace.py | 14 ++++++++++++++ kai/service/kai_application/kai_application.py | 9 +++++++++ 2 files changed, 23 insertions(+) diff --git a/kai/kai_trace.py b/kai/kai_trace.py index 372b7085..f0c34e18 100644 --- a/kai/kai_trace.py +++ b/kai/kai_trace.py @@ -138,6 +138,20 @@ def llm_result( with open(result_file_path, "w") as f: f.write(result.pretty_repr()) + @enabled_check + def llm_token_usage( + self, current_batch_count: int, retry_count: int, token_usage: dict + ): + token_usage_file_path = os.path.join( + self.trace_dir, + f"{current_batch_count}", + f"{retry_count}", + "token_usage.json", + ) + os.makedirs(os.path.dirname(token_usage_file_path), exist_ok=True) + with open(token_usage_file_path, "w") as f: + f.write(json.dumps(token_usage, indent=4, default=str)) + @enabled_check def exception( self, diff --git a/kai/service/kai_application/kai_application.py b/kai/service/kai_application/kai_application.py index d848ea41..52c961b3 100644 --- a/kai/service/kai_application/kai_application.py +++ b/kai/service/kai_application/kai_application.py @@ -180,6 +180,15 @@ def get_incident_solutions_for_file( ): llm_result = self.model_provider.llm.invoke(prompt) trace.llm_result(count, retry_attempt_count, llm_result) + try: + token_usage = llm_result.response_metadata["token_usage"] + trace.llm_token_usage( + count, retry_attempt_count, token_usage + ) + except KeyError as e: + KAI_LOG.warning( + f"Key does not exist in the dictionary: {e}" + ) content = parse_file_solution_content( src_file_language, str(llm_result.content)