Spaces:
Sleeping
Sleeping
Update app.py
Browse files
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, #
|
62 |
-
"TEST-CODE": test_code, #
|
63 |
-
"INTEGRATE-CODE": integrate_code, #
|
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 |
-
|
104 |
-
|
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()
|