Update app.py

app.py

@@ -1,246 +1,242 @@
 import streamlit as st
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import Chroma
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import torch
-import logging
+from transformers import pipeline
+import json
 from pathlib import Path
+import os
 
-class LocalWebDevAssistant:
+class FastWebGenerator:
     def __init__(self):
-        self.initialize_logging()
-        self.setup_model()
-        self.setup_embeddings()
-        self.initialize_vector_store()
-
-    def initialize_logging(self):
-        logging.basicConfig(level=logging.INFO)
-        self.logger = logging.getLogger(__name__)
-
-    def setup_model(self):
-        # Using a smaller, directly available model
-        model_name = "facebook/opt-350m"  # Smaller model that's good for code
-        
-        @st.cache_resource
-        def load_model_and_tokenizer(model_name):
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                torch_dtype=torch.float16,
-                low_cpu_mem_usage=True
-            )
-            return model, tokenizer
-
-        self.model, self.tokenizer = load_model_and_tokenizer(model_name)
+        # Load the smallest model possible for quick text completion
         self.generator = pipeline(
-            "text-generation",
-            model=self.model,
-            tokenizer=self.tokenizer,
-            max_length=1000
+            'text-generation',
+            model='distilgpt2',
+            device_map='auto'
         )
+        
+        # Load templates
+        self.templates = {
+            "react": {
+                "component": """
+import React, { useState, useEffect } from 'react';
 
-    def setup_embeddings(self):
-        self.embeddings = HuggingFaceEmbeddings(
-            model_name="all-MiniLM-L6-v2",
-            model_kwargs={'device': 'cpu'}
-        )
+const {component_name} = () => {
+    const [data, setData] = useState([]);
 
-    def initialize_vector_store(self):
-        if not Path("chroma_db").exists():
-            self.create_knowledge_base()
-        
-        self.vector_store = Chroma(
-            persist_directory="chroma_db",
-            embedding_function=self.embeddings
-        )
+    useEffect(() => {
+        // Fetch data
+        fetchData();
+    }, []);
 
-    def create_knowledge_base(self):
-        # Basic web development knowledge
-        documents = [
-            """React component structure:
-            import React from 'react';
-            
-            const Component = ({ props }) => {
-                return (
-                    <div>
-                        {/* Component content */}
-                    </div>
-                );
-            };
-            
-            export default Component;
-            """,
-            """FastAPI backend structure:
-            from fastapi import FastAPI
-            
-            app = FastAPI()
-            
-            @app.get("/")
-            async def root():
-                return {"message": "Hello World"}
-            """,
-            """MongoDB connection:
-            from pymongo import MongoClient
-            
-            client = MongoClient('mongodb://localhost:27017/')
-            db = client['database_name']
-            """
-        ]
-        
-        text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=500,
-            chunk_overlap=50
-        )
-        texts = text_splitter.split_text('\n\n'.join(documents))
-        
-        Chroma.from_texts(
-            texts,
-            self.embeddings,
-            persist_directory="chroma_db"
-        )
+    const fetchData = async () => {
+        try {
+            const response = await fetch('/api/data');
+            const jsonData = await response.json();
+            setData(jsonData);
+        } catch (error) {
+            console.error('Error:', error);
+        }
+    };
 
-    def generate_code(self, description, tech_stack, features):
-        try:
-            # Create prompt
-            prompt = f"""Generate code for a web application with:
-            Description: {description}
-            Technology Stack: {tech_stack}
-            Features: {features}
-            
-            Provide the code in sections:
-            """
-            
-            # Get relevant context from vector store
-            docs = self.vector_store.similarity_search(description, k=2)
-            context = "\n".join(doc.page_content for doc in docs)
-            
-            # Generate with context
-            full_prompt = f"{context}\n{prompt}"
-            
-            response = self.generator(
-                full_prompt,
-                max_length=1000,
-                num_return_sequences=1
-            )[0]['generated_text']
-            
-            return self.process_response(response)
-            
-        except Exception as e:
-            self.logger.error(f"Generation failed: {e}")
-            raise
-
-    def process_response(self, response):
-        # Extract different code sections
-        sections = {
-            "frontend": "",
-            "backend": "",
-            "database": "",
-            "instructions": ""
+    return (
+        <div className="container mx-auto p-4">
+            <h1 className="text-2xl font-bold mb-4">{title}</h1>
+            {content}
+        </div>
+    );
+};
+
+export default {component_name};
+""",
+                "api": """
+async function fetchData() {
+    const response = await fetch('/api/endpoint');
+    const data = await response.json();
+    return data;
+}
+
+async function postData(data) {
+    const response = await fetch('/api/endpoint', {
+        method: 'POST',
+        headers: {
+            'Content-Type': 'application/json',
+        },
+        body: JSON.stringify(data),
+    });
+    return response.json();
+}
+"""
+            },
+            "fastapi": {
+                "main": """
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List, Optional
+import uvicorn
+
+app = FastAPI()
+
+class ItemModel(BaseModel):
+    id: Optional[int]
+    name: str
+    description: Optional[str]
+
+@app.get("/")
+async def root():
+    return {"message": "API is running"}
+
+@app.get("/items")
+async def get_items():
+    return {"items": items}
+
+@app.post("/items")
+async def create_item(item: ItemModel):
+    items.append(item)
+    return item
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)
+""",
+                "crud": """
+@app.get("/items/{item_id}")
+async def get_item(item_id: int):
+    if item_id < 0 or item_id >= len(items):
+        raise HTTPException(status_code=404, detail="Item not found")
+    return items[item_id]
+
+@app.put("/items/{item_id}")
+async def update_item(item_id: int, item: ItemModel):
+    if item_id < 0 or item_id >= len(items):
+        raise HTTPException(status_code=404, detail="Item not found")
+    items[item_id] = item
+    return item
+
+@app.delete("/items/{item_id}")
+async def delete_item(item_id: int):
+    if item_id < 0 or item_id >= len(items):
+        raise HTTPException(status_code=404, detail="Item not found")
+    item = items.pop(item_id)
+    return item
+"""
+            },
+            "mongodb": """
+from motor.motor_asyncio import AsyncIOMotorClient
+
+client = AsyncIOMotorClient('mongodb://localhost:27017')
+db = client.database_name
+
+async def get_all():
+    cursor = db.collection.find({})
+    return await cursor.to_list(length=100)
+
+async def get_one(id: str):
+    return await db.collection.find_one({"_id": id})
+
+async def create_one(data: dict):
+    result = await db.collection.insert_one(data)
+    return result.inserted_id
+
+async def update_one(id: str, data: dict):
+    result = await db.collection.update_one(
+        {"_id": id},
+        {"$set": data}
+    )
+    return result.modified_count
+
+async def delete_one(id: str):
+    result = await db.collection.delete_one({"_id": id})
+    return result.deleted_count
+"""
         }
+
+    def customize_code(self, template, params):
+        # Quick text customization using the model
+        prompt = f"Customize this code: {template}\nWith these parameters: {params}\n"
+        result = self.generator(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']
+        return result
+
+    def generate_project(self, specs):
+        frontend_code = self.templates["react"]["component"].replace(
+            "{component_name}", specs["component_name"]
+        ).replace(
+            "{title}", specs["title"]
+        ).replace(
+            "{content}", specs["content"]
+        )
         
-        current_section = "frontend"
-        for line in response.split('\n'):
-            if "FRONTEND:" in line.upper():
-                current_section = "frontend"
-                continue
-            elif "BACKEND:" in line.upper():
-                current_section = "backend"
-                continue
-            elif "DATABASE:" in line.upper():
-                current_section = "database"
-                continue
-            elif "INSTRUCTIONS:" in line.upper():
-                current_section = "instructions"
-                continue
-                
-            sections[current_section] += line + '\n'
+        backend_code = self.templates["fastapi"]["main"]
+        if specs.get("crud", False):
+            backend_code += self.templates["fastapi"]["crud"]
             
-        return sections
+        db_code = self.templates["mongodb"] if specs["database"] == "MongoDB" else ""
+        
+        return {
+            "frontend": frontend_code,
+            "backend": backend_code,
+            "database": db_code
+        }
 
 def main():
-    st.set_page_config(page_title="Web Development Assistant", layout="wide")
+    st.set_page_config(page_title="Fast Code Generator", layout="wide")
     
-    st.title("🚀 Web Development Assistant")
-    st.write("Generate web application code using AI")
+    if 'generator' not in st.session_state:
+        with st.spinner("Loading generator..."):
+            st.session_state.generator = FastWebGenerator()
     
-    if 'assistant' not in st.session_state:
-        with st.spinner("Initializing... (this may take a minute)"):
-            st.session_state.assistant = LocalWebDevAssistant()
+    st.title("⚡ Fast Code Generator")
     
-    with st.form("project_details"):
-        description = st.text_area(
-            "Project Description",
-            placeholder="Describe your web application..."
-        )
-        
+    with st.form("generate_form"):
         col1, col2 = st.columns(2)
+        
         with col1:
-            frontend = st.selectbox(
-                "Frontend Framework",
-                ["React", "Vue", "Plain JavaScript"]
-            )
-            
-            backend = st.selectbox(
-                "Backend Framework",
-                ["FastAPI", "Express", "Flask"]
-            )
+            component_name = st.text_input("Component Name", "MyComponent")
+            title = st.text_input("Title", "My App")
+            content = st.text_area("Content", "<p>Your content here</p>")
         
         with col2:
-            database = st.selectbox(
-                "Database",
-                ["MongoDB", "PostgreSQL", "SQLite"]
-            )
+            database = st.selectbox("Database", ["MongoDB", "PostgreSQL"])
+            add_crud = st.checkbox("Add CRUD Operations", value=True)
+            add_auth = st.checkbox("Add Authentication", value=False)
+        
+        if st.form_submit_button("Generate"):
+            specs = {
+                "component_name": component_name,
+                "title": title,
+                "content": content,
+                "database": database,
+                "crud": add_crud,
+                "auth": add_auth
+            }
+            
+            result = st.session_state.generator.generate_project(specs)
             
-            features = st.multiselect(
-                "Features",
-                ["Authentication", "REST API", "Database CRUD", "Form Handling"]
+            # Display results
+            tabs = st.tabs(["Frontend", "Backend", "Database"])
+            
+            with tabs[0]:
+                st.code(result["frontend"], language="javascript")
+            with tabs[1]:
+                st.code(result["backend"], language="python")
+            with tabs[2]:
+                st.code(result["database"], language="python")
+            
+            # Add download button
+            combined_code = f"""
+            // Frontend Code
+            {result['frontend']}
+            
+            // Backend Code
+            {result['backend']}
+            
+            // Database Code
+            {result['database']}
+            """
+            
+            st.download_button(
+                "Download Code",
+                combined_code,
+                file_name="generated_code.txt",
+                mime="text/plain"
             )
-        
-        generate = st.form_submit_button("Generate Code")
-        
-        if generate:
-            try:
-                with st.spinner("Generating code..."):
-                    result = st.session_state.assistant.generate_code(
-                        description,
-                        {
-                            "frontend": frontend,
-                            "backend": backend,
-                            "database": database
-                        },
-                        features
-                    )
-                    
-                    # Display results in tabs
-                    tabs = st.tabs([
-                        "Frontend Code",
-                        "Backend Code",
-                        "Database Setup",
-                        "Instructions"
-                    ])
-                    
-                    with tabs[0]:
-                        st.code(result["frontend"], language="javascript")
-                    
-                    with tabs[1]:
-                        st.code(result["backend"], language="python")
-                    
-                    with tabs[2]:
-                        st.code(result["database"], language="sql")
-                    
-                    with tabs[3]:
-                        st.markdown(result["instructions"])
-                    
-                    # Add download button
-                    st.download_button(
-                        "Download Code",
-                        '\n\n'.join(result.values()),
-                        file_name="generated_code.txt"
-                    )
-                    
-            except Exception as e:
-                st.error(f"An error occurred: {str(e)}")
 
 if __name__ == "__main__":
     main()