Update app.py

app.py

@@ -2,114 +2,128 @@ import streamlit as st
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.llms import CTransformers
-from langchain.chains import RetrievalQA
-from langchain.prompts import PromptTemplate
-import os
-from pathlib import Path
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
 import logging
+from pathlib import Path
 
-class LocalWebDevRAG:
+class LocalWebDevAssistant:
     def __init__(self):
         self.initialize_logging()
+        self.setup_model()
         self.setup_embeddings()
-        self.setup_llm()
         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)
+        self.generator = pipeline(
+            "text-generation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            max_length=1000
+        )
+
     def setup_embeddings(self):
         self.embeddings = HuggingFaceEmbeddings(
             model_name="all-MiniLM-L6-v2",
             model_kwargs={'device': 'cpu'}
         )
-        
-    def setup_llm(self):
-        # Using CodeLlama local model
-        llm_config = {
-            'model': 'codellama-7b-instruct.ggmlv3.Q4_K_M.bin',
-            'model_type': 'llama',
-            'max_new_tokens': 2048,
-            'temperature': 0.7,
-            'context_length': 2048,
-        }
-        
-        self.llm = CTransformers(**llm_config)
-        
-        self.qa_prompt = PromptTemplate(
-            template="""You are an expert web developer. Based on the context and request, 
-            generate production-ready code.
-            
-            Context: {context}
-            Question: {question}
-            
-            Provide a detailed solution with explanations.""",
-            input_variables=["context", "question"]
-        )
 
     def initialize_vector_store(self):
-        try:
-            # Create or load vector store
-            if not Path("chroma_db").exists():
-                self.create_new_vector_store()
-            else:
-                self.vector_store = Chroma(
-                    persist_directory="chroma_db",
-                    embedding_function=self.embeddings
-                )
-                self.logger.info("Loaded existing vector store")
-                
-            self.qa_chain = RetrievalQA.from_chain_type(
-                llm=self.llm,
-                chain_type="stuff",
-                retriever=self.vector_store.as_retriever(),
-                chain_type_kwargs={"prompt": self.qa_prompt}
-            )
-            
-        except Exception as e:
-            self.logger.error(f"Vector store initialization failed: {e}")
-            raise
+        if not Path("chroma_db").exists():
+            self.create_knowledge_base()
+        
+        self.vector_store = Chroma(
+            persist_directory="chroma_db",
+            embedding_function=self.embeddings
+        )
 
-    def create_new_vector_store(self):
-        # Example code snippets and documentation
+    def create_knowledge_base(self):
+        # Basic web development knowledge
         documents = [
-            "React component best practices...",
-            "API security implementations...",
-            "Database schema designs...",
-            # Add more code examples and documentation
+            """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=1000,
-            chunk_overlap=200
+            chunk_size=500,
+            chunk_overlap=50
         )
         texts = text_splitter.split_text('\n\n'.join(documents))
         
-        self.vector_store = Chroma.from_texts(
+        Chroma.from_texts(
             texts,
             self.embeddings,
             persist_directory="chroma_db"
         )
-        self.logger.info("Created new vector store")
 
-    def generate_code(self, description, tech_stack, requirements):
+    def generate_code(self, description, tech_stack, features):
         try:
-            prompt = f"""
-            Create a web application with:
+            # Create prompt
+            prompt = f"""Generate code for a web application with:
             Description: {description}
-            Tech Stack: {tech_stack}
-            Requirements: {requirements}
+            Technology Stack: {tech_stack}
+            Features: {features}
             
-            Provide:
-            1. Frontend components
-            2. Backend API
-            3. Database schema
-            4. Setup instructions
+            Provide the code in sections:
             """
             
-            response = self.qa_chain.run(prompt)
+            # 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:
@@ -117,25 +131,44 @@ class LocalWebDevRAG:
             raise
 
     def process_response(self, response):
-        # Basic response processing
-        return {
-            "frontend": response.split("Frontend:")[1].split("Backend:")[0] if "Frontend:" in response else "",
-            "backend": response.split("Backend:")[1].split("Database:")[0] if "Backend:" in response else "",
-            "database": response.split("Database:")[1].split("Setup:")[0] if "Database:" in response else "",
-            "setup": response.split("Setup:")[1] if "Setup:" in response else response
+        # Extract different code sections
+        sections = {
+            "frontend": "",
+            "backend": "",
+            "database": "",
+            "instructions": ""
         }
+        
+        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'
+            
+        return sections
 
 def main():
-    st.set_page_config(page_title="Local Web Development AI", layout="wide")
+    st.set_page_config(page_title="Web Development Assistant", layout="wide")
     
-    st.title("🚀 Web Development AI Assistant")
-    st.write("Generate web applications using local AI - no API key required!")
-
-    if 'rag_system' not in st.session_state:
-        with st.spinner("Initializing AI system... (this may take a few minutes on first run)"):
-            st.session_state.rag_system = LocalWebDevRAG()
-
-    with st.form("project_specs"):
+    st.title("🚀 Web Development Assistant")
+    st.write("Generate web application code using AI")
+    
+    if 'assistant' not in st.session_state:
+        with st.spinner("Initializing... (this may take a minute)"):
+            st.session_state.assistant = LocalWebDevAssistant()
+    
+    with st.form("project_details"):
         description = st.text_area(
             "Project Description",
             placeholder="Describe your web application..."
@@ -144,30 +177,32 @@ def main():
         col1, col2 = st.columns(2)
         with col1:
             frontend = st.selectbox(
-                "Frontend",
-                ["React", "Vue", "Angular"]
+                "Frontend Framework",
+                ["React", "Vue", "Plain JavaScript"]
             )
-            database = st.selectbox(
-                "Database",
-                ["MongoDB", "PostgreSQL", "MySQL"]
+            
+            backend = st.selectbox(
+                "Backend Framework",
+                ["FastAPI", "Express", "Flask"]
             )
         
         with col2:
-            backend = st.selectbox(
-                "Backend",
-                ["Node.js", "Python/FastAPI", "Python/Django"]
+            database = st.selectbox(
+                "Database",
+                ["MongoDB", "PostgreSQL", "SQLite"]
             )
+            
             features = st.multiselect(
                 "Features",
-                ["Authentication", "REST API", "File Upload", "Real-time Updates"]
+                ["Authentication", "REST API", "Database CRUD", "Form Handling"]
             )
-
-        submitted = st.form_submit_button("Generate Code")
         
-        if submitted:
+        generate = st.form_submit_button("Generate Code")
+        
+        if generate:
             try:
-                with st.spinner("Generating your application..."):
-                    result = st.session_state.rag_system.generate_code(
+                with st.spinner("Generating code..."):
+                    result = st.session_state.assistant.generate_code(
                         description,
                         {
                             "frontend": frontend,
@@ -177,17 +212,32 @@ def main():
                         features
                     )
                     
-                    # Display results
-                    tabs = st.tabs(["Frontend", "Backend", "Database", "Setup"])
+                    # Display results in tabs
+                    tabs = st.tabs([
+                        "Frontend Code",
+                        "Backend Code",
+                        "Database Setup",
+                        "Instructions"
+                    ])
                     
                     with tabs[0]:
-                        st.code(result["frontend"])
+                        st.code(result["frontend"], language="javascript")
+                    
                     with tabs[1]:
-                        st.code(result["backend"])
+                        st.code(result["backend"], language="python")
+                    
                     with tabs[2]:
-                        st.code(result["database"])
+                        st.code(result["database"], language="sql")
+                    
                     with tabs[3]:
-                        st.markdown(result["setup"])
+                        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)}")