Team Ai
Apppublic

2008robocode-crypto/code-generation-system

sourceHugging Faceupdated 5mo agoView on Hugging Face
0likes
pipeline.py504 linesDownload Raw Back to src
1"""2Multi-stage generation pipeline orchestrator.3Implements the 4-stage compiler-like system for code generation.4"""5 6import json7import os8from typing import Any, Dict, List, Optional, Tuple9from datetime import datetime10import re11 12# Mock LLM calls for now - will be replaced with actual API calls13try:14    import anthropic15    HAS_ANTHROPIC = True16except ImportError:17    HAS_ANTHROPIC = False18 19from validator import Validator20from repair_engine import RepairEngine21 22 23class IntentExtractor:24    """Stage 1: Extract structured intent from natural language."""25    26    def __init__(self, use_llm: bool = True):27        self.use_llm = use_llm and HAS_ANTHROPIC28    29    def extract(self, user_prompt: str) -> Dict[str, Any]:30        """Extract structured intent from user prompt."""31        if self.use_llm:32            return self._extract_with_llm(user_prompt)33        else:34            return self._extract_pattern_based(user_prompt)35    36    def _extract_pattern_based(self, prompt: str) -> Dict[str, Any]:37        """Pattern-based intent extraction (fallback)."""38        intent = {39            "app_name": self._extract_app_name(prompt),40            "app_description": prompt[:200],41            "key_features": self._extract_features(prompt),42            "user_roles": self._extract_roles(prompt),43            "core_entities": self._extract_entities(prompt),44            "business_requirements": self._extract_requirements(prompt),45            "constraints": self._extract_constraints(prompt),46        }47        return intent48    49    def _extract_app_name(self, prompt: str) -> str:50        """Extract app name from prompt."""51        # Look for "Build a X" or "Create a X"52        match = re.search(r'(?:Build|Create|Make|Generate)\s+(?:a\s+)?([A-Z][a-zA-Z\s]+?)(?:\s+with|\s+that|\.|$)', prompt)53        if match:54            return match.group(1).strip().replace(" ", "")55        return "GeneratedApp"56    57    def _extract_features(self, prompt: str) -> List[str]:58        """Extract key features."""59        features = []60        61        # Common feature keywords62        feature_keywords = [63            "login", "authentication", "contacts", "dashboard", "analytics",64            "admin", "payments", "role-based", "access", "premium", "plan",65            "reports", "export", "import", "notifications", "search"66        ]67        68        for keyword in feature_keywords:69            if keyword.lower() in prompt.lower():70                features.append(keyword)71        72        return features or ["basic_crud"]73    74    def _extract_roles(self, prompt: str) -> List[str]:75        """Extract user roles."""76        roles = []77        role_keywords = {"admin": "admin", "user": "user", "guest": "guest", "customer": "user"}78        79        for keyword, role in role_keywords.items():80            if keyword.lower() in prompt.lower():81                roles.append(role)82        83        return roles or ["user"]84    85    def _extract_entities(self, prompt: str) -> List[str]:86        """Extract core data entities."""87        entities = []88        89        entity_keywords = {90            "contact": "Contact",91            "user": "User",92            "product": "Product",93            "order": "Order",94            "payment": "Payment",95            "report": "Report",96            "dashboard": "Dashboard",97        }98        99        for keyword, entity in entity_keywords.items():100            if keyword.lower() in prompt.lower():101                entities.append(entity)102        103        return entities or ["Item"]104    105    def _extract_requirements(self, prompt: str) -> List[str]:106        """Extract business requirements."""107        return [108            "User authentication and authorization",109            "Role-based access control",110            "Data persistence",111            "API endpoints for CRUD operations",112        ]113    114    def _extract_constraints(self, prompt: str) -> List[str]:115        """Extract constraints."""116        constraints = []117        118        if "premium" in prompt.lower():119            constraints.append("Payment processing required")120        if "real-time" in prompt.lower():121            constraints.append("Real-time synchronization needed")122        123        return constraints124    125    def _extract_with_llm(self, prompt: str) -> Dict[str, Any]:126        """Extract intent using Anthropic API."""127        try:128            client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))129            130            extraction_prompt = f"""Extract structured intent from this user prompt:131 132"{prompt}"133 134Return a JSON with these fields:135- app_name: string (extract or generate a name)136- app_description: string (2-3 sentences)137- key_features: list[string] (extracted features)138- user_roles: list[string] (roles mentioned)139- core_entities: list[string] (data models)140- business_requirements: list[string] (business rules)141- constraints: list[string] (any constraints mentioned)142 143Return ONLY valid JSON, no markdown formatting."""144            145            message = client.messages.create(146                model="claude-3-5-sonnet-20241022",147                max_tokens=1024,148                messages=[{"role": "user", "content": extraction_prompt}]149            )150            151            response_text = message.content[0].text152            return json.loads(response_text)153        except Exception as e:154            print(f"LLM extraction failed: {e}, falling back to pattern-based")155            return self._extract_pattern_based(prompt)156 157 158class SystemDesignLayer:159    """Stage 2: Convert intent to system design."""160    161    def __init__(self, use_llm: bool = True):162        self.use_llm = use_llm and HAS_ANTHROPIC163    164    def design(self, intent: Dict[str, Any]) -> Dict[str, Any]:165        """Generate system design from intent."""166        if self.use_llm:167            return self._design_with_llm(intent)168        else:169            return self._design_rule_based(intent)170    171    def _design_rule_based(self, intent: Dict[str, Any]) -> Dict[str, Any]:172        """Rule-based system design."""173        design = {174            "entities": self._generate_entities(intent),175            "user_flows": self._generate_flows(intent),176            "roles_and_permissions": self._generate_rbac(intent),177            "data_models": intent["core_entities"],178            "api_patterns": ["REST"],179            "ui_structure": self._generate_ui_structure(intent),180        }181        return design182    183    def _generate_entities(self, intent: Dict[str, Any]) -> Dict[str, List[str]]:184        """Generate entity definitions."""185        entities = {}186        187        for entity in intent["core_entities"]:188            if entity.lower() == "user":189                entities[entity] = ["id", "name", "email", "role", "created_at"]190            elif entity.lower() == "contact":191                entities[entity] = ["id", "name", "email", "phone", "owner_id"]192            elif entity.lower() == "product":193                entities[entity] = ["id", "name", "price", "description"]194            elif entity.lower() == "order":195                entities[entity] = ["id", "user_id", "total", "status", "created_at"]196            else:197                entities[entity] = ["id", "name", "created_at"]198        199        return entities200    201    def _generate_flows(self, intent: Dict[str, Any]) -> List[Dict[str, Any]]:202        """Generate user flows."""203        flows = [204            {"name": "Authentication", "steps": ["Login", "Verify", "Redirect to Dashboard"]},205            {"name": "CRUD Operations", "steps": ["View", "Create", "Update", "Delete"]},206        ]207        208        if "admin" in intent["user_roles"]:209            flows.append({"name": "Admin Panel", "steps": ["View Analytics", "Manage Users", "View Reports"]})210        211        return flows212    213    def _generate_rbac(self, intent: Dict[str, Any]) -> Dict[str, List[str]]:214        """Generate role-based access control."""215        rbac = {}216        217        for role in intent["user_roles"]:218            if role == "admin":219                rbac[role] = ["read_all", "write_all", "delete_all", "manage_users"]220            elif role == "user":221                rbac[role] = ["read_own", "write_own", "delete_own"]222            else:223                rbac[role] = ["read_public"]224        225        return rbac226    227    def _generate_ui_structure(self, intent: Dict[str, Any]) -> List[str]:228        """Generate UI page structure."""229        pages = ["/login", "/dashboard", "/profile"]230        231        if "contacts" in str(intent["key_features"]).lower():232            pages.append("/contacts")233        if "admin" in intent["user_roles"]:234            pages.append("/admin")235        if "analytics" in str(intent["key_features"]).lower():236            pages.append("/analytics")237        238        return pages239    240    def _design_with_llm(self, intent: Dict[str, Any]) -> Dict[str, Any]:241        """Generate system design using LLM."""242        try:243            client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))244            245            design_prompt = f"""Design a system architecture based on this intent:246 247{json.dumps(intent, indent=2)}248 249Return a JSON with these fields:250- entities: dict mapping entity names to attribute lists251- user_flows: list of flow objects with name and steps252- roles_and_permissions: dict mapping roles to permissions253- data_models: list of entity names254- api_patterns: list (e.g., ["REST", "GraphQL"])255- ui_structure: list of page paths256 257Return ONLY valid JSON."""258            259            message = client.messages.create(260                model="claude-3-5-sonnet-20241022",261                max_tokens=2048,262                messages=[{"role": "user", "content": design_prompt}]263            )264            265            response_text = message.content[0].text266            return json.loads(response_text)267        except Exception as e:268            print(f"LLM design failed: {e}, using rule-based")269            return self._design_rule_based(intent)270 271 272class SchemaGenerator:273    """Stage 3: Generate complete schemas (DB, API, UI, Auth)."""274    275    def __init__(self, use_llm: bool = True):276        self.use_llm = use_llm and HAS_ANTHROPIC277    278    def generate(self, design: Dict[str, Any], intent: Dict[str, Any]) -> Dict[str, Any]:279        """Generate complete schema from design."""280        if self.use_llm:281            return self._generate_with_llm(design, intent)282        else:283            return self._generate_rule_based(design, intent)284    285    def _generate_rule_based(self, design: Dict[str, Any], intent: Dict[str, Any]) -> Dict[str, Any]:286        """Rule-based schema generation."""287        schema = {288            "app_name": intent["app_name"],289            "app_description": intent["app_description"],290            "database_schema": self._generate_db_schema(design),291            "api_schema": self._generate_api_schema(design),292            "ui_schema": self._generate_ui_schema(design),293            "auth_config": self._generate_auth_config(design),294            "roles": self._generate_roles(design),295            "business_logic": self._generate_business_logic(intent),296        }297        return schema298    299    def _generate_db_schema(self, design: Dict[str, Any]) -> List[Dict[str, Any]]:300        """Generate database schema."""301        tables = []302        303        for entity, attributes in design.get("entities", {}).items():304            table = {305                "name": entity.lower() + "s",306                "fields": [307                    {"name": "id", "type": "string", "required": True},308                ] + [309                    {"name": attr, "type": "string", "required": True}310                    for attr in attributes if attr != "id"311                ],312                "primary_key": "id",313                "indexes": ["id"]314            }315            tables.append(table)316        317        return tables318    319    def _generate_api_schema(self, design: Dict[str, Any]) -> List[Dict[str, Any]]:320        """Generate API schema."""321        endpoints = []322        323        for entity in design.get("data_models", []):324            base_path = f"/api/{entity.lower()}s"325            326            endpoints.extend([327                {"path": base_path, "method": "GET", "description": f"List {entity}s"},328                {"path": f"{base_path}/{{id}}", "method": "GET", "description": f"Get {entity}"},329                {"path": base_path, "method": "POST", "description": f"Create {entity}"},330                {"path": f"{base_path}/{{id}}", "method": "PUT", "description": f"Update {entity}"},331                {"path": f"{base_path}/{{id}}", "method": "DELETE", "description": f"Delete {entity}"},332            ])333        334        return endpoints335    336    def _generate_ui_schema(self, design: Dict[str, Any]) -> List[Dict[str, Any]]:337        """Generate UI schema."""338        pages = []339        340        for path in design.get("ui_structure", []):341            page = {342                "path": path,343                "title": path.replace("/", " ").title(),344                "components": [345                    {"name": "header", "type": "header"},346                    {"name": "content", "type": "container"},347                    {"name": "footer", "type": "footer"},348                ]349            }350            pages.append(page)351        352        return pages353    354    def _generate_auth_config(self, design: Dict[str, Any]) -> Dict[str, Any]:355        """Generate authentication config."""356        return {357            "type": "jwt",358            "secret_key": "generated-secret",359            "expiry": 3600,360            "refresh_token_expiry": 86400,361        }362    363    def _generate_roles(self, design: Dict[str, Any]) -> List[Dict[str, Any]]:364        """Generate roles from RBAC."""365        roles = []366        367        for role_name, permissions in design.get("roles_and_permissions", {}).items():368            roles.append({369                "name": role_name,370                "permissions": permissions,371                "description": f"Role: {role_name}"372            })373        374        return roles375    376    def _generate_business_logic(self, intent: Dict[str, Any]) -> Dict[str, Any]:377        """Generate business logic rules."""378        logic = {379            "validation_rules": [380                "Email must be valid format",381                "Password must be at least 8 characters",382            ],383            "access_control": "Role-based access control enabled",384            "premium_features": "premium" in str(intent).lower(),385        }386        return logic387    388    def _generate_with_llm(self, design: Dict[str, Any], intent: Dict[str, Any]) -> Dict[str, Any]:389        """Generate schemas using LLM."""390        try:391            client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))392            393            schema_prompt = f"""Generate complete schemas from this design and intent:394 395Design: {json.dumps(design, indent=2)}396Intent: {json.dumps(intent, indent=2)}397 398Return a JSON with these fields:399- app_name: string400- app_description: string401- database_schema: list of tables (each with name, fields, primary_key)402- api_schema: list of endpoints (path, method, description)403- ui_schema: list of pages (path, title, components)404- auth_config: object with type, expiry, etc.405- roles: list of role objects (name, permissions, description)406- business_logic: object with business rules407 408All table fields must be objects with: name, type, required409Valid types: string, number, boolean, date, email, enum410All API endpoints must have valid HTTP methods: GET, POST, PUT, DELETE411 412Return ONLY valid JSON."""413            414            message = client.messages.create(415                model="claude-3-5-sonnet-20241022",416                max_tokens=4096,417                messages=[{"role": "user", "content": schema_prompt}]418            )419            420            response_text = message.content[0].text421            return json.loads(response_text)422        except Exception as e:423            print(f"LLM schema generation failed: {e}, using rule-based")424            return self._generate_rule_based(design, intent)425 426 427class RefinementLayer:428    """Stage 4: Refine and validate schemas across all layers."""429    430    def __init__(self):431        self.validator = Validator()432        self.repair_engine = RepairEngine()433    434    def refine(self, schema: Dict[str, Any], max_iterations: int = 3) -> Tuple[Dict[str, Any], Dict[str, Any]]:435        """Validate and repair schema iteratively."""436        metadata = {437            "iterations": 0,438            "validation_results": [],439            "repairs": [],440            "final_status": "unknown",441        }442        443        for i in range(max_iterations):444            metadata["iterations"] = i + 1445            446            # Validate447            result = self.validator.validate_complete(schema)448            metadata["validation_results"].append(result.to_dict())449            450            if result.is_valid:451                metadata["final_status"] = "valid"452                return schema, metadata453            454            # Repair455            schema, repairs = self.repair_engine.repair_config(schema)456            metadata["repairs"].extend(repairs)457        458        metadata["final_status"] = "repaired_with_warnings" if metadata["validation_results"][-1]["errors"] else "valid"459        return schema, metadata460 461 462class Pipeline:463    """Main orchestrator for the 4-stage pipeline."""464    465    def __init__(self, use_llm: bool = True):466        self.intent_extractor = IntentExtractor(use_llm=use_llm)467        self.system_design = SystemDesignLayer(use_llm=use_llm)468        self.schema_generator = SchemaGenerator(use_llm=use_llm)469        self.refinement = RefinementLayer()470        self.use_llm = use_llm471    472    def generate(self, user_prompt: str) -> Tuple[Dict[str, Any], Dict[str, Any]]:473        """Run complete pipeline: prompt → config."""474        execution_log = {475            "timestamp": datetime.now().isoformat(),476            "user_prompt": user_prompt[:500],477            "stages": {}478        }479        480        try:481            # Stage 1: Intent Extraction482            intent = self.intent_extractor.extract(user_prompt)483            execution_log["stages"]["intent_extraction"] = {"status": "completed"}484            485            # Stage 2: System Design486            design = self.system_design.design(intent)487            execution_log["stages"]["system_design"] = {"status": "completed"}488            489            # Stage 3: Schema Generation490            schema = self.schema_generator.generate(design, intent)491            execution_log["stages"]["schema_generation"] = {"status": "completed"}492            493            # Stage 4: Refinement494            refined_schema, refinement_metadata = self.refinement.refine(schema)495            execution_log["stages"]["refinement"] = refinement_metadata496            execution_log["final_status"] = "success"497            498            return refined_schema, execution_log499        500        except Exception as e:501            execution_log["final_status"] = "error"502            execution_log["error"] = str(e)503            return {}, execution_log504