triAGI-Coder

Runtime error

File size: 10,337 Bytes

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 = "<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

    import json
from json.decoder import JSONDecodeError

def generate_project_config(self, project_description: str) -> Optional[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:
        # Try to find and extract a JSON object from the response
        json_start = response.find('{')
        json_end = response.rfind('}') + 1
        if json_start != -1 and json_end != -1:
            json_str = response[json_start:json_end]
            config_dict = json.loads(json_str)
            return ProjectConfig(**config_dict)
        else:
            raise ValueError("No JSON object found in the response")
    except (JSONDecodeError, ValueError) as e:
        self.logger.error(f"Error parsing JSON from response: {str(e)}")
        self.logger.error(f"Full response from model: {response}")
        
        # Attempt to salvage partial information
        try:
            partial_config = self.extract_partial_config(response)
            if partial_config:
                self.logger.warning("Extracted partial config from malformed response")
                return partial_config
        except Exception as ex:
            self.logger.error(f"Failed to extract partial config: {str(ex)}")
        
        return None

def extract_partial_config(self, response: str) -> Optional[ProjectConfig]:
    """Attempt to extract partial configuration from a malformed response."""
    name = self.extract_field(response, "name")
    description = self.extract_field(response, "description")
    technologies = self.extract_list(response, "technologies")
    structure = self.extract_dict(response, "structure")
    
    if name and description:
        return ProjectConfig(
            name=name,
            description=description,
            technologies=technologies or [],
            structure=structure or {}
        )
    return None

def extract_field(self, text: str, field: str) -> Optional[str]:
    """Extract a simple field value from text."""
    match = re.search(rf'"{field}"\s*:\s*"([^"]*)"', text)
    return match.group(1) if match else None

def extract_list(self, text: str, field: str) -> Optional[List[str]]:
    """Extract a list from text."""
    match = re.search(rf'"{field}"\s*:\s*\[(.*?)\]', text, re.DOTALL)
    if match:
        items = re.findall(r'"([^"]*)"', match.group(1))
        return items
    return None

def extract_dict(self, text: str, field: str) -> Optional[Dict[str, List[str]]]:
    """Extract a dictionary from text."""
    match = re.search(rf'"{field}"\s*:\s*\{{(.*?)\}}', text, re.DOTALL)
    if match:
        dict_str = match.group(1)
        result = {}
        for item in re.finditer(r'"([^"]*)"\s*:\s*\[(.*?)\]', dict_str, re.DOTALL):
            key = item.group(1)
            values = re.findall(r'"([^"]*)"', item.group(2))
            result[key] = values
        return result
    return None
    
    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)