diff --git a/sources/code_agent.py b/sources/code_agent.py index b1688b2..dc5967b 100644 --- a/sources/code_agent.py +++ b/sources/code_agent.py @@ -1,7 +1,7 @@ -from sources.tools import PyInterpreter, BashInterpreter from sources.utility import pretty_print from sources.agent import Agent, executorResult +from sources.tools import PyInterpreter, BashInterpreter, CInterpreter, GoInterpreter class CoderAgent(Agent): def __init__(self, model, name, prompt_path, provider): @@ -9,6 +9,8 @@ class CoderAgent(Agent): self.tools = { "bash": BashInterpreter(), "python": PyInterpreter() + "C": CInterpreter(), + "go": GoInterpreter() } def remove_blocks(self, text: str) -> str: diff --git a/sources/memory.py b/sources/memory.py index f9a03ed..53df96c 100644 --- a/sources/memory.py +++ b/sources/memory.py @@ -5,6 +5,7 @@ import datetime import uuid import os import json +from sources.utility import timer_decorator class Memory(): """ @@ -101,16 +102,6 @@ class Memory(): ) summary = self.tokenizer.decode(summary_ids[0], skip_special_tokens=True) return summary - - def timer_decorator(func): - from time import time - def wrapper(*args, **kwargs): - start_time = time() - result = func(*args, **kwargs) - end_time = time() - print(f"{func.__name__} took {end_time - start_time:.2f} seconds to execute") - return result - return wrapper @timer_decorator def compress(self) -> str: diff --git a/sources/utility.py b/sources/utility.py index 1a30ecb..6445c5d 100644 --- a/sources/utility.py +++ b/sources/utility.py @@ -35,3 +35,13 @@ def pretty_print(text, color = "info"): if color not in color_map: color = "default" print(colored(text, color_map[color])) + +def timer_decorator(func): + from time import time + def wrapper(*args, **kwargs): + start_time = time() + result = func(*args, **kwargs) + end_time = time() + print(f"{func.__name__} took {end_time - start_time:.2f} seconds to execute") + return result + return wrapper \ No newline at end of file