Update supplemental.py

supplemental.py

@@ -1,12 +1,12 @@
 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 typing import List, Dict, Optional
+import re
+from huggingface_hub import HfApi, InferenceClient
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
+
 @dataclass
 class ProjectConfig:
     name: str
@@ -59,6 +59,25 @@ class JavaScriptGenerator(WebDevelopmentTool):
             js += "}\n\n"
         return js
 
+class ProjectConfig:
+    def __init__(self, name: str, description: str, technologies: List[str], structure: Dict[str, List[str]]):
+        self.name = name
+        self.description = description
+        self.technologies = technologies
+        self.structure = structure
+
+class HTMLGenerator:
+    def generate(self, content: str) -> str:
+        return f"<html><body>{content}</body></html>"
+
+class CSSGenerator:
+    def generate(self, styles: Dict[str, str]) -> str:
+        return "\n".join([f"{selector} {{ {'; '.join([f'{prop}: {value}' for prop, value in properties.items()])} }}" for selector, properties in styles.items()])
+
+class JavaScriptGenerator:
+    def generate(self, functionality: str) -> str:
+        return f"function main() {{ {functionality} }}"
+
 class EnhancedAIAgent:
     def __init__(self, name: str, description: str, skills: List[str], model_name: str):
         self.name = name
@@ -69,41 +88,28 @@ class EnhancedAIAgent:
         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.inference_client = InferenceClient(model=model_name, token=os.environ.get("HF_API_TOKEN"))
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, clean_up_tokenization_spaces=True)
+        self.model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+        self.text_generation = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer, clean_up_tokenization_spaces=True)
         self.logger = logging.getLogger(__name__)
+        self.logger.setLevel(logging.INFO)
+        handler = logging.StreamHandler()
+        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+        handler.setFormatter(formatter)
+        self.logger.addHandler(handler)
 
     def generate_agent_response(self, prompt: str) -> str:
         try:
-            response = self.inference_api(prompt)
-            return response[0]['generated_text']
+            response = self.inference_client.text_generation(prompt, max_new_tokens=100)
+            self.logger.info(f"Generated response for prompt: {prompt[:50]}...")
+            return response.generated_text
         except Exception as e:
-            self.logger.error(f"Error generating response: {str(e)}")
+            self.logger.error(f"Error generating response: {str(e)}", exc_info=True)
             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"""
+    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}
@@ -116,185 +122,97 @@ The ProjectConfig should include:
 
 Respond with a JSON object representing the ProjectConfig.
 """
-    response = self.generate_agent_response(prompt)
+        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)}")
+            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 (json.JSONDecodeError, ValueError) as e:
+            self.logger.error(f"Error parsing JSON from response: {str(e)}")
+            self.logger.error(f"Full response from model: {response}")
+            
+            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]:
+        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_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
+    def extract_field(self, text: str, field: str) -> Optional[str]:
+        match = re.search(rf'"{field}"\s*:\s*"([^"]*)"', text)
+        return match.group(1) if match else None
 
-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}
-```
+    def extract_list(self, text: str, field: str) -> Optional[List[str]]:
+        match = re.search(rf'"{field}"\s*:\s*\[(.*?)\]', text, re.DOTALL)
+        if match:
+            items = re.findall(r'"([^"]*)"', match.group(1))
+            return items
+        return None
 
-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
+    def extract_dict(self, text: str, field: str) -> Optional[Dict[str, List[str]]]:
+        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
 
-Explain in a way that would be understandable to a junior developer.
-"""
-        return self.generate_agent_response(prompt)
+    def generate_html(self, content: str) -> str:
+        return self.html_gen_tool.generate(content)
 
-    def suggest_refactoring(self, code: str) -> str:
-        prompt = f"""
-Please suggest refactoring improvements for the following code:
+    def generate_css(self, styles: Dict[str, str]) -> str:
+        return self.css_gen_tool.generate(styles)
 
-```
-{code}
-```
+    def generate_javascript(self, functionality: str) -> str:
+        return self.js_gen_tool.generate(functionality)
 
-Consider the following in your suggestions:
-1. Improving code readability
-2. Enhancing maintainability
-3. Applying design patterns where appropriate
-4. Optimizing performance (if applicable)
+    def create_project_files(self, config: ProjectConfig) -> Dict[str, str]:
+        files = {}
+        for directory, file_list in config.structure.items():
+            for file in file_list:
+                file_path = os.path.join(directory, file)
+                if file.endswith('.html'):
+                    files[file_path] = self.generate_html(f"Content for {file}")
+                elif file.endswith('.css'):
+                    files[file_path] = self.generate_css({"body": {"font-family": "Arial, sans-serif"}})
+                elif file.endswith('.js'):
+                    files[file_path] = self.generate_javascript(f"console.log('Script for {file}');")
+                else:
+                    files[file_path] = f"Content for {file}"
+        return files
 
-Provide specific suggestions and code examples.
-"""
-        return self.generate_agent_response(prompt)
+    def execute_project(self, project_description: str) -> Dict[str, str]:
+        config = self.generate_project_config(project_description)
+        if config:
+            return self.create_project_files(config)
+        else:
+            self.logger.error("Failed to generate project configuration")
+            return {}