import os import logging from typing import Dict, List, Optional from langchain_openai import ChatOpenAI from pydantic import BaseModel, Field from langchain_core.prompts import ChatPromptTemplate # Setup Logger logger = logging.getLogger(__name__) # --- Pydantic Models for Structured Output --- class CodeFile(BaseModel): filename: str = Field(description="Name of the file, e.g., 'NotificationFactory.java'") content: str = Field(description="Complete source code for the file") class LLDCodeResponse(BaseModel): title: str = Field(description="Title of the design pattern being demonstrated") subtitle: str = Field(description="Subtitle explaining the specific implementation") files: List[CodeFile] = Field(description="List of code files. Max 3-4 files.") execution_output: str = Field(description="Console output for the Main execution file.") typing_speed: int = Field(default=25, description="Typing animation speed in ms per char") auto_run: bool = Field(default=True, description="Whether to auto-run the animation") # Pattern-specific guidance PATTERN_GUIDANCE = { "singleton": { "description": "Ensures a class has only one instance and provides global access", "files_needed": ["Singleton class", "Main/Demo class"], "key_concepts": ["Private constructor", "Static instance", "Thread safety"], "example_domain": "Database connection, Logger, Configuration manager" }, "factory": { "description": "Creates objects without specifying exact class to create", "files_needed": ["Product interface", "Concrete products", "Factory class", "Main/Demo"], "key_concepts": ["Product interface", "Concrete implementations", "Factory method"], "example_domain": "Notification system (Email, SMS, Push), Vehicle factory, Payment processors" }, "observer": { "description": "Defines one-to-many dependency between objects", "files_needed": ["Subject interface", "Observer interface", "Concrete subject", "Concrete observers", "Main/Demo"], "key_concepts": ["Subject", "Observer", "Notify mechanism", "Subscription"], "example_domain": "Event system, Stock price updates, News feed, Weather station" }, "strategy": { "description": "Defines family of algorithms, encapsulates each, makes them interchangeable", "files_needed": ["Strategy interface", "Concrete strategies", "Context class", "Main/Demo"], "key_concepts": ["Strategy interface", "Concrete strategies", "Context switching"], "example_domain": "Payment methods, Sorting algorithms, Compression methods, Route calculation" }, "builder": { "description": "Separates construction of complex object from representation", "files_needed": ["Product class", "Builder interface/class", "Concrete builder", "Director (optional)", "Main/Demo"], "key_concepts": ["Step-by-step construction", "Fluent interface", "Immutable product"], "example_domain": "Building complex objects (House, Pizza, Document), Query builder" }, "decorator": { "description": "Attaches additional responsibilities to object dynamically", "files_needed": ["Component interface", "Concrete component", "Decorator base", "Concrete decorators", "Main/Demo"], "key_concepts": ["Component interface", "Wrapping", "Composition over inheritance"], "example_domain": "Coffee with addons, Text formatting, Stream wrappers" }, "adapter": { "description": "Converts interface of class into another interface clients expect", "files_needed": ["Target interface", "Adaptee class", "Adapter class", "Main/Demo"], "key_concepts": ["Target interface", "Adaptee", "Adapter wrapping"], "example_domain": "Legacy system integration, Third-party API wrapper" }, "command": { "description": "Encapsulates request as object, allowing parameterization", "files_needed": ["Command interface", "Concrete commands", "Receiver", "Invoker", "Main/Demo"], "key_concepts": ["Command interface", "Execute method", "Receiver", "Undo capability"], "example_domain": "Text editor operations, Remote control, Transaction system" }, "design_pattern": { "description": "Generic design pattern implementation", "files_needed": ["Interface", "Implementation", "Main/Demo"], "key_concepts": ["Clean code", "SOLID principles", "Best practices"], "example_domain": "Based on topic context" } } class LLDCodeGenerator: """ Generator for Low-Level Design code examples using LangChain for structured multi-file output. """ def __init__(self, client=None): api_key = os.getenv("OPENAI_API_KEY") if not api_key and client and hasattr(client, "api_key"): api_key = client.api_key self.llm = ChatOpenAI( model="gpt-4o", temperature=0.2, api_key=api_key ) self.logger = logging.getLogger(__name__) def generate_lld_code_content( self, slide_config: Dict, topic: str, architecture_description: str ) -> Dict: """ Generate multi-file LLD code examples with pattern-specific guidance. """ focus = slide_config.get('focus', 'Low-Level Design') language = slide_config.get('language', 'java') lld_pattern = slide_config.get('lld_pattern', 'design_pattern') self.logger.info(f"Generating LLD code for: {focus} in {language} using pattern: {lld_pattern}") # Get pattern-specific guidance pattern_info = PATTERN_GUIDANCE.get(lld_pattern.lower(), PATTERN_GUIDANCE['design_pattern']) # Define input model for type safety class PromptInput(BaseModel): language: str topic: str focus: str lld_pattern: str pattern_description: str files_needed: str key_concepts: str example_domain: str context: str # Prepare Prompt system_prompt = """You are an expert software architect creating clean, educational code examples for Low-Level Design patterns. Generate complete, working multi-file code that demonstrates design patterns and SOLID principles. CRITICAL REQUIREMENTS: 1. **Pattern Implementation**: - Implement the EXACT pattern specified: {lld_pattern} - Pattern description: {pattern_description} - Required files: {files_needed} - Key concepts to demonstrate: {key_concepts} - Use domain context: {example_domain} 2. **Multiple Files** (3-4 files): - Interface/Abstract class (if applicable) - Concrete implementation classes (2-3 variants) - Main/Demo class showing usage - Each file should be focused and follow Single Responsibility Principle 3. **REAL-WORLD NAMING**: - **NEVER** use generic names like `ConcreteProductA`, `Component1`, `Observer1`, `StrategyA` - **ALWAYS** use domain-specific names based on the TOPIC ({topic}) and FOCUS ({focus}) - Example for Notification System + Factory pattern: ✓ GOOD: `NotificationFactory`, `EmailNotification`, `SMSNotification`, `PushNotification` ✗ BAD: `Factory`, `ProductA`, `ProductB`, `ConcreteProduct1` - Example for Payment System + Strategy pattern: ✓ GOOD: `PaymentStrategy`, `CreditCardPayment`, `PayPalPayment`, `CryptoPayment` ✗ BAD: `Strategy`, `ConcreteStrategyA`, `ConcreteStrategyB` 4. **Clean Code**: - Use proper naming conventions for {language} - Add helpful comments explaining pattern-specific parts - COMPACT formatting: Use MINIMAL blank lines (max 1 blank line between methods/classes) - NO excessive spacing or empty lines at start of files - Keep code visually compact and professional 5. **Pattern-Specific Implementation**: - For **Factory**: Show polymorphic object creation without exposing instantiation logic - For **Singleton**: Include thread-safe implementation, private constructor, static instance - For **Observer**: Demonstrate subject-observer relationship, notify mechanism - For **Strategy**: Show algorithm family, context switching between strategies - For **Builder**: Implement step-by-step construction, fluent interface - For **Decorator**: Show component wrapping, responsibility addition - For **Adapter**: Demonstrate interface conversion between incompatible interfaces - For **Command**: Encapsulate requests as objects, show execute/undo 6. **Execution Output**: - Provide realistic console output showing the pattern in action - Output should demonstrate key pattern benefits - Show object creation, method calls, state changes 7. **Language**: {language} OUTPUT STRUCTURE: Return a JSON object with: - title: Pattern name + brief description - subtitle: Specific implementation focus - files: list of objects {{filename, content}} - execution_output: string containing the console output - typing_speed: 25-35 (ms per character) - auto_run: true """ user_prompt = """ TOPIC: {topic} FOCUS: {focus} PATTERN: {lld_pattern} PATTERN DESCRIPTION: {pattern_description} KEY CONCEPTS: {key_concepts} EXAMPLE DOMAIN: {example_domain} CONTEXT: {context} Generate the complete pattern implementation now with proper domain-specific names. """ try: structured_llm = self.llm.with_structured_output(LLDCodeResponse) prompt = ChatPromptTemplate.from_messages([ ("system", system_prompt), ("user", user_prompt) ]) chain = prompt | structured_llm # Create input with Pydantic model prompt_input = PromptInput( language=language, topic=topic, focus=focus, lld_pattern=lld_pattern, pattern_description=pattern_info['description'], files_needed=", ".join(pattern_info['files_needed']), key_concepts=", ".join(pattern_info['key_concepts']), example_domain=pattern_info['example_domain'], context=architecture_description[:800] ) result: LLDCodeResponse = chain.invoke(prompt_input.model_dump()) # Convert List[CodeFile] back to Dict[str, str] for frontend compatibility code_files_dict = {f.filename: f.content for f in result.files} # Construct execution output dict (Main file gets the output, others empty) exec_output_dict = {name: "" for name in code_files_dict} # Heuristic: Find Main file or use the last one main_file = next( (f for f in code_files_dict if "Main" in f or "Demo" in f or "App" in f or "Test" in f), list(code_files_dict.keys())[-1] ) if main_file: exec_output_dict[main_file] = result.execution_output self.logger.info(f"Generated {len(code_files_dict)} files for {lld_pattern} pattern") return { "title": result.title, "subtitle": result.subtitle, "code_files": code_files_dict, "execution_output": exec_output_dict, "typing_speed": result.typing_speed, "auto_run": result.auto_run, "pattern": lld_pattern } except Exception as e: self.logger.error(f"Error generating LLD code with LangChain: {e}") return self._get_fallback_lld_code(focus, language, lld_pattern) def _get_fallback_lld_code(self, focus: str, language: str, pattern: str) -> Dict: """Fallback LLD code if generation fails - pattern-specific""" self.logger.warning(f"Using fallback LLD code for {pattern}") if pattern.lower() == "singleton": return { "title": "Singleton Pattern", "subtitle": "Thread-Safe Database Connection", "code_files": { "DatabaseConnection.java": """public class DatabaseConnection { private static volatile DatabaseConnection instance; private String connectionString; private DatabaseConnection() { // Private constructor this.connectionString = "jdbc:mysql://localhost:3306/mydb"; System.out.println("Database connection initialized"); } public static DatabaseConnection getInstance() { if (instance == null) { synchronized (DatabaseConnection.class) { if (instance == null) { instance = new DatabaseConnection(); } } } return instance; } public void query(String sql) { System.out.println("Executing: " + sql); } }""", "Main.java": """public class Main { public static void main(String[] args) { // Only one instance created DatabaseConnection db1 = DatabaseConnection.getInstance(); DatabaseConnection db2 = DatabaseConnection.getInstance(); System.out.println("Same instance? " + (db1 == db2)); db1.query("SELECT * FROM users"); } }""" }, "execution_output": { "Main.java": """Database connection initialized Same instance? true Executing: SELECT * FROM users""" }, "typing_speed": 25, "auto_run": True, "pattern": "singleton" } # Generic fallback return { "title": focus, "subtitle": "Fallback Example", "code_files": { "Main.java": f"""public class Main {{ public static void main(String[] args) {{ System.out.println("Error generating {pattern} pattern code."); System.out.println("Focus: {focus}"); }} }}""" }, "execution_output": { "Main.java": f"Error generating {pattern} pattern code.\nFocus: {focus}" }, "typing_speed": 25, "auto_run": True, "pattern": pattern }