-
Notifications
You must be signed in to change notification settings - Fork 18
Expand file tree
/
Copy pathagents.py
More file actions
207 lines (183 loc) · 6.73 KB
/
agents.py
File metadata and controls
207 lines (183 loc) · 6.73 KB
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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
import sys
from abc import ABC, abstractmethod
from pathlib import Path
import logging
from aider.coders import Coder
from aider.models import Model
from aider.io import InputOutput
import re
import os
def handle_logging(logging_name: str, log_file: Path) -> None:
"""Handle logging for agent"""
logger = logging.getLogger(logging_name)
logger.setLevel(logging.INFO)
logger.propagate = False
logger_handler = logging.FileHandler(log_file)
logger_handler.setFormatter(
logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
)
logger.addHandler(logger_handler)
class AgentReturn(ABC):
def __init__(self, log_file: Path):
self.log_file = log_file
self.last_cost = 0.0
self.total_token_in = 0
self.total_token_out = 0
class Agents(ABC):
def __init__(self, max_iteration: int):
self.max_iteration = max_iteration
@abstractmethod
def run(self) -> AgentReturn:
"""Start agent"""
raise NotImplementedError
class AiderReturn(AgentReturn):
def __init__(self, log_file: Path):
super().__init__(log_file)
self.last_cost = self.get_money_cost()
self.total_token_in = self.get_total_token_in()
self.total_token_out = self.get_total_token_out()
def get_money_cost(self) -> float:
"""Get accumulated money cost from log file"""
last_cost = 0.0
with open(self.log_file, "r") as file:
for line in file:
if "Tokens:" in line and "Cost:" in line:
match = re.search(
r"Cost: \$\d+\.\d+ message, \$(\d+\.\d+) session", line
)
if match:
last_cost = float(match.group(1))
return last_cost
def get_total_token_in(self) -> int:
"""Get total token in from log file"""
total_tokens = 0
with open(self.log_file, "r") as file:
for line in file:
if "Tokens:" in line:
match = re.search(r"Tokens: ([\d.]+k?) sent", line)
if match:
token_str = match.group(1)
if token_str.endswith("k"):
total_tokens = int(float(token_str[:-1]) * 1000)
else:
total_tokens = int(float(token_str))
return total_tokens
def get_total_token_out(self) -> int:
"""Get total token out from log file"""
total_tokens = 0
with open(self.log_file, "r") as file:
for line in file:
if "Tokens:" in line:
match = re.search(r"(\d+) received", line)
if match:
total_str = match.group(1)
if total_str.endswith("k"):
total_tokens = int(float(total_str[:-1]) * 1000)
else:
total_tokens = int(float(total_str))
return total_tokens
class AiderAgents(Agents):
def __init__(self, max_iteration: int, model_name: str):
super().__init__(max_iteration)
self.model = Model(model_name)
# Check if API key is set for the model
if "openrouter" in model_name:
api_key = os.environ.get("OPENROUTER_API_KEY", None)
elif "gpt" in model_name:
api_key = os.environ.get("OPENAI_API_KEY", None)
elif "claude" in model_name:
api_key = os.environ.get("ANTHROPIC_API_KEY", None)
elif "gemini" in model_name:
api_key = os.environ.get("GEMINI_API_KEY", None)
elif "deepseek" in model_name:
api_key = os.environ.get("DEEPSEEK_API_KEY", None)
elif "mistral" in model_name:
api_key = os.environ.get("MISTRAL_API_KEY", None)
else:
raise ValueError(f"Unsupported model: {model_name}")
if not api_key:
raise ValueError(
"API Key Error: There is no API key associated with the model for this agent. "
"Edit model_name parameter in .agent.yaml, export API key for that model, and try again."
)
def run(
self,
message: str,
test_cmd: str,
lint_cmd: str,
fnames: list[str],
log_dir: Path,
test_first: bool = False,
lint_first: bool = False,
current_attempt: int = 0,
) -> AgentReturn:
"""Start aider agent"""
if test_cmd:
auto_test = True
else:
auto_test = False
if lint_cmd:
auto_lint = True
else:
auto_lint = False
log_dir = log_dir.resolve()
log_dir.mkdir(parents=True, exist_ok=True)
input_history_file = (
log_dir / ".aider.input.history"
if current_attempt == 0
else log_dir / f".aider_{current_attempt}.input.history"
)
chat_history_file = (
log_dir / ".aider.chat.history.md"
if current_attempt == 0
else log_dir / f".aider_{current_attempt}.chat.history.md"
)
# Set up logging
log_file = (
log_dir / "aider.log"
if current_attempt == 0
else log_dir / f"aider_{current_attempt}.log"
)
logging.basicConfig(
filename=log_file,
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
# Redirect print statements to the log file
sys.stdout = open(log_file, "a")
sys.stderr = open(log_file, "a")
# Configure httpx and backoff logging
handle_logging("httpx", log_file)
handle_logging("backoff", log_file)
io = InputOutput(
yes=True,
input_history_file=input_history_file,
chat_history_file=chat_history_file,
)
coder = Coder.create(
main_model=self.model,
fnames=fnames,
auto_lint=auto_lint,
auto_test=auto_test,
lint_cmds={"python": lint_cmd},
test_cmd=test_cmd,
io=io,
)
coder.max_reflections = self.max_iteration
coder.stream = False
# Run the agent
if test_first:
test_errors = coder.commands.cmd_test(test_cmd)
if test_errors:
coder.run(test_errors)
elif lint_first:
coder.commands.cmd_lint(fnames=fnames)
else:
coder.run(message)
# Close redirected stdout and stderr
sys.stdout.close()
sys.stderr.close()
# Restore original stdout and stderr
sys.stdout = sys.__stdout__
sys.stderr = sys.__stderr__
return AiderReturn(log_file)