@samitouri / QOSami_Agents_Zero / commits / b74d05d6

Docker + SSH, AgentConfig class

frdel committed Jul 5, 2024 at 22:58 UTC b74d05d6c2bf278537bd1388f399d4075fcdf5a6
27 files changed +422 -146
agent.py
+54 -48
@@ -1,58 +1,60 @@
1 +from dataclasses import dataclass, field
2 import time, importlib, inspect, os, json
3 import traceback
3 -from typing import Optional, Dict, TypedDict
4 -from tools.helpers import extract_tools, rate_limiter, files, errors
5 -from tools.helpers.print_style import PrintStyle
4 +from typing import Any, Optional, Dict, TypedDict
5 +from python.helpers import extract_tools, rate_limiter, files, errors
6 +from python.helpers.print_style import PrintStyle
7 from langchain.schema import AIMessage
8 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
9 from langchain_core.messages import HumanMessage, SystemMessage
10 from langchain_core.language_models.chat_models import BaseChatModel
11 from langchain_core.embeddings import Embeddings
11 -from tools.helpers.rate_limiter import RateLimiter
12 +from python.helpers.rate_limiter import RateLimiter
13
14 # rate_limit = rate_limiter.rate_limiter(30,160000) #TODO! implement properly
15
16 +@dataclass
17 +class AgentConfig:
18 + agent_number: int
19 + chat_model:BaseChatModel
20 + embeddings_model:Embeddings
21 + memory_subdir: str = ""
22 + auto_memory_count: int = 3
23 + auto_memory_skip: int = 2
24 + rate_limit_seconds: int = 60
25 + rate_limit_requests: int = 30
26 + rate_limit_input_tokens: int = 0
27 + rate_limit_output_tokens: int = 0
28 + msgs_keep_max: int = 25
29 + msgs_keep_start: int = 5
30 + msgs_keep_end: int = 10
31 + max_tool_response_length: int = 3000
32 + code_exec_docker_enabled: bool = True
33 + code_exec_docker_name: str = "agent-zero-exe"
34 + code_exec_docker_image: str = "frdel/agent-zero-exe:latest"
35 + code_exec_docker_ports: dict[str,int] = field(default_factory=lambda: {"22/tcp": 50022})
36 + code_exec_docker_volumes: dict[str, dict[str, str]] = field(default_factory=lambda: {files.get_abs_path("work_dir"): {"bind": "/root", "mode": "rw"}})
37 + code_exec_ssh_enabled: bool = True
38 + code_exec_ssh_addr: str = "localhost"
39 + code_exec_ssh_port: int = 50022
40 + code_exec_ssh_user: str = "root"
41 + code_exec_ssh_pass: str = "toor"
42 + additional: Dict[str, Any] = field(default_factory=dict)
43 +
44
45 class Agent:
46
47 paused=False
48 streaming_agent=None
49
21 - def __init__(self,
22 - agent_number: int,
23 - chat_model:BaseChatModel,
24 - embeddings_model:Embeddings,
25 - memory_subdir: str = "",
26 - auto_memory_count: int = 3,
27 - auto_memory_skip: int = 2,
28 - rate_limit_seconds: int = 60,
29 - rate_limit_requests: int = 30,
30 - rate_limit_input_tokens: int = 0,
31 - rate_limit_output_tokens: int = 0,
32 - msgs_keep_max: int = 25,
33 - msgs_keep_start: int = 5,
34 - msgs_keep_end: int = 10,
35 - max_tool_response_length: int = 3000,
36 - **kwargs):
37 -
38 - # agent config
39 - self.agent_number = agent_number
40 - self.chat_model = chat_model
41 - self.embeddings_model = embeddings_model
42 - self.memory_subdir = memory_subdir
43 - self.auto_memory_count = auto_memory_count
44 - self.auto_memory_skip = auto_memory_skip
45 - self.rate_limit_seconds = rate_limit_seconds
46 - self.rate_limit_requests = rate_limit_requests
47 - self.rate_limit_input_tokens = rate_limit_input_tokens
48 - self.rate_limit_output_tokens = rate_limit_output_tokens
49 - self.msgs_keep_max = msgs_keep_max
50 - self.msgs_keep_start = msgs_keep_start
51 - self.msgs_keep_end = msgs_keep_end
52 - self.max_tool_response_length = max_tool_response_length
50 + def __init__(self, number:int, config: AgentConfig):
51 +
52 + # agent config
53 + self.config = config
54
55 # non-config vars
55 - self.agent_name = f"Agent {self.agent_number}"
56 + self.number = number
57 + self.agent_name = f"Agent {self.config.agent_number}"
58
59 self.system_prompt = files.read_file("./prompts/agent.system.md").replace("{", "{{").replace("}", "}}")
60 self.tools_prompt = files.read_file("./prompts/agent.tools.md").replace("{", "{{").replace("}", "}}")
@@ -61,7 +63,7 @@ class Agent:
63 self.last_message = ""
64 self.intervention_message = ""
65 self.intervention_status = False
64 - self.rate_limiter = RateLimiter(max_calls=rate_limit_requests,max_input_tokens=rate_limit_input_tokens,max_output_tokens=rate_limit_output_tokens,window_seconds=rate_limit_seconds)
66 + self.rate_limiter = RateLimiter(max_calls=self.config.rate_limit_requests,max_input_tokens=self.config.rate_limit_input_tokens,max_output_tokens=self.config.rate_limit_output_tokens,window_seconds=self.config.rate_limit_seconds)
67 self.data = {} # free data object all the tools can use
68
69 os.chdir(files.get_abs_path("./work_dir")) #change CWD to work_dir
@@ -90,7 +92,7 @@ class Agent:
92 MessagesPlaceholder(variable_name="messages") ])
93
94 inputs = {"messages": self.history}
93 - chain = prompt | self.chat_model
95 + chain = prompt | self.config.chat_model
96
97 formatted_inputs = prompt.format(messages=self.history)
98 tokens = int(len(formatted_inputs)/4)
@@ -147,7 +149,7 @@ class Agent:
149 else:
150 new_message = HumanMessage(content=msg) if human else AIMessage(content=msg)
151 self.history.append(new_message)
150 - self.cleanup_history(self.msgs_keep_max, self.msgs_keep_start, self.msgs_keep_end)
152 + self.cleanup_history(self.config.msgs_keep_max, self.config.msgs_keep_start, self.config.msgs_keep_end)
153 if message_type=="ai":
154 self.last_message = msg
155
@@ -159,7 +161,7 @@ class Agent:
161 SystemMessage(content=system),
162 HumanMessage(content=msg)])
163
162 - chain = prompt | self.chat_model
164 + chain = prompt | self.config.chat_model
165 response = ""
166 printer = None
167
@@ -251,12 +253,12 @@ class Agent:
253
254
255 def get_tool(self, name: str, args: dict, message: str, **kwargs):
254 - from tools.unknown import Unknown
255 - from tools.helpers.tool import Tool
256 + from python.tools.unknown import Unknown
257 + from python.helpers.tool import Tool
258
259 tool_class = Unknown
258 - if files.exists("tools",f"{name}.py"):
259 - module = importlib.import_module("tools." + name) # Import the module
260 + if files.exists("python/tools",f"{name}.py"):
261 + module = importlib.import_module("python.tools." + name) # Import the module
262 class_list = inspect.getmembers(module, inspect.isclass) # Get all functions in the module
263
264 for cls in class_list:
@@ -267,14 +269,15 @@ class Agent:
269 return tool_class(agent=self, name=name, args=args, message=message, **kwargs)
270
271 def fetch_memories(self,reset_skip=False):
272 + if self.config.auto_memory_count<=0: return ""
273 if reset_skip: self.memory_skip_counter = 0
274
275 if self.memory_skip_counter > 0:
276 self.memory_skip_counter-=1
277 return ""
278 else:
276 - self.memory_skip_counter = self.auto_memory_skip
277 - from tools import memory_tool
279 + self.memory_skip_counter = self.config.auto_memory_skip
280 + from python.tools import memory_tool
281 messages = self.concat_messages(self.history)
282 memories = memory_tool.process_query(self,messages,"load")
283 input = {
@@ -283,4 +286,7 @@ class Agent:
286 }
287 cleanup_prompt = files.read_file("./prompts/msg.memory_cleanup.md").replace("{", "{{")
288 clean_memories = self.send_adhoc_message(cleanup_prompt,json.dumps(input), output_label="Memory injection")
286 - return clean_memories
\ No newline at end of file
289 + return clean_memories
290 +
291 + def call_extension(self, name: str, **kwargs) -> Any:
292 + pass
\ No newline at end of file
main.py
+48 -27
@@ -1,21 +1,20 @@
1 -import threading, sys, time, readline, models, os
1 +import threading, time, models, os
2 from ansio import application_keypad, mouse_input, raw_input
3 from ansio.input import InputEvent, get_input_event
4 -from agent import Agent
5 -from tools.helpers.print_style import PrintStyle
6 -from tools.helpers.files import read_file
4 +from agent import Agent, AgentConfig
5 +from python.helpers.print_style import PrintStyle
6 +from python.helpers.files import read_file
7 from pytimedinput import timedInput as timed_input
8 -from tools.helpers import files
8 +from python.helpers import files
9
10
11 input_lock = threading.Lock()
12
13 os.chdir(files.get_abs_path("./work_dir")) #change CWD to work_dir
14
15 -# Main conversation loop
16 -def chat():
15
18 - # chat model used for agents
16 +def initialize():
17 + # chat model used by agents
18 # chat_llm = models.get_groq_llama70b(temperature=0.2)
19 # chat_llm = models.get_groq_llama70b_json(temperature=0.2)
20 # chat_llm = models.get_groq_llama8b(temperature=0.2)
@@ -30,29 +29,51 @@ def chat():
29 # embedding model used for memory
30 # embedding_llm = models.get_embedding_openai()
31 embedding_llm = models.get_embedding_hf()
32 +
33 + # agent configuration
34 + config = AgentConfig(
35 + agent_number = 0,
36 + chat_model = chat_llm,
37 + embeddings_model = embedding_llm,
38 + # memory_subdir = "",
39 + auto_memory_count = 0,
40 + # auto_memory_skip = 2,
41 + # rate_limit_seconds = 60,
42 + # rate_limit_requests = 30,
43 + # rate_limit_input_tokens = 0,
44 + # rate_limit_output_tokens = 0,
45 + # msgs_keep_max = 25,
46 + # msgs_keep_start = 5,
47 + # msgs_keep_end = 10,
48 + # max_tool_response_length = 3000,
49 + code_exec_docker_enabled = True,
50 + # code_exec_docker_name = "agent-zero-exe",
51 + # code_exec_docker_image = "frdel/agent-zero-exe:latest",
52 + # code_exec_docker_ports = { "22/tcp": 50022 }
53 + # code_exec_docker_volumes = { files.get_abs_path("work_dir"): {"bind": "/root", "mode": "rw"} }
54 + code_exec_ssh_enabled = True,
55 + # code_exec_ssh_addr = "localhost",
56 + # code_exec_ssh_port = 50022,
57 + # code_exec_ssh_user = "root",
58 + # code_exec_ssh_pass = "toor",
59 + # additional = {},
60 + )
61
62 # create the first agent
35 - agent0 = Agent( agent_number=0,
36 - chat_model=chat_llm,
37 - embeddings_model=embedding_llm,
38 - # memory_subdir = "",
39 - # auto_memory_count = 3,
40 - # auto_memory_skip = 2,
41 - # rate_limit_seconds = 60,
42 - rate_limit_requests = 30,
43 - rate_limit_input_tokens = 160000,
44 - rate_limit_output_tokens = 8000,
45 - # msgs_keep_max = 25,
46 - # msgs_keep_start = 5,
47 - # msgs_keep_end = 10,
48 - # max_tool_response_length = 3000,
49 - )
63 + agent0 = Agent( number = 0, config = config )
64 +
65 + # start the chat loop
66 + chat(agent0)
67
68 +
69 +# Main conversation loop
70 +def chat(agent:Agent):
71 +
72 # start the conversation loop
73 while True:
74 # ask user for message
75 with input_lock:
55 - timeout = agent0.get_data("timeout") # how long the agent is willing to wait
76 + timeout = agent.get_data("timeout") # how long the agent is willing to wait
77 if not timeout: # if agent wants to wait for user input forever
78 PrintStyle(background_color="#6C3483", font_color="white", bold=True, padding=True).print(f"User message ('exit' to leave):")
79 user_input = input("> ")
@@ -77,10 +98,10 @@ def chat():
98 if user_input.lower() == 'exit': break
99
100 # send message to agent0,
80 - assistant_response = agent0.message_loop(user_input)
101 + assistant_response = agent.message_loop(user_input)
102
103 # print agent0 response
83 - PrintStyle(font_color="white",background_color="#1D8348", bold=True, padding=True).print(f"{agent0.agent_name}: reponse:")
104 + PrintStyle(font_color="white",background_color="#1D8348", bold=True, padding=True).print(f"{agent.agent_name}: reponse:")
105 PrintStyle(font_color="white").print(f"{assistant_response}")
106
107
@@ -123,4 +144,4 @@ if __name__ == "__main__":
144 threading.Thread(target=capture_keys, daemon=True).start()
145
146 # Start the chat
126 - chat()
\ No newline at end of file
147 + initialize()
\ No newline at end of file
python/__init__.py renamed
python/helpers/dirty_json.py renamed
python/helpers/docker.py new
+49
@@ -0,0 +1,49 @@
1 +import time
2 +import docker
3 +import atexit
4 +from typing import Dict, Optional
5 +from python.helpers.files import get_abs_path
6 +
7 +class DockerContainerManager:
8 + def __init__(self, image:str, name:str, ports: Optional[Dict[str, int]] = None, volumes: Optional[Dict[str, Dict[str, str]]] = None):
9 + self.client = docker.from_env()
10 + self.image = image
11 + self.name = name
12 + self.ports = ports
13 + self.volumes = volumes
14 + self.container = None
15 +
16 +
17 +
18 + def cleanup_container(self) -> None:
19 + if self.container:
20 + try:
21 + self.container.stop()
22 + self.container.remove()
23 + print(f"Stopped and removed the container: {self.container.id}")
24 + except Exception as e:
25 + print(f"Failed to stop and remove the container: {e}")
26 +
27 + def start_container(self) -> None:
28 + existing_container = None
29 + for container in self.client.containers.list():
30 + if container.name == self.name:
31 + existing_container = container
32 + break
33 +
34 + if existing_container:
35 + #print(f"Container with name '{self.name}' is already running with ID: {existing_container.id}")
36 + pass
37 + else:
38 + print(f"Initializing docker container {self.name}...")
39 + self.container = self.client.containers.run(
40 + self.image,
41 + detach=True,
42 + ports=self.ports,
43 + name=self.name,
44 + volumes=self.volumes,
45 + )
46 + atexit.register(self.cleanup_container)
47 + print(f"Started container with ID: {self.container.id}")
48 + time.sleep(1) # this helps to get SSH ready
49 +
python/helpers/errors.py renamed
python/helpers/extract_tools.py renamed
python/helpers/files.py renamed
python/helpers/messages.py renamed
python/helpers/perplexity_search.py renamed
python/helpers/print_style.py renamed
+5
@@ -1,4 +1,5 @@
1 import os, webcolors, html
2 +import sys
3 from datetime import datetime
4 from . import files
5
@@ -112,6 +113,10 @@ class PrintStyle:
113 self._log_html(html_text)
114 PrintStyle.last_endline = False
115
116 + def is_last_line_empty(self):
117 + lines = sys.stdin.readlines()
118 + return bool(lines) and not lines[-1].strip()
119 +
120 # Ensure HTML file is closed properly when the program exits
121 import atexit
122 atexit.register(PrintStyle._close_html_log)
python/helpers/rate_limiter.py renamed
python/helpers/shell_local.py new
+68
@@ -0,0 +1,68 @@
1 +import select
2 +import subprocess
3 +import time
4 +import sys
5 +from typing import Optional, Tuple
6 +
7 +class LocalInteractiveSession:
8 + def __init__(self):
9 + self.process = None
10 + self.full_output = ''
11 +
12 + def connect(self):
13 + # Start a new subprocess with the appropriate shell for the OS
14 + if sys.platform.startswith('win'):
15 + # Windows
16 + self.process = subprocess.Popen(
17 + ['cmd.exe'],
18 + stdin=subprocess.PIPE,
19 + stdout=subprocess.PIPE,
20 + stderr=subprocess.PIPE,
21 + text=True,
22 + bufsize=1
23 + )
24 + else:
25 + # macOS and Linux
26 + self.process = subprocess.Popen(
27 + ['/bin/bash'],
28 + stdin=subprocess.PIPE,
29 + stdout=subprocess.PIPE,
30 + stderr=subprocess.PIPE,
31 + text=True,
32 + bufsize=1
33 + )
34 +
35 + def close(self):
36 + if self.process:
37 + self.process.terminate()
38 + self.process.wait()
39 +
40 + def send_command(self, command: str):
41 + if not self.process:
42 + raise Exception("Shell not connected")
43 + self.full_output = ""
44 + self.process.stdin.write(command + '\n') # type: ignore
45 + self.process.stdin.flush() # type: ignore
46 +
47 + def read_output(self) -> Tuple[str, Optional[str]]:
48 + if not self.process:
49 + raise Exception("Shell not connected")
50 +
51 + partial_output = ''
52 + while True:
53 + rlist, _, _ = select.select([self.process.stdout], [], [], 0.1)
54 + if rlist:
55 + line = self.process.stdout.readline() # type: ignore
56 + if line:
57 + partial_output += line
58 + self.full_output += line
59 + time.sleep(0.1)
60 + else:
61 + break # No more output
62 + else:
63 + break # No data available
64 +
65 + if not partial_output:
66 + return self.full_output, None
67 +
68 + return self.full_output, partial_output
\ No newline at end of file
python/helpers/shell_ssh.py new
+71
@@ -0,0 +1,71 @@
1 +import paramiko
2 +import time
3 +import re
4 +from typing import Optional, Tuple
5 +
6 +class SSHInteractiveSession:
7 + def __init__(self, hostname: str, port: int, username: str, password: str):
8 + self.hostname = hostname
9 + self.port = port
10 + self.username = username
11 + self.password = password
12 + self.client = paramiko.SSHClient()
13 + self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
14 + self.shell = None
15 + self.full_output = ''
16 +
17 + def connect(self):
18 + # try 3 times with wait and then except
19 + errors = 0
20 + while True:
21 + try:
22 + self.client.connect(self.hostname, self.port, self.username, self.password)
23 + self.shell = self.client.invoke_shell()
24 + while True: # wait for end of initial output
25 + full, part = self.read_output()
26 + if full and not part: return
27 + time.sleep(0.1)
28 + except Exception as e:
29 + errors += 1
30 + if errors < 3:
31 + print(f"SSH Connection attempt {errors}...")
32 + time.sleep(5)
33 + else:
34 + raise e
35 +
36 + def close(self):
37 + if self.shell:
38 + self.shell.close()
39 + if self.client:
40 + self.client.close()
41 +
42 + def send_command(self, command: str):
43 + if not self.shell:
44 + raise Exception("Shell not connected")
45 + self.full_output = ""
46 + self.shell.send((command + '\n').encode())
47 +
48 + def read_output(self) -> Tuple[str, str]:
49 + if not self.shell:
50 + raise Exception("Shell not connected")
51 +
52 + partial_output = ''
53 + while self.shell.recv_ready():
54 + data = self.shell.recv(1024).decode('utf-8')
55 + data = self.clean_string(data)
56 + partial_output += data
57 + self.full_output += data
58 + time.sleep(0.1) # Prevent busy waiting
59 +
60 + self.full_output = self.clean_string(self.full_output)
61 + return self.full_output, partial_output
62 +
63 + def clean_string(self, input_string):
64 + # Remove ANSI escape codes
65 + ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
66 + cleaned = ansi_escape.sub('', input_string)
67 +
68 + # Replace '\r\n' with '\n'
69 + cleaned = cleaned.replace('\r\n', '\n')
70 +
71 + return cleaned
\ No newline at end of file
python/helpers/tool.py renamed
+3 -3
@@ -1,8 +1,8 @@
1 from abc import abstractmethod
2 from typing import TypedDict
3 from agent import Agent
4 -from tools.helpers.print_style import PrintStyle
5 -from tools.helpers import files, messages
4 +from python.helpers.print_style import PrintStyle
5 +from python.helpers import files, messages
6
7 class Response:
8 def __init__(self, message: str, break_loop: bool) -> None:
@@ -30,7 +30,7 @@ class Tool:
30 PrintStyle().print()
31
32 def after_execution(self, response: Response, **kwargs):
33 - text = messages.truncate_text(response.message.strip(), self.agent.max_tool_response_length)
33 + text = messages.truncate_text(response.message.strip(), self.agent.config.max_tool_response_length)
34 msg_response = files.read_file("./prompts/fw.tool_response.md", tool_name=self.name, tool_response=text)
35 self.agent.append_message(msg_response, human=True)
36 PrintStyle(font_color="#1B4F72", background_color="white", padding=True, bold=True).print(f"{self.agent.agent_name}: Response from tool '{self.name}':")
python/helpers/vector_db.py renamed
python/tools/call_subordinate.py renamed
+4 -4
@@ -1,7 +1,7 @@
1 from agent import Agent
2 -from tools.helpers.tool import Tool, Response
3 -from tools.helpers import files
4 -from tools.helpers.print_style import PrintStyle
2 +from python.helpers.tool import Tool, Response
3 +from python.helpers import files
4 +from python.helpers.print_style import PrintStyle
5
6 class Delegation(Tool):
7
@@ -10,7 +10,7 @@ class Delegation(Tool):
10 if self.agent.get_data("subordinate") is None or str(reset).lower().strip() == "true":
11 # subordinate = Agent(system_prompt=self.agent.system_prompt, tools_prompt=self.agent.tools_prompt, number=self.agent.number+1)
12 config = self.agent.__dict__.copy()
13 - config["agent_number"] = self.agent.agent_number+1
13 + config["agent_number"] = self.agent.number+1
14 subordinate = Agent(**config)
15 subordinate.set_data("superior", self.agent)
16 self.agent.set_data("subordinate", subordinate)
python/tools/code_execution_tool.py new
+93
@@ -0,0 +1,93 @@
1 +from dataclasses import dataclass
2 +import os, json, contextlib, subprocess, ast, shlex
3 +from io import StringIO
4 +import time
5 +from typing import Literal
6 +from python.helpers import files, messages
7 +from agent import Agent
8 +from python.helpers.tool import Tool, Response
9 +from python.helpers import files
10 +from python.helpers.print_style import PrintStyle
11 +from python.helpers.shell_local import LocalInteractiveSession
12 +from python.helpers.shell_ssh import SSHInteractiveSession
13 +from python.helpers.docker import DockerContainerManager
14 +
15 +@dataclass
16 +class State:
17 + shell: LocalInteractiveSession | SSHInteractiveSession
18 + docker: DockerContainerManager | None
19 +
20 +
21 +class CodeExecution(Tool):
22 +
23 + def execute(self,**kwargs):
24 + self.prepare_state()
25 +
26 + # os.chdir(files.get_abs_path("./work_dir")) #change CWD to work_dir
27 +
28 + runtime = self.args["runtime"].lower().strip()
29 + if runtime == "python":
30 + response = self.execute_python_code(self.args["code"])
31 + elif runtime == "nodejs":
32 + response = self.execute_nodejs_code(self.args["code"])
33 + elif runtime == "terminal":
34 + response = self.execute_terminal_command(self.args["code"])
35 + else:
36 + response = files.read_file("./prompts/fw.code_runtime_wrong.md", runtime=runtime)
37 +
38 + if not response: response = files.read_file("./prompts/fw.code_no_output.md")
39 + return Response(message=response, break_loop=False)
40 +
41 + def after_execution(self, response, **kwargs):
42 + msg_response = files.read_file("./prompts/fw.tool_response.md", tool_name=self.name, tool_response=response.message)
43 + self.agent.append_message(msg_response, human=True)
44 +
45 + def prepare_state(self):
46 + self.state = self.agent.get_data("cot_state")
47 + if not self.state:
48 +
49 + #initialize docker container if execution in docker is configured
50 + if self.agent.config.code_exec_docker_enabled:
51 + docker = DockerContainerManager(name=self.agent.config.code_exec_docker_name, image=self.agent.config.code_exec_docker_image, ports=self.agent.config.code_exec_docker_ports, volumes=self.agent.config.code_exec_docker_volumes)
52 + docker.start_container()
53 + else: docker = None
54 +
55 + #initialize local or remote interactive shell insterface
56 + if self.agent.config.code_exec_ssh_enabled:
57 + shell = SSHInteractiveSession(self.agent.config.code_exec_ssh_addr,self.agent.config.code_exec_ssh_port,self.agent.config.code_exec_ssh_user,self.agent.config.code_exec_ssh_pass)
58 + else: shell = LocalInteractiveSession()
59 +
60 + self.state = State(shell=shell,docker=docker)
61 + shell.connect()
62 + self.agent.set_data("cot_state", self.state)
63 +
64 + def execute_python_code(self, code):
65 + escaped_code = shlex.quote(code)
66 + command = f'python3 -c {escaped_code}'
67 + return self.terminal_session(command)
68 +
69 + def execute_nodejs_code(self, code):
70 + escaped_code = shlex.quote(code)
71 + command = f'node -e {escaped_code}'
72 + return self.terminal_session(command)
73 +
74 + def execute_terminal_command(self, command):
75 + return self.terminal_session(command)
76 +
77 + def terminal_session(self, command):
78 + self.state.shell.send_command(command)
79 +
80 + PrintStyle(background_color="white",font_color="#85C1E9",bold=True).print(f"{self.agent.agent_name} code execution output:")
81 +
82 + idle=0
83 + while True:
84 + time.sleep(0.1) # Wait for some output to be generated
85 + full_output, partial_output = self.state.shell.read_output()
86 +
87 + if partial_output:
88 + PrintStyle(font_color="#85C1E9").stream(partial_output)
89 + idle=0
90 + else:
91 + idle+=1
92 + if ( full_output and idle > 30 ) or ( not full_output and idle > 100 ): return full_output
93 +
\ No newline at end of file
python/tools/knowledge_tool.py renamed
+2 -2
@@ -5,8 +5,8 @@ import concurrent.futures
5
6
7
8 -from tools.helpers.tool import Tool, Response
9 -from tools.helpers import files
8 +from python.helpers.tool import Tool, Response
9 +from python.helpers import files
10
11 class Knowledge(Tool):
12 def execute(self, question="", **kwargs):
python/tools/memorize.py renamed
+3 -3
@@ -1,7 +1,7 @@
1 from agent import Agent
2 -from tools.helpers import files
3 -from tools.helpers.tool import Tool, Response
4 -from tools import memory_tool
2 +from python.helpers import files
3 +from python.helpers.tool import Tool, Response
4 +import memory_tool
5
6 class Memorize(Tool):
7 def execute(self,**kwargs):
python/tools/memory_tool.py renamed
+6 -7
@@ -1,16 +1,15 @@
1 from agent import Agent
2 -from tools.helpers.vector_db import VectorDB, Document
3 -from tools.helpers import files
2 +from python.helpers.vector_db import VectorDB, Document
3 +from python.helpers import files
4 import os, json
5 -from tools.helpers.tool import Tool, Response
6 -from tools.helpers.print_style import PrintStyle
5 +from python.helpers.tool import Tool, Response
6 +from python.helpers.print_style import PrintStyle
7
8 db: VectorDB | None = None
9
10 class Memory(Tool):
11 def execute(self,**kwargs):
12 - #TODO separate param for memory tool result count
13 - result = process_query(self.agent, self.args["memory"],self.args["action"], result_count=self.agent.auto_memory_count)
12 + result = process_query(self.agent, self.args["memory"],self.args["action"], result_count=self.agent.config.auto_memory_count)
13 return Response(message="\n\n".join(result), break_loop=False)
14
15
@@ -21,7 +20,7 @@ def initialize(embeddings_model, subdir=""):
20
21
22 def process_query(agent:Agent, message: str, action: str = "load", result_count: int = 3, **kwargs):
24 - if not db: initialize(agent.embeddings_model, subdir=agent.memory_subdir)
23 + if not db: initialize(agent.config.embeddings_model, subdir=agent.config.memory_subdir)
24
25 if action.strip().lower() == "save":
26 id = db.insert_document(str(message)) # type: ignore
python/tools/online_knowledge_tool.py renamed
+2 -2
@@ -1,6 +1,6 @@
1 from agent import Agent
2 -from tools.helpers import perplexity_search
3 -from tools.helpers.tool import Tool, Response
2 +from python.helpers import perplexity_search
3 +from python.helpers.tool import Tool, Response
4
5 class OnlineKnowledge(Tool):
6 def execute(self,**kwargs):
python/tools/response.py renamed
+5 -5
@@ -1,11 +1,11 @@
1 from agent import Agent
2 -from tools.helpers import files
3 -from tools.helpers.print_style import PrintStyle
2 +from python.helpers import files
3 +from python.helpers.print_style import PrintStyle
4
5 from agent import Agent
6 -from tools.helpers.tool import Tool, Response
7 -from tools.helpers import files
8 -from tools.helpers.print_style import PrintStyle
6 +from python.helpers.tool import Tool, Response
7 +from python.helpers import files
8 +from python.helpers.print_style import PrintStyle
9
10 class ResponseTool(Tool):
11
python/tools/task_done.py renamed
+5 -5
@@ -1,11 +1,11 @@
1 from agent import Agent
2 -from tools.helpers import files
3 -from tools.helpers.print_style import PrintStyle
2 +from python.helpers import files
3 +from python.helpers.print_style import PrintStyle
4
5 from agent import Agent
6 -from tools.helpers.tool import Tool, Response
7 -from tools.helpers import files
8 -from tools.helpers.print_style import PrintStyle
6 +from python.helpers.tool import Tool, Response
7 +from python.helpers import files
8 +from python.helpers.print_style import PrintStyle
9
10 class TaskDone(Tool):
11
python/tools/unknown.py renamed
+2 -2
@@ -1,5 +1,5 @@
1 -from tools.helpers.tool import Tool, Response
2 -from tools.helpers import files
1 +from python.helpers.tool import Tool, Response
2 +from python.helpers import files
3
4 class Unknown(Tool):
5 def execute(self, **kwargs):
requirements.txt
+2
@@ -9,3 +9,5 @@ langchain-chroma==0.1.1
9 webcolors==24.6.0
10 sentence-transformers==3.0.1
11 pytimedinput==2.0.1
12 +docker==7.1.0
13 +paramiko==3.4.0
tools/code_execution_tool.py deleted
-38
@@ -1,38 +0,0 @@
1 -import os, json, contextlib, subprocess, ast, shlex
2 -from io import StringIO
3 -from tools.helpers import files, messages
4 -from agent import Agent
5 -from tools.helpers.tool import Tool, Response
6 -from tools.helpers import files
7 -from tools.helpers.print_style import PrintStyle
8 -
9 -class CodeExecution(Tool):
10 -
11 - def execute(self,**kwargs):
12 -
13 - # os.chdir(files.get_abs_path("./work_dir")) #change CWD to work_dir
14 -
15 - runtime = self.args["runtime"].lower().strip()
16 - if runtime == "python":
17 - response = self.execute_python_code(self.args["code"])
18 - elif runtime == "nodejs":
19 - response = self.execute_nodejs_code(self.args["code"])
20 - elif runtime == "terminal":
21 - response = self.execute_terminal_command(self.args["code"])
22 - else:
23 - response = files.read_file("./prompts/fw.code_runtime_wrong.md", runtime=runtime)
24 -
25 - if not response: response = files.read_file("./prompts/fw.code_no_output.md")
26 - return Response(message=response, break_loop=False)
27 -
28 - def execute_python_code(self, code, input_data="y\n"):
29 - result = subprocess.run(['python', '-c', code], capture_output=True, text=True, input=input_data)
30 - return result.stdout + result.stderr
31 -
32 - def execute_nodejs_code(self, code, input_data="y\n"):
33 - result = subprocess.run(['node', '-e', code], capture_output=True, text=True, input=input_data)
34 - return result.stdout + result.stderr
35 -
36 - def execute_terminal_command(self, command, input_data="y\n"):
37 - result = subprocess.run(command, shell=True, capture_output=True, text=True, input=input_data)
38 - return result.stdout + result.stderr
\ No newline at end of file