Spaces:
Runtime error
Runtime error
| import os | |
| import json | |
| import logging | |
| from typing import List, Dict, Tuple, Optional, Any | |
| from dataclasses import dataclass | |
| from abc import ABC, abstractmethod | |
| from huggingface_hub import HfApi, InferenceApi | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline | |
| class ProjectConfig: | |
| name: str | |
| description: str | |
| technologies: List[str] | |
| structure: Dict[str, List[str]] | |
| class WebDevelopmentTool(ABC): | |
| def __init__(self, name: str, description: str): | |
| self.name = name | |
| self.description = description | |
| def generate_code(self, *args, **kwargs): | |
| pass | |
| class HTMLGenerator(WebDevelopmentTool): | |
| def __init__(self): | |
| super().__init__("HTML Generator", "Generates HTML code for web pages") | |
| def generate_code(self, structure: Dict[str, Any]) -> str: | |
| html = "<html><body>" | |
| for tag, content in structure.items(): | |
| html += f"<{tag}>{content}</{tag}>" | |
| html += "</body></html>" | |
| return html | |
| class CSSGenerator(WebDevelopmentTool): | |
| def __init__(self): | |
| super().__init__("CSS Generator", "Generates CSS code for styling web pages") | |
| def generate_code(self, styles: Dict[str, Dict[str, str]]) -> str: | |
| css = "" | |
| for selector, properties in styles.items(): | |
| css += f"{selector} {{\n" | |
| for prop, value in properties.items(): | |
| css += f" {prop}: {value};\n" | |
| css += "}\n" | |
| return css | |
| class JavaScriptGenerator(WebDevelopmentTool): | |
| def __init__(self): | |
| super().__init__("JavaScript Generator", "Generates JavaScript code for web functionality") | |
| def generate_code(self, functions: List[Dict[str, Any]]) -> str: | |
| js = "" | |
| for func in functions: | |
| js += f"function {func['name']}({', '.join(func['params'])}) {{\n" | |
| js += f" {func['body']}\n" | |
| js += "}\n\n" | |
| return js | |
| class EnhancedAIAgent: | |
| def __init__(self, name: str, description: str, skills: List[str], model_name: str): | |
| self.name = name | |
| self.description = description | |
| self.skills = skills | |
| self.model_name = model_name | |
| self.html_gen_tool = HTMLGenerator() | |
| self.css_gen_tool = CSSGenerator() | |
| self.js_gen_tool = JavaScriptGenerator() | |
| self.hf_api = HfApi() | |
| self.inference_api = InferenceApi(repo_id=model_name, token=os.environ.get("HF_API_TOKEN")) | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| self.model = AutoModelForCausalLM.from_pretrained(model_name) | |
| self.text_generation = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer) | |
| self.logger = logging.getLogger(__name__) | |
| def generate_agent_response(self, prompt: str) -> str: | |
| try: | |
| response = self.inference_api(prompt) | |
| return response[0]['generated_text'] | |
| except Exception as e: | |
| self.logger.error(f"Error generating response: {str(e)}") | |
| return f"Error: Unable to generate response. {str(e)}" | |
| def create_project_structure(self, project_config: ProjectConfig) -> Dict[str, str]: | |
| project_files = {} | |
| for directory, files in project_config.structure.items(): | |
| for file in files: | |
| file_path = os.path.join(directory, file) | |
| if file.endswith('.html'): | |
| content = self.html_gen_tool.generate_code({"body": f"<h1>{project_config.name}</h1>"}) | |
| elif file.endswith('.css'): | |
| content = self.css_gen_tool.generate_code({"body": {"font-family": "Arial, sans-serif"}}) | |
| elif file.endswith('.js'): | |
| content = self.js_gen_tool.generate_code([{"name": "init", "params": [], "body": "console.log('Initialized');"}]) | |
| else: | |
| content = f"// TODO: Implement {file}" | |
| project_files[file_path] = content | |
| return project_files | |
| def generate_project_config(self, project_description: str) -> ProjectConfig: | |
| prompt = f""" | |
| Based on the following project description, generate a ProjectConfig object: | |
| Description: {project_description} | |
| The ProjectConfig should include: | |
| - name: A short, descriptive name for the project | |
| - description: A brief summary of the project | |
| - technologies: A list of technologies to be used (e.g., ["HTML", "CSS", "JavaScript", "React"]) | |
| - structure: A dictionary representing the file structure, where keys are directories and values are lists of files | |
| Respond with a JSON object representing the ProjectConfig. | |
| """ | |
| response = self.generate_agent_response(prompt) | |
| config_dict = json.loads(response) | |
| return ProjectConfig(**config_dict) | |
| def implement_feature(self, feature_description: str, existing_code: Optional[str] = None) -> str: | |
| prompt = f""" | |
| Feature to implement: {feature_description} | |
| Existing code: | |
| ``` | |
| {existing_code if existing_code else 'No existing code provided.'} | |
| ``` | |
| Please implement the described feature, modifying the existing code if provided. | |
| Respond with only the code, no explanations. | |
| """ | |
| return self.generate_agent_response(prompt) | |
| def review_code(self, code: str) -> str: | |
| prompt = f""" | |
| Please review the following code and provide feedback: | |
| ``` | |
| {code} | |
| ``` | |
| Consider the following aspects in your review: | |
| 1. Code quality and readability | |
| 2. Potential bugs or errors | |
| 3. Adherence to best practices | |
| 4. Suggestions for improvement | |
| Provide your feedback in a structured format. | |
| """ | |
| return self.generate_agent_response(prompt) | |
| def optimize_code(self, code: str, optimization_goal: str) -> str: | |
| prompt = f""" | |
| Please optimize the following code with the goal of improving {optimization_goal}: | |
| ``` | |
| {code} | |
| ``` | |
| Provide only the optimized code in your response, no explanations. | |
| """ | |
| return self.generate_agent_response(prompt) | |
| def generate_documentation(self, code: str) -> str: | |
| prompt = f""" | |
| Please generate comprehensive documentation for the following code: | |
| ``` | |
| {code} | |
| ``` | |
| Include the following in your documentation: | |
| 1. Overview of the code's purpose | |
| 2. Description of functions/classes and their parameters | |
| 3. Usage examples | |
| 4. Any important notes or considerations | |
| Provide the documentation in Markdown format. | |
| """ | |
| return self.generate_agent_response(prompt) | |
| def suggest_tests(self, code: str) -> str: | |
| prompt = f""" | |
| Please suggest unit tests for the following code: | |
| ``` | |
| {code} | |
| ``` | |
| For each function or class, provide: | |
| 1. Test case description | |
| 2. Input values | |
| 3. Expected output or behavior | |
| Provide the suggestions in a structured format. | |
| """ | |
| return self.generate_agent_response(prompt) | |
| def explain_code(self, code: str) -> str: | |
| prompt = f""" | |
| Please provide a detailed explanation of the following code: | |
| ``` | |
| {code} | |
| ``` | |
| Include in your explanation: | |
| 1. Overall purpose of the code | |
| 2. Breakdown of each significant part | |
| 3. How different components interact | |
| 4. Any notable algorithms or design patterns used | |
| Explain in a way that would be understandable to a junior developer. | |
| """ | |
| return self.generate_agent_response(prompt) | |
| def suggest_refactoring(self, code: str) -> str: | |
| prompt = f""" | |
| Please suggest refactoring improvements for the following code: | |
| ``` | |
| {code} | |
| ``` | |
| Consider the following in your suggestions: | |
| 1. Improving code readability | |
| 2. Enhancing maintainability | |
| 3. Applying design patterns where appropriate | |
| 4. Optimizing performance (if applicable) | |
| Provide specific suggestions and explain the benefits of each. | |
| """ | |
| return self.generate_agent_response(prompt) |