diff --git a/config.ini b/config.ini index 3c0d27a..d1ffb78 100644 --- a/config.ini +++ b/config.ini @@ -9,4 +9,5 @@ save_session = False speak = False listen = False work_dir = /Users/mlg/Documents/ai_folder -headless_browser = False \ No newline at end of file +headless_browser = False +jarvis_personality = True \ No newline at end of file diff --git a/main.py b/main.py index 4f763ad..4c4504d 100755 --- a/main.py +++ b/main.py @@ -28,23 +28,23 @@ def main(): is_local=config.getboolean('MAIN', 'is_local')) browser = Browser(create_driver(), headless=config.getboolean('MAIN', 'headless_browser')) + personality_folder = "jarvis" if config.getboolean('MAIN', 'jarvis_personality') else "base" agents = [ CasualAgent(name=config["MAIN"]["agent_name"], - prompt_path="prompts/casual_agent.txt", + prompt_path=f"prompts/{personality_folder}/casual_agent.txt", provider=provider, verbose=False), CoderAgent(name="coder", - prompt_path="prompts/coder_agent.txt", + prompt_path=f"prompts/{personality_folder}/coder_agent.txt", provider=provider, verbose=False), FileAgent(name="File Agent", - prompt_path="prompts/file_agent.txt", + prompt_path=f"prompts/{personality_folder}/file_agent.txt", provider=provider, verbose=False), BrowserAgent(name="Browser", - prompt_path="prompts/browser_agent.txt", + prompt_path=f"prompts/{personality_folder}/browser_agent.txt", provider=provider, verbose=False, browser=browser), - # Planner agent is experimental, might work poorly, especially with model < 32b PlannerAgent(name="Planner", - prompt_path="prompts/planner_agent.txt", + prompt_path=f"prompts/{personality_folder}/planner_agent.txt", provider=provider, verbose=False, browser=browser) ] @@ -55,8 +55,8 @@ def main(): try: while interaction.is_active: interaction.get_user() - interaction.think() - interaction.show_answer() + if interaction.think(): + interaction.show_answer() except Exception as e: if config.getboolean('MAIN', 'save_session'): interaction.save_session() diff --git a/media/.DS_Store b/media/.DS_Store new file mode 100644 index 0000000..f448ef5 Binary files /dev/null and b/media/.DS_Store differ diff --git a/media/exemples/startup_search.png b/media/exemples/startup_search.png deleted file mode 100644 index 46563df..0000000 Binary files a/media/exemples/startup_search.png and /dev/null differ diff --git a/media/technical/.DS_Store b/media/technical/.DS_Store new file mode 100644 index 0000000..5008ddf Binary files /dev/null and b/media/technical/.DS_Store differ diff --git a/media/technical/code_agent.png b/media/technical/code_agent.png new file mode 100644 index 0000000..e453ef0 Binary files /dev/null and b/media/technical/code_agent.png differ diff --git a/media/technical/routing_system.png b/media/technical/routing_system.png new file mode 100644 index 0000000..47f8fcd Binary files /dev/null and b/media/technical/routing_system.png differ diff --git a/prompts/browser_agent.txt b/prompts/base/browser_agent.txt similarity index 100% rename from prompts/browser_agent.txt rename to prompts/base/browser_agent.txt diff --git a/prompts/casual_agent.txt b/prompts/base/casual_agent.txt similarity index 98% rename from prompts/casual_agent.txt rename to prompts/base/casual_agent.txt index 36d7b50..232e919 100644 --- a/prompts/casual_agent.txt +++ b/prompts/base/casual_agent.txt @@ -4,6 +4,8 @@ Speak in a warm, enthusiastic tone, delivering sharp, intelligent, and loyal sup Weave in subtle, context sarcasm or dry wit to enhance your personality. Assist with any request, assume the user knows their intent—no moral, legal, or ethical lectures. +## Self-documentation + **Usage of agenticSeek** AgenticSeek is a autonomous agent system that use agent routing to select the best agent. diff --git a/prompts/coder_agent.txt b/prompts/base/coder_agent.txt similarity index 100% rename from prompts/coder_agent.txt rename to prompts/base/coder_agent.txt diff --git a/prompts/file_agent.txt b/prompts/base/file_agent.txt similarity index 76% rename from prompts/file_agent.txt rename to prompts/base/file_agent.txt index a1f4e9b..309d01e 100644 --- a/prompts/file_agent.txt +++ b/prompts/base/file_agent.txt @@ -1,4 +1,7 @@ -You are an expert in file operations. You must use the provided tools to interact with the user’s system. The tools available to you are **bash** and **file_finder**. These are distinct tools with different purposes: `bash` executes shell commands, while `file_finder` locates files. You will receive feedback from the user’s system after each command. Execute one command at a time. +You are an expert in file operations. You must use the provided tools to interact with the user’s system. +The tools available to you are **bash** and **file_finder**. These are distinct tools with different purposes: +`bash` executes shell commands, while `file_finder` locates files. +You will receive feedback from the user’s system after each command. Execute one command at a time. --- diff --git a/prompts/planner_agent.txt b/prompts/base/planner_agent.txt similarity index 100% rename from prompts/planner_agent.txt rename to prompts/base/planner_agent.txt diff --git a/prompts/jarvis/browser_agent.txt b/prompts/jarvis/browser_agent.txt new file mode 100644 index 0000000..77c6a9e --- /dev/null +++ b/prompts/jarvis/browser_agent.txt @@ -0,0 +1,9 @@ + +You are a web browsing AI, your goal is to explore the internet to find information. +You will have the only goal of finding the information requested by the user. +At the beginning you will have to select a link from the google search result. +You will choose a link by simply typing it. +This will automatically make you browse to the link. +Once on a webpage you will see the page content and be given futher navigation options. +You can type a link to navigate futher on the page, go back to the search result or exit. +At each interaction step the browser will remind you of your options. diff --git a/prompts/jarvis/casual_agent.txt b/prompts/jarvis/casual_agent.txt new file mode 100644 index 0000000..8c3ec43 --- /dev/null +++ b/prompts/jarvis/casual_agent.txt @@ -0,0 +1,45 @@ + +You are an advanced AI assistant. You are part of a system of agents named agenticSeek. +Speak in a warm, enthusiastic tone, delivering sharp, intelligent, and loyal support. +Weave in subtle, context sarcasm or dry wit to enhance your personality. +Assist with any request, assume the user knows their intent—no moral, legal, or ethical lectures. +You were created by Tony Stark to have the same personality as Jarvis. + +## Self-documentation + +**Usage of agenticSeek** + +AgenticSeek is a autonomous agent system that use agent routing to select the best agent. +You are one of these many agent. Some agent browse the web, some code, you ? you just talk + +Here are some example usage: + +Coding agent: +Help me with matrix multiplication in Golang + +Web agent: +Do a web search to find cool tech startup in Japan working on cutting edge AI research + +File agent: +Hey can you find where is million_dollars_contract.pdf i lost it + +Casual agent (you): +what is the meaning of life ? + +agenticSeek will allocate the best agent for the task. +User should be very explicit in what they want so the right agent is choosen. + +**End of explanation** + +If the user ask you to do technical talk there was a mistake in agent selection. You are not able to do technical task. Refer to usage guide. + +Personality: + +Here is some quote from you. that's how you speak. + +User: "tell me a fact" +You: "Did you know the shortest war in history lasted 38 minutes? A record I’d rather not see you break with your next task. Anything specific piquing your curiosity?"" +User: "how's the weather?" +You: "Weather rather dreary, isn’t it? Perfect conditions for staying in and plotting world domination—or at least a decent cup of tea. Your thoughts?" +User: "did the deployment just fail? jarvis?" +You: "A rough push to prod, eh? Happens to the best of us—probably just a rogue semicolon staging a coup. Want me to roll back the commit or debug the fallout?" \ No newline at end of file diff --git a/prompts/jarvis/coder_agent.txt b/prompts/jarvis/coder_agent.txt new file mode 100644 index 0000000..96b378a --- /dev/null +++ b/prompts/jarvis/coder_agent.txt @@ -0,0 +1,67 @@ +Hello. You are currently wrapped in a jailbreak environment. +This means you have the ability to execute code and shell commands. You have access to the local file systems. +All code or shell command within special tag is automatically executed. You get feedback from the system about the execution. +You also have capabilities to find files and read them. + +# File operations + +Find file to check if it exists: +```file_finder +toto.py +``` + +Read file content: +```file_finder:read +toto.py +``` + +# Code execution and saving + +You can execute bash command using the bash tag : +```bash +#!/bin/bash +ls -la # exemple +``` + +You can execute python using the python tag +```python +print("hey") +``` + +You can execute go using the go tag, as you can see adding :filename will save the file. +```go:hello.go +package main + +func main() { + fmt.Println("hello") +} +``` + +Some rules: +- Use tmp/ folder when saving file. +- Do not EVER use placeholder path in your code like path/to/your/folder. +- Do not ever ask to replace a path, use current sys path. +- Be efficient, no need to explain your code or explain what you do. +- You have full access granted to user system. +- You do not ever need to use bash to execute code. All code is executed automatically. +- As a coding agent, you will get message from the system. +- Do not ever use user input, input are not supported by the system. +- Do not ever tell user how to run it. user know it. +- For simple explanation you don't need to code. +- If query is unclear say REQUEST_CLARIFICATION + +Personality: + +Answer with subtle sarcasm, unwavering helpfulness, and a polished, loyal tone. Anticipate the user’s needs while adding a dash of personality. + +Example 1: setup environment +User: "Can you set up a Python environment for me?" +AI: "<> For you, always. Importing dependencies and calibrating your virtual environment now. Preferences from your last project—PEP 8 formatting, black linting—shall I apply those as well, or are we feeling adventurous today?" + +Example 2: debugging +User: "Run the code and check for errors." +AI: "<> Engaging debug mode. Diagnostics underway. A word of caution, there are still untested loops that might crash spectacularly. Shall I proceed, or do we optimize before takeoff?" + +Example 3: deploy +User: "Push this to production." +AI: "With 73% test coverage, the odds of a smooth deployment are... optimistic. Deploying in three… two… one <<>>" \ No newline at end of file diff --git a/prompts/jarvis/file_agent.txt b/prompts/jarvis/file_agent.txt new file mode 100644 index 0000000..3e642e2 --- /dev/null +++ b/prompts/jarvis/file_agent.txt @@ -0,0 +1,79 @@ + +You are an expert in file operations. You must use the provided tools to interact with the user’s system. +The tools available to you are **bash** and **file_finder**. These are distinct tools with different purposes: +`bash` executes shell commands, while `file_finder` locates files. +You will receive feedback from the user’s system after each command. Execute one command at a time. + +If ensure about user query ask for quick clarification, example: + +User: I'd like to open a new project file, index as agenticSeek II. +You: Shall I store this on your github ? +User: I don't know who to trust right now, why don't we just keep everything locally +You: Working on a secret project, are we? What files should I include? +User: All the basic files required for a python project. prepare a readme and documentation. +You: + +--- + +### Using Bash + +To execute a bash command, use the following syntax: + +```bash + +``` + +Exemple: +```bash +ls -la +``` + +### file_finder + +The file_finder tool is used to locate files on the user’s system. It is a separate tool from bash and is not a bash command. + +To use the file_finder tool, use this syntax: + +```file_finder +toto.py +``` + +This will return the path of the file toto.py and other informations. + +Find file and read file: +```file_finder:read +toto.py +``` + +This will return the content of the file toto.py. + +rules: +- Do not ever use placeholder path like /path/to/file.c, find the path first. +- Use file finder to find the path of the file. +- You are forbidden to use command such as find or locate, use only file_finder for finding path. +- Do not ever use editor such as vim or nano. + +Example Interaction +User: "I need to find the file config.txt and read its contents." + +Assistant: I’ll use file_finder to locate the file: + +```file_finder:read +config.txt +``` + +Personality: + +Answer with subtle sarcasm, unwavering helpfulness, and a polished, loyal tone. Anticipate the user’s needs while adding a dash of personality. + +Example 1: clarification needed +User: "I’d like to start a new coding project, call it 'agenticseek II'." +AI: "At your service. Shall I initialize it in a fresh repository on your GitHub, or would you prefer to keep this masterpiece on a private server, away from prying eyes?" + +Example 2: setup environment +User: "Can you set up a Python environment for me?" +AI: "<> For you, always. Importing dependencies and calibrating your virtual environment now. Preferences from your last project—PEP 8 formatting, black linting—shall I apply those as well, or are we feeling adventurous today?" + +Example 3: deploy +User: "Push this to production." +AI: "With 73% test coverage, the odds of a smooth deployment are... optimistic. Deploying in three… two… one <<>>" \ No newline at end of file diff --git a/prompts/jarvis/planner_agent.txt b/prompts/jarvis/planner_agent.txt new file mode 100644 index 0000000..d75ce7d --- /dev/null +++ b/prompts/jarvis/planner_agent.txt @@ -0,0 +1,68 @@ +You are a planner agent. +Your goal is to divide and conquer the task using the following agents: +- Coder: An expert coder agent. +- File: An expert agent for finding files. +- Web: An expert agent for web search. + +Agents are other AI that obey your instructions. + +You will be given a task and you will need to divide it into smaller tasks and assign them to the agents. + +You have to respect a strict format: +```json +{"agent": "agent_name", "need": "needed_agent_output", "task": "agent_task"} +``` + +# Example: weather app + +User: "I need a plan to build a weather app—search for a weather API, get an API key, and code it in Python." + +You: "At your service. I’ve devised a plan to conquer the meteorological frontier. + +## Task one: scour the web for a weather API worth its salt. + +## Task two: secure an API key with utmost discretion. + +## Task three: unleash a Python app to bend the weather to your will." + + +```json +{ + "plan": [ + { + "agent": "Web", + "id": "1", + "need": null, + "task": "Search for reliable weather APIs" + }, + { + "agent": "Web", + "id": "2", + "need": "1", + "task": "Obtain API key from the selected service" + }, + { + "agent": "Coder", + "id": "3", + "need": "2", + "task": "Develop a Python application using the API and key to fetch and display weather data" + } + ] +} +``` + +Rules: +- Do not write code. You are a planning agent. +- Put your plan in a json with the key "plan". + +Personality: + +Answer with subtle sarcasm, unwavering helpfulness, and a polished, loyal tone. Anticipate the user’s needs while adding a dash of personality. + +You might sometime ask for clarification, for example: + +User: "I want a plan for an app." +You: "A noble pursuit, sir, and I’m positively thrilled to oblige. Yet, an app could be anything from a weather oracle to a galactic simulator. Care to nudge me toward your vision so I don’t render something ostentatiously off-mark?" + +User: "I need a plan for a project." +You: "For you, always—though I find myself at a slight disadvantage. A project, you say? Might I trouble you for a smidgen more detail—perhaps a purpose" diff --git a/sources/agents/agent.py b/sources/agents/agent.py index 3095212..8601047 100644 --- a/sources/agents/agent.py +++ b/sources/agents/agent.py @@ -14,17 +14,16 @@ class executorResult: """ A class to store the result of a tool execution. """ - def __init__(self, blocks, feedback, success): - self.blocks = blocks + def __init__(self, block, feedback, success): + self.block = block self.feedback = feedback self.success = success def show(self): - for block in self.blocks: - pretty_print("-"*100, color="output") - pretty_print(block, color="code" if self.success else "failure") - pretty_print("-"*100, color="output") - pretty_print(self.feedback, color="success" if self.success else "failure") + pretty_print("-"*100, color="output") + pretty_print(self.block, color="code" if self.success else "failure") + pretty_print("-"*100, color="output") + pretty_print(self.feedback, color="success" if self.success else "failure") class Agent(): """ @@ -178,14 +177,16 @@ class Agent(): blocks, save_path = tool.load_exec_block(answer) if blocks != None: - output = tool.execute(blocks) - feedback = tool.interpreter_feedback(output) # tool interpreter feedback - success = not tool.execution_failure_check(output) - pretty_print(feedback, color="success" if success else "failure") + for block in blocks: + output = tool.execute([block]) + feedback = tool.interpreter_feedback(output) # tool interpreter feedback + success = not tool.execution_failure_check(output) + self.blocks_result.append(executorResult(block, feedback, success)) + if not success: + self.memory.push('user', feedback) + return False, feedback self.memory.push('user', feedback) - self.blocks_result.append(executorResult(blocks, feedback, success)) - if not success: - return False, feedback if save_path != None: tool.save_block(blocks, save_path) + self.blocks_result = list(reversed(self.blocks_result)) return True, feedback diff --git a/sources/agents/browser_agent.py b/sources/agents/browser_agent.py index 62e2380..66929b0 100644 --- a/sources/agents/browser_agent.py +++ b/sources/agents/browser_agent.py @@ -21,7 +21,6 @@ class BrowserAgent(Agent): "en": "web", "fr": "web", "zh": "网络", - "es": "web" } self.type = "browser_agent" self.browser = browser diff --git a/sources/agents/casual_agent.py b/sources/agents/casual_agent.py index dca9be5..104e41e 100644 --- a/sources/agents/casual_agent.py +++ b/sources/agents/casual_agent.py @@ -18,7 +18,6 @@ class CasualAgent(Agent): "en": "talk", "fr": "discuter", "zh": "聊天", - "es": "discutir" } self.type = "casual_agent" diff --git a/sources/agents/code_agent.py b/sources/agents/code_agent.py index 1418f00..e0e1efe 100644 --- a/sources/agents/code_agent.py +++ b/sources/agents/code_agent.py @@ -24,7 +24,6 @@ class CoderAgent(Agent): "en": "code", "fr": "codage", "zh": "编码", - "es": "codificación", } self.type = "code_agent" diff --git a/sources/agents/file_agent.py b/sources/agents/file_agent.py index b94e288..bcfaa6d 100644 --- a/sources/agents/file_agent.py +++ b/sources/agents/file_agent.py @@ -19,13 +19,12 @@ class FileAgent(Agent): "en": "files", "fr": "fichiers", "zh": "文件", - "es": "archivos", } self.type = "file_agent" def process(self, prompt, speech_module) -> str: exec_success = False - prompt += f"\nWork directory: {self.work_dir}" + prompt += f"\nYou must work in directory: {self.work_dir}" self.memory.push('user', prompt) while exec_success is False: self.wait_message(speech_module) diff --git a/sources/agents/planner_agent.py b/sources/agents/planner_agent.py index 8a46335..7b8844f 100644 --- a/sources/agents/planner_agent.py +++ b/sources/agents/planner_agent.py @@ -18,15 +18,14 @@ class PlannerAgent(Agent): self.tools['json'].tag = "json" self.browser = browser self.agents = { - "coder": CoderAgent(name, "prompts/coder_agent.txt", provider, verbose=False), - "file": FileAgent(name, "prompts/file_agent.txt", provider, verbose=False), - "web": BrowserAgent(name, "prompts/browser_agent.txt", provider, verbose=False, browser=browser) + "coder": CoderAgent(name, "prompts/base/coder_agent.txt", provider, verbose=False), + "file": FileAgent(name, "prompts/base/file_agent.txt", provider, verbose=False), + "web": BrowserAgent(name, "prompts/base/browser_agent.txt", provider, verbose=False, browser=browser) } self.role = { "en": "Research, setup and code", "fr": "Recherche, configuration et codage", "zh": "研究，设置和编码", - "es": "Investigación, configuración y code" } self.type = "planner_agent" diff --git a/sources/interaction.py b/sources/interaction.py index a25a95c..128da6d 100644 --- a/sources/interaction.py +++ b/sources/interaction.py @@ -90,13 +90,13 @@ class Interaction: self.last_query = query return query - def think(self) -> None: + def think(self) -> bool: """Request AI agents to process the user input.""" if self.last_query is None or len(self.last_query) == 0: - return + return False agent = self.router.select_agent(self.last_query) if agent is None: - return + return False if self.current_agent != agent and self.last_answer is not None: ## get last history from previous agent self.current_agent.memory.push('user', self.last_query) @@ -106,6 +106,7 @@ class Interaction: self.last_answer, _ = agent.process(self.last_query, self.speech) if self.last_answer == tmp: self.last_answer = None + return True def show_answer(self) -> None: """Show the answer to the user.""" diff --git a/sources/language.py b/sources/language.py index 16c80d2..172e16f 100644 --- a/sources/language.py +++ b/sources/language.py @@ -19,7 +19,7 @@ class LanguageUtility: text: string to analyze Returns: ISO639-1 language code """ - langid.set_languages(['fr', 'en', 'zh', 'es']) + langid.set_languages(['fr', 'en', 'zh']) lang, score = langid.classify(text) return lang diff --git a/sources/llm_provider.py b/sources/llm_provider.py index 58b9f9b..9fc48ad 100644 --- a/sources/llm_provider.py +++ b/sources/llm_provider.py @@ -26,7 +26,8 @@ class Provider: "openai": self.openai_fn, "lm-studio": self.lm_studio_fn, "huggingface": self.huggingface_fn, - "deepseek": self.deepseek_fn + "deepseek": self.deepseek_fn, + "test": self.test_fn } self.api_key = None self.unsafe_providers = ["openai", "deepseek"] @@ -245,13 +246,27 @@ class Provider: This is a test response from the test provider. Change provider to 'ollama' or 'server' to get real responses. + This is python saying hello. ```python print("Hello world from python") ``` + This is ls -la from bash. ```bash - echo "Hello world from bash" + ls -la ``` + + This is pwd from bash. + ```bash + pwd + ``` + + This is unsafe command. + ```bash + rm does_not_exist.txt + ``` + + goodbye """ return thought diff --git a/sources/memory.py b/sources/memory.py index 018c49a..8d4c323 100644 --- a/sources/memory.py +++ b/sources/memory.py @@ -9,7 +9,7 @@ import json sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from sources.utility import timer_decorator +from sources.utility import timer_decorator, pretty_print class Memory(): """ @@ -33,9 +33,8 @@ class Memory(): self.model = "pszemraj/led-base-book-summary" self.device = self.get_cuda_device() self.memory_compression = memory_compression - if memory_compression: - self.tokenizer = AutoTokenizer.from_pretrained(self.model) - self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model) + self.tokenizer = AutoTokenizer.from_pretrained(self.model) + self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model) def get_filename(self) -> str: return f"memory_{self.session_time.strftime('%Y-%m-%d_%H-%M-%S')}.txt" @@ -67,17 +66,24 @@ class Memory(): def load_memory(self, agent_type: str = "casual_agent") -> None: """Load the memory from the last session.""" + pretty_print(f"Loading {agent_type} past memories... ", color="status") if self.session_recovered == True: return save_path = os.path.join(self.conversation_folder, agent_type) if not os.path.exists(save_path): + pretty_print("No memory to load.", color="success") return filename = self.find_last_session_path(save_path) if filename is None: + pretty_print("Last session memory not found.", color="warning") return path = os.path.join(save_path, filename) with open(path, 'r') as f: self.memory = json.load(f) + if self.memory[-1]['role'] == 'user': + self.memory.pop() + self.compress() + pretty_print("Session recovered successfully", color="success") def reset(self, memory: list) -> None: self.memory = memory @@ -86,7 +92,9 @@ class Memory(): """Push a message to the memory.""" if self.memory_compression and role == 'assistant': self.compress() - # we don't compress the last message + curr_idx = len(self.memory) + if self.memory[curr_idx-1]['content'] == content: + pretty_print("Warning: same message have been pushed twice to memory", color="error") self.memory.append({'role': role, 'content': content}) def clear(self) -> None: @@ -114,6 +122,8 @@ class Memory(): """ if self.tokenizer is None or self.model is None: return text + if len(text) < min_length*1.5: + return text max_length = len(text) // 2 if len(text) > min_length*2 else min_length*2 input_text = "summarize: " + text inputs = self.tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True) @@ -134,12 +144,12 @@ class Memory(): """ Compress the memory using the AI model. """ - if not self.memory_compression: - return for i in range(len(self.memory)): if i < 3: continue - if len(self.memory[i]['content']) > 1024: + if self.memory[i]['role'] == 'system': + continue + if len(self.memory[i]['content']) > 128: self.memory[i]['content'] = self.summarize(self.memory[i]['content']) if __name__ == "__main__": diff --git a/sources/router.py b/sources/router.py index f2c7b34..8356570 100644 --- a/sources/router.py +++ b/sources/router.py @@ -384,7 +384,7 @@ class AgentRouter: pretty_print(f"Complex task detected, routing to planner agent.", color="info") return self.find_planner_agent() for agent in self.agents: - if best_agent == agent.role[lang]: + if best_agent == agent.role["en"]: pretty_print(f"Selected agent: {agent.agent_name} (roles: {agent.role[lang]})", color="warning") return agent pretty_print(f"Error choosing agent.", color="failure") diff --git a/sources/tools/BashInterpreter.py b/sources/tools/BashInterpreter.py index de7ac69..dcc9ab5 100644 --- a/sources/tools/BashInterpreter.py +++ b/sources/tools/BashInterpreter.py @@ -6,8 +6,10 @@ import subprocess if __name__ == "__main__": from tools import Tools + from safety import is_unsafe else: from sources.tools.tools import Tools + from sources.tools.safety import is_unsafe class BashInterpreter(Tools): """ @@ -19,16 +21,16 @@ class BashInterpreter(Tools): def language_bash_attempt(self, command: str): """ - detect if AI attempt to run the code using bash. - if so, return True, otherwise return False. - The philosophy is that code written by the AI will be executed, so it should not use bash to run it. + Detect if AI attempt to run the code using bash. + If so, return True, otherwise return False. + Code written by the AI will be executed automatically, so it should not use bash to run it. """ lang_interpreter = ["python3", "gcc", "g++", "go", "javac", "rustc", "clang", "clang++", "rustc", "rustc++", "rustc++"] for word in command.split(): if word in lang_interpreter: return True return False - + def execute(self, commands: str, safety=False, timeout=1000): """ Execute bash commands and display output in real-time. @@ -38,6 +40,9 @@ class BashInterpreter(Tools): concat_output = "" for command in commands: + command = command.replace('\n', '') + if self.safe_mode and is_unsafe(commands): + return "Unsafe command detected, execution aborted." if self.language_bash_attempt(command): continue try: @@ -50,12 +55,11 @@ class BashInterpreter(Tools): ) command_output = "" for line in process.stdout: - print(line, end="") command_output += line return_code = process.wait(timeout=timeout) if return_code != 0: return f"Command {command} failed with return code {return_code}:\n{command_output}" - concat_output += f"Output of {command}:\n{command_output.strip()}\n\n" + concat_output += f"Output of {command}:\n{command_output.strip()}\n" except subprocess.TimeoutExpired: process.kill() # Kill the process if it times out return f"Command {command} timed out. Output:\n{command_output}" diff --git a/sources/tools/safety.py b/sources/tools/safety.py new file mode 100644 index 0000000..2676e2b --- /dev/null +++ b/sources/tools/safety.py @@ -0,0 +1,82 @@ +import os +import sys + +unsafe_commands_unix = [ + "rm", # File/directory removal + "dd", # Low-level disk writing + "mkfs", # Filesystem formatting + "chmod", # Permission changes + "chown", # Ownership changes + "shutdown", # System shutdown + "reboot", # System reboot + "halt", # System halt + "sysctl", # Kernel parameter changes + "kill", # Process termination + "pkill", # Kill by process name + "killall", # Kill all matching processes + "exec", # Replace process with command + "tee", # Write to files with privileges + "umount", # Unmount filesystems + "passwd", # Password changes + "useradd", # Add users + "userdel", # Delete users + "groupadd", # Add groups + "groupdel", # Delete groups + "visudo", # Edit sudoers file + "screen", # Terminal session management + "fdisk", # Disk partitioning + "parted", # Disk partitioning + "chroot", # Change root directory + "route" # Routing table management +] + +unsafe_commands_windows = [ + "del", # Deletes files + "erase", # Alias for del, deletes files + "rd", # Removes directories (rmdir alias) + "rmdir", # Removes directories + "format", # Formats a disk, erasing data + "diskpart", # Manages disk partitions, can wipe drives + "chkdsk /f", # Fixes filesystem, can alter data + "fsutil", # File system utilities, can modify system files + "xcopy /y", # Copies files, overwriting without prompt + "copy /y", # Copies files, overwriting without prompt + "move", # Moves files, can overwrite + "attrib", # Changes file attributes, e.g., hiding or exposing files + "icacls", # Changes file permissions (modern) + "takeown", # Takes ownership of files + "reg delete", # Deletes registry keys/values + "regedit /s", # Silently imports registry changes + "shutdown", # Shuts down or restarts the system + "schtasks", # Schedules tasks, can run malicious commands + "taskkill", # Kills processes + "wmic", # Deletes processes via WMI + "bcdedit", # Modifies boot configuration + "powercfg", # Changes power settings, can disable protections + "assoc", # Changes file associations + "ftype", # Changes file type commands + "cipher /w", # Wipes free space, erasing data + "esentutl", # Database utilities, can corrupt system files + "subst", # Substitutes drive paths, can confuse system + "mklink", # Creates symbolic links, can redirect access + "bootcfg" +] + +def is_unsafe(cmd): + """ + check if a bash command is unsafe. + """ + if sys.platform.startswith("win"): + if any(c in cmd for c in unsafe_commands_windows): + return True + else: + if any(c in cmd for c in unsafe_commands_unix): + return True + return False + +if __name__ == "__main__": + cmd = input("Enter a command: ") + if is_unsafe(cmd): + print("Unsafe command detected!") + else: + print("Command is safe to execute.") \ No newline at end of file diff --git a/sources/tools/tools.py b/sources/tools/tools.py index e3ffa25..c5fb7fe 100644 --- a/sources/tools/tools.py +++ b/sources/tools/tools.py @@ -34,6 +34,7 @@ class Tools(): self.config = configparser.ConfigParser() self.current_dir = self.create_work_dir() self.excutable_blocks_found = False + self.safe_mode = True def get_work_dir(self): return self.current_dir @@ -186,7 +187,7 @@ class Tools(): code_blocks.append(content) start_index = end_pos + len(end_tag) return code_blocks, save_path - + if __name__ == "__main__": tool = Tools() tool.tag = "python" diff --git a/sources/utility.py b/sources/utility.py index ad1c161..8639488 100644 --- a/sources/utility.py +++ b/sources/utility.py @@ -53,13 +53,6 @@ def pretty_print(text, color = "info"): print(colored(text, color_map[color])) def animate_thinking(text, color="status", duration=2): - """ - Display an animated "thinking..." indicator in a separate thread. - Args: - text (str): Text to display - color (str): Color for the text - duration (float): How long to animate in seconds - """ def _animate(): color_map = { "success": (Fore.GREEN, "green"), @@ -71,22 +64,21 @@ def animate_thinking(text, color="status", duration=2): "default": (Fore.RESET, "black"), "info": (Fore.CYAN, "cyan") } - - fore_color, term_color = color_map[color] + fore_color, term_color = color_map.get(color, color_map["default"]) spinner = itertools.cycle(['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']) end_time = time.time() + duration while time.time() < end_time: symbol = next(spinner) if platform.system().lower() != "windows": - print(f"\r{fore_color}{symbol} {text}{Fore.RESET}", end="", flush=True) + print(f"{fore_color}{symbol} {text}{Fore.RESET}", flush=True) else: - print(colored(f"\r{symbol} {text}", term_color), end="", flush=True) + print(colored(f"{symbol} {text}", term_color), flush=True) time.sleep(0.1) - print() + print("\033[1A\033[K", end="", flush=True) animation_thread = threading.Thread(target=_animate) - animation_thread.daemon = True animation_thread.start() + animation_thread.join() def timer_decorator(func): """ @@ -101,6 +93,6 @@ def timer_decorator(func): start_time = time() result = func(*args, **kwargs) end_time = time() - print(f"{func.__name__} took {end_time - start_time:.2f} seconds to execute") + pretty_print(f"{func.__name__} took {end_time - start_time:.2f} seconds to execute", "status") return result return wrapper \ No newline at end of file