File size: 2,938 Bytes
702d20f
 
 
 
 
 
 
 
 
 
 
 
 
e5bf02b
702d20f
fc85cfc
702d20f
 
fc85cfc
702d20f
 
 
 
 
 
 
 
 
 
 
 
 
35a0e84
702d20f
 
35a0e84
702d20f
 
9c146b2
fc85cfc
 
702d20f
 
 
 
d9b59c4
fc85cfc
702d20f
 
 
 
 
 
 
 
 
 
 
 
 
 
35a0e84
702d20f
 
 
 
 
 
 
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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import subprocess
import webbrowser
from typing import Dict, Any
from flows.base_flows.atomic import AtomicFlow
import os


class PlanFileEditAtomicFlow(AtomicFlow):
    def _generate_content(self, plan_file_location, plan_str) -> str:
        content = (
                "The below plan will be written to " +
                plan_file_location + "\n"
                "Edit the plan directly or provide your thoughts down below if you have any suggestions.\n"
                "When you are done editing plan and providing feedback, save file and close the current VSCode session to continue.\n"
                "###########\n"
                "Plan:\n" +
                plan_str +
                "\n############\n"
                "Thoughts:"
        )
        return content

    def _generate_temp_file_location(self, plan_file_location):
        directory = os.path.dirname(plan_file_location)
        ret = os.path.join(directory, 'temp_plan.txt')
        return ret

    def _write_plan_content_to_file(self, file_location, content: str):
        try:
            with open(file_location, "w") as file:
                file.write(content)

            return True, f"Plan written to {file_location}", file_location

        except Exception as e:
            return False, str(e), file_location

    def _check_input(self, input_data: Dict[str, Any]):
        assert any(item in input_data for item in ["plan", "new_plan"]), "plan or new_plan is not passed to PlanFileEditAtomicFlow"
        assert "plan_file_location" in input_data, "plan_file_location not passed to PlanFileEditAtomicFlow"
        plan_file_loc = input_data["plan_file_location"]
        assert os.path.exists(plan_file_loc), f"{plan_file_loc} does not exist"
        assert os.path.isfile(plan_file_loc), f"{plan_file_loc} is not a file"

    def _generate_input_to_writer(self, input_data: Dict[str, Any]):
        plan_str = input_data['plan'] if "plan" in input_data else input_data['new_plan']
        plan_file_location = input_data["plan_file_location"]
        content_to_write = self._generate_content(plan_file_location, plan_str)
        file_location_to_write = self._generate_temp_file_location(plan_file_location)
        return content_to_write, file_location_to_write

    def run(
            self,
            input_data: Dict[str, Any]
    ):
        self._check_input(input_data)

        # ~~~ Getting input data to the file editor ~~~
        content_to_write, file_location_to_write = self._generate_input_to_writer(input_data)

        # ~~~ Calling the writer function ~~~
        result, plan_editor_output, temp_file_location = self._write_plan_content_to_file(
            file_location_to_write, content_to_write)

        # ~~~ Generating return variables ~~~
        response = {}
        response["plan_editor_output"] = plan_editor_output
        response["temp_plan_file_location"] = temp_file_location
        return response