Dev_Agent/modify_code.py at main · Anish-CodeDev/Dev_Agent · GitHub

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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
from google.genai import Client,types
import os
from dotenv import load_dotenv
from rest import RestApi
load_dotenv()
gemini = Client()


class ModifyCodeFuncs:
    def __init__(self,instruction,folder_name):
        self.instruction = instruction
        self.model = "gemma-4-26b-a4b-it"
        self.folder_name = folder_name
    def generate_summary_for_code(self,filename:str):
        with open(filename,"r") as f:
            code = f.read()
        res = gemini.models.generate_content(
            model=self.model,
            contents=f"""
            You are a code analysis tool.
            Analyze the following code and provide a summary of what it does.
            Code:
            {code}
            """,
            config=types.GenerateContentConfig(
                response_mime_type="application/json",
                response_schema={
                    "type": "object",
                    "properties": {
                        "summary": {
                            "type": "string",
                            "description": "Summary of the code"
                        }
                    },
                    "required": ["summary"]
                }
            )
        )
        return res.text
    def compile_descriptions_of_all_files(self):
        res = []
        for filename in os.listdir(self.folder_name):
            if not filename.endswith('.py'):
                continue
            if filename == '.git' or filename == 'agent.py':
                continue
            with open(os.path.join(self.folder_name, filename), "r") as f:
                code = f.read()
            # Use Gemma model to determine if the file is relevant to the instruction
            relevance_res = gemini.models.generate_content(
                model=self.model,
                contents=f"""
                You are a code analysis tool.
                Given the following code from file "{filename}" and the user's task: "{self.instruction}",
                determine whether this file is relevant to the task and would need to be modified.

                Code:
                {code}
                """,
                config=types.GenerateContentConfig(
                    response_mime_type="application/json",
                    response_schema={
                        "type": "object",
                        "properties": {
                            "is_relevant": {
                                "type": "string",
                                "description": "yes if the file is relevant to the task, no otherwise"
                            }
                        },
                        "required": ["is_relevant"]
                    }
                )
            )
            relevance = eval(relevance_res.text)
            if relevance.get("is_relevant", "no").lower() == "yes":
                print(filename)
                res.append({"filename": filename, "summary": eval(self.generate_summary_for_code(os.path.join(self.folder_name, filename)))["summary"]})
                print(res)
        return res
    def generate_task_for_each_file_based_on_summary(self,summary:str):
        res = gemini.models.generate_content(
            model=self.model,
            contents=f"""
            You are a code analysis tool.
            Analyze the following code and provide a task for each file based on the summary with respect to the primary task of the user: {self.instruction}.
            If you feel it has no relevance with respect to the primary task of the user: {self.instruction}, set the task to None
            Code:
            {summary}
            """,
            config=types.GenerateContentConfig(
                response_mime_type="application/json",
                response_schema={
                    "type": "object",
                    "properties": {
                        "filename":{
                            "type": "string"
                        },
                        "task":{
                            "type": "string"
                        }
                    },
                    "required": ["filename","task"]
                }
            )
        )
        return eval(res.text)
    def generate_final_tasks(self):
        descriptions = self.compile_descriptions_of_all_files()
        res = []
        for description in descriptions:
            r = self.generate_task_for_each_file_based_on_summary(description["summary"])
            if r['task'] != "None":
                res.append(r)
        return res
'''
obj = ModifyCodeFuncs("I want you to add a new feature to the agent.py file which will help in modifying the code")
print(obj.generate_final_tasks())

obj = ModifyCodeFuncs("I want you to add a new feature to the agent.py file which will help in modifying the code")
print(obj.compile_descriptions_of_all_files())
'''