Update app.py

app.py

@@ -12,6 +12,8 @@ from datetime import datetime
 from huggingface_hub import hf_hub_download, HfApi, CommitOperationAdd
 from pathlib import Path
 import tempfile
+import spaces  # for ZeroGPU
+import requests  # for IP geolocation
 
 # Load environment variables and initialize clients
 load_dotenv()
@@ -74,9 +76,19 @@ else:
     # Load from Hugging Face dataset
     db = load_chroma_db()
 
+@spaces.GPU(memory="40g")  # Add GPU decorator for initialize_system
 def initialize_system():
     """Initialize the system components"""
     
+    # Add GPU diagnostics
+    print("\n=== GPU Diagnostics ===")
+    print(f"CUDA available: {torch.cuda.is_available()}")
+    if torch.cuda.is_available():
+        print(f"Current CUDA device: {torch.cuda.current_device()}")
+        print(f"Device name: {torch.cuda.get_device_name()}")
+        print(f"Device memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
+    print("=====================\n")
+    
     # Use the same ChromaDB client that was loaded from HF
     chroma_client = db  # Use the global db instance we created
     
@@ -91,12 +103,16 @@ def initialize_system():
     collection = chroma_client.get_collection(name="website_content", embedding_function=embedding_function)
     print(f"Found {collection.count()} documents in collection")
     
-    # Initialize the reranker
+    # Initialize the reranker and explicitly move to GPU if available
     print("\nInitialising Cross-Encoder...")
     reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', device=DEVICE)
+    if torch.cuda.is_available():
+        reranker.model.to('cuda')  # Ensure model is on GPU
+        print("Reranker moved to GPU")
     
     return chroma_client, collection, reranker
 
+@spaces.GPU(memory="40g")  # Add GPU decorator for get_context
 def get_context(message):
     results = collection.query(
         query_texts=[message],
@@ -165,8 +181,28 @@ def get_context(message):
     print(f"\nFinal context length: {total_chars} characters")
     return context
 
+def get_ip_info(ip_address):
+    """Get geolocation info for an IP address"""
+    if ip_address in ['127.0.0.1', 'localhost', '0.0.0.0']:
+        return {"country": "Local", "city": "Local"}
+    try:
+        response = requests.get(f'https://ipapi.co/{ip_address}/json/')
+        if response.status_code == 200:
+            data = response.json()
+            return {
+                "country": data.get("country_name", "Unknown"),
+                "city": data.get("city", "Unknown"),
+                "region": data.get("region", "Unknown")
+            }
+    except Exception as e:
+        print(f"Error getting IP info: {str(e)}")
+    return {"country": "Unknown", "city": "Unknown"}
+
 def log_conversation(timestamp, user_message, assistant_response, model_name, context, error=None, client_ip=None):
     """Log conversation details to JSON file - local directory or HuggingFace Dataset repository"""
+    # Get IP geolocation
+    ip_info = get_ip_info(client_ip) if client_ip else {"country": "Unknown", "city": "Unknown"}
+    
     # Create a log entry
     log_entry = {
         "timestamp": timestamp,
@@ -175,7 +211,8 @@ def log_conversation(timestamp, user_message, assistant_response, model_name, co
         "assistant_response": assistant_response,
         "context": context,
         "error": str(error) if error else None,
-        "client_ip": client_ip
+        "client_ip": client_ip,
+        "location": ip_info
     }
     
     # Check if running on Hugging Face Spaces
@@ -285,9 +322,7 @@ def chat_response(message, history, model_name, request: gr.Request):
         # Add history first without context
         if history:
             for h in history:
-                messages.append({"role": "user", "content": str(h[0])})
-                if h[1]:  # If there's a response
-                    messages.append({"role": "assistant", "content": str(h[1])})
+                messages.append({"role": h["role"], "content": str(h["content"])})
         
         # Add current message
         messages.append({"role": "user", "content": str(message)})
@@ -408,7 +443,7 @@ if __name__ == "__main__":
                 info="Please try out the other AI models to use for responses (all LLMs are running on [GroqCloud](https://groq.com/groqrack/)) 😊"
             )
             
-            chatbot = gr.Chatbot(height=600)
+            chatbot = gr.Chatbot(height=600, type="messages")
             with gr.Row(equal_height=True):
                 msg = gr.Textbox(
                     placeholder="What would you like to know? Or choose an example question...❓",
@@ -431,13 +466,17 @@ if __name__ == "__main__":
             )
             
             def user(user_message, history):
-                return "", history + [[user_message, None]]
+                history = history or []
+                return "", history + [{"role": "user", "content": user_message}]
             
             def bot(history, model_name, request: gr.Request):
-                if history and history[-1][1] is None:
-                    for response in chat_response(history[-1][0], history[:-1], model_name, request):
-                        history[-1][1] = response
-                        yield history
+                history = history or []
+                if history and history[-1]["role"] == "user":
+                    user_message = history[-1]["content"]
+                    history_without_last = history[:-1]
+                    for response in chat_response(user_message, history_without_last, model_name, request):
+                        history_with_response = history + [{"role": "assistant", "content": response}]
+                        yield history_with_response
             
             msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
                 bot, [chatbot, model_selector], chatbot
@@ -446,7 +485,7 @@ if __name__ == "__main__":
                 bot, [chatbot, model_selector], chatbot
             )
             
-            clear.click(lambda: None, None, chatbot, queue=False)
+            clear.click(lambda: [], None, chatbot, queue=False)  # Updated to return empty list
             gr.Markdown("<footer style='text-align: center; margin-top: 5px;'>🤖 AI-generated content; while the Chat Assistant strives for accuracy, errors may occur; please thoroughly check critical information 🤖<br>⚠️ <strong><u>Disclaimer: This system was not produced by the British Antarctic Survey (BAS) and AI generated output does not reflect the views or opinions of BAS</u></strong> ⚠️ <br>(just a bit of fun :D)</footer>")
         demo.launch(
             server_name="0.0.0.0",