Upload 2 files

agents/software_engineer_agent.py

@@ -1,6 +1,8 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 from langchain_core.messages import AIMessage
+import asyncio
+from typing import Generator, Dict, Any
 
 MODEL_REPO = "Rahul-8799/software_engineer_mellum"
 
@@ -11,79 +13,79 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
-def run(state: dict) -> dict:
-    """Software Engineer generates clean, modern UI code using best practices"""
+async def stream_inference(prompt: str) -> Generator[str, None, None]:
+    """Stream the model's output token by token"""
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
+    
+    for _ in range(100):
+        output_ids = model.generate(
+            input_ids,
+            max_new_tokens=1,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        
+        new_token = output_ids[0][-1]
+        if new_token == tokenizer.eos_token_id:
+            break
+            
+        token_text = tokenizer.decode([new_token])
+        yield token_text
+        
+        input_ids = output_ids
+        await asyncio.sleep(0.05)
+
+async def run(state: Dict[str, Any]) -> Dict[str, Any]:
+    """Software Engineer generates responsive and interactive UI code"""
     messages = state["messages"]
     prompt = messages[-1].content
     
-    # Enhance the prompt with UI implementation guidelines
+    # Enhance the prompt with modern web development requirements
     enhanced_prompt = f"""
-Objective
-
-Generate modern, responsive, and accessible UI code that is visually appealing and adheres to current frontend development best practices.
-
-1. Styling Framework: Tailwind CSS
-	•	Use Tailwind CSS utility classes for styling all elements.
-	•	Apply spacing, typography, sizing, and layout using Tailwind classes.
-	•	Follow a mobile-first design approach.
-	•	Use Tailwind’s built-in responsive breakpoints (sm, md, lg, xl, 2xl) to adapt layouts for different screen sizes.
-
-2. Layout Techniques
-	•	Use CSS Grid for complex, multi-column or two-dimensional layouts.
-	•	Use Flexbox for flexible alignment of components like navigation bars, cards, buttons, and modals.
-	•	Maintain consistent spacing with utility classes such as gap, space-x, space-y, p-*, and m-*.
-
-3. Semantic HTML
-	•	Use semantic HTML tags appropriately: <header>, <nav>, <main>, <section>, <article>, <footer>, etc.
-	•	Avoid unnecessary <div> elements to prevent cluttered and unstructured markup.
-	•	Ensure proper nesting and hierarchy of elements.
-
-4. Accessibility
-	•	Add ARIA labels, role attributes, and alt text where needed for screen reader support.
-	•	Ensure keyboard accessibility with tabindex, proper focus states, and interactive elements being navigable.
-	•	Use <label> elements properly linked to form fields via the for attribute.
-
-5. Responsive Design
-	•	Use Tailwind’s responsive utilities to adjust layouts across various screen sizes.
-	•	Design components to be fully usable on both desktop and mobile devices.
-	•	Use collapsible or toggleable UI patterns (e.g., hamburger menus) for smaller viewports.
-
-6. Theming and Styling Consistency
-	•	Define and use CSS variables (--primary-color, --font-family, etc.) for theme consistency across components.
-	•	Maintain a clear visual hierarchy with consistent font sizes, weights, and colors.
-	•	Customize Tailwind’s theme configuration if needed for project-specific design tokens.
-
-7. JavaScript and Interactivity
-	•	Add interactivity using plain JavaScript, Alpine.js, or React if specified.
-	•	Implement common UI components such as modals, dropdowns, tooltips, accordions with appropriate open/close behavior.
-	•	Provide user feedback through form validations, dynamic updates, and transitions.
-
-8. Loading and Error States
-	•	Implement loading states using spinners, skeleton screens, or placeholders while data is being fetched or actions are processing.
-	•	Show error states using alerts, banners, or toast messages when applicable.
-	•	Use conditional rendering or state flags to handle visibility and transitions between states.
-
-9. Component Structure and Reusability
-	•	Break down the UI into modular, reusable components (e.g., Button, Card, Modal, Form).
-	•	Each component should:
-	•	Be self-contained with a clear purpose.
-	•	Accept inputs or props when necessary.
-	•	Maintain responsive and accessible markup by default.
-
-10. Code Quality Standards
-	•	Write clean, readable, and maintainable code.
-	•	Remove unused classes, scripts, or markup.
-	•	Follow consistent naming conventions and indentation rules.
-	•	Add comments only when necessary for clarity.
-
-
+    Generate modern, responsive, and interactive UI code following these requirements:
+    1. Use Tailwind CSS for responsive design
+    2. Implement JavaScript for interactivity
+    3. Add smooth animations and transitions
+    4. Ensure mobile-first approach
+    5. Include proper error handling
+    6. Add loading states and feedback
+    
     Original requirements: {prompt}
+    
+    Generate the following files:
+    
+    1. index.html - Main HTML structure
+    2. styles.css - Custom styles (if needed beyond Tailwind)
+    3. script.js - Interactive features
+    4. tailwind.config.js - Tailwind configuration
+    
+    Format the output as:
+    
+    ## HTML Structure
+    ```html
+    [HTML code]
+    ```
+    
+    ## CSS Styles
+    ```css
+    [CSS code]
+    ```
+    
+    ## JavaScript
+    ```javascript
+    [JavaScript code]
+    ```
+    
+    ## Tailwind Config
+    ```javascript
+    [Tailwind configuration]
+    ```
     """
     
-    input_ids = tokenizer(enhanced_prompt, return_tensors="pt").input_ids.to(model.device)
-    output_ids = model.generate(input_ids, max_new_tokens=3000)
-    output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
+    # Stream the output
+    output = ""
+    async for token in stream_inference(enhanced_prompt):
+        output += token
+    
     return {
         "messages": [AIMessage(content=output)],
         "chat_log": state["chat_log"] + [{"role": "Software Engineer", "content": output}],

agents/ui_designer_agent.py

@@ -1,8 +1,10 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 from langchain_core.messages import AIMessage
+import asyncio
+from typing import Generator, Dict, Any
 
-MODEL_REPO = "Rahul-8799/ui_designer_mistral"  # You'll need to fine-tune this model
+MODEL_REPO = "Rahul-8799/ui_designer_mistral"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
@@ -11,30 +13,64 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
-def run(state: dict) -> dict:
-    """UI Designer creates beautiful and structured UI designs with proper spacing and layout"""
+async def stream_inference(prompt: str) -> Generator[str, None, None]:
+    """Stream the model's output token by token"""
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
+    
+    for _ in range(100):
+        output_ids = model.generate(
+            input_ids,
+            max_new_tokens=1,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        
+        new_token = output_ids[0][-1]
+        if new_token == tokenizer.eos_token_id:
+            break
+            
+        token_text = tokenizer.decode([new_token])
+        yield token_text
+        
+        input_ids = output_ids
+        await asyncio.sleep(0.05)
+
+async def run(state: Dict[str, Any]) -> Dict[str, Any]:
+    """UI Designer creates responsive and interactive UI designs"""
     messages = state["messages"]
     prompt = messages[-1].content
     
-    # Enhance the prompt with UI design principles
+    # Enhance the prompt with responsive design requirements
     enhanced_prompt = f"""
-    Create a beautiful and well-structured UI design following these principles:
-    1. Use proper spacing and padding (recommended: 1rem/16px for padding, 2rem/32px for margins)
-    2. Implement a consistent color scheme
-    3. Ensure proper hierarchy with clear headings
-    4. Use responsive design principles
-    5. Implement proper grid system
-    6. Add smooth transitions and hover effects
-    7. Ensure proper contrast and readability
-    8. Use modern UI components and patterns
+    Create a modern, responsive UI design following these requirements:
+    1. Mobile-first approach with responsive breakpoints
+    2. Modern CSS features (Flexbox, Grid, CSS Variables)
+    3. Interactive elements with JavaScript
+    4. Smooth animations and transitions
+    5. Accessibility features
+    6. Cross-browser compatibility
     
     Original requirements: {prompt}
+    
+    Provide the design in this format:
+    
+    ## Responsive Layout
+    [Describe the responsive layout structure]
+    
+    ## CSS Framework
+    [Specify CSS framework and custom styles]
+    
+    ## JavaScript Features
+    [List interactive features and animations]
+    
+    ## Component Structure
+    [Describe component hierarchy and relationships]
     """
     
-    input_ids = tokenizer(enhanced_prompt, return_tensors="pt").input_ids.to(model.device)
-    output_ids = model.generate(input_ids, max_new_tokens=3000)
-    output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
+    # Stream the output
+    output = ""
+    async for token in stream_inference(enhanced_prompt):
+        output += token
+    
     return {
         "messages": [AIMessage(content=output)],
         "chat_log": state["chat_log"] + [{"role": "UI Designer", "content": output}],