Spaces:
Sleeping
Sleeping
| 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) | |