Update app.py

app.py

@@ -5,6 +5,9 @@ import pandas as pd
 import numpy as np
 from datetime import datetime, timedelta
 from typing import Dict, List, Any
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import streamlit as st
 
 # --- Data Processing Class ---
 class DataProcessor:
@@ -356,27 +359,95 @@ def render_brainstorm_page():
         else:
             st.info("No products yet. Create one to get started!")
 
+
+class LLaMAAssistant:
+    def __init__(self, model_name="meta-llama/Llama-2-7b-chat-hf"):
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_name, 
+                torch_dtype=torch.float16, 
+                device_map="auto"
+            )
+        except Exception as e:
+            st.error(f"Model loading error: {e}")
+            self.model = None
+            self.tokenizer = None
+
+    def generate_response(self, prompt: str, context: list = None) -> str:
+        if not self.model or not self.tokenizer:
+            return "LLM not initialized. Please check model configuration."
+        
+        # Prepare conversation context
+        if context is None:
+            context = []
+        
+        # Create full prompt with conversation history
+        full_prompt = "".join([f"{msg['role']}: {msg['content']}\n" for msg in context])
+        full_prompt += f"user: {prompt}\nassistant: "
+        
+        # Tokenize input
+        input_ids = self.tokenizer(full_prompt, return_tensors="pt").input_ids.to(self.model.device)
+        
+        # Generate response
+        try:
+            output = self.model.generate(
+                input_ids, 
+                max_length=500,
+                num_return_sequences=1,
+                no_repeat_ngram_size=2,
+                temperature=0.7,
+                top_p=0.9
+            )
+            
+            # Decode response
+            response = self.tokenizer.decode(output[0], skip_special_tokens=True)
+            
+            # Extract only the new part of the response
+            response = response[len(full_prompt):].strip()
+            
+            return response
+        except Exception as e:
+            return f"Response generation error: {e}"
+
 def render_chat():
-    st.header("💬 Business Assistant")
+    st.header("💬 Business AI Assistant")
     
+    # Initialize LLaMA model (only once)
+    if 'llama_assistant' not in st.session_state:
+        st.session_state.llama_assistant = LLaMAAssistant()
+    
+    # Initialize message history
     if "messages" not in st.session_state:
         st.session_state.messages = []
     
+    # Display chat history
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
     
+    # User input
     if prompt := st.chat_input("Ask about your business..."):
+        # Add user message to chat history
         st.session_state.messages.append({"role": "user", "content": prompt})
+        
+        # Display user message
         with st.chat_message("user"):
             st.markdown(prompt)
         
-        response = f"Thank you for your question about '{prompt}'. The LLM integration will be implemented soon."
-        
+        # Generate AI response
         with st.chat_message("assistant"):
-            st.markdown(response)
+            with st.spinner("Generating response..."):
+                response = st.session_state.llama_assistant.generate_response(
+                    prompt, 
+                    st.session_state.messages[:-1]  # Exclude current message
+                )
+                st.markdown(response)
+        
+        # Add assistant response to history
         st.session_state.messages.append({"role": "assistant", "content": response})
 
+# Replace the existing render_chat() function in your main.py with this implementation
 def main():
     st.set_page_config(
         page_title="Prospira",