detection_documentation_analysis/llm_document_detection_analysis.py at main · gjGatos/detection_documentation_analysis · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

import json
import nest_asyncio
import yaml
import os

# Graph imports
from langchain_core.rate_limiters import InMemoryRateLimiter
from langchain_openai import ChatOpenAI
from langchain_community.callbacks import get_openai_callback

import dda_graph_definitions as graph_defs
from schemas.detection_state import DetectionState

import logging
logger = logging.getLogger(__name__)
logging.basicConfig(filename='./logs/detection_analysis.log', level=logging.INFO)

class LLMDetectionParsingEngine:

	def __init__(self):
		self.reset_llm_costs();
		self.compiled_graph = graph_defs.define_graph()


	def perform_investigation(self, repo_path, local_path, data_sensor_name=None, use_fake_data=False):
		self.reset_llm_costs()

		results = self.execute_graph(self.compiled_graph, repo_path=repo_path, local_file_path=local_path, data_sensor_name=data_sensor_name, graph_name="get_detections", use_fake_data=use_fake_data)

		self.print_costs()
		self.save_results(results, repo_path, local_path, './results')
		return results, self.llm_costs

	def save_results(self, results, repo_path, local_filename, results_path):

		result_info = {}
		result_info['header'] = { 'repo_path': repo_path, 'filename': local_filename }
		llm_is = []

		for r in results['investigative_steps']:
			llm_is.append({ 'name': r.name, 'description': r.description })
		result_info['llm_generated_investigative_steps'] = llm_is

		if 'comparison_results' in results and results['comparison_results'] != None:
			unique_steps_from_investigation_1 = []
			unique_steps_from_investigation_2 = []
			same_steps = []
			for r in results['comparison_results'].unique_steps_from_investigation_1:
				unique_steps_from_investigation_1.append({ 'name': r.name, 'description': r.description, 'is_generic': r.generic })
			for r in results['comparison_results'].unique_steps_from_investigation_2:
				unique_steps_from_investigation_2.append({ 'name': r.name, 'description': r.description, 'is_generic': r.generic })
			for r in results['comparison_results'].same_steps:
				same_steps.append({ 'name': r.name, 'description': r.description, 'is_generic': r.generic })
			result_info['llm_unique_steps_from_investigation'] = unique_steps_from_investigation_1
			result_info['human_unique_steps_from_investigation'] = unique_steps_from_investigation_2
			result_info['same_steps'] = same_steps

			result_info['metrics'] = {}
			result_info['metrics']['llm_metrics'] = self.get_metrics(results['comparison_results'].unique_steps_from_investigation_1)
			result_info['metrics']['human_metrics'] = self.get_metrics(results['comparison_results'].unique_steps_from_investigation_2)
			result_info['metrics']['same_steps'] = self.get_metrics(results['comparison_results'].same_steps)

			output_filename = os.path.join(results_path, 'results_' + local_filename)

			with open(output_filename.replace(".toml", ".json"), 'w') as yaml_file:
				json.dump(result_info, yaml_file, indent=2)

	def get_metrics(self, response):
		result = { 'count': len(response) }
		generic_count = 0
		for r in response:
			if r.generic:
				generic_count += 1
		result['generic_count'] = generic_count
		return result


	def execute_graph(self, compiled_graph, repo_path="", local_file_path="", data_sensor_name=None, procedures=[], graph_name="undefined", use_fake_data=False):
		results = {}
		rate_limiter = InMemoryRateLimiter(
			requests_per_second=2.0,
			check_every_n_seconds=0.05, # Wake up every 50 ms to check whether allowed to make a request,
			max_bucket_size=50,  		# Controls the maximum burst size.
		)

		# This line is needed when running in jupyter notebooks
		nest_asyncio.apply()

		llm = ChatOpenAI(model="gpt-4o-2024-08-06", temperature=0, rate_limiter=rate_limiter)
		config = {"configurable": {"thread_id": "1", "llm": llm, \
		"use_fake_data": use_fake_data, "data_sensor_name": data_sensor_name, "min_note_size": 2000 }}

		with get_openai_callback() as cb:
			path = os.path.join(repo_path, local_file_path)

			results = compiled_graph.invoke({"file_path":path, "investigative_steps":[], "raw_data": []}, config)

			self.save_llm_costs(cb)


		return results


	def reset_llm_costs(self):
		self.llm_costs = {
			'total_tokens': 0,
			'prompt_tokens': 0,
			'completion_tokens': 0,
			'total_cost': 0
		}

	def save_llm_costs(self, cb):
		self.llm_costs['total_tokens'] += cb.total_tokens
		self.llm_costs['prompt_tokens'] += cb.prompt_tokens
		self.llm_costs['completion_tokens'] += cb.completion_tokens
		self.llm_costs['total_cost'] += cb.total_cost

	def print_costs(self):
		print(f"\nTotal Tokens: {self.llm_costs['total_tokens']}")
		print(f"Prompt Tokens: {self.llm_costs['prompt_tokens']}")
		print(f"Completion Tokens: {self.llm_costs['completion_tokens']}")
		print(f"Total Cost (USD): ${self.llm_costs['total_cost']}")