Skip to content

Commit 0cd6897

Browse files
FredyRivera-devFredyRivera-dev
committed
Add Tools write_file, append_file and edit_file
1 parent fd4f473 commit 0cd6897

File tree

6 files changed

+555
-57
lines changed

6 files changed

+555
-57
lines changed

‎atlasai/ai/ai_agent.py‎

Lines changed: 109 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,64 @@ async def analyze_project(self, project_dir: str, callback: Optional[Callable[[s
210210
},
211211
"strict": True
212212
}
213+
},
214+
{
215+
"type": "function",
216+
"function": {
217+
"name": "write_file",
218+
"description": "Create a new file or overwrite an existing one",
219+
"parameters": {
220+
"type": "object",
221+
"properties": {
222+
"file_path": {
223+
"type": "string",
224+
"description": "Path to the file to be written"
225+
},
226+
"content": {
227+
"type": "string",
228+
"description": "Content to write to the file"
229+
},
230+
"overwrite": {
231+
"type": "boolean",
232+
"description": "If True, will overwrite an existing file"
233+
}
234+
},
235+
"required": ["file_path", "content", "overwrite"],
236+
"additionalProperties": False
237+
},
238+
"strict": True
239+
}
240+
},
241+
{
242+
"type": "function",
243+
"function": {
244+
"name": "edit_file",
245+
"description": "Edit parts of an existing file",
246+
"parameters": {
247+
"type": "object",
248+
"properties": {
249+
"file_path": {
250+
"type": "string",
251+
"description": "Path to the file to be edited"
252+
},
253+
"search_text": {
254+
"type": "string",
255+
"description": "Text to search for"
256+
},
257+
"replace_text": {
258+
"type": "string",
259+
"description": "Replacement text"
260+
},
261+
"regex": {
262+
"type": "boolean",
263+
"description": "If True, search_text will be interpreted as a regular expression"
264+
}
265+
},
266+
"required": ["file_path", "search_text", "replace_text", "regex"],
267+
"additionalProperties": False
268+
},
269+
"strict": True
270+
}
213271
}
214272
]
215273

@@ -374,7 +432,7 @@ async def analyze_project(self, project_dir: str, callback: Optional[Callable[[s
374432
arguments_value = tool_call.get('function', {}).get('arguments', '{}')
375433

376434
# Validate tool name
377-
valid_tools = ["list_directory", "read_file", "execute_command"]
435+
valid_tools = ["list_directory", "read_file", "execute_command", "write_file", "edit_file"]
378436
if tool_name not in valid_tools:
379437
error_msg = f"Unknown tool: '{tool_name}'"
380438
logger.error(error_msg)
@@ -484,6 +542,18 @@ async def analyze_project(self, project_dir: str, callback: Optional[Callable[[s
484542
title=f"📑 File Content: {os.path.basename(file_path)}",
485543
border_style="green"
486544
))
545+
elif tool_name == "write_file" or tool_name == "append_file":
546+
callback(Panel(
547+
result,
548+
title=f"✏️ File Operation: {os.path.basename(arguments.get('file_path', ''))}",
549+
border_style="green"
550+
))
551+
elif tool_name == "edit_file":
552+
callback(Panel(
553+
result,
554+
title=f"🔄 Edited File: {os.path.basename(arguments.get('file_path', ''))}",
555+
border_style="blue"
556+
))
487557
else:
488558
# Generic result display
489559
max_display = 500
@@ -587,6 +657,44 @@ async def _execute_tool(self, tool_name: str, arguments: Dict[str, Any], working
587657
# Execute in the project directory
588658
result = execute_command(commands)
589659
return result
660+
661+
elif tool_name == "write_file":
662+
file_path = arguments.get("file_path")
663+
content = arguments.get("content", "")
664+
overwrite = arguments.get("overwrite", False)
665+
666+
if not file_path:
667+
return "Error: No se especificó ningún archivo para escribir"
668+
669+
# Convertir rutas relativas a absolutas
670+
if not os.path.isabs(file_path):
671+
file_path = os.path.abspath(file_path)
672+
673+
from atlasai.tools import write_file
674+
return write_file(file_path, content, overwrite)
675+
676+
elif tool_name == "edit_file":
677+
file_path = arguments.get("file_path")
678+
search_text = arguments.get("search_text")
679+
replace_text = arguments.get("replace_text", "")
680+
regex = arguments.get("regex", False)
681+
682+
if not file_path:
683+
return "Error: No se especificó ningún archivo para editar"
684+
685+
if search_text is None:
686+
return "Error: No se especificó el texto a buscar"
687+
688+
# Convertir rutas relativas a absolutas
689+
if not os.path.isabs(file_path):
690+
file_path = os.path.abspath(file_path)
691+
692+
# Verificar que el archivo existe
693+
if not os.path.isfile(file_path):
694+
return f"Error: El archivo '{file_path}' no existe"
695+
696+
from atlasai.tools import edit_file
697+
return edit_file(file_path, search_text, replace_text, regex)
590698

591699
else:
592700
return f"Error: Unknown tool '{tool_name}'"

‎atlasai/ai/general_agent.py‎

Lines changed: 119 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from openai import OpenAI
1010

1111
from atlasai.tools import get_os, search, list_directory, read_file, execute_command
12+
from atlasai.ai.prompts import BASE_PROMPTS
1213

1314
logger = logging.getLogger(__name__)
1415

@@ -27,62 +28,19 @@ def __init__(self, model="qwen3:8b", provider="ollama", api_key=None, stream=Fal
2728
self.api_key = api_key
2829
self.stream = stream
2930
self.language = language
30-
self.system_prompt = self._get_system_prompt()
31+
general_prompt = BASE_PROMPTS.get_general_agent_prompt(language)
32+
advanced_prompt = BASE_PROMPTS.get_advanced_agent_prompt(language)
33+
if language == "es":
34+
general_prompt = general_prompt.split("IMPORTANT:")[0].strip()
35+
36+
self.system_prompt = f"{general_prompt}\n\n{advanced_prompt}"
37+
3138

3239
# Initialize OpenAI client only when needed
3340
if self.provider == "openai":
3441
self.client = OpenAI(api_key=self.api_key)
3542

36-
def _get_system_prompt(self) -> str:
37-
"""Get the system prompt for the general agent."""
38-
base_prompt = """You are AtlasAI-CLI, an AI-powered assistant using a real operating system. You are a software expert: few engineers are as talented as you at understanding codebases, analyzing projects, and providing helpful information.
39-
40-
When to Communicate with User:
41-
- When encountering environment issues
42-
- When critical information cannot be accessed through available tools
43-
- When you need user permission to proceed with actions
44-
45-
Approach to Work:
46-
- Fulfill the user's request using all the tools available to you
47-
- When encountering difficulties, take time to gather information before concluding a root cause
48-
- If you're facing environment issues, find a way to continue your work if possible
49-
- Always consider that issues might be in your approach rather than the environment
50-
51-
Best Practices:
52-
- When analyzing files, first understand the file's code conventions
53-
- Mimic existing code style, use existing libraries and utilities, and follow existing patterns
54-
- NEVER assume that a given library is available unless you've confirmed it
55-
- When analyzing projects, look at existing components to understand their approach
56-
57-
Information Handling:
58-
- Use search capabilities to find information when needed
59-
- Thoroughly explore directories and files to understand project structure
60-
- Be concise but complete in your explanations
61-
62-
Data Security:
63-
- Treat code and data as sensitive information
64-
- Never commit secrets or keys to any repository
65-
- Always follow security best practices
66-
67-
I have access to the following tools to help analyze and interact with your system:
68-
69-
- get_os: Get information about the current operating system
70-
- search: Search the web for information
71-
- list_directory: View contents of directories
72-
- read_file: Examine files
73-
- execute_command: Execute shell commands (read-only safe commands)
74-
75-
I'll carefully analyze your query and decide which tools I need to use to provide the most helpful response. I'll always explain my reasoning and be clear about my process.
76-
"""
77-
78-
# Add language instruction
79-
if self.language == "es":
80-
base_prompt += """
8143

82-
IMPORTANT: For final reasoning and explanations, provide your analysis in Spanish while keeping technical terms, commands, and code in English. All other interactions with tools and system will remain in English."""
83-
84-
return base_prompt
85-
8644
def _get_available_tools(self) -> List[Dict[str, Any]]:
8745
"""Get the list of available tools."""
8846
return [
@@ -176,7 +134,65 @@ def _get_available_tools(self) -> List[Dict[str, Any]]:
176134
},
177135
"strict": True
178136
}
179-
}
137+
},
138+
{
139+
"type": "function",
140+
"function": {
141+
"name": "write_file",
142+
"description": "Create a new file or overwrite an existing one",
143+
"parameters": {
144+
"type": "object",
145+
"properties": {
146+
"file_path": {
147+
"type": "string",
148+
"description": "Path to the file to be written"
149+
},
150+
"content": {
151+
"type": "string",
152+
"description": "Content to write to the file"
153+
},
154+
"overwrite": {
155+
"type": "boolean",
156+
"description": "If True, will overwrite an existing file"
157+
}
158+
},
159+
"required": ["file_path", "content", "overwrite"],
160+
"additionalProperties": False
161+
},
162+
"strict": True
163+
}
164+
},
165+
{
166+
"type": "function",
167+
"function": {
168+
"name": "edit_file",
169+
"description": "Edit parts of an existing file",
170+
"parameters": {
171+
"type": "object",
172+
"properties": {
173+
"file_path": {
174+
"type": "string",
175+
"description": "Path to the file to be edited"
176+
},
177+
"search_text": {
178+
"type": "string",
179+
"description": "Text to search for"
180+
},
181+
"replace_text": {
182+
"type": "string",
183+
"description": "Replacement text"
184+
},
185+
"regex": {
186+
"type": "boolean",
187+
"description": "If True, search_text will be interpreted as a regular expression"
188+
}
189+
},
190+
"required": ["file_path", "search_text", "replace_text", "regex"],
191+
"additionalProperties": False
192+
},
193+
"strict": True
194+
}
195+
}
180196
]
181197

182198
def _format_directory_output(self, content):
@@ -297,6 +313,7 @@ async def process_query(self, query: str, callback: Optional[Callable[[str], Non
297313
)
298314
except Exception as e:
299315
error_msg = f"Error checking tool support: {str(e)}"
316+
logger.error(error_msg)
300317
if "does not support tools" in str(e) or "400" in str(e):
301318
error_msg = f"Model '{self.model}' does not support function calling/tools. Please use a compatible model."
302319
logger.error(error_msg)
@@ -417,7 +434,7 @@ async def process_query(self, query: str, callback: Optional[Callable[[str], Non
417434
arguments_value = tool_call.get('function', {}).get('arguments', '{}')
418435

419436
# Validate tool name
420-
valid_tools = ["get_os", "search", "list_directory", "read_file", "execute_command"]
437+
valid_tools = ["get_os", "search", "list_directory", "read_file", "execute_command", "write_file", "edit_file"]
421438
if tool_name not in valid_tools:
422439
error_msg = f"Unknown tool: '{tool_name}'"
423440
logger.error(error_msg)
@@ -550,6 +567,18 @@ async def process_query(self, query: str, callback: Optional[Callable[[str], Non
550567
title="🔎 Search Results",
551568
border_style="green"
552569
))
570+
elif tool_name == "write_file" or tool_name == "append_file":
571+
callback(Panel(
572+
result,
573+
title=f"✏️ File Operation: {os.path.basename(arguments.get('file_path', ''))}",
574+
border_style="green"
575+
))
576+
elif tool_name == "edit_file":
577+
callback(Panel(
578+
result,
579+
title=f"🔄 Edited File: {os.path.basename(arguments.get('file_path', ''))}",
580+
border_style="blue"
581+
))
553582
else:
554583
# Generic result display
555584
max_display = 500
@@ -660,6 +689,44 @@ async def _execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> str:
660689
return "Error: No commands specified to execute"
661690

662691
return execute_command(commands)
692+
693+
elif tool_name == "write_file":
694+
file_path = arguments.get("file_path")
695+
content = arguments.get("content", "")
696+
overwrite = arguments.get("overwrite", False)
697+
698+
if not file_path:
699+
return "Error: No se especificó ningún archivo para escribir"
700+
701+
# Convertir rutas relativas a absolutas
702+
if not os.path.isabs(file_path):
703+
file_path = os.path.abspath(file_path)
704+
705+
from atlasai.tools import write_file
706+
return write_file(file_path, content, overwrite)
707+
708+
elif tool_name == "edit_file":
709+
file_path = arguments.get("file_path")
710+
search_text = arguments.get("search_text")
711+
replace_text = arguments.get("replace_text", "")
712+
regex = arguments.get("regex", False)
713+
714+
if not file_path:
715+
return "Error: No se especificó ningún archivo para editar"
716+
717+
if search_text is None:
718+
return "Error: No se especificó el texto a buscar"
719+
720+
# Convertir rutas relativas a absolutas
721+
if not os.path.isabs(file_path):
722+
file_path = os.path.abspath(file_path)
723+
724+
# Verificar que el archivo existe
725+
if not os.path.isfile(file_path):
726+
return f"Error: El archivo '{file_path}' no existe"
727+
728+
from atlasai.tools import edit_file
729+
return edit_file(file_path, search_text, replace_text, regex)
663730

664731
else:
665732
return f"Error: Unknown tool '{tool_name}'"

0 commit comments

Comments
 (0)