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 @dataclass 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 @abstractmethod 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 = "" for tag, content in structure.items(): html += f"<{tag}>{content}" 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"

{project_config.name}

"}) 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) try: config_dict = json.loads(response) return ProjectConfig(**config_dict) except JSONDecodeError as e: self.logger.error(f"Error decoding JSON: {str(e)}") self.logger.error(f"Response from model: {response}") return None # Or handle the error differently 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 code examples. """ return self.generate_agent_response(prompt)