| from typing import Dict, Any | |
| import os | |
| from flows.base_flows import AtomicFlow | |
| from flows.utils.general_helpers import get_pyfile_functions_metadata_from_file | |
| class MemoryReadingAtomicFlow(AtomicFlow): | |
| """A flow to read memory from given files. | |
| Any composite flow that uses this flow should have | |
| memory_files: Dict[str, str] which maps memory name to its memory file location in the flow_state | |
| *Input Interface*: - `memory_files` : name of the Dict which maps the memory name to its file location e.g. | |
| {"plan": "examples/JARVIS/plan.txt"} | |
| """ | |
| def __init__(self, **kwargs): | |
| super().__init__(**kwargs) | |
| self.supported_mem_name = ["plan", "logs", "code_library"] | |
| def _check_input_data(self, input_data: Dict[str, Any]): | |
| """input data sanity check""" | |
| assert "memory_files" in input_data, "memory_files not passed to MemoryReadingAtomicFlow" | |
| for mem_name, mem_path in input_data["memory_files"].items(): | |
| assert mem_name in self.supported_mem_name, (f"{mem_name} is not supported in MemoryReadingAtomicFlow, " | |
| f"supported names are: {self.supported_mem_name}") | |
| assert os.path.exists(mem_path), f"{mem_path} does not exist." | |
| assert os.path.isfile(mem_path), f"{mem_path} is not a file." | |
| def _read_text(self, file_location): | |
| with open(file_location, 'r', encoding='utf-8') as file: | |
| content = file.read() | |
| return content | |
| def _read_py_code_library(self, file_location): | |
| metadata = get_pyfile_functions_metadata_from_file(file_location) | |
| if len(metadata) == 0: | |
| return "No functions yet." | |
| descriptions = [] | |
| for function in metadata: | |
| description = f"Function Name: {function['name']}\n" | |
| description += f"Documentation: {function.get('docstring', 'No documentation')}\n" | |
| description += f"Parameters: {', '.join(function.get('parameters', []))}\n" | |
| description += f"Default Values: {function.get('defaults', 'None')}" | |
| descriptions.append(description) | |
| return "\n\n".join(descriptions) | |
| def run( | |
| self, | |
| input_data: Dict[str, Any]): | |
| self._check_input_data(input_data) | |
| response = {} | |
| for mem_name, mem_path in input_data["memory_files"].items(): | |
| if mem_name in ['plan', 'logs']: | |
| response[mem_name] = self._read_text(mem_path) | |
| elif mem_name == 'code_library' and mem_path.endswith('.py'): | |
| response[mem_name] = self._read_py_code_library(mem_path) | |
| return response | |