Update app.py

app.py

@@ -1,262 +1,196 @@
 import streamlit as st
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.document_loaders import TextLoader, DirectoryLoader
+from langchain_community.llms import CTransformers
 from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
-from langchain.chat_models import ChatOpenAI
 import os
 from pathlib import Path
-import json
 import logging
-from typing import Dict, List, Optional
 
-class WebDevRAG:
-    def __init__(self, api_key: str, model_name: str = "gpt-4"):
-        self.api_key = api_key
-        self.model_name = model_name
-        self.embeddings = HuggingFaceEmbeddings(
-            model_name="sentence-transformers/all-MiniLM-L6-v2"
-        )
+class LocalWebDevRAG:
+    def __init__(self):
         self.initialize_logging()
-        self.initialize_vector_store()
+        self.setup_embeddings()
         self.setup_llm()
-        
+        self.initialize_vector_store()
+
     def initialize_logging(self):
-        logging.basicConfig(
-            level=logging.INFO,
-            format='%(asctime)s - %(levelname)s - %(message)s'
-        )
+        logging.basicConfig(level=logging.INFO)
         self.logger = logging.getLogger(__name__)
 
+    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):
-        """Initialize FAISS vector store with web development knowledge"""
         try:
-            # Load from existing if available
-            if Path("vector_store").exists():
-                self.vector_store = FAISS.load_local(
-                    "vector_store",
-                    self.embeddings
+            # 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")
-            else:
-                self.create_new_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
 
     def create_new_vector_store(self):
-        """Create new vector store from knowledge base"""
-        loader = DirectoryLoader(
-            "./knowledge_base",
-            glob="**/*.txt",
-            loader_cls=TextLoader
-        )
-        documents = loader.load()
+        # Example code snippets and documentation
+        documents = [
+            "React component best practices...",
+            "API security implementations...",
+            "Database schema designs...",
+            # Add more code examples and documentation
+        ]
+        
         text_splitter = RecursiveCharacterTextSplitter(
             chunk_size=1000,
             chunk_overlap=200
         )
-        texts = text_splitter.split_documents(documents)
-        self.vector_store = FAISS.from_documents(texts, self.embeddings)
-        self.vector_store.save_local("vector_store")
-        self.logger.info("Created new vector store")
-
-    def setup_llm(self):
-        """Setup LLM with custom prompts"""
-        self.llm = ChatOpenAI(
-            model_name=self.model_name,
-            temperature=0.7,
-            api_key=self.api_key
-        )
-        
-        self.qa_prompt = PromptTemplate(
-            template="""You are an expert web developer AI assistant. Using the context provided, 
-            generate high-quality, production-ready code for the user's request. Include best practices, 
-            security considerations, and modern development patterns.
-
-            Context: {context}
-            
-            Question: {question}
-            
-            Provide a detailed, professional solution with explanations for key decisions.""",
-            input_variables=["context", "question"]
-        )
+        texts = text_splitter.split_text('\n\n'.join(documents))
         
-        self.qa_chain = RetrievalQA.from_chain_type(
-            llm=self.llm,
-            chain_type="stuff",
-            retriever=self.vector_store.as_retriever(
-                search_kwargs={"k": 5}
-            ),
-            chain_type_kwargs={"prompt": self.qa_prompt}
+        self.vector_store = Chroma.from_texts(
+            texts,
+            self.embeddings,
+            persist_directory="chroma_db"
         )
+        self.logger.info("Created new vector store")
 
-    def generate_code(self, 
-                     description: str, 
-                     tech_stack: Dict[str, str],
-                     requirements: List[str]) -> Dict:
-        """
-        Generate web application code based on description and requirements
-        
-        Args:
-            description: Project description
-            tech_stack: Dictionary of technology choices
-            requirements: List of specific requirements
-            
-        Returns:
-            Dictionary containing generated code and documentation
-        """
+    def generate_code(self, description, tech_stack, requirements):
         try:
-            # Construct detailed prompt
             prompt = f"""
-            Project Description: {description}
-            
-            Technology Stack:
-            {json.dumps(tech_stack, indent=2)}
-            
-            Requirements:
-            {json.dumps(requirements, indent=2)}
+            Create a web application with:
+            Description: {description}
+            Tech Stack: {tech_stack}
+            Requirements: {requirements}
             
-            Generate a complete solution including:
-            1. Frontend code (component architecture)
-            2. Backend API design
+            Provide:
+            1. Frontend components
+            2. Backend API
             3. Database schema
-            4. Security implementations
-            5. Deployment considerations
+            4. Setup instructions
             """
             
-            # Get RAG-enhanced response
             response = self.qa_chain.run(prompt)
-            
-            # Process and structure the response
             return self.process_response(response)
             
         except Exception as e:
-            self.logger.error(f"Code generation failed: {e}")
+            self.logger.error(f"Generation failed: {e}")
             raise
 
-    def process_response(self, response: str) -> Dict:
-        """Process and structure the LLM response"""
-        # Add processing logic here
+    def process_response(self, response):
+        # Basic response processing
         return {
-            "frontend": self.extract_frontend_code(response),
-            "backend": self.extract_backend_code(response),
-            "database": self.extract_database_schema(response),
-            "documentation": self.extract_documentation(response)
+            "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
         }
 
-    def extract_frontend_code(self, response: str) -> Dict:
-        # Add extraction logic
-        pass
-
-    def extract_backend_code(self, response: str) -> Dict:
-        # Add extraction logic
-        pass
-
-    def extract_database_schema(self, response: str) -> Dict:
-        # Add extraction logic
-        pass
-
-    def extract_documentation(self, response: str) -> str:
-        # Add extraction logic
-        pass
-
-def create_streamlit_interface():
-    st.set_page_config(page_title="AI Web Developer", layout="wide")
-    
-    # API key handling
-    api_key = st.secrets["OPENAI_API_KEY"]
+def main():
+    st.set_page_config(page_title="Local Web Development AI", layout="wide")
     
-    if "rag_system" not in st.session_state:
-        st.session_state.rag_system = WebDevRAG(api_key)
+    st.title("🚀 Web Development AI Assistant")
+    st.write("Generate web applications using local AI - no API key required!")
 
-    st.title("AI Web Development Assistant")
-    
-    # Project Configuration
-    with st.form("project_config"):
+    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"):
         description = st.text_area(
             "Project Description",
-            "Describe your web application in detail..."
+            placeholder="Describe your web application..."
         )
         
-        # Technology Stack Selection
         col1, col2 = st.columns(2)
         with col1:
             frontend = st.selectbox(
-                "Frontend Framework",
-                ["React", "Vue", "Angular", "Next.js"]
+                "Frontend",
+                ["React", "Vue", "Angular"]
             )
             database = st.selectbox(
                 "Database",
-                ["PostgreSQL", "MongoDB", "MySQL"]
+                ["MongoDB", "PostgreSQL", "MySQL"]
             )
-            
+        
         with col2:
             backend = st.selectbox(
-                "Backend Framework",
-                ["Node.js/Express", "Django", "FastAPI"]
+                "Backend",
+                ["Node.js", "Python/FastAPI", "Python/Django"]
             )
-            auth = st.selectbox(
-                "Authentication",
-                ["JWT", "OAuth2", "Session-based"]
+            features = st.multiselect(
+                "Features",
+                ["Authentication", "REST API", "File Upload", "Real-time Updates"]
             )
-        
-        # Additional Requirements
-        requirements = st.multiselect(
-            "Additional Requirements",
-            [
-                "REST API",
-                "GraphQL",
-                "Real-time updates",
-                "File upload",
-                "Email integration",
-                "Payment processing",
-                "Admin dashboard",
-                "Analytics",
-                "SEO optimization",
-                "Mobile responsiveness",
-                "Automated testing",
-                "CI/CD pipeline"
-            ]
-        )
-        
-        submitted = st.form_submit_button("Generate Solution")
+
+        submitted = st.form_submit_button("Generate Code")
         
         if submitted:
-            with st.spinner("Generating comprehensive solution..."):
-                tech_stack = {
-                    "frontend": frontend,
-                    "backend": backend,
-                    "database": database,
-                    "authentication": auth
-                }
-                
-                try:
+            try:
+                with st.spinner("Generating your application..."):
                     result = st.session_state.rag_system.generate_code(
                         description,
-                        tech_stack,
-                        requirements
+                        {
+                            "frontend": frontend,
+                            "backend": backend,
+                            "database": database
+                        },
+                        features
                     )
                     
-                    # Display results in organized tabs
-                    tabs = st.tabs(["Frontend", "Backend", "Database", "Documentation"])
+                    # Display results
+                    tabs = st.tabs(["Frontend", "Backend", "Database", "Setup"])
                     
                     with tabs[0]:
-                        st.code(result["frontend"], language="javascript")
-                        
+                        st.code(result["frontend"])
                     with tabs[1]:
-                        st.code(result["backend"], language="python")
-                        
+                        st.code(result["backend"])
                     with tabs[2]:
-                        st.code(result["database"], language="sql")
-                        
+                        st.code(result["database"])
                     with tabs[3]:
-                        st.markdown(result["documentation"])
-                        
-                except Exception as e:
-                    st.error(f"An error occurred: {str(e)}")
+                        st.markdown(result["setup"])
+                    
+            except Exception as e:
+                st.error(f"An error occurred: {str(e)}")
 
 if __name__ == "__main__":
-    create_streamlit_interface()
+    main()