geolocation + session tracking to log

.gitignore

@@ -1,8 +1,13 @@
-.DS_store
+.DS_Store
 .env
 /testing/
 /logs/
 logging_config.py
 /data/
 app_interactions.jsonl
-auditqa/__pycache__/
+
+# Python cache files
+__pycache__/
+*.py[cod]
+*$py.class
+.pytest_cache/

app.py

@@ -4,7 +4,7 @@ import logging
 import asyncio
 import os
 from uuid import uuid4
-from datetime import datetime
+from datetime import datetime, timedelta
 from pathlib import Path
 from huggingface_hub import CommitScheduler
 from auditqa.sample_questions import QUESTIONS
@@ -12,17 +12,20 @@ from auditqa.reports import files, report_list, new_files, new_report_list
 from auditqa.process_chunks import load_chunks, getconfig, get_local_qdrant, load_new_chunks
 from auditqa.retriever import get_context
 from auditqa.reader import nvidia_client, dedicated_endpoint
-from auditqa.utils import make_html_source, parse_output_llm_with_sources, save_logs, get_message_template
+from auditqa.utils import make_html_source, parse_output_llm_with_sources, save_logs, get_message_template, get_client_location, get_client_ip
 from dotenv import load_dotenv
 from threading import Lock
-import json
-from functools import partial
+# import json
+# from functools import partial
+# import time
+from gradio.routes import Request
 
 # TESTING DEBUG LOG
 from auditqa.logging_config import setup_logging
 setup_logging()
 import logging
 logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)  # Ensure debug logging is enabled
 
 load_dotenv()
 
@@ -99,11 +102,57 @@ def submit_feedback(feedback, logs_data):
         # Still need to return the expected outputs even on error
         return gr.update(visible=False), gr.update(visible=True)
 
-async def chat(query,history,sources,reports,subtype,year):
-    """taking a query and a message history, use a pipeline (reformulation, retriever, answering) 
-       to yield a tuple of:(messages in gradio format/messages in langchain format, source documents)
-    """
-
+# Session Manager added (track session duration & location)
+class SessionManager:
+    def __init__(self):
+        self.sessions = {}
+        
+    def create_session(self, client_ip):
+        session_id = str(uuid4())
+        self.sessions[session_id] = {
+            'start_time': datetime.now(),
+            'last_activity': datetime.now(),
+            'client_ip': client_ip,
+            'location_info': get_client_location(client_ip)
+        }
+        return session_id
+    
+    def update_session(self, session_id):
+        if session_id in self.sessions:
+            self.sessions[session_id]['last_activity'] = datetime.now()
+    
+    def get_session_duration(self, session_id):
+        if session_id in self.sessions:
+            start = self.sessions[session_id]['start_time']
+            last = self.sessions[session_id]['last_activity']
+            return (last - start).total_seconds()
+        return 0
+    
+    def get_session_data(self, session_id):
+        return self.sessions.get(session_id)
+
+# Initialize session manager
+session_manager = SessionManager()
+
+async def chat(query, history, sources, reports, subtype, year, client_ip=None, session_id=None):
+    """Update chat function to handle session data"""
+    # TESTING: DEBUG LOG
+    logger.debug(f"Chat function called with query: {query}")
+    logger.debug(f"Client IP: {client_ip}")
+    logger.debug(f"Session ID: {session_id}")
+    
+    if not session_id: # Session managment
+        session_id = session_manager.create_session(client_ip)
+        logger.debug(f"Created new session: {session_id}")
+    else:
+        session_manager.update_session(session_id)
+        logger.debug(f"Updated existing session: {session_id}")
+    
+    # Get session data
+    session_data = session_manager.get_session_data(session_id)
+    session_duration = session_manager.get_session_duration(session_id)
+    logger.debug(f"Session duration: {session_duration}")
+    
     print(f">> NEW QUESTION : {query}")
     print(f"history:{history}")
     print(f"sources:{sources}")
@@ -158,9 +207,12 @@ async def chat(query,history,sources,reports,subtype,year):
 
     ##-----------------------get answer from endpoints------------------------------
     answer_yet = ""
-    # Create logs data structure at the beginning (so that feedback can be saved after streaming
+    # Logs strcuture updated for feedback + session data (moved up here because: feedback)
     timestamp = str(datetime.now().timestamp())
     logs_data = {
+        "session_id": session_id,
+        "client_location": session_data['location_info'],
+        "session_duration_seconds": session_duration,
         "system_prompt": SYSTEM_PROMPT,
         "sources": sources,
         "reports": reports,
@@ -171,7 +223,7 @@ async def chat(query,history,sources,reports,subtype,year):
         "endpoint_type": model_config.get('reader','TYPE'),
         "reader": model_config.get('reader','NVIDIA_MODEL'),
         "docs": [doc.page_content for doc in context_retrieved],
-        "answer": "",  # Updated after streaming
+        "answer": "",
         "time": timestamp,
     }
 
@@ -196,14 +248,14 @@ async def chat(query,history,sources,reports,subtype,year):
                     history[-1] = (query, parsed_answer)
                     # Update logs_data with current answer
                     logs_data["answer"] = parsed_answer
-                    yield [tuple(x) for x in history], docs_html, logs_data
+                    yield [tuple(x) for x in history], docs_html, logs_data, session_id
 
         # Stream the response updates
         async for update in process_stream():
             yield update
 
     else:
-        chat_model = dedicated_endpoint() # TESTING: ADAPTED FOR HF INFERENCE API
+        chat_model = dedicated_endpoint() # TESTING: ADAPTED FOR HF INFERENCE API (needs to be reverted for production version)
         async def process_stream():
             nonlocal answer_yet
             try:
@@ -228,7 +280,7 @@ async def chat(query,history,sources,reports,subtype,year):
                     history[-1] = (query, parsed_answer)
                     # Update logs_data with current answer
                     logs_data["answer"] = parsed_answer
-                    yield [tuple(x) for x in history], docs_html, logs_data
+                    yield [tuple(x) for x in history], docs_html, logs_data, session_id
                     await asyncio.sleep(0.05)
                 
             except Exception as e:
@@ -482,24 +534,36 @@ with gr.Blocks(title="Audit Q&A", css= "style.css", theme=theme,elem_id = "main-
 
     #-------------------- Gradio voodoo continued -------------------------
 
-    # Using event listeners for 1. query box 2. click on example question
-    # https://www.gradio.app/docs/gradio/textbox#event-listeners-arguments
+    # Add these state components at the top level of the Blocks
+    session_id = gr.State(None)
+    client_ip = gr.State(None)
+    
+    @demo.load(api_name="get_client_ip")
+    def get_client_ip_handler(dummy_input="", request: gr.Request = None):
+        """Handler for getting client IP in Gradio context"""
+        return get_client_ip(request)
+    
+    # Update the event handlers
     (textbox
-        .submit(start_chat, [textbox, chatbot], [textbox, tabs, chatbot], queue=False, api_name="start_chat_textbox")
-        # queue must be set as False (default) so the process is not waiting for another to be finished
-        .then(chat, [textbox, chatbot, dropdown_sources, dropdown_reports, dropdown_category, dropdown_year], [chatbot, sources_textbox, feedback_state], queue=True, concurrency_limit=8, api_name="chat_textbox")
+        .submit(get_client_ip_handler, [textbox], [client_ip], api_name="get_ip_textbox")
+        .then(start_chat, [textbox, chatbot], [textbox, tabs, chatbot], queue=False, api_name="start_chat_textbox")
+        .then(chat, 
+            [textbox, chatbot, dropdown_sources, dropdown_reports, dropdown_category, dropdown_year, client_ip, session_id], 
+            [chatbot, sources_textbox, feedback_state, session_id], 
+            queue=True, concurrency_limit=8, api_name="chat_textbox")
         .then(show_feedback, [feedback_state], [feedback_row, feedback_thanks, feedback_state], api_name="show_feedback_textbox")
         .then(finish_chat, None, [textbox], api_name="finish_chat_textbox"))
 
     (examples_hidden
         .change(start_chat, [examples_hidden, chatbot], [textbox, tabs, chatbot], queue=False, api_name="start_chat_examples")
-        # queue must be set as False (default) so the process is not waiting for another to be finished
-        .then(chat, [examples_hidden, chatbot, dropdown_sources, dropdown_reports, dropdown_category, dropdown_year], [chatbot, sources_textbox, feedback_state], queue=True, concurrency_limit=8, api_name="chat_examples")
+        .then(get_client_ip_handler, [examples_hidden], [client_ip], api_name="get_ip_examples")
+        .then(chat, 
+            [examples_hidden, chatbot, dropdown_sources, dropdown_reports, dropdown_category, dropdown_year, client_ip, session_id], 
+            [chatbot, sources_textbox, feedback_state, session_id], 
+            concurrency_limit=8, api_name="chat_examples")
         .then(show_feedback, [feedback_state], [feedback_row, feedback_thanks, feedback_state], api_name="show_feedback_examples")
         .then(finish_chat, None, [textbox], api_name="finish_chat_examples"))
 
-
-
     demo.queue()
 
 demo.launch()

auditqa/utils.py

@@ -5,6 +5,12 @@ from langchain.schema import (
     HumanMessage,
     SystemMessage,
 )
+import requests
+from datetime import datetime
+from uuid import uuid4
+
+
+
 
 def save_logs(scheduler, JSON_DATASET_PATH, logs, feedback=None) -> None:
     """ Every interaction with app saves the log of question and answer, 
@@ -73,4 +79,55 @@ def parse_output_llm_with_sources(output):
         else:
             parts.append(part)
     content_parts = "".join(parts)
-    return content_parts
+    return content_parts
+
+def get_client_ip(request=None):
+    """Get the client IP address from the request context"""
+    try:
+        if request:
+            # Try different headers that might contain the real IP
+            ip = request.client.host
+            # Check for proxy headers
+            forwarded_for = request.headers.get('X-Forwarded-For')
+            if forwarded_for:
+                # X-Forwarded-For can contain multiple IPs - first one is the client
+                ip = forwarded_for.split(',')[0].strip()
+            
+            logging.debug(f"Client IP detected: {ip}")
+            return ip
+    except Exception as e:
+        logging.error(f"Error getting client IP: {e}")
+    return "127.0.0.1"
+
+
+def get_client_location(ip_address) -> dict | None:
+    """Get geolocation info using ipapi.co"""
+    # Add headers so we don't get blocked...
+    headers = {
+        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
+    }
+    try:
+        response = requests.get(
+            f'https://ipapi.co/{ip_address}/json/',
+            headers=headers,
+            timeout=5
+        )
+        if response.status_code == 200:
+            data = response.json()
+            return {
+                'city': data.get('city'),
+                'region': data.get('region'),
+                'country': data.get('country_name'),
+                'latitude': data.get('latitude'),
+                'longitude': data.get('longitude')
+            }
+        elif response.status_code == 429:
+            logging.warning(f"Rate limit exceeded. Response: {response.text}")
+            return None
+        else:
+            logging.error(f"Error: Status code {response.status_code}. Response: {response.text}")
+            return None
+            
+    except requests.exceptions.RequestException as e:
+        logging.error(f"Request failed: {str(e)}")
+        return None