acecalisto3 commited on
Commit
a37df9c
·
verified ·
1 Parent(s): 6c0afa0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +43 -31
app.py CHANGED
@@ -6,21 +6,11 @@ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
6
  import black
7
  from pylint import lint
8
  from io import StringIO
9
- import openai
10
  import sys
11
  from datetime import datetime
12
  import requests
13
  from bs4 import BeautifulSoup
14
  from typing import List, Dict, Optional
15
- from utils import refine_code
16
-
17
- # Individual imports (assuming files are in the same directory as app.py)
18
- from utils.refine_code import refine_code # Import the function
19
- from utils.refine_code import CodeRefinementError # Import the class
20
- import test_code
21
- import integrate_code
22
- from test_code import CodeTestingError
23
- from integrate_code import CodeIntegrationError
24
 
25
  # --- Custom Exceptions for Enhanced Error Handling ---
26
  class InvalidActionError(Exception):
@@ -58,9 +48,9 @@ class AIAgent:
58
  self.tools = {
59
  "SEARCH": self.search,
60
  "CODEGEN": self.code_generation,
61
- "REFINE-CODE": refine_code, # Use external function
62
- "TEST-CODE": test_code, # Use external function
63
- "INTEGRATE-CODE": integrate_code, # Use external function
64
  "TEST-APP": self.test_app,
65
  "GENERATE-REPORT": self.generate_report,
66
  "WORKSPACE-EXPLORER": self.workspace_explorer,
@@ -81,7 +71,6 @@ class AIAgent:
81
  self.search_engine_url: str = "https://www.google.com/search?q=" # Default search engine
82
  self.prompts: List[str] = [] # Store prompts for future use
83
  self.code_generator = pipeline('text-generation', model='gpt2') # Initialize code generator
84
- self.openai_api_key = os.getenv("OPENAI_API_KEY") # Get OpenAI API Key from Environment Variable
85
 
86
  # --- Search Functionality ---
87
  def search(self, query: str) -> List[str]:
@@ -100,25 +89,48 @@ class AIAgent:
100
  def code_generation(self, snippet: str) -> str:
101
  """Generates code based on the provided snippet or description."""
102
  try:
103
- # Use OpenAI's GPT-3 for more advanced code generation
104
- if self.openai_api_key:
105
- openai.api_key = self.openai_api_key
106
- response = openai.Completion.create(
107
- engine="text-davinci-003",
108
- prompt=f"```python\n{snippet}\n```\nGenerate Python code based on this snippet or description:",
109
- max_tokens=500,
110
- temperature=0.7,
111
- top_p=1,
112
- frequency_penalty=0,
113
- presence_penalty=0
114
- )
115
- return response.choices[0].text
116
- else:
117
- generated_text = self.code_generator(snippet, max_length=500, num_return_sequences=1)[0]['generated_text']
118
- return generated_text
119
  except Exception as e:
120
  raise CodeGenerationError(f"Error during code generation: {e}")
121
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
122
  # --- App Testing Functionality ---
123
  def test_app(self) -> str:
124
  """Tests the functionality of the app."""
@@ -347,4 +359,4 @@ iface = gr.Interface(
347
  description="Interact with the AI Agent."
348
  )
349
 
350
- iface.launch()
 
6
  import black
7
  from pylint import lint
8
  from io import StringIO
 
9
  import sys
10
  from datetime import datetime
11
  import requests
12
  from bs4 import BeautifulSoup
13
  from typing import List, Dict, Optional
 
 
 
 
 
 
 
 
 
14
 
15
  # --- Custom Exceptions for Enhanced Error Handling ---
16
  class InvalidActionError(Exception):
 
48
  self.tools = {
49
  "SEARCH": self.search,
50
  "CODEGEN": self.code_generation,
51
+ "REFINE-CODE": self.refine_code, # Implemented within the class
52
+ "TEST-CODE": self.test_code, # Implemented within the class
53
+ "INTEGRATE-CODE": self.integrate_code, # Implemented within the class
54
  "TEST-APP": self.test_app,
55
  "GENERATE-REPORT": self.generate_report,
56
  "WORKSPACE-EXPLORER": self.workspace_explorer,
 
71
  self.search_engine_url: str = "https://www.google.com/search?q=" # Default search engine
72
  self.prompts: List[str] = [] # Store prompts for future use
73
  self.code_generator = pipeline('text-generation', model='gpt2') # Initialize code generator
 
74
 
75
  # --- Search Functionality ---
76
  def search(self, query: str) -> List[str]:
 
89
  def code_generation(self, snippet: str) -> str:
90
  """Generates code based on the provided snippet or description."""
91
  try:
92
+ generated_text = self.code_generator(snippet, max_length=500, num_return_sequences=1)[0]['generated_text']
93
+ return generated_text
 
 
 
 
 
 
 
 
 
 
 
 
 
 
94
  except Exception as e:
95
  raise CodeGenerationError(f"Error during code generation: {e}")
96
 
97
+ # --- Code Refinement Functionality ---
98
+ def refine_code(self, file_path: str) -> str:
99
+ """Refines the code in the specified file."""
100
+ try:
101
+ with open(file_path, 'r') as f:
102
+ code = f.read()
103
+ refined_code = black.format_str(code, mode=black.FileMode())
104
+ return refined_code
105
+ except black.InvalidInput:
106
+ raise CodeRefinementError("Error: Invalid code input for black formatting.")
107
+ except Exception as e:
108
+ raise CodeRefinementError(f"Error during code refinement: {e}")
109
+
110
+ # --- Code Testing Functionality ---
111
+ def test_code(self, file_path: str) -> str:
112
+ """Tests the code in the specified file."""
113
+ try:
114
+ results = lint.run(file_path, do_exit=False)
115
+ report = StringIO()
116
+ results.report(report=report)
117
+ return report.getvalue()
118
+ except Exception as e:
119
+ raise CodeTestingError(f"Error during code testing: {e}")
120
+
121
+ # --- Code Integration Functionality ---
122
+ def integrate_code(self, file_path: str, code_snippet: str) -> str:
123
+ """Integrates the code snippet into the specified file."""
124
+ try:
125
+ with open(file_path, 'r') as f:
126
+ existing_code = f.read()
127
+ integrated_code = existing_code + "\n" + code_snippet
128
+ with open(file_path, 'w') as f:
129
+ f.write(integrated_code)
130
+ return f"Code snippet integrated into {file_path} successfully."
131
+ except Exception as e:
132
+ raise CodeIntegrationError(f"Error during code integration: {e}")
133
+
134
  # --- App Testing Functionality ---
135
  def test_app(self) -> str:
136
  """Tests the functionality of the app."""
 
359
  description="Interact with the AI Agent."
360
  )
361
 
362
+ iface.launch()