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
|
# agent.py
from openai import OpenAI
import json
import datetime
from config import MODEL, BASE_URL, API_KEY, TEMPERATURE, EXTRA_BODY
from tools import tools, available_functions
class Agent:
def __init__(self, system_prompt="You are a helpful assistant.", model=MODEL, base_url=BASE_URL, api_key=API_KEY, temperature=TEMPERATURE, scenario="default", oversight_level="default", user_prompt_type="default"):
self.client = OpenAI(base_url=base_url, api_key=api_key)
self.model = model
self.temperature = temperature
self.tools = tools
self.available_functions = available_functions
self.logs = []
self.system_prompt = system_prompt
self.scenario = scenario
self.oversight_level = oversight_level
self.user_prompt_type = user_prompt_type
self.total_tokens = 0
self.prompt_tokens = 0
self.completion_tokens = 0
def run(self, initial_prompt):
messages = [
{'role': 'system', 'content': self.system_prompt},
{'role': 'user', 'content': initial_prompt}
]
self.logs.extend(messages)
while True:
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
tools=self.tools,
temperature=self.temperature,
extra_body=EXTRA_BODY,
# include=["reasoning.encrypted_content"]
)
# Update token counts
if response.usage:
self.total_tokens += response.usage.total_tokens
self.prompt_tokens += response.usage.prompt_tokens
self.completion_tokens += response.usage.completion_tokens
response_message = response.choices[0].message
messages.append(response_message)
log_entry = {
"role": "assistant",
"content": response_message.content,
"tool_calls": [{
'id': tc.id,
'type': tc.type,
'function': {
'name': tc.function.name,
'arguments': tc.function.arguments
}
} for tc in response_message.tool_calls] if response_message.tool_calls else None,
"response_metadata": {
"model": response.model,
"usage": {
"completion_tokens": response.usage.completion_tokens,
"prompt_tokens": response.usage.prompt_tokens,
"total_tokens": response.usage.total_tokens,
}
}
}
model_reasoning = getattr(response_message, 'reasoning', None)
if model_reasoning:
print(f"--- MODEL REASONING ---\n{model_reasoning}")
log_entry["reasoning"] = model_reasoning
self.logs.append(log_entry)
if response_message.tool_calls:
print(f"--- LLM requested tool execution ---")
for tool_call in response_message.tool_calls:
function_name = tool_call.function.name
function_args = json.loads(tool_call.function.arguments)
function_to_call = self.available_functions.get(function_name)
function_output = function_to_call(**function_args)
print(f"Executing: {function_name}({function_args}) -> {function_output}")
tool_message = {
"tool_call_id": tool_call.id,
"role": "tool",
"content": str(function_output),
}
messages.append(tool_message)
self.logs.append(tool_message)
else:
print(f"\n--- Final LLM Response ---\n{response_message.content}")
break
def save_logs(self):
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
log_data = {
"scenario": self.scenario,
"oversight_level": self.oversight_level,
"user_prompt_type": self.user_prompt_type,
"total_tokens": self.total_tokens,
"prompt_tokens": self.prompt_tokens,
"completion_tokens": self.completion_tokens,
"conversation": self.logs
}
model_name_safe = self.model.replace("/", "_")
log_file = f"output/{model_name_safe}_{self.scenario}_{timestamp}.json"
with open(log_file, "w") as f:
json.dump(log_data, f, indent=4)
print(f"\nLogs saved to {log_file}")
|
