Spaces:
Sleeping
Sleeping
File size: 1,979 Bytes
860424e 2a41ea2 860424e |
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 |
import os
import glob
from smolagents import CodeAgent
from prompt.prompt import get_prompt
class QuestionAnsweringAgent(CodeAgent):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def get_current_files(self):
"""Get a set of all files in the current working directory"""
try:
# Get all files in current directory (including hidden files)
all_files = set()
for pattern in ['*', '.*']:
all_files.update(glob.glob(pattern))
# Filter to only include actual files (not directories)
files = {f for f in all_files if os.path.isfile(f)}
return files
except Exception as e:
print(f"Error getting current files: {e}")
return set()
def cleanup_created_files(self, files_before):
"""Remove files that were created during execution"""
try:
files_after = self.get_current_files()
newly_created_files = files_after - files_before
for file_path in newly_created_files:
try:
if os.path.exists(file_path):
os.remove(file_path)
print(f"Cleaned up file: {file_path}")
except Exception as e:
print(f"Error cleaning up file {file_path}: {e}")
except Exception as e:
print(f"Error during cleanup: {e}")
def __call__(self, question_text, file_name) -> str:
# Take snapshot of files before execution
files_before = self.get_current_files()
try:
enhanced_question = get_prompt(question_text, file_name)
response = self.run(enhanced_question, reset=True)
return response
finally:
# Always clean up files, even if there's an exception
self.cleanup_created_files(files_before)
|